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

@@ -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
 

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

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)

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.

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.
 

src/anemoi/graphs/nodes/__init__.py

@@ -21,13 +21,15 @@
 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",
     "NPZFileNodes",
     "TriNodes",
     "HexNodes",
     "HEALPixNodes",
+    "LatLonNodes",
     "LimitedAreaHEALPixNodes",
     "LimitedAreaNPZFileNodes",
     "LimitedAreaTriNodes",

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)

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]