chore: Merge Release 0.4.2 pull request #97 from ecmwf/develop

Release 0.4.2

.pre-commit-config.yaml

@@ -40,7 +40,7 @@ repos:
     - --force-single-line-imports
     - --profile black
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.7.2
+  rev: v0.8.1
   hooks:
   - id: ruff
     args:
@@ -64,7 +64,7 @@ repos:
   hooks:
   - id: pyproject-fmt
 -   repo: https://github.com/jshwi/docsig # Check docstrings against function sig
-    rev: v0.64.0
+    rev: v0.65.0
     hooks:
     -   id: docsig
         args:
@@ -74,6 +74,5 @@ repos:
         - --check-protected  # Check protected methods
         - --check-class      # Check class docstrings
         - --disable=E113     # Disable empty docstrings
-        - --summary          # Print a summary
 ci:
   autoupdate_schedule: monthly

CHANGELOG.md

@@ -10,6 +10,16 @@ Keep it human-readable, your future self will thank you!
 
 ## [Unreleased](https://github.com/ecmwf/anemoi-graphs/compare/0.4.1...HEAD)
 
+### 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
 
 ### Added

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/generate/tri_icosahedron.py

@@ -183,9 +183,8 @@ def add_edges_to_nx_graph(
         node_neighbours = get_neighbours_within_hops(r_sphere, x_hops, valid_nodes=valid_nodes)
 
         _, vertex_mapping_index = tree.query(r_vertices_rad, k=1)
-        for idx_node, idx_neighbours in node_neighbours.items():
-            graph = add_neigbours_edges(graph, idx_node, idx_neighbours, vertex_mapping_index=vertex_mapping_index)
-
+        neighbour_pairs = create_node_neighbours_list(graph, node_neighbours, vertex_mapping_index)
+        graph.add_edges_from(neighbour_pairs)
     return graph
 
 
@@ -270,3 +269,42 @@ def add_neigbours_edges(
             graph.add_edge(node_neighbour, node)
 
     return graph
+
+
+def create_node_neighbours_list(
+    graph: nx.Graph,
+    node_neighbours: dict[int, set[int]],
+    vertex_mapping_index: np.ndarray | None = None,
+    self_loops: bool = False,
+) -> list[tuple]:
+    """Preprocesses the dict of node neighbours.
+
+    Parameters:
+    -----------
+    graph: nx.Graph
+        The graph.
+    node_neighbours: dict[int, set[int]]
+        dictionairy with key: node index and value: set of neighbour node indices
+    vertex_mapping_index: np.ndarry
+        Index to map the vertices from the refined sphere to the original one, by default None.
+    self_loops: bool
+        Whether is supported to add self-loops, by default False.
+
+    Returns:
+    --------
+    list: tuple
+        A list with containing node neighbour pairs in tuples
+    """
+    graph_nodes_idx = list(sorted(graph.nodes))
+
+    if vertex_mapping_index is None:
+        vertex_mapping_index = np.arange(len(graph.nodes)).reshape(len(graph.nodes), 1)
+
+    neighbour_pairs = [
+        (graph_nodes_idx[vertex_mapping_index[node_neighbour][0]], graph_nodes_idx[vertex_mapping_index[node][0]])
+        for node, neighbours in node_neighbours.items()
+        for node_neighbour in neighbours
+        if node != node_neighbour or (self_loops and node == node_neighbour)
+    ]
+
+    return neighbour_pairs

src/anemoi/graphs/nodes/__init__.py

@@ -9,6 +9,7 @@
 
 from .builders.from_file import LimitedAreaNPZFileNodes
 from .builders.from_file import NPZFileNodes
+from .builders.from_file import TextNodes
 from .builders.from_file import ZarrDatasetNodes
 from .builders.from_healpix import HEALPixNodes
 from .builders.from_healpix import LimitedAreaHEALPixNodes
@@ -20,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",
@@ -35,4 +38,5 @@
     "ICONMultimeshNodes",
     "ICONCellGridNodes",
     "ICONNodes",
+    "TextNodes",
 ]

src/anemoi/graphs/nodes/attributes.py

@@ -18,7 +18,9 @@
 import numpy as np
 import torch
 from anemoi.datasets import open_dataset
+from scipy.spatial import ConvexHull
 from scipy.spatial import SphericalVoronoi
+from scipy.spatial import Voronoi
 from torch_geometric.data import HeteroData
 from torch_geometric.data.storage import NodeStorage
 
@@ -101,6 +103,68 @@ def get_raw_values(self, nodes: NodeStorage, **kwargs) -> np.ndarray:
 class AreaWeights(BaseNodeAttribute):
     """Implements the area of the nodes as the weights.
 
+    Attributes
+    ----------
+    flat: bool
+        If True, the area is computed in 2D, otherwise in 3D.
+    **other: Any
+        Additional keyword arguments, see PlanarAreaWeights and SphericalAreaWeights
+        for details.
+
+    Methods
+    -------
+    compute(self, graph, nodes_name)
+        Compute the area attributes for each node.
+    """
+
+    def __new__(cls, flat: bool = False, **kwargs):
+        logging.warning(
+            "Creating %s with flat=%s and kwargs=%s. In a future release, AreaWeights will be deprecated: please use directly PlanarAreaWeights or SphericalAreaWeights.",
+            cls.__name__,
+            flat,
+            kwargs,
+        )
+        if flat:
+            return PlanarAreaWeights(**kwargs)
+        return SphericalAreaWeights(**kwargs)
+
+
+class PlanarAreaWeights(BaseNodeAttribute):
+    """Implements the 2D area of the nodes as the weights.
+
+    Attributes
+    ----------
+    norm : str
+        Normalisation of the weights.
+
+    Methods
+    -------
+    compute(self, graph, nodes_name)
+        Compute the area attributes for each node.
+    """
+
+    def __init__(
+        self,
+        norm: str | None = None,
+        dtype: str = "float32",
+    ) -> None:
+        super().__init__(norm, dtype)
+
+    def get_raw_values(self, nodes: NodeStorage, **kwargs) -> np.ndarray:
+        latitudes, longitudes = nodes.x[:, 0], nodes.x[:, 1]
+        points = np.stack([latitudes, longitudes], -1)
+        v = Voronoi(points, qhull_options="QJ Pp")
+        areas = []
+        for r in v.regions:
+            area = ConvexHull(v.vertices[r, :]).volume
+            areas.append(area)
+        result = np.asarray(areas)
+        return result
+
+
+class SphericalAreaWeights(BaseNodeAttribute):
+    """Implements the 3D area of the nodes as the weights.
+
     Attributes
     ----------
     norm : str
@@ -132,22 +196,6 @@ def __init__(
         self.fill_value = fill_value
 
     def get_raw_values(self, nodes: NodeStorage, **kwargs) -> np.ndarray:
-        """Compute the area associated to each node.
-
-        It uses Voronoi diagrams to compute the area of each node.
-
-        Parameters
-        ----------
-        nodes : NodeStorage
-            Nodes of the graph.
-        kwargs : dict
-            Additional keyword arguments.
-
-        Returns
-        -------
-        np.ndarray
-            Attributes.
-        """
         latitudes, longitudes = nodes.x[:, 0], nodes.x[:, 1]
         points = latlon_rad_to_cartesian((np.asarray(latitudes), np.asarray(longitudes)))
         sv = SphericalVoronoi(points, self.radius, self.centre)

src/anemoi/graphs/nodes/builders/from_file.py

@@ -63,6 +63,37 @@ def get_coordinates(self) -> torch.Tensor:
         return self.reshape_coords(dataset.latitudes, dataset.longitudes)
 
 
+class TextNodes(BaseNodeBuilder):
+    """Nodes from text file.
+
+    Attributes
+    ----------
+    dataset : str | DictConfig
+        The path to txt file containing the coordinates of the nodes.
+    idx_lon : int
+        The index of the longitude in the dataset.
+    idx_lat : int
+        The index of the latitude in the dataset.
+    """
+
+    def __init__(self, dataset, name: str, idx_lon: int = 0, idx_lat: int = 1) -> None:
+        LOGGER.info("Reading the dataset from %s.", dataset)
+        self.dataset = np.loadtxt(dataset)
+        self.idx_lon = idx_lon
+        self.idx_lat = idx_lat
+        super().__init__(name)
+
+    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.dataset[self.idx_lat, :], self.dataset[self.idx_lon, :])
+
+
 class NPZFileNodes(BaseNodeBuilder):
     """Nodes from NPZ defined grids.
 
@@ -146,7 +177,10 @@ def get_coordinates(self) -> np.ndarray:
         )
         area_mask = self.area_mask_builder.get_mask(coords)
 
-        LOGGER.info("Dropping %d nodes from the processor mesh.", len(area_mask) - area_mask.sum())
+        LOGGER.info(
+            "Dropping %d nodes from the processor mesh.",
+            len(area_mask) - area_mask.sum(),
+        )
         coords = coords[area_mask]
 
         return coords