feature: support icon graphs (#53)

* feat: topology-based encoder/processor/decoder graphs derived from an ICON grid file. * docs: attempt to set a correct license header. * doc: changed the docstring style to the Numpy-style instead of Google-style. * refactor: changed variable name `verts` into `vertices`. * fix: changed float division to double slash operator. * refactor: removed unnecessary `else` branch. * refactor: more descriptive variable name in for loop. * refactor: renamed variable (ignoring the minus sign of `phi`). * refactor: changed name of temporary variable. * refactor: remove unnecessary `else´branch. * refactor: added type annotation for completeness. * refactor: remove default in function argument. * refactor: change argument name to a more understandable name. * refactor: remove redundant code. * refactor: more Pythonic if-else statement. * refactor: more Pythonic if-else statement. * refactor: use more appropriate `LOGGER.debug` instead of verbosity flag. * refactor: more appropriate variable name. * refactor: more appropriate variable name. * refactor: unified the three ICON Edgebuilders and the two Nodebuilders. * refactor: more verbose but also clearer names for variables in mesh construction algorithm. * remove: removed obsolete function `set_constant_edge_id`. * refactor: replaced the sequential ID counter by a UUID. * revert change of copyright notice * [fix] add encoder & processor edges to the __all__ variable * [refactor] move auxiliary functions to utils.py in graphs/generate/ * [doc] added empty torso for class documentation. * [fix] fixed interfaces (masks). * [refactor] remove edge attribute calculation from this PR. * [chore] adjust copyright notice. * [doc] elaborate on icon mesh classes in rst file. * Add Icon tests * update Change log * Add __future__ annotations * fix change log --------- Co-authored-by: Marek Jacob <[email protected]> Co-authored-by: Florian Prill <[email protected]>
ecmwf · Nov 13, 2024 · 6cfcade · 6cfcade
1 parent 30f039f
commit 6cfcade
Show file tree

Hide file tree

Showing 10 changed files with 968 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ Please add your functional changes to the appropriate section in the PR.
 Keep it human-readable, your future self will thank you!
 
 ## [Unreleased](https://github.com/ecmwf/anemoi-graphs/compare/0.4.0...HEAD)
+- feat: Define node sets and edges based on an ICON icosahedral mesh (#53)
 
 ## [0.4.0 - LAM and stretched graphs](https://github.com/ecmwf/anemoi-graphs/compare/0.3.0...0.4.0) - 2024-11-08
 

diff --git a/docs/graphs/node_coordinates/icon_mesh.rst b/docs/graphs/node_coordinates/icon_mesh.rst
@@ -0,0 +1,64 @@
+####################################
+ Triangular Mesh with ICON Topology
+####################################
+
+The classes `ICONMultimeshNodes` and `ICONCellGridNodes` define node
+sets based on an ICON icosahedral mesh:
+
+-  class `ICONCellGridNodes`: data grid, representing cell circumcenters
+-  class `ICONMultimeshNodes`: hidden mesh, representing the vertices of
+   a grid hierarchy
+
+Both classes, together with the corresponding edge builders
+
+-  class `ICONTopologicalProcessorEdges`
+-  class `ICONTopologicalEncoderEdges`
+-  class `ICONTopologicalDecoderEdges`
+
+are based on the mesh hierarchy that is reconstructed from an ICON mesh
+file in NetCDF format, making use of the `refinement_level_v` and
+`refinement_level_c` property contained therein.
+
+-  `refinement_level_v[vertex] = 0,1,2, ...`,
+      where 0 denotes the vertices of the base grid, ie. the icosahedron
+      including the step of root subdivision RXXB00.
+
+-  `refinement_level_c[cell]`: cell refinement level index such that
+   value 0 denotes the cells of the base grid, ie. the icosahedron
+   including the step of root subdivision RXXB00.
+
+To avoid multiple runs of the reconstruction algorithm, a separate
+`ICONNodes` instance is created and used by the builders, see the
+following YAML example:
+
+.. code:: yaml
+
+   nodes:
+     # ICON mesh
+     icon_mesh:
+       node_builder:
+         _target_: anemoi.graphs.nodes.ICONNodes
+         name: "icon_grid_0026_R03B07_G"
+         grid_filename: "icon_grid_0026_R03B07_G.nc"
+         max_level_multimesh: 3
+         max_level_dataset: 3
+     # Data nodes
+     data:
+       node_builder:
+         _target_: anemoi.graphs.nodes.ICONCellGridNodes
+         icon_mesh: "icon_mesh"
+       attributes: ${graph.attributes.nodes}
+     # Hidden nodes
+     hidden:
+       node_builder:
+         _target_: anemoi.graphs.nodes.ICONMultimeshNodes
+         icon_mesh: "icon_mesh"
+
+   edges:
+     # Processor configuration
+     - source_name: ${graph.hidden}
+       target_name: ${graph.hidden}
+       edge_builder:
+         _target_: anemoi.graphs.edges.ICONTopologicalProcessorEdges
+         icon_mesh: "icon_mesh"
+       attributes: ${graph.attributes.edges}
diff --git a/pyproject.toml b/pyproject.toml
@@ -45,11 +45,13 @@ dependencies = [
   "healpy>=1.17",
   "hydra-core>=1.3",
   "matplotlib>=3.4",
+  "netcdf4",
   "networkx>=3.1",
   "plotly>=5.19",
   "torch>=2.2",
   "torch-geometric>=2.3.1,<2.5",
   "trimesh>=4.1",
+  "typeguard",
 ]
 
 optional-dependencies.all = [  ]

diff --git a/src/anemoi/graphs/edges/__init__.py b/src/anemoi/graphs/edges/__init__.py
@@ -8,7 +8,17 @@
 # nor does it submit to any jurisdiction.
 
 from .builder import CutOffEdges
+from .builder import ICONTopologicalDecoderEdges
+from .builder import ICONTopologicalEncoderEdges
+from .builder import ICONTopologicalProcessorEdges
 from .builder import KNNEdges
 from .builder import MultiScaleEdges
 
-__all__ = ["KNNEdges", "CutOffEdges", "MultiScaleEdges"]
+__all__ = [
+    "KNNEdges",
+    "CutOffEdges",
+    "MultiScaleEdges",
+    "ICONTopologicalProcessorEdges",
+    "ICONTopologicalEncoderEdges",
+    "ICONTopologicalDecoderEdges",
+]
diff --git a/src/anemoi/graphs/edges/builder.py b/src/anemoi/graphs/edges/builder.py
@@ -15,6 +15,7 @@
 
 import networkx as nx
 import numpy as np
+import scipy
 import torch
 from anemoi.utils.config import DotDict
 from hydra.utils import instantiate
@@ -80,10 +81,8 @@ def get_edge_index(self, graph: HeteroData) -> torch.Tensor:
         source_nodes, target_nodes = self.prepare_node_data(graph)
 
         adjmat = self.get_adjacency_matrix(source_nodes, target_nodes)
-
         # Get source & target indices of the edges
         edge_index = np.stack([adjmat.col, adjmat.row], axis=0)
-
         return torch.from_numpy(edge_index).to(torch.int32)
 
     def register_edges(self, graph: HeteroData) -> HeteroData:
@@ -381,7 +380,13 @@ class MultiScaleEdges(BaseEdgeBuilder):
         Update the graph with the edges.
     """
 
-    VALID_NODES = [TriNodes, HexNodes, LimitedAreaTriNodes, LimitedAreaHexNodes, StretchedTriNodes]
+    VALID_NODES = [
+        TriNodes,
+        HexNodes,
+        LimitedAreaTriNodes,
+        LimitedAreaHexNodes,
+        StretchedTriNodes,
+    ]
 
     def __init__(self, source_name: str, target_name: str, x_hops: int, **kwargs):
         super().__init__(source_name, target_name)
@@ -444,3 +449,129 @@ def update_graph(self, graph: HeteroData, attrs_config: DotDict | None = None) -
         ), f"{self.__class__.__name__} requires {','.join(valid_node_names)} nodes."
 
         return super().update_graph(graph, attrs_config)
+
+
+class ICONTopologicalBaseEdgeBuilder(BaseEdgeBuilder):
+    """Base class for computing edges based on ICON grid topology.
+
+    Attributes
+    ----------
+    source_name : str
+        The name of the source nodes.
+    target_name : str
+        The name of the target nodes.
+    icon_mesh   : str
+        The name of the ICON mesh (defines both the processor mesh and the data)
+    """
+
+    def __init__(
+        self,
+        source_name: str,
+        target_name: str,
+        icon_mesh: str,
+        source_mask_attr_name: str | None = None,
+        target_mask_attr_name: str | None = None,
+    ):
+        self.icon_mesh = icon_mesh
+        super().__init__(source_name, target_name, source_mask_attr_name, target_mask_attr_name)
+
+    def update_graph(self, graph: HeteroData, attrs_config: DotDict = None) -> HeteroData:
+        """Update the graph with the edges."""
+        assert self.icon_mesh is not None, f"{self.__class__.__name__} requires initialized icon_mesh."
+        self.icon_sub_graph = graph[self.icon_mesh][self.sub_graph_address]
+        return super().update_graph(graph, attrs_config)
+
+    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage):
+        """Parameters
+        ----------
+        source_nodes : NodeStorage
+            The source nodes.
+        target_nodes : NodeStorage
+            The target nodes.
+        """
+        LOGGER.info(f"Using ICON topology {self.source_name}>{self.target_name}")
+        nrows = self.icon_sub_graph.num_edges
+        adj_matrix = scipy.sparse.coo_matrix(
+            (
+                np.ones(nrows),
+                (
+                    self.icon_sub_graph.edge_vertices[:, self.vertex_index[0]],
+                    self.icon_sub_graph.edge_vertices[:, self.vertex_index[1]],
+                ),
+            )
+        )
+        return adj_matrix
+
+
+class ICONTopologicalProcessorEdges(ICONTopologicalBaseEdgeBuilder):
+    """Computes edges based on ICON grid topology: processor grid built
+    from ICON grid vertices.
+    """
+
+    def __init__(
+        self,
+        source_name: str,
+        target_name: str,
+        icon_mesh: str,
+        source_mask_attr_name: str | None = None,
+        target_mask_attr_name: str | None = None,
+    ):
+        self.sub_graph_address = "_multi_mesh"
+        self.vertex_index = (1, 0)
+        super().__init__(
+            source_name,
+            target_name,
+            icon_mesh,
+            source_mask_attr_name,
+            target_mask_attr_name,
+        )
+
+
+class ICONTopologicalEncoderEdges(ICONTopologicalBaseEdgeBuilder):
+    """Computes encoder edges based on ICON grid topology: ICON cell
+    circumcenters for mapped onto processor grid built from ICON grid
+    vertices.
+    """
+
+    def __init__(
+        self,
+        source_name: str,
+        target_name: str,
+        icon_mesh: str,
+        source_mask_attr_name: str | None = None,
+        target_mask_attr_name: str | None = None,
+    ):
+        self.sub_graph_address = "_cell_grid"
+        self.vertex_index = (1, 0)
+        super().__init__(
+            source_name,
+            target_name,
+            icon_mesh,
+            source_mask_attr_name,
+            target_mask_attr_name,
+        )
+
+
+class ICONTopologicalDecoderEdges(ICONTopologicalBaseEdgeBuilder):
+    """Computes encoder edges based on ICON grid topology: mapping from
+    processor grid built from ICON grid vertices onto ICON cell
+    circumcenters.
+    """
+
+    def __init__(
+        self,
+        source_name: str,
+        target_name: str,
+        icon_mesh: str,
+        source_mask_attr_name: str | None = None,
+        target_mask_attr_name: str | None = None,
+    ):
+        self.sub_graph_address = "_cell_grid"
+        self.vertex_index = (0, 1)
+        super().__init__(
+            source_name,
+            target_name,
+            icon_mesh,
+            source_mask_attr_name,
+            target_mask_attr_name,
+        )