From cd46c2b1a7dab56a6e876859b7e9d2e11db2a133 Mon Sep 17 00:00:00 2001 From: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com> Date: Thu, 19 Dec 2024 09:39:37 +0100 Subject: [PATCH] feat: support for defining nodes from lat-lon arrays (#98) * feat: add transform and LatLonNodes * docs: add docstrings * feat(tests): add test for LatLonNodes * fix: update changelog * fix: change uppercase file extensions to lowercase --- CHANGELOG.md | 11 ++- docs/graphs/node_coordinates.rst | 2 + .../graphs/node_coordinates/latlon_arrays.rst | 18 +++++ docs/graphs/node_coordinates/text_file.rst | 20 ++++++ src/anemoi/graphs/generate/transforms.py | 36 ++++++++++ src/anemoi/graphs/nodes/__init__.py | 2 + .../graphs/nodes/builders/from_vectors.py | 67 +++++++++++++++++++ tests/nodes/test_arrays.py | 64 ++++++++++++++++++ 8 files changed, 217 insertions(+), 3 deletions(-) create mode 100644 docs/graphs/node_coordinates/latlon_arrays.rst create mode 100644 docs/graphs/node_coordinates/text_file.rst create mode 100644 src/anemoi/graphs/nodes/builders/from_vectors.py create mode 100644 tests/nodes/test_arrays.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f61f7cf..02a96e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,14 @@ Keep it human-readable, your future self will thank you! ## [Unreleased](https://github.com/ecmwf/anemoi-graphs/compare/0.4.1...HEAD) -# Changed +### Added + +- feat: Support for providing lon/lat coordinates from a text file (loaded with numpy loadtxt method) to build the graph `TextNodes` (#93) +- feat: Build 2D graphs with `Voronoi` in case `SphericalVoronoi` does not work well/is an overkill (LAM). Set `flat=true` in the nodes attributes to compute area weight using Voronoi with a qhull options preventing the empty region creation (#93)  +- feat: Support for defining nodes from lat& lon NumPy arrays (#98) +- feat: new transform functions to map from sin&cos values to latlon (#98) + +### Changed - fix: faster edge builder for tri icosahedron. (#92) ## [0.4.1 - ICON graphs, multiple edge builders and post processors](https://github.com/ecmwf/anemoi-graphs/compare/0.4.0...0.4.1) - 2024-11-26 @@ -22,8 +29,6 @@ Keep it human-readable, your future self will thank you! - feat: Add `RemoveUnconnectedNodes` post processor to clean unconnected nodes in LAM. (#71) - feat: Define node sets and edges based on an ICON icosahedral mesh (#53) - feat: Support for multiple edge builders between two sets of nodes (#70) -- feat: Support for providing lon/lat coordinates from a text file (loaded with numpy loadtxt method) to build the graph `TextNodes` (#93) -- feat: Build 2D graphs with `Voronoi` in case `SphericalVoronoi` does not work well/is an overkill (LAM). Set `flat=true` in the nodes attributes to compute area weight using Voronoi with a qhull options preventing the empty region creation (#93) # Changed diff --git a/docs/graphs/node_coordinates.rst b/docs/graphs/node_coordinates.rst index 76973d2..d120a87 100644 --- a/docs/graphs/node_coordinates.rst +++ b/docs/graphs/node_coordinates.rst @@ -25,6 +25,8 @@ a file: node_coordinates/zarr_dataset node_coordinates/npz_file node_coordinates/icon_mesh + node_coordinates/text_file + node_coordinates/latlon_arrays or based on other algorithms. A commonn approach is to use an icosahedron to project the earth's surface, and refine it iteratively to diff --git a/docs/graphs/node_coordinates/latlon_arrays.rst b/docs/graphs/node_coordinates/latlon_arrays.rst new file mode 100644 index 0000000..d21d036 --- /dev/null +++ b/docs/graphs/node_coordinates/latlon_arrays.rst @@ -0,0 +1,18 @@ +####################################### + From latitude & longitude coordinates +####################################### + +Nodes can also be created directly using latitude and longitude +coordinates. Below is an example demonstrating how to add these nodes to +a graph: + +.. code:: python + + from anemoi.graphs.nodes import LatLonNodes + + ... + + lats = np.array([45.0, 45.0, 40.0, 40.0]) + lons = np.array([5.0, 10.0, 10.0, 5.0]) + + graph = LatLonNodes(latitudes=lats, longitudes=lons, name="my_nodes").update_graph(graph) diff --git a/docs/graphs/node_coordinates/text_file.rst b/docs/graphs/node_coordinates/text_file.rst new file mode 100644 index 0000000..c83710c --- /dev/null +++ b/docs/graphs/node_coordinates/text_file.rst @@ -0,0 +1,20 @@ +################ + From text file +################ + +To define the `node coordinates` based on a `.txt` file, you can +configure the `.yaml` as follows: + +.. code:: yaml + + nodes: + data: # name of the nodes + node_builder: + _target_: anemoi.graphs.nodes.TextNodes + dataset: my_file.txt + idx_lon: 0 + idx_lat: 1 + +Here, dataset refers to the path of the `.txt` file that contains the +latitude and longitude values in the columns specified by `idx_lat` and +`idx_lon`, respectively. diff --git a/src/anemoi/graphs/generate/transforms.py b/src/anemoi/graphs/generate/transforms.py index 9392241..094f9c9 100644 --- a/src/anemoi/graphs/generate/transforms.py +++ b/src/anemoi/graphs/generate/transforms.py @@ -52,6 +52,42 @@ def cartesian_to_latlon_rad(xyz: np.ndarray) -> np.ndarray: return np.array((lat, lon), dtype=np.float32).transpose() +def sincos_to_latlon_rad(sincos: np.ndarray) -> np.ndarray: + """Sine & cosine components to lat-lon coordinates. + + Parameters + ---------- + sincos : np.ndarray + The sine and cosine componenets of the latitude and longitude. Shape: (N, 4). + The dimensions correspond to: sin(lat), cos(lat), sin(lon) and cos(lon). + + Returns + ------- + np.ndarray + A 2D array of the coordinates of shape (N, 2) in radians. + """ + latitudes = np.arctan2(sincos[:, 0], sincos[:, 1]) + longitudes = np.arctan2(sincos[:, 2], sincos[:, 3]) + return np.stack([latitudes, longitudes], axis=-1) + + +def sincos_to_latlon_degrees(sincos: np.ndarray) -> np.ndarray: + """Sine & cosine components to lat-lon coordinates. + + Parameters + ---------- + sincos : np.ndarray + The sine and cosine componenets of the latitude and longitude. Shape: (N, 4). + The dimensions correspond to: sin(lat), cos(lat), sin(lon) and cos(lon). + + Returns + ------- + np.ndarray + A 2D array of the coordinates of shape (N, 2) in degrees. + """ + return np.rad2deg(sincos_to_latlon_rad(sincos)) + + def latlon_rad_to_cartesian(loc: tuple[np.ndarray, np.ndarray], radius: float = 1) -> np.ndarray: """Convert planar coordinates to 3D coordinates in a sphere. diff --git a/src/anemoi/graphs/nodes/__init__.py b/src/anemoi/graphs/nodes/__init__.py index 6fc429b..0a9eaf9 100644 --- a/src/anemoi/graphs/nodes/__init__.py +++ b/src/anemoi/graphs/nodes/__init__.py @@ -21,6 +21,7 @@ from .builders.from_refined_icosahedron import LimitedAreaTriNodes from .builders.from_refined_icosahedron import StretchedTriNodes from .builders.from_refined_icosahedron import TriNodes +from .builders.from_vectors import LatLonNodes __all__ = [ "ZarrDatasetNodes", @@ -28,6 +29,7 @@ "TriNodes", "HexNodes", "HEALPixNodes", + "LatLonNodes", "LimitedAreaHEALPixNodes", "LimitedAreaNPZFileNodes", "LimitedAreaTriNodes", diff --git a/src/anemoi/graphs/nodes/builders/from_vectors.py b/src/anemoi/graphs/nodes/builders/from_vectors.py new file mode 100644 index 0000000..f58767c --- /dev/null +++ b/src/anemoi/graphs/nodes/builders/from_vectors.py @@ -0,0 +1,67 @@ +# (C) Copyright 2024 Anemoi contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +from __future__ import annotations + +import logging + +import numpy as np +import torch + +from anemoi.graphs.nodes.builders.base import BaseNodeBuilder + +LOGGER = logging.getLogger(__name__) + + +class LatLonNodes(BaseNodeBuilder): + """Nodes from its latitude and longitude positions (in numpy arrays). + + Attributes + ---------- + latitudes : list | np.ndarray + The latitude of the nodes, in degrees. + longitudes : list | np.ndarray + The longitude of the nodes, in degrees. + + Methods + ------- + get_coordinates() + Get the lat-lon coordinates of the nodes. + register_nodes(graph, name) + Register the nodes in the graph. + register_attributes(graph, name, config) + Register the attributes in the nodes of the graph specified. + update_graph(graph, name, attrs_config) + Update the graph with new nodes and attributes. + """ + + def __init__(self, latitudes: list[float] | np.ndarray, longitudes: list[float] | np.ndarray, name: str) -> None: + super().__init__(name) + self.latitudes = latitudes if isinstance(latitudes, np.ndarray) else np.array(latitudes) + self.longitudes = longitudes if isinstance(longitudes, np.ndarray) else np.array(longitudes) + + assert len(self.latitudes) == len( + self.longitudes + ), f"Lenght of latitudes and longitudes must match but {len(self.latitudes)}!={len(self.longitudes)}." + assert self.latitudes.ndim == 1 or ( + self.latitudes.ndim == 2 and self.latitudes.shape[1] == 1 + ), "latitudes must have shape (N, ) or (N, 1)." + assert self.longitudes.ndim == 1 or ( + self.longitudes.ndim == 2 and self.longitudes.shape[1] == 1 + ), "longitudes must have shape (N, ) or (N, 1)." + + def get_coordinates(self) -> torch.Tensor: + """Get the coordinates of the nodes. + + Returns + ------- + torch.Tensor of shape (num_nodes, 2) + A 2D tensor with the coordinates, in radians. + """ + return self.reshape_coords(self.latitudes, self.longitudes) diff --git a/tests/nodes/test_arrays.py b/tests/nodes/test_arrays.py new file mode 100644 index 0000000..77e6f20 --- /dev/null +++ b/tests/nodes/test_arrays.py @@ -0,0 +1,64 @@ +# (C) Copyright 2024 Anemoi contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +import pytest +import torch +from torch_geometric.data import HeteroData + +from anemoi.graphs.nodes.attributes import AreaWeights +from anemoi.graphs.nodes.attributes import UniformWeights +from anemoi.graphs.nodes.builders.from_vectors import LatLonNodes + +lats = [45.0, 45.0, 40.0, 40.0] +lons = [5.0, 10.0, 10.0, 5.0] + + +def test_init(): + """Test LatLonNodes initialization.""" + node_builder = LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + assert isinstance(node_builder, LatLonNodes) + + +def test_fail_init_length_mismatch(): + """Test LatLonNodes initialization with invalid argument.""" + lons = [5.0, 10.0, 10.0, 5.0, 5.0] + + with pytest.raises(AssertionError): + LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + + +def test_fail_init_missing_argument(): + """Test NPZFileNodes initialization with missing argument.""" + with pytest.raises(TypeError): + LatLonNodes(name="test_nodes") + + +def test_register_nodes(): + """Test LatLonNodes register correctly the nodes.""" + graph = HeteroData() + node_builder = LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + graph = node_builder.register_nodes(graph) + + assert graph["test_nodes"].x is not None + assert isinstance(graph["test_nodes"].x, torch.Tensor) + assert graph["test_nodes"].x.shape == (len(lats), 2) + assert graph["test_nodes"].node_type == "LatLonNodes" + + +@pytest.mark.parametrize("attr_class", [UniformWeights, AreaWeights]) +def test_register_attributes(graph_with_nodes: HeteroData, attr_class): + """Test LatLonNodes register correctly the weights.""" + node_builder = LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + config = {"test_attr": {"_target_": f"anemoi.graphs.nodes.attributes.{attr_class.__name__}"}} + + graph = node_builder.register_attributes(graph_with_nodes, config) + + assert graph["test_nodes"]["test_attr"] is not None + assert isinstance(graph["test_nodes"]["test_attr"], torch.Tensor) + assert graph["test_nodes"]["test_attr"].shape[0] == graph["test_nodes"].x.shape[0]