From 100c251f87409848ac6deb4118ff9d683a6f2f69 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Thu, 26 Sep 2024 07:54:23 +0000
Subject: [PATCH 01/35] feat: topology-based encoder/processor/decoder graphs
 derived from an ICON grid file.

---
 pyproject.toml                          |   2 +
 src/anemoi/graphs/edges/__init__.py     |   5 +-
 src/anemoi/graphs/edges/builder.py      | 117 ++++-
 src/anemoi/graphs/generate/icon_mesh.py | 556 ++++++++++++++++++++++++
 src/anemoi/graphs/nodes/__init__.py     |   5 +-
 src/anemoi/graphs/nodes/builder.py      |  73 +++-
 6 files changed, 753 insertions(+), 5 deletions(-)
 create mode 100644 src/anemoi/graphs/generate/icon_mesh.py

diff --git a/pyproject.toml b/pyproject.toml
index a6148ef..128a110 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -47,11 +47,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
index 281860d..2424982 100644
--- a/src/anemoi/graphs/edges/__init__.py
+++ b/src/anemoi/graphs/edges/__init__.py
@@ -1,5 +1,8 @@
 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", "ICONTopologicalDecoderEdges"]
diff --git a/src/anemoi/graphs/edges/builder.py b/src/anemoi/graphs/edges/builder.py
index 7f34390..094815d 100644
--- a/src/anemoi/graphs/edges/builder.py
+++ b/src/anemoi/graphs/edges/builder.py
@@ -6,6 +6,7 @@
 
 import networkx as nx
 import numpy as np
+import scipy
 import torch
 from anemoi.utils.config import DotDict
 from hydra.utils import instantiate
@@ -58,10 +59,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:
@@ -352,3 +351,117 @@ def update_graph(self, graph: HeteroData, attrs_config: DotDict | None = None) -
         self.node_type = graph[self.source_name].node_type
 
         return super().update_graph(graph, attrs_config)
+
+
+class ICONTopologicalProcessorEdges(BaseEdgeBuilder):
+    """Computes 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):
+        self.icon_mesh = icon_mesh
+        super().__init__(source_name, target_name)
+
+    def update_graph(self, graph: HeteroData, attrs_config: DotDict = None) -> HeteroData:
+        """Update the graph with the edges."""
+        self.multi_mesh = graph[self.icon_mesh]["_multi_mesh"]
+        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.multi_mesh.num_edges
+        adj_matrix = scipy.sparse.coo_matrix(
+            (np.ones(nrows), (self.multi_mesh.edge_vertices[:, 1], self.multi_mesh.edge_vertices[:, 0]))
+        )
+        return adj_matrix
+
+
+class ICONTopologicalEncoderEdges(BaseEdgeBuilder):
+    """Computes 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):
+        self.icon_mesh = icon_mesh
+        super().__init__(source_name, target_name)
+
+    def update_graph(self, graph: HeteroData, attrs_config: DotDict = None) -> HeteroData:
+        """Update the graph with the edges."""
+        self.cell_data_grid = graph[self.icon_mesh]["_cell_grid"]
+        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.cell_data_grid.num_edges
+        adj_matrix = scipy.sparse.coo_matrix(
+            (np.ones(nrows), (self.cell_data_grid.edge_vertices[:, 1], self.cell_data_grid.edge_vertices[:, 0]))
+        )
+        return adj_matrix
+
+
+class ICONTopologicalDecoderEdges(BaseEdgeBuilder):
+    """Computes 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):
+        self.icon_mesh = icon_mesh
+        super().__init__(source_name, target_name)
+
+    def update_graph(self, graph: HeteroData, attrs_config: DotDict = None) -> HeteroData:
+        """Update the graph with the edges."""
+        self.cell_data_grid = graph[self.icon_mesh]["_cell_grid"]
+        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.cell_data_grid.num_edges
+        adj_matrix = scipy.sparse.coo_matrix(
+            (np.ones(nrows), (self.cell_data_grid.edge_vertices[:, 0], self.cell_data_grid.edge_vertices[:, 1]))
+        )
+        return adj_matrix
diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
new file mode 100644
index 0000000..8c2b809
--- /dev/null
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -0,0 +1,556 @@
+# AICON-GRAPH
+#
+# ---------------------------------------------------------------
+# Copyright (C) 2004-2024, DWD, MPI-M, DKRZ, KIT, ETH, MeteoSwiss
+# Contact information: icon-model.org
+#
+# See AUTHORS.md for a list of authors
+# See LICENSES/ for license information
+# SPDX-License-Identifier: BSD-3-Clause
+# ---------------------------------------------------------------
+#
+
+import itertools
+import logging
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Optional
+
+import netCDF4
+import numpy as np
+import scipy
+import torch
+import torch_geometric
+from typeguard import typechecked
+from typing_extensions import Self
+
+LOGGER = logging.getLogger(__name__)
+
+
+@typechecked
+class NodeSet:
+    """Stores nodes on the unit sphere."""
+
+    id_iter: int = itertools.count()  # unique ID for each object
+    gc_vertices: np.ndarray  # geographical (lat/lon) coordinates [rad], shape [:,2]
+
+    def __init__(self, lon: np.ndarray, lat: np.ndarray):
+        self.gc_vertices = np.column_stack((lon, lat))
+        self.id = next(self.id_iter)
+
+    @property
+    def num_verts(self) -> int:
+        return self.gc_vertices.shape[0]
+
+    @cached_property
+    def cc_vertices(self):
+        """Cartesian coordinates [rad], shape [:,3]."""
+        return self._gc_to_cartesian()
+
+    def __add__(self, other: Self) -> Self:
+        """concatenates two node sets."""
+        gc_vertices = np.concatenate((self.gc_vertices, other.gc_vertices))
+        return NodeSet(gc_vertices[:, 0], gc_vertices[:, 1])
+
+    def __eq__(self, other: Self) -> bool:
+        """Compares two node sets."""
+        return self.id == other.id
+
+    def _gc_to_cartesian(self, radius: float = 1.0) -> np.ndarray:
+        """Returns Cartesian coordinates of the node set, shape [:,3]."""
+        xyz = (
+            radius * np.cos(lat_rad := self.gc_vertices[:, 1]) * np.cos(lon_rad := self.gc_vertices[:, 0]),
+            radius * np.cos(lat_rad) * np.sin(lon_rad),
+            radius * np.sin(lat_rad),
+        )
+        return np.stack(xyz, axis=-1)
+
+
+@typechecked
+@dataclass
+class EdgeID:
+    """Stores additional categorical data for each edge (IDs for heterogeneous input)."""
+
+    edge_id: np.ndarray
+    num_classes: int
+
+    def __add__(self, other: Self):
+        """Concatenates two edge ID datasets."""
+        assert self.num_classes == other.num_classes
+        return EdgeID(edge_id=np.concatenate((self.edge_id, other.edge_id)), num_classes=self.num_classes)
+
+
+@typechecked
+class GeneralGraph:
+    """Stores edges for a given node set, calculates edge features."""
+
+    nodeset: NodeSet  # graph nodes
+    edge_vertices: np.ndarray  # vertex indices for each edge, shape [:,2]
+
+    def __init__(self, nodeset: NodeSet, bidirectional: bool, edge_vertices: np.ndarray):
+        self.nodeset = nodeset
+        # (optional) duplicate edges (bi-directional):
+        if bidirectional:
+            self.edge_vertices = np.concatenate([edge_vertices, np.fliplr(edge_vertices)])
+        else:
+            self.edge_vertices = edge_vertices
+
+    @property
+    def num_verts(self) -> int:
+        return self.nodeset.num_verts
+
+    @property
+    def num_edges(self) -> int:
+        return self.edge_vertices.shape[0]
+
+    @cached_property
+    def edge_features(self) -> torch.tensor:
+        return torch.tensor(
+            get_edge_attributes(
+                self.nodeset.gc_vertices[self.edge_vertices[:, 1]],
+                self.nodeset.cc_vertices[self.edge_vertices[:, 1]],
+                self.nodeset.gc_vertices[self.edge_vertices[:, 0]],
+                self.nodeset.cc_vertices[self.edge_vertices[:, 0]],
+            ),
+            dtype=torch.float32,
+        )
+
+
+@typechecked
+class BipartiteGraph:
+    """Graph defined on a pair of NodeSets."""
+
+    nodeset: tuple[NodeSet, NodeSet]  # source and target node set
+    edge_vertices: np.ndarray  # vertex indices for each edge, shape [:,2]
+    edge_id: np.ndarray  # additional ID for each edge (markers for heterogeneous input)
+
+    def __init__(self, nodeset: tuple[NodeSet, NodeSet], edge_vertices: np.ndarray, edge_id: Optional[EdgeID] = None):
+        self.nodeset = nodeset
+        self.edge_vertices = edge_vertices
+        if edge_id is None:
+            self.edge_id = None
+        else:
+            self.edge_id = edge_id
+
+    @property
+    def num_edges(self) -> int:
+        return self.edge_vertices.shape[0]
+
+    @cached_property
+    def edge_features(self) -> torch.tensor:
+        edge_feature_arr = torch.tensor(
+            get_edge_attributes(
+                self.nodeset[0].gc_vertices[self.edge_vertices[:, 0]],
+                self.nodeset[0].cc_vertices[self.edge_vertices[:, 0]],
+                self.nodeset[1].gc_vertices[self.edge_vertices[:, 1]],
+                self.nodeset[1].cc_vertices[self.edge_vertices[:, 1]],
+            ),
+            dtype=torch.float32,
+        )
+        if self.edge_id is None:
+            return edge_feature_arr
+        else:
+            # one-hot-encoded categorical data (`edge_id`)
+            one_hot = torch_geometric.utils.one_hot(
+                index=torch.tensor(self.edge_id.edge_id), num_classes=self.edge_id.num_classes
+            )
+            return torch.concatenate((one_hot, edge_feature_arr), dim=1)
+
+    def set_constant_edge_id(self, edge_id: int, num_classes: int):
+        self.edge_id = EdgeID(edge_id=np.full(self.num_edges, fill_value=edge_id), num_classes=num_classes)
+
+    def __add__(self, other: "BipartiteGraph"):
+        """Concatenates two bipartite graphs that share a common target node set.
+        Shifts the node indices of the second bipartite graph.
+        """
+
+        if not self.nodeset[1] == other.nodeset[1]:
+            raise ValueError("Only bipartite graphs with common target node set can be merged.")
+        shifted_edge_vertices = other.edge_vertices
+        shifted_edge_vertices[:, 0] += self.nodeset[0].num_verts
+        # (Optional:) merge one-hot-encoded categorical data (`edge_id`)
+        if None not in (self.edge_id, other.edge_id):
+            edge_id = self.edge_id + other.edge_id
+        else:
+            edge_id = None
+        return BipartiteGraph(
+            nodeset=(self.nodeset[0] + other.nodeset[0], self.nodeset[1]),
+            edge_vertices=np.concatenate((self.edge_vertices, shifted_edge_vertices)),
+            edge_id=edge_id,
+        )
+
+
+@typechecked
+class BipartiteGraphProximity(BipartiteGraph):
+    """Graph defined on a pair of NodeSets, where the definition of the
+    graph edges is based on geographical proximity (Euclidean distance
+    in Cartesian coordinates).
+    """
+
+    def __init__(self, nodeset: tuple[NodeSet, NodeSet], radius: float, edge_id: Optional[EdgeID] = None):
+
+        src, tgt = nodeset
+        point_tree = scipy.spatial.KDTree(tgt.cc_vertices)
+        neighbor_list = point_tree.query_ball_point(src.cc_vertices, r=radius)
+        # turn `neighbor_list` into a list of pairs:
+        nb_list = [(idx, x) for idx, x in enumerate(neighbor_list)]
+        nb_list = np.asarray([(x, k) for x, y in nb_list for k in y])
+        super().__init__(nodeset=nodeset, edge_id=edge_id, edge_vertices=nb_list)
+
+
+@typechecked
+class ICONMultiMesh(GeneralGraph):
+    """Reads vertices and topology from an ICON grid file; creates multi-mesh."""
+
+    uuidOfHGrid: str
+    max_level: int
+    nodeset: NodeSet  # set of ICON grid vertices
+    cell_vertices: np.ndarray
+
+    def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None, iverbosity: int = 0):
+
+        # open file, representing the finest level
+        if iverbosity > 0:
+            LOGGER.info(f"{type(self).__name__}: read ICON grid file '{icon_grid_filename}'")
+        with netCDF4.Dataset(icon_grid_filename, "r") as ncfile:
+            # read vertex coordinates
+            vlon = read_coordinate_array(ncfile, "vlon", "vertex")
+            vlat = read_coordinate_array(ncfile, "vlat", "vertex")
+
+            edge_verts_fine = np.asarray(ncfile.variables["edge_vertices"][:] - 1, dtype=np.int64).transpose()
+            assert ncfile.variables["edge_vertices"].dimensions == ("nc", "edge")
+
+            cell_verts_fine = np.asarray(ncfile.variables["vertex_of_cell"][:] - 1, dtype=np.int64).transpose()
+            assert ncfile.variables["vertex_of_cell"].dimensions == ("nv", "cell")
+
+            reflvl_vertex = ncfile.variables["refinement_level_v"][:]
+            assert ncfile.variables["refinement_level_v"].dimensions == ("vertex",)
+
+            self.uuidOfHGrid = ncfile.uuidOfHGrid
+
+        if max_level is not None:
+            self.max_level = max_level
+        else:
+            self.max_level = reflvl_vertex.max()
+
+        # generate edge-vertex relations for coarser levels:
+        (edge_vertices, cell_vertices) = self._get_hierarchy_of_icon_edge_graphs(
+            edge_verts_fine=edge_verts_fine,
+            cell_verts_fine=cell_verts_fine,
+            reflvl_vertex=reflvl_vertex,
+            iverbosity=iverbosity,
+        )
+        # restrict edge-vertex list to multi_mesh level "max_level":
+        if self.max_level < len(edge_vertices):
+            (self.edge_vertices, self.cell_vertices, vlon, vlat) = self._restrict_multi_mesh_level(
+                edge_vertices, cell_vertices, reflvl_vertex=reflvl_vertex, vlon=vlon, vlat=vlat
+            )
+        # store vertices as a `NodeSet`:
+        self.nodeset = NodeSet(vlon, vlat)
+        # concatenate edge-vertex lists (= edges of the multi-level mesh):
+        multi_mesh_edges = np.concatenate([edges for edges in self.edge_vertices], axis=0)
+        # generate multi-mesh graph data structure:
+        super().__init__(nodeset=self.nodeset, bidirectional=True, edge_vertices=multi_mesh_edges)
+
+    def _restrict_multi_mesh_level(
+        self,
+        edge_vertices: list[np.ndarray],
+        cell_vertices: np.ndarray,
+        reflvl_vertex: np.ndarray,
+        vlon: np.ndarray,
+        vlat: np.ndarray,
+    ) -> tuple[list[np.ndarray], np.ndarray, np.ndarray, np.ndarray]:
+        """Creates a new mesh with only the vertices at the desired level."""
+
+        num_verts = reflvl_vertex.shape[0]
+        vertex_idx = reflvl_vertex <= self.max_level
+        vertex_glb2loc = np.zeros(num_verts, dtype=int)
+        vertex_glb2loc[vertex_idx] = np.arange(vertex_idx.sum())
+        return (
+            [vertex_glb2loc[M] for M in edge_vertices[: self.max_level + 1]],
+            # cell_vertices: preserve negative indices (incomplete cells)
+            np.where(cell_vertices >= 0, vertex_glb2loc[cell_vertices], cell_vertices),
+            vlon[vertex_idx],
+            vlat[vertex_idx],
+        )
+
+    def _get_hierarchy_of_icon_edge_graphs(
+        self, edge_verts_fine: np.ndarray, cell_verts_fine: np.ndarray, reflvl_vertex: np.ndarray, iverbosity: int = 1
+    ) -> tuple[list[np.ndarray], np.ndarray]:
+        """Returns a list of edge-vertex relations (coarsest to finest level)."""
+
+        edge_vertices = [edge_verts_fine]  # list of edge-vertex relations (coarsest to finest level).
+
+        num_verts = reflvl_vertex.shape[0]
+        # edge-to-vertex adjacency matrix with 2 non-zero entries per row:
+        e2v_matrix = convert_list_to_adjacency_matrix(edge_verts_fine, num_verts)
+        # cell-to-vertex adjacency matrix with 3 non-zero entries per row:
+        c2v_matrix = convert_list_to_adjacency_matrix(cell_verts_fine, num_verts)
+        v2v_matrix = e2v_matrix.transpose() * e2v_matrix
+        v2v_matrix.setdiag(np.ones(num_verts))  # vertices are self-connected
+
+        s_v_coarse = scipy.sparse.diags(np.ones(num_verts), dtype=bool)
+
+        # coarsen edge-vertex list from level `ilevel -> ilevel - 1`:
+        for ilevel in reversed(range(1, reflvl_vertex.max() + 1)):
+            if iverbosity > 0:
+                LOGGER.info(f"  edges[{ilevel}] = {edge_vertices[0].shape[0] : >9}")
+
+            # define edge selection matrix (selecting only edges of which have
+            # exactly one coarse vertex):
+            #
+            # get a boolean mask, matching all edges where one of its vertices
+            # has refinement level index `ilevel`:
+            ref_level = reflvl_vertex[edge_vertices[0]] == ilevel
+            edges_coarse = np.logical_xor(ref_level[:, 0], ref_level[:, 1])  # = bisected coarse edges
+            idx_e2e = np.argwhere(edges_coarse).flatten()
+            s_edges = selection_matrix(idx_e2e, edges_coarse.shape[0])
+
+            # define vertex selection matrix selecting only vertices of
+            # level `ilevel`:
+            idx_v_fine = np.argwhere(reflvl_vertex == ilevel).flatten()
+            s_v_fine = selection_matrix(idx_v_fine, num_verts)
+            # define vertex selection matrix, selecting only vertices of
+            # level < `ilevel`, by successively removing `s_fine` from an identity matrix.
+            s_v_coarse.data[0][idx_v_fine] = False
+
+            # create an adjacency matrix which links each fine level
+            # vertex to its two coarser neighbor vertices:
+            v2v_fine2coarse = s_v_fine * v2v_matrix * s_v_coarse
+            # remove rows that have only one non-zero entry
+            # (corresponding to incomplete parent triangles in LAM grids):
+            csum = v2v_fine2coarse * np.ones((v2v_fine2coarse.shape[1], 1))
+            s_v2v = selection_matrix(np.argwhere(csum == 2).flatten(), v2v_fine2coarse.shape[0])
+            v2v_fine2coarse = s_v2v * v2v_fine2coarse
+
+            # then construct the edges-to-parent-vertex adjacency matrix:
+            parent_edge_verts = s_edges * e2v_matrix * v2v_fine2coarse
+            # again, we have might have selected edges within
+            # `s_edges` which are part of an incomplete parent edge
+            # (LAM case). We filter these here:
+            csum = parent_edge_verts * np.ones((parent_edge_verts.shape[1], 1))
+            s_e2e = selection_matrix(np.argwhere(csum == 2).flatten(), parent_edge_verts.shape[0])
+            parent_edge_verts = s_e2e * parent_edge_verts
+
+            # note: the edges-vertex adjacency matrix still has duplicate
+            # rows, since two child edges have the same parent edge.
+            edge_verts_coarse = convert_adjacency_matrix_to_list(parent_edge_verts, ncols_per_row=2)
+            edge_vertices.insert(0, edge_verts_coarse)
+
+            # store cell-to-vert adjacency matrix
+            if ilevel > self.max_level:
+                c2v_matrix = c2v_matrix * v2v_fine2coarse
+                # similar to the treatment above, we need to handle
+                # coarse LAM cells which are incomplete.
+                csum = c2v_matrix * np.ones((c2v_matrix.shape[1], 1))
+                s_c2c = selection_matrix(np.argwhere(csum == 3).flatten(), c2v_matrix.shape[0])
+                c2v_matrix = s_c2c * c2v_matrix
+
+            # replace edge-to-vertex and vert-to-vert adjacency matrices (for next level):
+            if ilevel > 1:
+                v2v_matrix = s_v_coarse * v2v_matrix * v2v_fine2coarse
+                e2v_matrix = convert_list_to_adjacency_matrix(edge_verts_coarse, num_verts)
+
+        # Fine-level cells outside of multi-mesh (LAM boundary)
+        # correspond to empty rows in the adjacency matrix. We
+        # substitute these by three additional, non-existent vertices:
+        csum = 3 - c2v_matrix * np.ones((c2v_matrix.shape[1], 1))
+        nvmax = c2v_matrix.shape[1]
+        c2v_matrix = scipy.sparse.csr_matrix(scipy.sparse.hstack((c2v_matrix, csum, csum, csum)))
+
+        # build a list of cell-vertices [1..num_cells,1..3] for all
+        # fine-level cells:
+        cell_vertices = convert_adjacency_matrix_to_list(c2v_matrix, remove_duplicates=False, ncols_per_row=3)
+
+        # finally, translate non-existent vertices into "-1" indices:
+        cell_vertices = np.where(cell_vertices >= nvmax, -np.ones(cell_vertices.shape, dtype=np.int32), cell_vertices)
+
+        return (edge_vertices, cell_vertices)
+
+
+@typechecked
+class ICONCellDataGrid(BipartiteGraph):
+    """Reads cell locations from an ICON grid file; builds grid-to-mesh edges based on ICON topology."""
+
+    uuidOfHGrid: str
+    nodeset: NodeSet  # set of ICON cell circumcenters
+    max_level: int
+    select_c: np.ndarray
+
+    def __init__(
+        self,
+        icon_grid_filename: str,
+        multi_mesh: Optional[ICONMultiMesh] = None,
+        max_level: Optional[int] = None,
+        iverbosity: int = 0,
+    ):
+        # open file, representing the finest level
+        if iverbosity > 0:
+            LOGGER.info(f"{type(self).__name__}: read ICON grid file '{icon_grid_filename}'")
+        with netCDF4.Dataset(icon_grid_filename, "r") as ncfile:
+            # read cell circumcenter coordinates
+            clon = read_coordinate_array(ncfile, "clon", "cell")
+            clat = read_coordinate_array(ncfile, "clat", "cell")
+
+            reflvl_cell = ncfile.variables["refinement_level_c"][:]
+            assert ncfile.variables["refinement_level_c"].dimensions == ("cell",)
+
+            self.uuidOfHGrid = ncfile.uuidOfHGrid
+
+        if max_level is not None:
+            self.max_level = max_level
+        else:
+            self.max_level = reflvl_cell.max()
+
+        # restrict to level `max_level`:
+        self.select_c = np.argwhere(reflvl_cell <= self.max_level)
+        # generate source grid node set:
+        self.nodeset = NodeSet(clon[self.select_c], clat[self.select_c])
+
+        if multi_mesh is not None:
+            # generate edges between source grid nodes and multi-mesh nodes:
+            edge_vertices = self._get_grid2mesh_edges(self.select_c, multi_mesh=multi_mesh)
+            super().__init__((self.nodeset, multi_mesh.nodeset), edge_vertices)
+
+    def _get_grid2mesh_edges(self, select_c: np.ndarray, multi_mesh: ICONMultiMesh) -> np.ndarray:
+        """Create "grid-to-mesh" edges, ie. edges from (clat,clon) to the
+        vertices of the multi-mesh.
+        """
+
+        num_cells = select_c.shape[0]
+        num_verts_per_cell = multi_mesh.cell_vertices.shape[1]
+        src_list = np.kron(np.arange(num_cells), np.ones((1, num_verts_per_cell), dtype=np.int64)).flatten()
+        dst_list = multi_mesh.cell_vertices[select_c[:, 0], :].flatten()
+        edge_vertices = np.stack((src_list, dst_list), axis=1, dtype=np.int64)
+        return edge_vertices
+
+
+# -------------------------------------------------------------
+
+
+@typechecked
+def get_icon_mesh_and_grid(
+    iverbosity: int, grid_file: str, max_level_multimesh: int, max_level_dataset: int
+) -> tuple[ICONMultiMesh, ICONCellDataGrid]:
+    """Factory function, creating an ICON multi-mesh and an ICON cell-grid."""
+    return (
+        multi_mesh := ICONMultiMesh(grid_file, max_level=max_level_multimesh, iverbosity=iverbosity),
+        ICONCellDataGrid(grid_file, multi_mesh, max_level=max_level_dataset, iverbosity=iverbosity),
+    )
+
+
+@typechecked
+def get_edge_attributes(
+    gc_send: np.ndarray, cc_send: np.ndarray, gc_recv: np.ndarray, cc_recv: np.ndarray
+) -> np.ndarray:
+    """Build edge features using the position on the unit sphere of the mesh
+    nodes.
+
+    The features are (4 input features in total):
+
+    - length of the edge
+    - vector difference between the 3d positions of the sender node and the
+      receiver node computed in a local coordinate system of the receiver.
+
+    The local coordinate system of the receiver is computed by applying two elemental
+    rotations:
+
+    - the first rotation changes the azimuthal angle until that receiver
+      node lies at longitude 0,
+    - the second rotation changes the polar angle until that receiver
+      node lies at latitude 0.
+
+    Literature: Appdx. A2.4 in
+       Lam, R. et al. (2022). GraphCast: Learning skillful medium-range
+       global weather forecasting. 10.48550/arXiv.2212.12794.
+
+    """
+
+    phi_theta = gc_recv  # precession and nutation angle
+    theta__phi = np.fliplr(phi_theta)
+    theta__phi[:, 1] *= -1.0  # angles = [theta, -phi]
+    # rotation matrix:
+    R = scipy.spatial.transform.Rotation.from_euler(seq="YZ", angles=theta__phi)
+    edge_length = np.array([arc_length(cc_recv, cc_send)])
+    distance = R.apply(cc_send) - [1.0, 0.0, 0.0]  # subtract the rotated position of sender
+    return np.concatenate((edge_length.transpose(), distance), axis=1)
+
+
+@typechecked
+def arc_length(v1, v2) -> np.ndarray:
+    """Calculate length of great circle arc on the unit sphere."""
+    return np.arccos(np.clip(np.einsum("ij,ij->i", v1, v2), a_min=0.0, a_max=1.0))
+
+
+@typechecked
+def read_coordinate_array(ncfile, arrname: str, dimname: str) -> np.ndarray:
+    """Auxiliary routine, reading a coordinate array, checking consistency."""
+    arr = ncfile.variables[arrname][:]
+    assert ncfile.variables[arrname].dimensions == (dimname,)
+    assert ncfile.variables[arrname].units == "radian"
+    # netCDF4 returns all variables as numpy.ma.core.MaskedArray.
+    # -> convert to regular arrays
+    assert not arr.mask.any(), f"There are missing values in {arrname}"
+    return arr.data
+
+
+@typechecked
+def convert_list_to_adjacency_matrix(list_matrix: np.ndarray, ncols: int = 0) -> scipy.sparse.csr_matrix:
+    """Convert an edge list into an adjacency matrix.
+
+    Args:
+      list_matrix : np.ndarray
+        boolean matrix given by list of column indices for each row.
+      ncols : int
+        number of columns in result matrix.
+    Returns:
+      Returns sparse matrix [nrows, ncols]
+    """
+    nrows, ncols_per_row = list_matrix.shape
+    indptr = np.arange(ncols_per_row * (nrows + 1), step=ncols_per_row)
+    indices = list_matrix.ravel()
+    return scipy.sparse.csr_matrix((np.ones(nrows * ncols_per_row), indices, indptr), dtype=bool, shape=(nrows, ncols))
+
+
+@typechecked
+def convert_adjacency_matrix_to_list(
+    adj_matrix: scipy.sparse.csr_matrix,
+    ncols_per_row: int,
+    remove_duplicates: bool = True,
+) -> np.ndarray:
+    """Convert an adjacency matrix into an edge list.
+
+    Args:
+      adj_matrix: scipy.sparse.csr_matrix
+          sparse (boolean) adjacency matrix
+      ncols_per_row: int
+          number of nonzero entries per row
+      remove_duplicates: bool
+          logical flag: remove duplicate rows.
+    Returns:
+      boolean matrix given by list of column indices for each row.
+    """
+    if remove_duplicates:
+        # The edges-vertex adjacency matrix may have duplicate rows, remove
+        # them by selecting the rows that are unique:
+        nrows = int(adj_matrix.nnz / ncols_per_row)
+        mat = adj_matrix.indices.reshape((nrows, ncols_per_row))
+        return np.unique(mat, axis=0)
+    else:
+        nrows = adj_matrix.shape[0]
+        return adj_matrix.indices.reshape((nrows, ncols_per_row))
+
+
+@typechecked
+def selection_matrix(idx: np.ndarray, isize: int = 0) -> scipy.sparse.csr_matrix:
+    """Create a diagonal selection matrix.
+
+    Args:
+        idx : np.ndarray
+           integer array of indices
+        isize: int
+           size of (square) selection matrix
+    Returns:
+        diagonal matrix with ones at selected indices (idx,idx).
+    """
+    return scipy.sparse.csr_matrix((np.ones((len(idx))), (idx, idx)), dtype=bool, shape=(isize, isize))
diff --git a/src/anemoi/graphs/nodes/__init__.py b/src/anemoi/graphs/nodes/__init__.py
index 7b6b149..348dc1c 100644
--- a/src/anemoi/graphs/nodes/__init__.py
+++ b/src/anemoi/graphs/nodes/__init__.py
@@ -1,6 +1,9 @@
 from .builder import HexNodes
+from .builder import ICONCellGridNodes
+from .builder import ICONMultimeshNodes
+from .builder import ICONNodes
 from .builder import NPZFileNodes
 from .builder import TriNodes
 from .builder import ZarrDatasetNodes
 
-__all__ = ["ZarrDatasetNodes", "NPZFileNodes", "TriNodes", "HexNodes"]
+__all__ = ["ZarrDatasetNodes", "NPZFileNodes", "TriNodes", "HexNodes", "ICONMultimeshNodes", "ICONCellGridNodes"]
diff --git a/src/anemoi/graphs/nodes/builder.py b/src/anemoi/graphs/nodes/builder.py
index 12c818c..837005b 100644
--- a/src/anemoi/graphs/nodes/builder.py
+++ b/src/anemoi/graphs/nodes/builder.py
@@ -13,6 +13,7 @@
 from torch_geometric.data import HeteroData
 
 from anemoi.graphs.generate.hexagonal import create_hexagonal_nodes
+from anemoi.graphs.generate.icon_mesh import get_icon_mesh_and_grid
 from anemoi.graphs.generate.icosahedral import create_icosahedral_nodes
 
 LOGGER = logging.getLogger(__name__)
@@ -104,7 +105,6 @@ def update_graph(self, graph: HeteroData, attr_config: DotDict | None = None) ->
 
         if attr_config is None:
             return graph
-
         graph = self.register_attributes(graph, attr_config)
 
         return graph
@@ -282,6 +282,77 @@ def create_nodes(self) -> np.ndarray:
         return create_hexagonal_nodes(self.resolutions)
 
 
+class ICONNodes(BaseNodeBuilder):
+    """ICON grid (cell and vertex locations)."""
+
+    def __init__(self, name: str, grid_filename: str, max_level_multimesh: int, max_level_dataset: int) -> None:
+        self.grid_filename = grid_filename
+        self.multi_mesh, self.cell_grid = get_icon_mesh_and_grid(
+            iverbosity=1,
+            grid_file=self.grid_filename,
+            max_level_multimesh=max_level_multimesh,
+            max_level_dataset=max_level_dataset,
+        )
+        super().__init__(name)
+
+    def get_coordinates(self) -> torch.Tensor:
+        return torch.from_numpy(self.multi_mesh.nodeset.gc_vertices.astype(np.float32)).fliplr()
+
+    def register_attributes(self, graph: HeteroData, config: DotDict) -> HeteroData:
+        graph[self.name]["_grid_filename"] = self.grid_filename
+        graph[self.name]["_multi_mesh"] = self.multi_mesh
+        graph[self.name]["_cell_grid"] = self.cell_grid
+        return super().register_attributes(graph, config)
+
+
+class ICONMultimeshNodes(BaseNodeBuilder):
+    """Processor mesh based on an ICON grid.
+
+    Parameters
+    ----------
+    name : str
+        key for the nodes in the HeteroData graph object.
+    icon_mesh : str
+        key corresponding to the ICON mesh (cells and vertices).
+    """
+
+    def __init__(self, name: str, icon_mesh: str) -> None:
+        self.icon_mesh = icon_mesh
+        super().__init__(name)
+
+    def update_graph(self, graph: HeteroData, attr_config: DotDict | None = None) -> HeteroData:
+        """Update the graph with new nodes."""
+        self.multi_mesh = graph[self.icon_mesh]["_multi_mesh"]
+        return super().update_graph(graph, attr_config)
+
+    def get_coordinates(self) -> torch.Tensor:
+        return torch.from_numpy(self.multi_mesh.nodeset.gc_vertices.astype(np.float32)).fliplr()
+
+
+class ICONCellGridNodes(BaseNodeBuilder):
+    """Data mesh based on an ICON grid.
+
+    Parameters
+    ----------
+    name : str
+        key for the nodes in the HeteroData graph object.
+    icon_mesh : str
+        key corresponding to the ICON mesh (cells and vertices).
+    """
+
+    def __init__(self, name: str, icon_mesh: str) -> None:
+        self.icon_mesh = icon_mesh
+        super().__init__(name)
+
+    def update_graph(self, graph: HeteroData, attr_config: DotDict | None = None) -> HeteroData:
+        """Update the graph with new nodes."""
+        self.cell_grid = graph[self.icon_mesh]["_cell_grid"]
+        return super().update_graph(graph, attr_config)
+
+    def get_coordinates(self) -> torch.Tensor:
+        return torch.from_numpy(self.cell_grid.nodeset[0].gc_vertices.astype(np.float32)).fliplr()
+
+
 class HEALPixNodes(BaseNodeBuilder):
     """Nodes from HEALPix grid.
 

From 9848cf0a62dafd8952f81a61212d47e35cf4fbc7 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Thu, 26 Sep 2024 15:10:01 +0000
Subject: [PATCH 02/35] docs: attempt to set a correct license header.

---
 src/anemoi/graphs/commands/__init__.py  |  8 ++++----
 src/anemoi/graphs/generate/icon_mesh.py | 15 ++++++---------
 2 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/src/anemoi/graphs/commands/__init__.py b/src/anemoi/graphs/commands/__init__.py
index cebb539..30e21bc 100644
--- a/src/anemoi/graphs/commands/__init__.py
+++ b/src/anemoi/graphs/commands/__init__.py
@@ -1,11 +1,11 @@
 #!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 ECMWF, DWD.
 #
 # 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.
+# In applying this licence, the above institution do 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 os
diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 8c2b809..69b5d7f 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -1,13 +1,10 @@
-# AICON-GRAPH
+# (C) Copyright 2024 ECMWF, DWD.
 #
-# ---------------------------------------------------------------
-# Copyright (C) 2004-2024, DWD, MPI-M, DKRZ, KIT, ETH, MeteoSwiss
-# Contact information: icon-model.org
-#
-# See AUTHORS.md for a list of authors
-# See LICENSES/ for license information
-# SPDX-License-Identifier: BSD-3-Clause
-# ---------------------------------------------------------------
+# 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, the above institution do 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 itertools

From 69ae2fa84e88f9ec86381edb52064708dcf2748a Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 06:54:04 +0000
Subject: [PATCH 03/35] doc: changed the docstring style to the Numpy-style
 instead of Google-style.

---
 src/anemoi/graphs/generate/icon_mesh.py | 54 +++++++++++++++----------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 69b5d7f..69507b4 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -2,7 +2,7 @@
 #
 # 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, the above institution do not waive the privileges
+# In applying this licence, the above institutions do 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.
 #
@@ -495,13 +495,17 @@ def read_coordinate_array(ncfile, arrname: str, dimname: str) -> np.ndarray:
 def convert_list_to_adjacency_matrix(list_matrix: np.ndarray, ncols: int = 0) -> scipy.sparse.csr_matrix:
     """Convert an edge list into an adjacency matrix.
 
-    Args:
-      list_matrix : np.ndarray
+    Parameters
+    ----------
+    list_matrix : np.ndarray
         boolean matrix given by list of column indices for each row.
-      ncols : int
+    ncols : int
         number of columns in result matrix.
-    Returns:
-      Returns sparse matrix [nrows, ncols]
+
+    Returns
+    -------
+    scipy.sparse.csr_matrix
+        sparse matrix [nrows, ncols]
     """
     nrows, ncols_per_row = list_matrix.shape
     indptr = np.arange(ncols_per_row * (nrows + 1), step=ncols_per_row)
@@ -517,15 +521,19 @@ def convert_adjacency_matrix_to_list(
 ) -> np.ndarray:
     """Convert an adjacency matrix into an edge list.
 
-    Args:
-      adj_matrix: scipy.sparse.csr_matrix
-          sparse (boolean) adjacency matrix
-      ncols_per_row: int
-          number of nonzero entries per row
-      remove_duplicates: bool
-          logical flag: remove duplicate rows.
-    Returns:
-      boolean matrix given by list of column indices for each row.
+    Parameters
+    ----------
+    adj_matrix : scipy.sparse.csr_matrix
+        sparse (boolean) adjacency matrix
+    ncols_per_row : int
+        number of nonzero entries per row
+    remove_duplicates : bool
+        logical flag: remove duplicate rows.
+
+    Returns
+    -------
+    np.ndarray
+        boolean matrix given by list of column indices for each row.
     """
     if remove_duplicates:
         # The edges-vertex adjacency matrix may have duplicate rows, remove
@@ -542,12 +550,16 @@ def convert_adjacency_matrix_to_list(
 def selection_matrix(idx: np.ndarray, isize: int = 0) -> scipy.sparse.csr_matrix:
     """Create a diagonal selection matrix.
 
-    Args:
-        idx : np.ndarray
-           integer array of indices
-        isize: int
-           size of (square) selection matrix
-    Returns:
+    Parameters
+    ----------
+    idx : np.ndarray
+        integer array of indices
+    isize : int
+        size of (square) selection matrix
+
+    Returns
+    -------
+    scipy.sparse.csr_matrix
         diagonal matrix with ones at selected indices (idx,idx).
     """
     return scipy.sparse.csr_matrix((np.ones((len(idx))), (idx, idx)), dtype=bool, shape=(isize, isize))

From d858d0e1298772e0b46b9f10fdccc809eb9f6c1f Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 06:56:04 +0000
Subject: [PATCH 04/35] refactor: changed variable name `verts` into
 `vertices`.

---
 src/anemoi/graphs/generate/icon_mesh.py | 58 +++++++++++++------------
 1 file changed, 31 insertions(+), 27 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 69507b4..cf10d12 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -36,7 +36,7 @@ def __init__(self, lon: np.ndarray, lat: np.ndarray):
         self.id = next(self.id_iter)
 
     @property
-    def num_verts(self) -> int:
+    def num_vertices(self) -> int:
         return self.gc_vertices.shape[0]
 
     @cached_property
@@ -93,8 +93,8 @@ def __init__(self, nodeset: NodeSet, bidirectional: bool, edge_vertices: np.ndar
             self.edge_vertices = edge_vertices
 
     @property
-    def num_verts(self) -> int:
-        return self.nodeset.num_verts
+    def num_vertices(self) -> int:
+        return self.nodeset.num_vertices
 
     @property
     def num_edges(self) -> int:
@@ -164,7 +164,7 @@ def __add__(self, other: "BipartiteGraph"):
         if not self.nodeset[1] == other.nodeset[1]:
             raise ValueError("Only bipartite graphs with common target node set can be merged.")
         shifted_edge_vertices = other.edge_vertices
-        shifted_edge_vertices[:, 0] += self.nodeset[0].num_verts
+        shifted_edge_vertices[:, 0] += self.nodeset[0].num_vertices
         # (Optional:) merge one-hot-encoded categorical data (`edge_id`)
         if None not in (self.edge_id, other.edge_id):
             edge_id = self.edge_id + other.edge_id
@@ -214,10 +214,10 @@ def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None, ive
             vlon = read_coordinate_array(ncfile, "vlon", "vertex")
             vlat = read_coordinate_array(ncfile, "vlat", "vertex")
 
-            edge_verts_fine = np.asarray(ncfile.variables["edge_vertices"][:] - 1, dtype=np.int64).transpose()
+            edge_vertices_fine = np.asarray(ncfile.variables["edge_vertices"][:] - 1, dtype=np.int64).transpose()
             assert ncfile.variables["edge_vertices"].dimensions == ("nc", "edge")
 
-            cell_verts_fine = np.asarray(ncfile.variables["vertex_of_cell"][:] - 1, dtype=np.int64).transpose()
+            cell_vertices_fine = np.asarray(ncfile.variables["vertex_of_cell"][:] - 1, dtype=np.int64).transpose()
             assert ncfile.variables["vertex_of_cell"].dimensions == ("nv", "cell")
 
             reflvl_vertex = ncfile.variables["refinement_level_v"][:]
@@ -232,8 +232,8 @@ def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None, ive
 
         # generate edge-vertex relations for coarser levels:
         (edge_vertices, cell_vertices) = self._get_hierarchy_of_icon_edge_graphs(
-            edge_verts_fine=edge_verts_fine,
-            cell_verts_fine=cell_verts_fine,
+            edge_vertices_fine=edge_vertices_fine,
+            cell_vertices_fine=cell_vertices_fine,
             reflvl_vertex=reflvl_vertex,
             iverbosity=iverbosity,
         )
@@ -259,9 +259,9 @@ def _restrict_multi_mesh_level(
     ) -> tuple[list[np.ndarray], np.ndarray, np.ndarray, np.ndarray]:
         """Creates a new mesh with only the vertices at the desired level."""
 
-        num_verts = reflvl_vertex.shape[0]
+        num_vertices = reflvl_vertex.shape[0]
         vertex_idx = reflvl_vertex <= self.max_level
-        vertex_glb2loc = np.zeros(num_verts, dtype=int)
+        vertex_glb2loc = np.zeros(num_vertices, dtype=int)
         vertex_glb2loc[vertex_idx] = np.arange(vertex_idx.sum())
         return (
             [vertex_glb2loc[M] for M in edge_vertices[: self.max_level + 1]],
@@ -272,21 +272,25 @@ def _restrict_multi_mesh_level(
         )
 
     def _get_hierarchy_of_icon_edge_graphs(
-        self, edge_verts_fine: np.ndarray, cell_verts_fine: np.ndarray, reflvl_vertex: np.ndarray, iverbosity: int = 1
+        self,
+        edge_vertices_fine: np.ndarray,
+        cell_vertices_fine: np.ndarray,
+        reflvl_vertex: np.ndarray,
+        iverbosity: int = 1,
     ) -> tuple[list[np.ndarray], np.ndarray]:
         """Returns a list of edge-vertex relations (coarsest to finest level)."""
 
-        edge_vertices = [edge_verts_fine]  # list of edge-vertex relations (coarsest to finest level).
+        edge_vertices = [edge_vertices_fine]  # list of edge-vertex relations (coarsest to finest level).
 
-        num_verts = reflvl_vertex.shape[0]
+        num_vertices = reflvl_vertex.shape[0]
         # edge-to-vertex adjacency matrix with 2 non-zero entries per row:
-        e2v_matrix = convert_list_to_adjacency_matrix(edge_verts_fine, num_verts)
+        e2v_matrix = convert_list_to_adjacency_matrix(edge_vertices_fine, num_vertices)
         # cell-to-vertex adjacency matrix with 3 non-zero entries per row:
-        c2v_matrix = convert_list_to_adjacency_matrix(cell_verts_fine, num_verts)
+        c2v_matrix = convert_list_to_adjacency_matrix(cell_vertices_fine, num_vertices)
         v2v_matrix = e2v_matrix.transpose() * e2v_matrix
-        v2v_matrix.setdiag(np.ones(num_verts))  # vertices are self-connected
+        v2v_matrix.setdiag(np.ones(num_vertices))  # vertices are self-connected
 
-        s_v_coarse = scipy.sparse.diags(np.ones(num_verts), dtype=bool)
+        s_v_coarse = scipy.sparse.diags(np.ones(num_vertices), dtype=bool)
 
         # coarsen edge-vertex list from level `ilevel -> ilevel - 1`:
         for ilevel in reversed(range(1, reflvl_vertex.max() + 1)):
@@ -306,7 +310,7 @@ def _get_hierarchy_of_icon_edge_graphs(
             # define vertex selection matrix selecting only vertices of
             # level `ilevel`:
             idx_v_fine = np.argwhere(reflvl_vertex == ilevel).flatten()
-            s_v_fine = selection_matrix(idx_v_fine, num_verts)
+            s_v_fine = selection_matrix(idx_v_fine, num_vertices)
             # define vertex selection matrix, selecting only vertices of
             # level < `ilevel`, by successively removing `s_fine` from an identity matrix.
             s_v_coarse.data[0][idx_v_fine] = False
@@ -321,18 +325,18 @@ def _get_hierarchy_of_icon_edge_graphs(
             v2v_fine2coarse = s_v2v * v2v_fine2coarse
 
             # then construct the edges-to-parent-vertex adjacency matrix:
-            parent_edge_verts = s_edges * e2v_matrix * v2v_fine2coarse
+            parent_edge_vertices = s_edges * e2v_matrix * v2v_fine2coarse
             # again, we have might have selected edges within
             # `s_edges` which are part of an incomplete parent edge
             # (LAM case). We filter these here:
-            csum = parent_edge_verts * np.ones((parent_edge_verts.shape[1], 1))
-            s_e2e = selection_matrix(np.argwhere(csum == 2).flatten(), parent_edge_verts.shape[0])
-            parent_edge_verts = s_e2e * parent_edge_verts
+            csum = parent_edge_vertices * np.ones((parent_edge_vertices.shape[1], 1))
+            s_e2e = selection_matrix(np.argwhere(csum == 2).flatten(), parent_edge_vertices.shape[0])
+            parent_edge_vertices = s_e2e * parent_edge_vertices
 
             # note: the edges-vertex adjacency matrix still has duplicate
             # rows, since two child edges have the same parent edge.
-            edge_verts_coarse = convert_adjacency_matrix_to_list(parent_edge_verts, ncols_per_row=2)
-            edge_vertices.insert(0, edge_verts_coarse)
+            edge_vertices_coarse = convert_adjacency_matrix_to_list(parent_edge_vertices, ncols_per_row=2)
+            edge_vertices.insert(0, edge_vertices_coarse)
 
             # store cell-to-vert adjacency matrix
             if ilevel > self.max_level:
@@ -346,7 +350,7 @@ def _get_hierarchy_of_icon_edge_graphs(
             # replace edge-to-vertex and vert-to-vert adjacency matrices (for next level):
             if ilevel > 1:
                 v2v_matrix = s_v_coarse * v2v_matrix * v2v_fine2coarse
-                e2v_matrix = convert_list_to_adjacency_matrix(edge_verts_coarse, num_verts)
+                e2v_matrix = convert_list_to_adjacency_matrix(edge_vertices_coarse, num_vertices)
 
         # Fine-level cells outside of multi-mesh (LAM boundary)
         # correspond to empty rows in the adjacency matrix. We
@@ -415,8 +419,8 @@ def _get_grid2mesh_edges(self, select_c: np.ndarray, multi_mesh: ICONMultiMesh)
         """
 
         num_cells = select_c.shape[0]
-        num_verts_per_cell = multi_mesh.cell_vertices.shape[1]
-        src_list = np.kron(np.arange(num_cells), np.ones((1, num_verts_per_cell), dtype=np.int64)).flatten()
+        num_vertices_per_cell = multi_mesh.cell_vertices.shape[1]
+        src_list = np.kron(np.arange(num_cells), np.ones((1, num_vertices_per_cell), dtype=np.int64)).flatten()
         dst_list = multi_mesh.cell_vertices[select_c[:, 0], :].flatten()
         edge_vertices = np.stack((src_list, dst_list), axis=1, dtype=np.int64)
         return edge_vertices

From 5e21ff5ce9309130ecb339ad7ee3f7115cfb7fcc Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 06:58:57 +0000
Subject: [PATCH 05/35] fix: changed float division to double slash operator.

---
 src/anemoi/graphs/generate/icon_mesh.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index cf10d12..73c6291 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -542,7 +542,7 @@ def convert_adjacency_matrix_to_list(
     if remove_duplicates:
         # The edges-vertex adjacency matrix may have duplicate rows, remove
         # them by selecting the rows that are unique:
-        nrows = int(adj_matrix.nnz / ncols_per_row)
+        nrows = int(adj_matrix.nnz // ncols_per_row)
         mat = adj_matrix.indices.reshape((nrows, ncols_per_row))
         return np.unique(mat, axis=0)
     else:

From 3d662d171e5f5f13e113b0bc469c4362624836d3 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:00:27 +0000
Subject: [PATCH 06/35] refactor: removed unnecessary `else` branch.

---
 src/anemoi/graphs/generate/icon_mesh.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 73c6291..2dab380 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -146,12 +146,12 @@ def edge_features(self) -> torch.tensor:
         )
         if self.edge_id is None:
             return edge_feature_arr
-        else:
-            # one-hot-encoded categorical data (`edge_id`)
-            one_hot = torch_geometric.utils.one_hot(
-                index=torch.tensor(self.edge_id.edge_id), num_classes=self.edge_id.num_classes
-            )
-            return torch.concatenate((one_hot, edge_feature_arr), dim=1)
+
+        # one-hot-encoded categorical data (`edge_id`)
+        one_hot = torch_geometric.utils.one_hot(
+            index=torch.tensor(self.edge_id.edge_id), num_classes=self.edge_id.num_classes
+        )
+        return torch.concatenate((one_hot, edge_feature_arr), dim=1)
 
     def set_constant_edge_id(self, edge_id: int, num_classes: int):
         self.edge_id = EdgeID(edge_id=np.full(self.num_edges, fill_value=edge_id), num_classes=num_classes)

From cfd82a9ebca14336428cddf03294e1d2d836a698 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:03:32 +0000
Subject: [PATCH 07/35] refactor: more descriptive variable name in for loop.

---
 src/anemoi/graphs/generate/icon_mesh.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 2dab380..7e8af26 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -264,7 +264,7 @@ def _restrict_multi_mesh_level(
         vertex_glb2loc = np.zeros(num_vertices, dtype=int)
         vertex_glb2loc[vertex_idx] = np.arange(vertex_idx.sum())
         return (
-            [vertex_glb2loc[M] for M in edge_vertices[: self.max_level + 1]],
+            [vertex_glb2loc[vertices] for vertices in edge_vertices[: self.max_level + 1]],
             # cell_vertices: preserve negative indices (incomplete cells)
             np.where(cell_vertices >= 0, vertex_glb2loc[cell_vertices], cell_vertices),
             vlon[vertex_idx],

From 48230ec1c197fd1d7ffb4883c19ad3c98d621d74 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:04:51 +0000
Subject: [PATCH 08/35] refactor: renamed variable (ignoring the minus sign of
 `phi`).

---
 src/anemoi/graphs/generate/icon_mesh.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 7e8af26..a4a8031 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -468,10 +468,10 @@ def get_edge_attributes(
     """
 
     phi_theta = gc_recv  # precession and nutation angle
-    theta__phi = np.fliplr(phi_theta)
-    theta__phi[:, 1] *= -1.0  # angles = [theta, -phi]
+    theta_phi = np.fliplr(phi_theta)
+    theta_phi[:, 1] *= -1.0  # angles = [theta, -phi]
     # rotation matrix:
-    R = scipy.spatial.transform.Rotation.from_euler(seq="YZ", angles=theta__phi)
+    R = scipy.spatial.transform.Rotation.from_euler(seq="YZ", angles=theta_phi)
     edge_length = np.array([arc_length(cc_recv, cc_send)])
     distance = R.apply(cc_send) - [1.0, 0.0, 0.0]  # subtract the rotated position of sender
     return np.concatenate((edge_length.transpose(), distance), axis=1)

From 6ad883f073da15d9a70ac80d2c670dfeb454523e Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:05:59 +0000
Subject: [PATCH 09/35] refactor: changed name of temporary variable.

---
 src/anemoi/graphs/generate/icon_mesh.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index a4a8031..c47cfde 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -471,9 +471,9 @@ def get_edge_attributes(
     theta_phi = np.fliplr(phi_theta)
     theta_phi[:, 1] *= -1.0  # angles = [theta, -phi]
     # rotation matrix:
-    R = scipy.spatial.transform.Rotation.from_euler(seq="YZ", angles=theta_phi)
+    rotation = scipy.spatial.transform.Rotation.from_euler(seq="YZ", angles=theta_phi)
     edge_length = np.array([arc_length(cc_recv, cc_send)])
-    distance = R.apply(cc_send) - [1.0, 0.0, 0.0]  # subtract the rotated position of sender
+    distance = rotation.apply(cc_send) - [1.0, 0.0, 0.0]  # subtract the rotated position of sender
     return np.concatenate((edge_length.transpose(), distance), axis=1)
 
 

From 1862a66ff1df8c5fc47426ede036147f3da4dc94 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:06:53 +0000
Subject: [PATCH 10/35] =?UTF-8?q?refactor:=20remove=20unnecessary=20`else?=
 =?UTF-8?q?=C2=B4branch.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/anemoi/graphs/generate/icon_mesh.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index c47cfde..abdc8a5 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -545,9 +545,9 @@ def convert_adjacency_matrix_to_list(
         nrows = int(adj_matrix.nnz // ncols_per_row)
         mat = adj_matrix.indices.reshape((nrows, ncols_per_row))
         return np.unique(mat, axis=0)
-    else:
-        nrows = adj_matrix.shape[0]
-        return adj_matrix.indices.reshape((nrows, ncols_per_row))
+
+    nrows = adj_matrix.shape[0]
+    return adj_matrix.indices.reshape((nrows, ncols_per_row))
 
 
 @typechecked

From 6f6e50143818b25ecd1681a0f82eabd23b292320 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:08:49 +0000
Subject: [PATCH 11/35] refactor: added type annotation for completeness.

---
 src/anemoi/graphs/generate/icon_mesh.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index abdc8a5..f03497e 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -478,7 +478,7 @@ def get_edge_attributes(
 
 
 @typechecked
-def arc_length(v1, v2) -> np.ndarray:
+def arc_length(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
     """Calculate length of great circle arc on the unit sphere."""
     return np.arccos(np.clip(np.einsum("ij,ij->i", v1, v2), a_min=0.0, a_max=1.0))
 

From 9bc7ae0767be3928037c1eb31e79131b56f2421e Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:16:01 +0000
Subject: [PATCH 12/35] refactor: remove default in function argument.

---
 src/anemoi/graphs/generate/icon_mesh.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index f03497e..c811dfb 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -551,7 +551,7 @@ def convert_adjacency_matrix_to_list(
 
 
 @typechecked
-def selection_matrix(idx: np.ndarray, isize: int = 0) -> scipy.sparse.csr_matrix:
+def selection_matrix(idx: np.ndarray, isize: int) -> scipy.sparse.csr_matrix:
     """Create a diagonal selection matrix.
 
     Parameters

From 0fae14d06573d0c195ecab0c8db5a87b1d350fcb Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:19:12 +0000
Subject: [PATCH 13/35] refactor: change argument name to a more understandable
 name.

---
 src/anemoi/graphs/generate/icon_mesh.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index c811dfb..6c65656 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -551,14 +551,14 @@ def convert_adjacency_matrix_to_list(
 
 
 @typechecked
-def selection_matrix(idx: np.ndarray, isize: int) -> scipy.sparse.csr_matrix:
+def selection_matrix(idx: np.ndarray, num_diagonals: int) -> scipy.sparse.csr_matrix:
     """Create a diagonal selection matrix.
 
     Parameters
     ----------
     idx : np.ndarray
         integer array of indices
-    isize : int
+    num_diagonals : int
         size of (square) selection matrix
 
     Returns
@@ -566,4 +566,4 @@ def selection_matrix(idx: np.ndarray, isize: int) -> scipy.sparse.csr_matrix:
     scipy.sparse.csr_matrix
         diagonal matrix with ones at selected indices (idx,idx).
     """
-    return scipy.sparse.csr_matrix((np.ones((len(idx))), (idx, idx)), dtype=bool, shape=(isize, isize))
+    return scipy.sparse.csr_matrix((np.ones((len(idx))), (idx, idx)), dtype=bool, shape=(num_diagonals, num_diagonals))

From a44fb01b38bd83b1151e709b1b46c38f660f8f50 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:20:48 +0000
Subject: [PATCH 14/35] refactor: remove redundant code.

---
 src/anemoi/graphs/generate/icon_mesh.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 6c65656..e1ac260 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -124,10 +124,7 @@ class BipartiteGraph:
     def __init__(self, nodeset: tuple[NodeSet, NodeSet], edge_vertices: np.ndarray, edge_id: Optional[EdgeID] = None):
         self.nodeset = nodeset
         self.edge_vertices = edge_vertices
-        if edge_id is None:
-            self.edge_id = None
-        else:
-            self.edge_id = edge_id
+        self.edge_id = edge_id
 
     @property
     def num_edges(self) -> int:

From 25630cca7056cd86badbf7a538e6d6de5e573a91 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:23:15 +0000
Subject: [PATCH 15/35] refactor: more Pythonic if-else statement.

---
 src/anemoi/graphs/generate/icon_mesh.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index e1ac260..de862d7 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -163,10 +163,8 @@ def __add__(self, other: "BipartiteGraph"):
         shifted_edge_vertices = other.edge_vertices
         shifted_edge_vertices[:, 0] += self.nodeset[0].num_vertices
         # (Optional:) merge one-hot-encoded categorical data (`edge_id`)
-        if None not in (self.edge_id, other.edge_id):
-            edge_id = self.edge_id + other.edge_id
-        else:
-            edge_id = None
+        edge_id = None if None in (self.edge_id, other.edge_id) else self.edge_id + other.edge_id
+
         return BipartiteGraph(
             nodeset=(self.nodeset[0] + other.nodeset[0], self.nodeset[1]),
             edge_vertices=np.concatenate((self.edge_vertices, shifted_edge_vertices)),

From 612afe5a68ded4364a435e5f3ffb32dfafa85817 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:26:04 +0000
Subject: [PATCH 16/35] refactor: more Pythonic if-else statement.

---
 src/anemoi/graphs/generate/icon_mesh.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index de862d7..c4f5de3 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -220,10 +220,7 @@ def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None, ive
 
             self.uuidOfHGrid = ncfile.uuidOfHGrid
 
-        if max_level is not None:
-            self.max_level = max_level
-        else:
-            self.max_level = reflvl_vertex.max()
+        self.max_level = max_level if max_level is not None else reflvl_vertex.max()
 
         # generate edge-vertex relations for coarser levels:
         (edge_vertices, cell_vertices) = self._get_hierarchy_of_icon_edge_graphs(

From 368427ca0356d82117592f19c05f77ce2a0b2905 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:27:59 +0000
Subject: [PATCH 17/35] refactor: use more appropriate `LOGGER.debug` instead
 of verbosity flag.

---
 src/anemoi/graphs/generate/icon_mesh.py | 20 +++++++-------------
 src/anemoi/graphs/nodes/builder.py      |  1 -
 2 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index c4f5de3..646fb9b 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -199,11 +199,10 @@ class ICONMultiMesh(GeneralGraph):
     nodeset: NodeSet  # set of ICON grid vertices
     cell_vertices: np.ndarray
 
-    def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None, iverbosity: int = 0):
+    def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None):
 
         # open file, representing the finest level
-        if iverbosity > 0:
-            LOGGER.info(f"{type(self).__name__}: read ICON grid file '{icon_grid_filename}'")
+        LOGGER.debug(f"{type(self).__name__}: read ICON grid file '{icon_grid_filename}'")
         with netCDF4.Dataset(icon_grid_filename, "r") as ncfile:
             # read vertex coordinates
             vlon = read_coordinate_array(ncfile, "vlon", "vertex")
@@ -227,7 +226,6 @@ def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None, ive
             edge_vertices_fine=edge_vertices_fine,
             cell_vertices_fine=cell_vertices_fine,
             reflvl_vertex=reflvl_vertex,
-            iverbosity=iverbosity,
         )
         # restrict edge-vertex list to multi_mesh level "max_level":
         if self.max_level < len(edge_vertices):
@@ -268,7 +266,6 @@ def _get_hierarchy_of_icon_edge_graphs(
         edge_vertices_fine: np.ndarray,
         cell_vertices_fine: np.ndarray,
         reflvl_vertex: np.ndarray,
-        iverbosity: int = 1,
     ) -> tuple[list[np.ndarray], np.ndarray]:
         """Returns a list of edge-vertex relations (coarsest to finest level)."""
 
@@ -286,8 +283,7 @@ def _get_hierarchy_of_icon_edge_graphs(
 
         # coarsen edge-vertex list from level `ilevel -> ilevel - 1`:
         for ilevel in reversed(range(1, reflvl_vertex.max() + 1)):
-            if iverbosity > 0:
-                LOGGER.info(f"  edges[{ilevel}] = {edge_vertices[0].shape[0] : >9}")
+            LOGGER.debug(f"  edges[{ilevel}] = {edge_vertices[0].shape[0] : >9}")
 
             # define edge selection matrix (selecting only edges of which have
             # exactly one coarse vertex):
@@ -375,11 +371,9 @@ def __init__(
         icon_grid_filename: str,
         multi_mesh: Optional[ICONMultiMesh] = None,
         max_level: Optional[int] = None,
-        iverbosity: int = 0,
     ):
         # open file, representing the finest level
-        if iverbosity > 0:
-            LOGGER.info(f"{type(self).__name__}: read ICON grid file '{icon_grid_filename}'")
+        LOGGER.debug(f"{type(self).__name__}: read ICON grid file '{icon_grid_filename}'")
         with netCDF4.Dataset(icon_grid_filename, "r") as ncfile:
             # read cell circumcenter coordinates
             clon = read_coordinate_array(ncfile, "clon", "cell")
@@ -423,12 +417,12 @@ def _get_grid2mesh_edges(self, select_c: np.ndarray, multi_mesh: ICONMultiMesh)
 
 @typechecked
 def get_icon_mesh_and_grid(
-    iverbosity: int, grid_file: str, max_level_multimesh: int, max_level_dataset: int
+    grid_file: str, max_level_multimesh: int, max_level_dataset: int
 ) -> tuple[ICONMultiMesh, ICONCellDataGrid]:
     """Factory function, creating an ICON multi-mesh and an ICON cell-grid."""
     return (
-        multi_mesh := ICONMultiMesh(grid_file, max_level=max_level_multimesh, iverbosity=iverbosity),
-        ICONCellDataGrid(grid_file, multi_mesh, max_level=max_level_dataset, iverbosity=iverbosity),
+        multi_mesh := ICONMultiMesh(grid_file, max_level=max_level_multimesh),
+        ICONCellDataGrid(grid_file, multi_mesh, max_level=max_level_dataset),
     )
 
 
diff --git a/src/anemoi/graphs/nodes/builder.py b/src/anemoi/graphs/nodes/builder.py
index 837005b..26582a4 100644
--- a/src/anemoi/graphs/nodes/builder.py
+++ b/src/anemoi/graphs/nodes/builder.py
@@ -288,7 +288,6 @@ class ICONNodes(BaseNodeBuilder):
     def __init__(self, name: str, grid_filename: str, max_level_multimesh: int, max_level_dataset: int) -> None:
         self.grid_filename = grid_filename
         self.multi_mesh, self.cell_grid = get_icon_mesh_and_grid(
-            iverbosity=1,
             grid_file=self.grid_filename,
             max_level_multimesh=max_level_multimesh,
             max_level_dataset=max_level_dataset,

From 5c76e10c5f573572b1aa4bab33ad858304930920 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:29:33 +0000
Subject: [PATCH 18/35] refactor: more appropriate variable name.

---
 src/anemoi/graphs/generate/icon_mesh.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 646fb9b..211fbf6 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -250,15 +250,15 @@ def _restrict_multi_mesh_level(
         """Creates a new mesh with only the vertices at the desired level."""
 
         num_vertices = reflvl_vertex.shape[0]
-        vertex_idx = reflvl_vertex <= self.max_level
+        vertex_mask = reflvl_vertex <= self.max_level
         vertex_glb2loc = np.zeros(num_vertices, dtype=int)
-        vertex_glb2loc[vertex_idx] = np.arange(vertex_idx.sum())
+        vertex_glb2loc[vertex_mask] = np.arange(vertex_mask.sum())
         return (
             [vertex_glb2loc[vertices] for vertices in edge_vertices[: self.max_level + 1]],
             # cell_vertices: preserve negative indices (incomplete cells)
             np.where(cell_vertices >= 0, vertex_glb2loc[cell_vertices], cell_vertices),
-            vlon[vertex_idx],
-            vlat[vertex_idx],
+            vlon[vertex_mask],
+            vlat[vertex_mask],
         )
 
     def _get_hierarchy_of_icon_edge_graphs(

From f6d63f0056376a93f585ce565bc22360185750d2 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 07:30:48 +0000
Subject: [PATCH 19/35] refactor: more appropriate variable name.

---
 src/anemoi/graphs/generate/icon_mesh.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 211fbf6..4eed667 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -290,8 +290,8 @@ def _get_hierarchy_of_icon_edge_graphs(
             #
             # get a boolean mask, matching all edges where one of its vertices
             # has refinement level index `ilevel`:
-            ref_level = reflvl_vertex[edge_vertices[0]] == ilevel
-            edges_coarse = np.logical_xor(ref_level[:, 0], ref_level[:, 1])  # = bisected coarse edges
+            ref_level_mask = reflvl_vertex[edge_vertices[0]] == ilevel
+            edges_coarse = np.logical_xor(ref_level_mask[:, 0], ref_level_mask[:, 1])  # = bisected coarse edges
             idx_e2e = np.argwhere(edges_coarse).flatten()
             s_edges = selection_matrix(idx_e2e, edges_coarse.shape[0])
 

From 8efd2c6975163104bc71b4e7e2576a8166a88b7e Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 09:50:17 +0000
Subject: [PATCH 20/35] refactor: unified the three ICON Edgebuilders and the
 two Nodebuilders.

---
 src/anemoi/graphs/edges/builder.py | 101 ++++++++++-------------------
 src/anemoi/graphs/nodes/builder.py |  39 +++++------
 2 files changed, 50 insertions(+), 90 deletions(-)

diff --git a/src/anemoi/graphs/edges/builder.py b/src/anemoi/graphs/edges/builder.py
index 094815d..1038b14 100644
--- a/src/anemoi/graphs/edges/builder.py
+++ b/src/anemoi/graphs/edges/builder.py
@@ -353,8 +353,8 @@ def update_graph(self, graph: HeteroData, attrs_config: DotDict | None = None) -
         return super().update_graph(graph, attrs_config)
 
 
-class ICONTopologicalProcessorEdges(BaseEdgeBuilder):
-    """Computes edges based on ICON grid topology.
+class ICONTopologicalBaseEdgeBuilder(BaseEdgeBuilder):
+    """Base class for computing edges based on ICON grid topology.
 
     Attributes
     ----------
@@ -372,7 +372,7 @@ def __init__(self, source_name: str, target_name: str, icon_mesh: str):
 
     def update_graph(self, graph: HeteroData, attrs_config: DotDict = None) -> HeteroData:
         """Update the graph with the edges."""
-        self.multi_mesh = graph[self.icon_mesh]["_multi_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):
@@ -384,84 +384,49 @@ def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStor
             The target nodes.
         """
         LOGGER.info(f"Using ICON topology {self.source_name}>{self.target_name}")
-        nrows = self.multi_mesh.num_edges
+        nrows = self.icon_sub_graph.num_edges
         adj_matrix = scipy.sparse.coo_matrix(
-            (np.ones(nrows), (self.multi_mesh.edge_vertices[:, 1], self.multi_mesh.edge_vertices[:, 0]))
+            (
+                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 ICONTopologicalEncoderEdges(BaseEdgeBuilder):
-    """Computes 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)
+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):
-        self.icon_mesh = icon_mesh
-        super().__init__(source_name, target_name)
+        self.sub_graph_address = "_multi_mesh"
+        self.vertex_index = (1, 0)
+        super().__init__(source_name, target_name, icon_mesh)
 
-    def update_graph(self, graph: HeteroData, attrs_config: DotDict = None) -> HeteroData:
-        """Update the graph with the edges."""
-        self.cell_data_grid = graph[self.icon_mesh]["_cell_grid"]
-        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.cell_data_grid.num_edges
-        adj_matrix = scipy.sparse.coo_matrix(
-            (np.ones(nrows), (self.cell_data_grid.edge_vertices[:, 1], self.cell_data_grid.edge_vertices[:, 0]))
-        )
-        return adj_matrix
+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):
+        self.sub_graph_address = "_cell_grid"
+        self.vertex_index = (1, 0)
+        super().__init__(source_name, target_name, icon_mesh)
 
-class ICONTopologicalDecoderEdges(BaseEdgeBuilder):
-    """Computes 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)
+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):
-        self.icon_mesh = icon_mesh
-        super().__init__(source_name, target_name)
-
-    def update_graph(self, graph: HeteroData, attrs_config: DotDict = None) -> HeteroData:
-        """Update the graph with the edges."""
-        self.cell_data_grid = graph[self.icon_mesh]["_cell_grid"]
-        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.cell_data_grid.num_edges
-        adj_matrix = scipy.sparse.coo_matrix(
-            (np.ones(nrows), (self.cell_data_grid.edge_vertices[:, 0], self.cell_data_grid.edge_vertices[:, 1]))
-        )
-        return adj_matrix
+        self.sub_graph_address = "_cell_grid"
+        self.vertex_index = (0, 1)
+        super().__init__(source_name, target_name, icon_mesh)
diff --git a/src/anemoi/graphs/nodes/builder.py b/src/anemoi/graphs/nodes/builder.py
index 26582a4..9191b7c 100644
--- a/src/anemoi/graphs/nodes/builder.py
+++ b/src/anemoi/graphs/nodes/builder.py
@@ -304,8 +304,8 @@ def register_attributes(self, graph: HeteroData, config: DotDict) -> HeteroData:
         return super().register_attributes(graph, config)
 
 
-class ICONMultimeshNodes(BaseNodeBuilder):
-    """Processor mesh based on an ICON grid.
+class ICONTopologicalBaseNodeBuilder(BaseNodeBuilder):
+    """Base class for data mesh or processor mesh based on an ICON grid.
 
     Parameters
     ----------
@@ -321,35 +321,30 @@ def __init__(self, name: str, icon_mesh: str) -> None:
 
     def update_graph(self, graph: HeteroData, attr_config: DotDict | None = None) -> HeteroData:
         """Update the graph with new nodes."""
-        self.multi_mesh = graph[self.icon_mesh]["_multi_mesh"]
+        self.icon_sub_graph = graph[self.icon_mesh][self.sub_graph_address]
         return super().update_graph(graph, attr_config)
 
-    def get_coordinates(self) -> torch.Tensor:
-        return torch.from_numpy(self.multi_mesh.nodeset.gc_vertices.astype(np.float32)).fliplr()
 
+class ICONMultimeshNodes(ICONTopologicalBaseNodeBuilder):
+    """Processor mesh based on an ICON grid."""
 
-class ICONCellGridNodes(BaseNodeBuilder):
-    """Data mesh based on an ICON grid.
+    def __init__(self, name: str, icon_mesh: str) -> None:
+        self.sub_graph_address = "_multi_mesh"
+        super().__init__(name, icon_mesh)
 
-    Parameters
-    ----------
-    name : str
-        key for the nodes in the HeteroData graph object.
-    icon_mesh : str
-        key corresponding to the ICON mesh (cells and vertices).
-    """
+    def get_coordinates(self) -> torch.Tensor:
+        return torch.from_numpy(self.icon_sub_graph.nodeset.gc_vertices.astype(np.float32)).fliplr()
 
-    def __init__(self, name: str, icon_mesh: str) -> None:
-        self.icon_mesh = icon_mesh
-        super().__init__(name)
 
-    def update_graph(self, graph: HeteroData, attr_config: DotDict | None = None) -> HeteroData:
-        """Update the graph with new nodes."""
-        self.cell_grid = graph[self.icon_mesh]["_cell_grid"]
-        return super().update_graph(graph, attr_config)
+class ICONCellGridNodes(ICONTopologicalBaseNodeBuilder):
+    """Data mesh based on an ICON grid."""
+
+    def __init__(self, name: str, icon_mesh: str) -> None:
+        self.sub_graph_address = "_cell_grid"
+        super().__init__(name, icon_mesh)
 
     def get_coordinates(self) -> torch.Tensor:
-        return torch.from_numpy(self.cell_grid.nodeset[0].gc_vertices.astype(np.float32)).fliplr()
+        return torch.from_numpy(self.icon_sub_graph.nodeset[0].gc_vertices.astype(np.float32)).fliplr()
 
 
 class HEALPixNodes(BaseNodeBuilder):

From 1961150cfc22e62f71a4ab2e26cea7552ed5d79d Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 10:17:18 +0000
Subject: [PATCH 21/35] refactor: more verbose but also clearer names for
 variables in mesh construction algorithm.

---
 src/anemoi/graphs/generate/icon_mesh.py | 56 +++++++++++++------------
 1 file changed, 29 insertions(+), 27 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 4eed667..eea36a9 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -273,13 +273,13 @@ def _get_hierarchy_of_icon_edge_graphs(
 
         num_vertices = reflvl_vertex.shape[0]
         # edge-to-vertex adjacency matrix with 2 non-zero entries per row:
-        e2v_matrix = convert_list_to_adjacency_matrix(edge_vertices_fine, num_vertices)
+        edge2vertex_matrix = convert_list_to_adjacency_matrix(edge_vertices_fine, num_vertices)
         # cell-to-vertex adjacency matrix with 3 non-zero entries per row:
-        c2v_matrix = convert_list_to_adjacency_matrix(cell_vertices_fine, num_vertices)
-        v2v_matrix = e2v_matrix.transpose() * e2v_matrix
-        v2v_matrix.setdiag(np.ones(num_vertices))  # vertices are self-connected
+        cell2vertex_matrix = convert_list_to_adjacency_matrix(cell_vertices_fine, num_vertices)
+        vertex2vertex_matrix = edge2vertex_matrix.transpose() * edge2vertex_matrix
+        vertex2vertex_matrix.setdiag(np.ones(num_vertices))  # vertices are self-connected
 
-        s_v_coarse = scipy.sparse.diags(np.ones(num_vertices), dtype=bool)
+        selected_vertex_coarse = scipy.sparse.diags(np.ones(num_vertices), dtype=bool)
 
         # coarsen edge-vertex list from level `ilevel -> ilevel - 1`:
         for ilevel in reversed(range(1, reflvl_vertex.max() + 1)):
@@ -292,34 +292,36 @@ def _get_hierarchy_of_icon_edge_graphs(
             # has refinement level index `ilevel`:
             ref_level_mask = reflvl_vertex[edge_vertices[0]] == ilevel
             edges_coarse = np.logical_xor(ref_level_mask[:, 0], ref_level_mask[:, 1])  # = bisected coarse edges
-            idx_e2e = np.argwhere(edges_coarse).flatten()
-            s_edges = selection_matrix(idx_e2e, edges_coarse.shape[0])
+            idx_edge2edge = np.argwhere(edges_coarse).flatten()
+            selected_edges = selection_matrix(idx_edge2edge, edges_coarse.shape[0])
 
             # define vertex selection matrix selecting only vertices of
             # level `ilevel`:
             idx_v_fine = np.argwhere(reflvl_vertex == ilevel).flatten()
-            s_v_fine = selection_matrix(idx_v_fine, num_vertices)
+            selected_vertex_fine = selection_matrix(idx_v_fine, num_vertices)
             # define vertex selection matrix, selecting only vertices of
             # level < `ilevel`, by successively removing `s_fine` from an identity matrix.
-            s_v_coarse.data[0][idx_v_fine] = False
+            selected_vertex_coarse.data[0][idx_v_fine] = False
 
             # create an adjacency matrix which links each fine level
             # vertex to its two coarser neighbor vertices:
-            v2v_fine2coarse = s_v_fine * v2v_matrix * s_v_coarse
+            vertex2vertex_fine2coarse = selected_vertex_fine * vertex2vertex_matrix * selected_vertex_coarse
             # remove rows that have only one non-zero entry
             # (corresponding to incomplete parent triangles in LAM grids):
-            csum = v2v_fine2coarse * np.ones((v2v_fine2coarse.shape[1], 1))
-            s_v2v = selection_matrix(np.argwhere(csum == 2).flatten(), v2v_fine2coarse.shape[0])
-            v2v_fine2coarse = s_v2v * v2v_fine2coarse
+            csum = vertex2vertex_fine2coarse * np.ones((vertex2vertex_fine2coarse.shape[1], 1))
+            selected_vertex2vertex = selection_matrix(
+                np.argwhere(csum == 2).flatten(), vertex2vertex_fine2coarse.shape[0]
+            )
+            vertex2vertex_fine2coarse = selected_vertex2vertex * vertex2vertex_fine2coarse
 
             # then construct the edges-to-parent-vertex adjacency matrix:
-            parent_edge_vertices = s_edges * e2v_matrix * v2v_fine2coarse
+            parent_edge_vertices = selected_edges * edge2vertex_matrix * vertex2vertex_fine2coarse
             # again, we have might have selected edges within
-            # `s_edges` which are part of an incomplete parent edge
+            # `selected_edges` which are part of an incomplete parent edge
             # (LAM case). We filter these here:
             csum = parent_edge_vertices * np.ones((parent_edge_vertices.shape[1], 1))
-            s_e2e = selection_matrix(np.argwhere(csum == 2).flatten(), parent_edge_vertices.shape[0])
-            parent_edge_vertices = s_e2e * parent_edge_vertices
+            selected_edge2edge = selection_matrix(np.argwhere(csum == 2).flatten(), parent_edge_vertices.shape[0])
+            parent_edge_vertices = selected_edge2edge * parent_edge_vertices
 
             # note: the edges-vertex adjacency matrix still has duplicate
             # rows, since two child edges have the same parent edge.
@@ -328,28 +330,28 @@ def _get_hierarchy_of_icon_edge_graphs(
 
             # store cell-to-vert adjacency matrix
             if ilevel > self.max_level:
-                c2v_matrix = c2v_matrix * v2v_fine2coarse
+                cell2vertex_matrix = cell2vertex_matrix * vertex2vertex_fine2coarse
                 # similar to the treatment above, we need to handle
                 # coarse LAM cells which are incomplete.
-                csum = c2v_matrix * np.ones((c2v_matrix.shape[1], 1))
-                s_c2c = selection_matrix(np.argwhere(csum == 3).flatten(), c2v_matrix.shape[0])
-                c2v_matrix = s_c2c * c2v_matrix
+                csum = cell2vertex_matrix * np.ones((cell2vertex_matrix.shape[1], 1))
+                selected_cell2cell = selection_matrix(np.argwhere(csum == 3).flatten(), cell2vertex_matrix.shape[0])
+                cell2vertex_matrix = selected_cell2cell * cell2vertex_matrix
 
             # replace edge-to-vertex and vert-to-vert adjacency matrices (for next level):
             if ilevel > 1:
-                v2v_matrix = s_v_coarse * v2v_matrix * v2v_fine2coarse
-                e2v_matrix = convert_list_to_adjacency_matrix(edge_vertices_coarse, num_vertices)
+                vertex2vertex_matrix = selected_vertex_coarse * vertex2vertex_matrix * vertex2vertex_fine2coarse
+                edge2vertex_matrix = convert_list_to_adjacency_matrix(edge_vertices_coarse, num_vertices)
 
         # Fine-level cells outside of multi-mesh (LAM boundary)
         # correspond to empty rows in the adjacency matrix. We
         # substitute these by three additional, non-existent vertices:
-        csum = 3 - c2v_matrix * np.ones((c2v_matrix.shape[1], 1))
-        nvmax = c2v_matrix.shape[1]
-        c2v_matrix = scipy.sparse.csr_matrix(scipy.sparse.hstack((c2v_matrix, csum, csum, csum)))
+        csum = 3 - cell2vertex_matrix * np.ones((cell2vertex_matrix.shape[1], 1))
+        nvmax = cell2vertex_matrix.shape[1]
+        cell2vertex_matrix = scipy.sparse.csr_matrix(scipy.sparse.hstack((cell2vertex_matrix, csum, csum, csum)))
 
         # build a list of cell-vertices [1..num_cells,1..3] for all
         # fine-level cells:
-        cell_vertices = convert_adjacency_matrix_to_list(c2v_matrix, remove_duplicates=False, ncols_per_row=3)
+        cell_vertices = convert_adjacency_matrix_to_list(cell2vertex_matrix, remove_duplicates=False, ncols_per_row=3)
 
         # finally, translate non-existent vertices into "-1" indices:
         cell_vertices = np.where(cell_vertices >= nvmax, -np.ones(cell_vertices.shape, dtype=np.int32), cell_vertices)

From 30dee38eea8d74924562606bd2d38c77ea9dd3e3 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 10:24:57 +0000
Subject: [PATCH 22/35] remove: removed obsolete function
 `set_constant_edge_id`.

---
 src/anemoi/graphs/generate/icon_mesh.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index eea36a9..de30094 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -150,9 +150,6 @@ def edge_features(self) -> torch.tensor:
         )
         return torch.concatenate((one_hot, edge_feature_arr), dim=1)
 
-    def set_constant_edge_id(self, edge_id: int, num_classes: int):
-        self.edge_id = EdgeID(edge_id=np.full(self.num_edges, fill_value=edge_id), num_classes=num_classes)
-
     def __add__(self, other: "BipartiteGraph"):
         """Concatenates two bipartite graphs that share a common target node set.
         Shifts the node indices of the second bipartite graph.

From 92cb023c2b4d27ef1db83eb76330dd99152c58e0 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Mon, 30 Sep 2024 11:29:03 +0000
Subject: [PATCH 23/35] refactor: replaced the sequential ID counter by a UUID.

---
 src/anemoi/graphs/generate/icon_mesh.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index de30094..d3e1e14 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -9,6 +9,7 @@
 
 import itertools
 import logging
+import uuid
 from dataclasses import dataclass
 from functools import cached_property
 from typing import Optional
@@ -33,7 +34,7 @@ class NodeSet:
 
     def __init__(self, lon: np.ndarray, lat: np.ndarray):
         self.gc_vertices = np.column_stack((lon, lat))
-        self.id = next(self.id_iter)
+        self.id = uuid.uuid4()
 
     @property
     def num_vertices(self) -> int:

From 9937a610d67c263d22fe2e574749e6dfba1318e7 Mon Sep 17 00:00:00 2001
From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de>
Date: Wed, 16 Oct 2024 12:39:42 +0200
Subject: [PATCH 24/35] revert change of copyright notice

---
 src/anemoi/graphs/commands/__init__.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/anemoi/graphs/commands/__init__.py b/src/anemoi/graphs/commands/__init__.py
index 30e21bc..cebb539 100644
--- a/src/anemoi/graphs/commands/__init__.py
+++ b/src/anemoi/graphs/commands/__init__.py
@@ -1,11 +1,11 @@
 #!/usr/bin/env python
-# (C) Copyright 2024 ECMWF, DWD.
+# (C) Copyright 2024 ECMWF.
 #
 # 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, the above institution do 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.
+# 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 os

From 75057c372286e55f10a3f7d409034c31072b1f2b Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Wed, 6 Nov 2024 10:12:47 +0100
Subject: [PATCH 25/35] [fix]  add encoder & processor edges to the __all__
 variable

---
 src/anemoi/graphs/edges/__init__.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/anemoi/graphs/edges/__init__.py b/src/anemoi/graphs/edges/__init__.py
index 2424982..f5737f2 100644
--- a/src/anemoi/graphs/edges/__init__.py
+++ b/src/anemoi/graphs/edges/__init__.py
@@ -5,4 +5,11 @@
 from .builder import KNNEdges
 from .builder import MultiScaleEdges
 
-__all__ = ["KNNEdges", "CutOffEdges", "MultiScaleEdges", "ICONTopologicalDecoderEdges"]
+__all__ = [
+    "KNNEdges",
+    "CutOffEdges",
+    "MultiScaleEdges",
+    "ICONTopologicalProcessorEdges",
+    "ICONTopologicalEncoderEdges",
+    "ICONTopologicalDecoderEdges",
+]

From dcff399876836630158c0f77889ea7e9bb6fe135 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Wed, 6 Nov 2024 10:17:39 +0100
Subject: [PATCH 26/35] [refactor]  move auxiliary functions to utils.py in
 graphs/generate/

---
 src/anemoi/graphs/generate/icon_mesh.py | 118 ++++++++----------------
 src/anemoi/graphs/generate/utils.py     |  76 +++++++++++++++
 2 files changed, 113 insertions(+), 81 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index d3e1e14..72e10b0 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -22,6 +22,10 @@
 from typeguard import typechecked
 from typing_extensions import Self
 
+from anemoi.graphs.generate.utils import convert_adjacency_matrix_to_list
+from anemoi.graphs.generate.utils import convert_list_to_adjacency_matrix
+from anemoi.graphs.generate.utils import selection_matrix
+
 LOGGER = logging.getLogger(__name__)
 
 
@@ -75,7 +79,10 @@ class EdgeID:
     def __add__(self, other: Self):
         """Concatenates two edge ID datasets."""
         assert self.num_classes == other.num_classes
-        return EdgeID(edge_id=np.concatenate((self.edge_id, other.edge_id)), num_classes=self.num_classes)
+        return EdgeID(
+            edge_id=np.concatenate((self.edge_id, other.edge_id)),
+            num_classes=self.num_classes,
+        )
 
 
 @typechecked
@@ -122,7 +129,12 @@ class BipartiteGraph:
     edge_vertices: np.ndarray  # vertex indices for each edge, shape [:,2]
     edge_id: np.ndarray  # additional ID for each edge (markers for heterogeneous input)
 
-    def __init__(self, nodeset: tuple[NodeSet, NodeSet], edge_vertices: np.ndarray, edge_id: Optional[EdgeID] = None):
+    def __init__(
+        self,
+        nodeset: tuple[NodeSet, NodeSet],
+        edge_vertices: np.ndarray,
+        edge_id: Optional[EdgeID] = None,
+    ):
         self.nodeset = nodeset
         self.edge_vertices = edge_vertices
         self.edge_id = edge_id
@@ -147,7 +159,8 @@ def edge_features(self) -> torch.tensor:
 
         # one-hot-encoded categorical data (`edge_id`)
         one_hot = torch_geometric.utils.one_hot(
-            index=torch.tensor(self.edge_id.edge_id), num_classes=self.edge_id.num_classes
+            index=torch.tensor(self.edge_id.edge_id),
+            num_classes=self.edge_id.num_classes,
         )
         return torch.concatenate((one_hot, edge_feature_arr), dim=1)
 
@@ -177,7 +190,12 @@ class BipartiteGraphProximity(BipartiteGraph):
     in Cartesian coordinates).
     """
 
-    def __init__(self, nodeset: tuple[NodeSet, NodeSet], radius: float, edge_id: Optional[EdgeID] = None):
+    def __init__(
+        self,
+        nodeset: tuple[NodeSet, NodeSet],
+        radius: float,
+        edge_id: Optional[EdgeID] = None,
+    ):
 
         src, tgt = nodeset
         point_tree = scipy.spatial.KDTree(tgt.cc_vertices)
@@ -228,7 +246,11 @@ def __init__(self, icon_grid_filename: str, max_level: Optional[int] = None):
         # restrict edge-vertex list to multi_mesh level "max_level":
         if self.max_level < len(edge_vertices):
             (self.edge_vertices, self.cell_vertices, vlon, vlat) = self._restrict_multi_mesh_level(
-                edge_vertices, cell_vertices, reflvl_vertex=reflvl_vertex, vlon=vlon, vlat=vlat
+                edge_vertices,
+                cell_vertices,
+                reflvl_vertex=reflvl_vertex,
+                vlon=vlon,
+                vlat=vlat,
             )
         # store vertices as a `NodeSet`:
         self.nodeset = NodeSet(vlon, vlat)
@@ -352,7 +374,11 @@ def _get_hierarchy_of_icon_edge_graphs(
         cell_vertices = convert_adjacency_matrix_to_list(cell2vertex_matrix, remove_duplicates=False, ncols_per_row=3)
 
         # finally, translate non-existent vertices into "-1" indices:
-        cell_vertices = np.where(cell_vertices >= nvmax, -np.ones(cell_vertices.shape, dtype=np.int32), cell_vertices)
+        cell_vertices = np.where(
+            cell_vertices >= nvmax,
+            -np.ones(cell_vertices.shape, dtype=np.int32),
+            cell_vertices,
+        )
 
         return (edge_vertices, cell_vertices)
 
@@ -459,7 +485,11 @@ def get_edge_attributes(
     # rotation matrix:
     rotation = scipy.spatial.transform.Rotation.from_euler(seq="YZ", angles=theta_phi)
     edge_length = np.array([arc_length(cc_recv, cc_send)])
-    distance = rotation.apply(cc_send) - [1.0, 0.0, 0.0]  # subtract the rotated position of sender
+    distance = rotation.apply(cc_send) - [
+        1.0,
+        0.0,
+        0.0,
+    ]  # subtract the rotated position of sender
     return np.concatenate((edge_length.transpose(), distance), axis=1)
 
 
@@ -479,77 +509,3 @@ def read_coordinate_array(ncfile, arrname: str, dimname: str) -> np.ndarray:
     # -> convert to regular arrays
     assert not arr.mask.any(), f"There are missing values in {arrname}"
     return arr.data
-
-
-@typechecked
-def convert_list_to_adjacency_matrix(list_matrix: np.ndarray, ncols: int = 0) -> scipy.sparse.csr_matrix:
-    """Convert an edge list into an adjacency matrix.
-
-    Parameters
-    ----------
-    list_matrix : np.ndarray
-        boolean matrix given by list of column indices for each row.
-    ncols : int
-        number of columns in result matrix.
-
-    Returns
-    -------
-    scipy.sparse.csr_matrix
-        sparse matrix [nrows, ncols]
-    """
-    nrows, ncols_per_row = list_matrix.shape
-    indptr = np.arange(ncols_per_row * (nrows + 1), step=ncols_per_row)
-    indices = list_matrix.ravel()
-    return scipy.sparse.csr_matrix((np.ones(nrows * ncols_per_row), indices, indptr), dtype=bool, shape=(nrows, ncols))
-
-
-@typechecked
-def convert_adjacency_matrix_to_list(
-    adj_matrix: scipy.sparse.csr_matrix,
-    ncols_per_row: int,
-    remove_duplicates: bool = True,
-) -> np.ndarray:
-    """Convert an adjacency matrix into an edge list.
-
-    Parameters
-    ----------
-    adj_matrix : scipy.sparse.csr_matrix
-        sparse (boolean) adjacency matrix
-    ncols_per_row : int
-        number of nonzero entries per row
-    remove_duplicates : bool
-        logical flag: remove duplicate rows.
-
-    Returns
-    -------
-    np.ndarray
-        boolean matrix given by list of column indices for each row.
-    """
-    if remove_duplicates:
-        # The edges-vertex adjacency matrix may have duplicate rows, remove
-        # them by selecting the rows that are unique:
-        nrows = int(adj_matrix.nnz // ncols_per_row)
-        mat = adj_matrix.indices.reshape((nrows, ncols_per_row))
-        return np.unique(mat, axis=0)
-
-    nrows = adj_matrix.shape[0]
-    return adj_matrix.indices.reshape((nrows, ncols_per_row))
-
-
-@typechecked
-def selection_matrix(idx: np.ndarray, num_diagonals: int) -> scipy.sparse.csr_matrix:
-    """Create a diagonal selection matrix.
-
-    Parameters
-    ----------
-    idx : np.ndarray
-        integer array of indices
-    num_diagonals : int
-        size of (square) selection matrix
-
-    Returns
-    -------
-    scipy.sparse.csr_matrix
-        diagonal matrix with ones at selected indices (idx,idx).
-    """
-    return scipy.sparse.csr_matrix((np.ones((len(idx))), (idx, idx)), dtype=bool, shape=(num_diagonals, num_diagonals))
diff --git a/src/anemoi/graphs/generate/utils.py b/src/anemoi/graphs/generate/utils.py
index 2dfc320..ccb2c0c 100644
--- a/src/anemoi/graphs/generate/utils.py
+++ b/src/anemoi/graphs/generate/utils.py
@@ -8,6 +8,8 @@
 # nor does it submit to any jurisdiction.
 
 import numpy as np
+import scipy
+from typeguard import typechecked
 
 
 def get_coordinates_ordering(coords: np.ndarray) -> np.ndarray:
@@ -29,3 +31,77 @@ def get_coordinates_ordering(coords: np.ndarray) -> np.ndarray:
     index_longitude = np.argsort(coords[index_latitude][:, 0])[::-1]
     node_ordering = np.arange(coords.shape[0])[index_latitude][index_longitude]
     return node_ordering
+
+
+@typechecked
+def convert_list_to_adjacency_matrix(list_matrix: np.ndarray, ncols: int = 0) -> scipy.sparse.csr_matrix:
+    """Convert an edge list into an adjacency matrix.
+
+    Parameters
+    ----------
+    list_matrix : np.ndarray
+        boolean matrix given by list of column indices for each row.
+    ncols : int
+        number of columns in result matrix.
+
+    Returns
+    -------
+    scipy.sparse.csr_matrix
+        sparse matrix [nrows, ncols]
+    """
+    nrows, ncols_per_row = list_matrix.shape
+    indptr = np.arange(ncols_per_row * (nrows + 1), step=ncols_per_row)
+    indices = list_matrix.ravel()
+    return scipy.sparse.csr_matrix((np.ones(nrows * ncols_per_row), indices, indptr), dtype=bool, shape=(nrows, ncols))
+
+
+@typechecked
+def convert_adjacency_matrix_to_list(
+    adj_matrix: scipy.sparse.csr_matrix,
+    ncols_per_row: int,
+    remove_duplicates: bool = True,
+) -> np.ndarray:
+    """Convert an adjacency matrix into an edge list.
+
+    Parameters
+    ----------
+    adj_matrix : scipy.sparse.csr_matrix
+        sparse (boolean) adjacency matrix
+    ncols_per_row : int
+        number of nonzero entries per row
+    remove_duplicates : bool
+        logical flag: remove duplicate rows.
+
+    Returns
+    -------
+    np.ndarray
+        boolean matrix given by list of column indices for each row.
+    """
+    if remove_duplicates:
+        # The edges-vertex adjacency matrix may have duplicate rows, remove
+        # them by selecting the rows that are unique:
+        nrows = int(adj_matrix.nnz // ncols_per_row)
+        mat = adj_matrix.indices.reshape((nrows, ncols_per_row))
+        return np.unique(mat, axis=0)
+
+    nrows = adj_matrix.shape[0]
+    return adj_matrix.indices.reshape((nrows, ncols_per_row))
+
+
+@typechecked
+def selection_matrix(idx: np.ndarray, num_diagonals: int) -> scipy.sparse.csr_matrix:
+    """Create a diagonal selection matrix.
+
+    Parameters
+    ----------
+    idx : np.ndarray
+        integer array of indices
+    num_diagonals : int
+        size of (square) selection matrix
+
+    Returns
+    -------
+    scipy.sparse.csr_matrix
+        diagonal matrix with ones at selected indices (idx,idx).
+    """
+    return scipy.sparse.csr_matrix((np.ones((len(idx))), (idx, idx)), dtype=bool, shape=(num_diagonals, num_diagonals))

From 9b19fade5418e2f5f782026cef7a5ce589660925 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Wed, 6 Nov 2024 11:11:35 +0100
Subject: [PATCH 27/35] [doc]  added empty torso for class documentation.

---
 docs/graphs/node_coordinates/icon_mesh.rst | 8 ++++++++
 1 file changed, 8 insertions(+)
 create mode 100644 docs/graphs/node_coordinates/icon_mesh.rst

diff --git a/docs/graphs/node_coordinates/icon_mesh.rst b/docs/graphs/node_coordinates/icon_mesh.rst
new file mode 100644
index 0000000..9bb01f2
--- /dev/null
+++ b/docs/graphs/node_coordinates/icon_mesh.rst
@@ -0,0 +1,8 @@
+####################################
+ Triangular Mesh with ICON Topology
+####################################
+
+This class defines nodes based on the `refinement_level` property
+attached to ICON grid files, where `0` denotes the vertices of the base
+grid, ie. the icosahedron including the step of root subdivision
+`RXXB00`.

From 611e264898bce97074d7fa8bb97e296a34882b59 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Fri, 8 Nov 2024 15:04:01 +0100
Subject: [PATCH 28/35] [fix] fixed interfaces (masks).

---
 src/anemoi/graphs/edges/builder.py | 71 ++++++++++++++++++++++++++----
 1 file changed, 62 insertions(+), 9 deletions(-)

diff --git a/src/anemoi/graphs/edges/builder.py b/src/anemoi/graphs/edges/builder.py
index af3bf32..92c3003 100644
--- a/src/anemoi/graphs/edges/builder.py
+++ b/src/anemoi/graphs/edges/builder.py
@@ -380,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)
@@ -458,12 +464,20 @@ class ICONTopologicalBaseEdgeBuilder(BaseEdgeBuilder):
         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):
+    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)
+        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)
 
@@ -494,10 +508,23 @@ class ICONTopologicalProcessorEdges(ICONTopologicalBaseEdgeBuilder):
     from ICON grid vertices.
     """
 
-    def __init__(self, source_name: str, target_name: str, icon_mesh: str):
+    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)
+        super().__init__(
+            source_name,
+            target_name,
+            icon_mesh,
+            source_mask_attr_name,
+            target_mask_attr_name,
+        )
 
 
 class ICONTopologicalEncoderEdges(ICONTopologicalBaseEdgeBuilder):
@@ -506,10 +533,23 @@ class ICONTopologicalEncoderEdges(ICONTopologicalBaseEdgeBuilder):
     vertices.
     """
 
-    def __init__(self, source_name: str, target_name: str, icon_mesh: str):
+    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)
+        super().__init__(
+            source_name,
+            target_name,
+            icon_mesh,
+            source_mask_attr_name,
+            target_mask_attr_name,
+        )
 
 
 class ICONTopologicalDecoderEdges(ICONTopologicalBaseEdgeBuilder):
@@ -518,7 +558,20 @@ class ICONTopologicalDecoderEdges(ICONTopologicalBaseEdgeBuilder):
     circumcenters.
     """
 
-    def __init__(self, source_name: str, target_name: str, icon_mesh: str):
+    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)
+        super().__init__(
+            source_name,
+            target_name,
+            icon_mesh,
+            source_mask_attr_name,
+            target_mask_attr_name,
+        )

From 797a6a21f2794caf0017bb404b4d9f444a072610 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Fri, 8 Nov 2024 15:23:56 +0100
Subject: [PATCH 29/35] [refactor] remove edge attribute calculation from this
 PR.

---
 src/anemoi/graphs/generate/icon_mesh.py       | 118 +-----------------
 src/anemoi/graphs/nodes/builders/from_icon.py |  12 +-
 2 files changed, 7 insertions(+), 123 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index 72e10b0..ddf26fe 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -17,8 +17,6 @@
 import netCDF4
 import numpy as np
 import scipy
-import torch
-import torch_geometric
 from typeguard import typechecked
 from typing_extensions import Self
 
@@ -87,7 +85,7 @@ def __add__(self, other: Self):
 
 @typechecked
 class GeneralGraph:
-    """Stores edges for a given node set, calculates edge features."""
+    """Stores edges for a given node set."""
 
     nodeset: NodeSet  # graph nodes
     edge_vertices: np.ndarray  # vertex indices for each edge, shape [:,2]
@@ -108,18 +106,6 @@ def num_vertices(self) -> int:
     def num_edges(self) -> int:
         return self.edge_vertices.shape[0]
 
-    @cached_property
-    def edge_features(self) -> torch.tensor:
-        return torch.tensor(
-            get_edge_attributes(
-                self.nodeset.gc_vertices[self.edge_vertices[:, 1]],
-                self.nodeset.cc_vertices[self.edge_vertices[:, 1]],
-                self.nodeset.gc_vertices[self.edge_vertices[:, 0]],
-                self.nodeset.cc_vertices[self.edge_vertices[:, 0]],
-            ),
-            dtype=torch.float32,
-        )
-
 
 @typechecked
 class BipartiteGraph:
@@ -143,27 +129,6 @@ def __init__(
     def num_edges(self) -> int:
         return self.edge_vertices.shape[0]
 
-    @cached_property
-    def edge_features(self) -> torch.tensor:
-        edge_feature_arr = torch.tensor(
-            get_edge_attributes(
-                self.nodeset[0].gc_vertices[self.edge_vertices[:, 0]],
-                self.nodeset[0].cc_vertices[self.edge_vertices[:, 0]],
-                self.nodeset[1].gc_vertices[self.edge_vertices[:, 1]],
-                self.nodeset[1].cc_vertices[self.edge_vertices[:, 1]],
-            ),
-            dtype=torch.float32,
-        )
-        if self.edge_id is None:
-            return edge_feature_arr
-
-        # one-hot-encoded categorical data (`edge_id`)
-        one_hot = torch_geometric.utils.one_hot(
-            index=torch.tensor(self.edge_id.edge_id),
-            num_classes=self.edge_id.num_classes,
-        )
-        return torch.concatenate((one_hot, edge_feature_arr), dim=1)
-
     def __add__(self, other: "BipartiteGraph"):
         """Concatenates two bipartite graphs that share a common target node set.
         Shifts the node indices of the second bipartite graph.
@@ -183,29 +148,6 @@ def __add__(self, other: "BipartiteGraph"):
         )
 
 
-@typechecked
-class BipartiteGraphProximity(BipartiteGraph):
-    """Graph defined on a pair of NodeSets, where the definition of the
-    graph edges is based on geographical proximity (Euclidean distance
-    in Cartesian coordinates).
-    """
-
-    def __init__(
-        self,
-        nodeset: tuple[NodeSet, NodeSet],
-        radius: float,
-        edge_id: Optional[EdgeID] = None,
-    ):
-
-        src, tgt = nodeset
-        point_tree = scipy.spatial.KDTree(tgt.cc_vertices)
-        neighbor_list = point_tree.query_ball_point(src.cc_vertices, r=radius)
-        # turn `neighbor_list` into a list of pairs:
-        nb_list = [(idx, x) for idx, x in enumerate(neighbor_list)]
-        nb_list = np.asarray([(x, k) for x, y in nb_list for k in y])
-        super().__init__(nodeset=nodeset, edge_id=edge_id, edge_vertices=nb_list)
-
-
 @typechecked
 class ICONMultiMesh(GeneralGraph):
     """Reads vertices and topology from an ICON grid file; creates multi-mesh."""
@@ -441,64 +383,6 @@ def _get_grid2mesh_edges(self, select_c: np.ndarray, multi_mesh: ICONMultiMesh)
 # -------------------------------------------------------------
 
 
-@typechecked
-def get_icon_mesh_and_grid(
-    grid_file: str, max_level_multimesh: int, max_level_dataset: int
-) -> tuple[ICONMultiMesh, ICONCellDataGrid]:
-    """Factory function, creating an ICON multi-mesh and an ICON cell-grid."""
-    return (
-        multi_mesh := ICONMultiMesh(grid_file, max_level=max_level_multimesh),
-        ICONCellDataGrid(grid_file, multi_mesh, max_level=max_level_dataset),
-    )
-
-
-@typechecked
-def get_edge_attributes(
-    gc_send: np.ndarray, cc_send: np.ndarray, gc_recv: np.ndarray, cc_recv: np.ndarray
-) -> np.ndarray:
-    """Build edge features using the position on the unit sphere of the mesh
-    nodes.
-
-    The features are (4 input features in total):
-
-    - length of the edge
-    - vector difference between the 3d positions of the sender node and the
-      receiver node computed in a local coordinate system of the receiver.
-
-    The local coordinate system of the receiver is computed by applying two elemental
-    rotations:
-
-    - the first rotation changes the azimuthal angle until that receiver
-      node lies at longitude 0,
-    - the second rotation changes the polar angle until that receiver
-      node lies at latitude 0.
-
-    Literature: Appdx. A2.4 in
-       Lam, R. et al. (2022). GraphCast: Learning skillful medium-range
-       global weather forecasting. 10.48550/arXiv.2212.12794.
-
-    """
-
-    phi_theta = gc_recv  # precession and nutation angle
-    theta_phi = np.fliplr(phi_theta)
-    theta_phi[:, 1] *= -1.0  # angles = [theta, -phi]
-    # rotation matrix:
-    rotation = scipy.spatial.transform.Rotation.from_euler(seq="YZ", angles=theta_phi)
-    edge_length = np.array([arc_length(cc_recv, cc_send)])
-    distance = rotation.apply(cc_send) - [
-        1.0,
-        0.0,
-        0.0,
-    ]  # subtract the rotated position of sender
-    return np.concatenate((edge_length.transpose(), distance), axis=1)
-
-
-@typechecked
-def arc_length(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
-    """Calculate length of great circle arc on the unit sphere."""
-    return np.arccos(np.clip(np.einsum("ij,ij->i", v1, v2), a_min=0.0, a_max=1.0))
-
-
 @typechecked
 def read_coordinate_array(ncfile, arrname: str, dimname: str) -> np.ndarray:
     """Auxiliary routine, reading a coordinate array, checking consistency."""
diff --git a/src/anemoi/graphs/nodes/builders/from_icon.py b/src/anemoi/graphs/nodes/builders/from_icon.py
index 000f469..a056e6b 100644
--- a/src/anemoi/graphs/nodes/builders/from_icon.py
+++ b/src/anemoi/graphs/nodes/builders/from_icon.py
@@ -12,7 +12,8 @@
 from anemoi.utils.config import DotDict
 from torch_geometric.data import HeteroData
 
-from anemoi.graphs.generate.icon_mesh import get_icon_mesh_and_grid
+from anemoi.graphs.generate.icon_mesh import ICONCellDataGrid
+from anemoi.graphs.generate.icon_mesh import ICONMultiMesh
 from anemoi.graphs.nodes.builders.base import BaseNodeBuilder
 
 
@@ -21,11 +22,10 @@ class ICONNodes(BaseNodeBuilder):
 
     def __init__(self, name: str, grid_filename: str, max_level_multimesh: int, max_level_dataset: int) -> None:
         self.grid_filename = grid_filename
-        self.multi_mesh, self.cell_grid = get_icon_mesh_and_grid(
-            grid_file=self.grid_filename,
-            max_level_multimesh=max_level_multimesh,
-            max_level_dataset=max_level_dataset,
-        )
+
+        self.multi_mesh = ICONMultiMesh(self.grid_filename, max_level=max_level_multimesh)
+        self.cell_grid = ICONCellDataGrid(self.grid_filename, self.multi_mesh, max_level=max_level_dataset)
+
         super().__init__(name)
 
     def get_coordinates(self) -> torch.Tensor:

From 37cb87a806089193706b5b9190b65d967711ecb1 Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Fri, 8 Nov 2024 15:26:03 +0100
Subject: [PATCH 30/35] [chore] adjust copyright notice.

---
 src/anemoi/graphs/generate/icon_mesh.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/anemoi/graphs/generate/icon_mesh.py b/src/anemoi/graphs/generate/icon_mesh.py
index ddf26fe..bf0a851 100644
--- a/src/anemoi/graphs/generate/icon_mesh.py
+++ b/src/anemoi/graphs/generate/icon_mesh.py
@@ -1,11 +1,11 @@
-# (C) Copyright 2024 ECMWF, DWD.
+# (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, the above institutions do 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.
 #
+# 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 itertools
 import logging

From 90a17213e9cf7ca256669d8bcf0b88cbaaea2e0f Mon Sep 17 00:00:00 2001
From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de>
Date: Fri, 8 Nov 2024 16:43:59 +0100
Subject: [PATCH 31/35] [doc] elaborate on icon mesh classes in rst file.

---
 docs/graphs/node_coordinates/icon_mesh.rst | 64 ++++++++++++++++++++--
 1 file changed, 60 insertions(+), 4 deletions(-)

diff --git a/docs/graphs/node_coordinates/icon_mesh.rst b/docs/graphs/node_coordinates/icon_mesh.rst
index 9bb01f2..dc306ed 100644
--- a/docs/graphs/node_coordinates/icon_mesh.rst
+++ b/docs/graphs/node_coordinates/icon_mesh.rst
@@ -2,7 +2,63 @@
  Triangular Mesh with ICON Topology
 ####################################
 
-This class defines nodes based on the `refinement_level` property
-attached to ICON grid files, where `0` denotes the vertices of the base
-grid, ie. the icosahedron including the step of root subdivision
-`RXXB00`.
+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}

From b45e21ecabba5163b6a96fa8693494a6a96fd022 Mon Sep 17 00:00:00 2001
From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de>
Date: Mon, 11 Nov 2024 16:59:11 +0100
Subject: [PATCH 32/35] Add Icon tests

---
 tests/nodes/test_icon.py | 196 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 196 insertions(+)
 create mode 100644 tests/nodes/test_icon.py

diff --git a/tests/nodes/test_icon.py b/tests/nodes/test_icon.py
new file mode 100644
index 0000000..528222a
--- /dev/null
+++ b/tests/nodes/test_icon.py
@@ -0,0 +1,196 @@
+# (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 netCDF4
+import numpy as np
+import pytest
+import torch
+from torch_geometric.data import HeteroData
+
+from anemoi.graphs.edges import ICONTopologicalDecoderEdges
+from anemoi.graphs.edges import ICONTopologicalEncoderEdges
+from anemoi.graphs.edges import ICONTopologicalProcessorEdges
+from anemoi.graphs.generate.icon_mesh import ICONCellDataGrid
+from anemoi.graphs.generate.icon_mesh import ICONMultiMesh
+from anemoi.graphs.nodes import ICONCellGridNodes
+from anemoi.graphs.nodes import ICONMultimeshNodes
+from anemoi.graphs.nodes import ICONNodes
+from anemoi.graphs.nodes.builders.base import BaseNodeBuilder
+
+
+class DatasetMock:
+    """This datasets emulates the most primitive unstructured grid with
+    refinement.
+
+    Enumeration of cells , edges and vertices in netCDF file is 1 based.
+    C: cell
+    E: edge
+    V: vertex
+
+    Cell C2 with its additional vertex V4 and edges E4 and E4 were added as
+    a first refinement.
+
+    [V1: 0, 1]🢀-E3--[V3: 1, 1]
+      🢁      ╲             🢁
+      |       ╲ [C1: ⅔, ⅔] |
+      |        ╲           |
+      E5        E1         E2
+      |          ╲         |
+      |           ╲        |
+      | [C2: ⅓, ⅓] ╲       |
+      |             🢆     |
+    [V4: 0, 1]🢀-E4--[V2: 1, 1]
+
+    Note: Triangular refinement does not actually work like this. This grid
+    mock serves testing purposes only.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+
+        class MockVariable:
+            def __init__(self, data, units, dimensions):
+                self.data = np.ma.asarray(data)
+                self.shape = data.shape
+                self.units = units
+                self.dimensions = dimensions
+
+            def __getitem__(self, key):
+                return self.data[key]
+
+        self.variables = {
+            "vlon": MockVariable(np.array([0, 1, 1, 0]), "radian", ("vertex",)),
+            "vlat": MockVariable(np.array([1, 0, 1, 0]), "radian", ("vertex",)),
+            "clon": MockVariable(np.array([0.66, 0.33]), "radian", ("cell",)),
+            "clat": MockVariable(np.array([0.66, 0.33]), "radian", ("cell",)),
+            "edge_vertices": MockVariable(np.array([[1, 2], [2, 3], [3, 1], [2, 4], [4, 1]]).T, "", ("nc", "edge")),
+            "vertex_of_cell": MockVariable(np.array([[1, 2, 3], [1, 2, 4]]).T, "", ("nv", "cell")),
+            "refinement_level_v": MockVariable(np.array([0, 0, 0, 1]), "", ("vertex",)),
+            "refinement_level_c": MockVariable(np.array([0, 1]), "", ("cell",)),
+        }
+        """common array dimensions:
+            nc: 2, # constant
+            nv: 3, # constant
+            vertex: 4,
+            edge: 5,
+            cell: 2,
+        """
+        self.uuidOfHGrid = "__test_data__"
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+@pytest.mark.parametrize("max_level_multimesh,max_level_dataset", [(0, 0), (0, 1), (1, 1)])
+def test_init(monkeypatch, max_level_multimesh: int, max_level_dataset: int):
+    """Test ICONNodes initialization."""
+
+    monkeypatch.setattr(netCDF4, "Dataset", DatasetMock)
+    node_builder = ICONNodes(
+        name="test_nodes",
+        grid_filename="test.nc",
+        max_level_multimesh=max_level_multimesh,
+        max_level_dataset=max_level_dataset,
+    )
+    assert isinstance(node_builder, BaseNodeBuilder)
+    assert isinstance(node_builder, ICONNodes)
+
+
+@pytest.mark.parametrize("Node_builder", [ICONCellGridNodes, ICONMultimeshNodes])
+def test_Node_builder_dependencies(monkeypatch, Node_builder):
+    """Test that the `Node_builder` depends on the presence of ICONNodes."""
+
+    monkeypatch.setattr(netCDF4, "Dataset", DatasetMock)
+    nodes = ICONNodes("test_icon_nodes", "test.nc", 0, 0)
+    data_nodes = Node_builder("data_nodes", "test_icon_nodes")
+    config = {}
+    graph = HeteroData()
+    graph = nodes.register_attributes(graph, config)
+
+    data_nodes.update_graph(graph)
+
+    data_nodes = ICONCellGridNodes("data_nodes2", "missing_icon_nodes")
+    with pytest.raises(KeyError):
+        data_nodes.update_graph(graph)
+
+
+class Test_Edge_builder_dependencies:
+    @pytest.fixture()
+    def icon_graph(self, monkeypatch) -> HeteroData:
+        """Return a HeteroData object with ICONNodes nodes."""
+        graph = HeteroData()
+        monkeypatch.setattr(netCDF4, "Dataset", DatasetMock)
+        nodes = ICONNodes("test_icon_nodes", "test.nc", 1, 0)
+
+        graph = nodes.update_graph(graph, {})
+
+        data_nodes = ICONCellGridNodes("data", "test_icon_nodes")
+        graph = data_nodes.register_attributes(graph, {})
+
+        return graph
+
+    @pytest.mark.parametrize(
+        "Edge_builder", [ICONTopologicalProcessorEdges, ICONTopologicalEncoderEdges, ICONTopologicalDecoderEdges]
+    )
+    def test_ProcessorEdges_dependencies(self, icon_graph, Edge_builder):
+        """Test that the `Edge_builder` depends on the presence of ICONNodes."""
+
+        edges = Edge_builder(
+            source_name="data",
+            target_name="data",
+            icon_mesh="test_icon_nodes",
+        )
+        edges.update_graph(icon_graph)
+
+        edges2 = Edge_builder(
+            source_name="data",
+            target_name="data",
+            icon_mesh="missing_icon_nodes",
+        )
+        with pytest.raises(KeyError):
+            edges2.update_graph(icon_graph)
+
+
+def test_register_nodes(monkeypatch):
+    """Test ICONNodes register correctly the nodes."""
+    monkeypatch.setattr(netCDF4, "Dataset", DatasetMock)
+    nodes = ICONNodes("test_icon_nodes", "test.nc", 0, 0)
+    graph = HeteroData()
+
+    graph = nodes.register_nodes(graph)
+
+    assert graph["test_icon_nodes"].x is not None
+    assert isinstance(graph["test_icon_nodes"].x, torch.Tensor)
+    assert graph["test_icon_nodes"].x.shape[1] == 2
+    assert graph["test_icon_nodes"].x.shape[0] == 3, "number of vertices at refinement_level_v == 0"
+    assert graph["test_icon_nodes"].node_type == "ICONNodes"
+
+    nodes2 = ICONNodes("test_icon_nodes", "test.nc", 1, 0)
+    graph = nodes2.register_nodes(graph)
+    assert graph["test_icon_nodes"].x.shape[0] == 4, "number of vertices at refinement_level_v == 1"
+
+
+def test_register_attributes(
+    monkeypatch,
+    graph_with_nodes: HeteroData,
+):
+    """Test ICONNodes register correctly the weights."""
+    monkeypatch.setattr(netCDF4, "Dataset", DatasetMock)
+    nodes = ICONNodes("test_icon_nodes", "test.nc", 0, 0)
+    config = {"test_attr": {"_target_": "anemoi.graphs.nodes.attributes.UniformWeights"}}
+    graph = HeteroData()
+
+    graph = nodes.register_attributes(graph, config)
+
+    assert graph["test_icon_nodes"]["_grid_filename"] is not None
+    assert isinstance(graph["test_icon_nodes"]["_multi_mesh"], ICONMultiMesh)
+    assert isinstance(graph["test_icon_nodes"]["_cell_grid"], ICONCellDataGrid)

From 63ea1b020408baef53e9e8f861331c8d1316560e Mon Sep 17 00:00:00 2001
From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de>
Date: Mon, 11 Nov 2024 17:24:52 +0100
Subject: [PATCH 33/35] update Change log

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9307b58..fef75d1 100644
--- 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)
+- 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
 

From 4aaa9ff0d50adb93cf0ae5dc91dedba2bebcd404 Mon Sep 17 00:00:00 2001
From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de>
Date: Mon, 11 Nov 2024 17:47:22 +0100
Subject: [PATCH 34/35] Add __future__ annotations

---
 src/anemoi/graphs/nodes/builders/from_icon.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/anemoi/graphs/nodes/builders/from_icon.py b/src/anemoi/graphs/nodes/builders/from_icon.py
index a056e6b..f0473e9 100644
--- a/src/anemoi/graphs/nodes/builders/from_icon.py
+++ b/src/anemoi/graphs/nodes/builders/from_icon.py
@@ -7,6 +7,8 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from __future__ import annotations
+
 import numpy as np
 import torch
 from anemoi.utils.config import DotDict

From 01a94f0ef330bb68884bc00caafc162cdd3c70b3 Mon Sep 17 00:00:00 2001
From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de>
Date: Tue, 12 Nov 2024 08:59:26 +0100
Subject: [PATCH 35/35] fix change log

---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fef75d1..7743f30 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +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)
-- define node sets and edges based on an ICON icosahedral mesh (#53)
+- 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