introduce new TriNodesEdgebuilder-strategy

Author: mpvginde

graphs/src/anemoi/graphs/edges/builders/multi_scale.py

@@ -38,6 +38,8 @@ class MultiScaleEdges(BaseEdgeBuilder):
         Defines the refinement levels at which edges are computed. If an integer is provided, edges are computed for all
         levels up to and including that level. For instance, `scale_resolutions=4` includes edges at levels 1 through 4,
         whereas `scale_resolutions=[4]` only includes edges at level 4.
+    new_method: bool, optional
+        Use the new edge building strategy for x_hops=1. Default is False
 
     Methods
     -------
@@ -59,6 +61,7 @@ def __init__(
         target_name: str,
         x_hops: int,
         scale_resolutions: int | list[int] | None = None,
+        new_method: bool = False,
         **kwargs,
     ):
         super().__init__(source_name, target_name)
@@ -74,6 +77,7 @@ def __init__(
             scale_resolutions is None or min(scale_resolutions) > 0
         ), "The scale_resolutions argument only supports positive integers."
         self.scale_resolutions = scale_resolutions
+        self.new_method = new_method
 
     @staticmethod
     def get_edge_builder_class(node_type: str) -> type[BaseIcosahedronEdgeStrategy]:
@@ -102,10 +106,14 @@ def compute_edge_index(self, source_nodes: NodeStorage, _target_nodes: NodeStora
             )
 
         # Add edges
-        source_nodes = edge_builder_cls().add_edges(source_nodes, self.x_hops, scale_resolutions=scale_resolutions)
-        adjmat = nx.to_scipy_sparse_array(source_nodes["_nx_graph"], format="coo")
-
-        # Get source & target indices of the edges
-        edge_index = np.stack([adjmat.col, adjmat.row], axis=0)
+        source_nodes = edge_builder_cls().add_edges(source_nodes, self.x_hops, scale_resolutions=scale_resolutions, new_method=self.new_method)
+        if self.new_method:
+            # If the new method is used, the edges are already computed and stored in the node storage
+            edge_index = source_nodes["_multiscale_edges"]
+        else:
+            adjmat = nx.to_scipy_sparse_array(source_nodes["_nx_graph"], format="coo")
+            # 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)

graphs/src/anemoi/graphs/generate/multi_scale_edges.py

@@ -12,7 +12,8 @@
 from abc import abstractmethod
 
 from torch_geometric.data.storage import NodeStorage
-
+import logging
+LOGGER = logging.getLogger(__name__)
 
 class BaseIcosahedronEdgeStrategy(ABC):
     """Abstract base class for different edge-building strategies."""
@@ -24,16 +25,28 @@ def add_edges(self, nodes: NodeStorage, x_hops: int, scale_resolutions: list[int
 class TriNodesEdgeBuilder(BaseIcosahedronEdgeStrategy):
     """Edge builder for TriNodes and LimitedAreaTriNodes."""
 
-    def add_edges(self, nodes: NodeStorage, x_hops: int, scale_resolutions: list[int]) -> NodeStorage:
+    def add_edges(self, nodes: NodeStorage, x_hops: int, scale_resolutions: list[int], new_method: bool = False ) -> NodeStorage:
         from anemoi.graphs.generate import tri_icosahedron
 
-        nodes["_nx_graph"] = tri_icosahedron.add_edges_to_nx_graph(
-            nodes["_nx_graph"],
-            resolutions=scale_resolutions,
-            x_hops=x_hops,
-            area_mask_builder=nodes.get("_area_mask_builder", None),
-        )
-        return nodes
+        if new_method:
+            assert x_hops == 1, "New strategy currently only supports x_hops=1."
+            LOGGER.info("Using new strategy for x_hops=1 multiscale-edge building.")
+            # Compute the multiscale edges directly and store them in the node storage
+            multiscale_edges = tri_icosahedron.add_edges_hop_1(
+                nodes_coords_rad = nodes["x"],
+                resolutions=scale_resolutions,
+                node_ordering=nodes["_node_ordering"],
+            )
+            nodes["_multiscale_edges"] = multiscale_edges
+        else:
+            LOGGER.info("Using existing strategy for multiscale-edge building.")
+            nodes["_nx_graph"] = tri_icosahedron.add_edges_to_nx_graph(
+                nodes["_nx_graph"],
+                resolutions=scale_resolutions,
+                x_hops=x_hops,
+                area_mask_builder=nodes.get("_area_mask_builder", None),
+            )
+            return nodes
 
 
 class HexNodesEdgeBuilder(BaseIcosahedronEdgeStrategy):

graphs/src/anemoi/graphs/generate/tri_icosahedron.py

@@ -14,11 +14,13 @@
 import numpy as np
 import trimesh
 from sklearn.neighbors import BallTree
+import logging 
 
 from anemoi.graphs.generate.masks import KNNAreaMaskBuilder
 from anemoi.graphs.generate.transforms import cartesian_to_latlon_rad
 from anemoi.graphs.generate.utils import get_coordinates_ordering
 
+LOGGER = logging.getLogger(__name__)
 
 def create_tri_nodes(
     resolution: int, area_mask_builder: KNNAreaMaskBuilder | None = None
@@ -133,6 +135,29 @@ def create_nx_graph_from_tri_coords(coords_rad: np.ndarray, node_ordering: np.nd
     assert graph.number_of_nodes() == len(node_ordering), "The number of nodes must be the same."
     return graph
 
+def add_edges_hop_1(nodes_coords_rad, resolutions: list[int], node_ordering: list[int]) -> np.ndarray:
+    """Adds edges for x_hops = 1 relying on trimesh only."""
+    
+    hop_1_edges = []
+    for subdivisions in resolutions:
+        sphere = trimesh.creation.icosphere(subdivisions=subdivisions, radius=1.0)
+        coords_rad = cartesian_to_latlon_rad(sphere.vertices)
+        LOGGER.debug("Adding %d 1-hop edges for resolution %d", sphere.edges.shape[0], subdivisions)
+        hop_1_edges.append(sphere.edges)
+    # Check that the node coordinates from the node storage match those were generated
+    diff = np.sum(np.abs(nodes_coords_rad.numpy()[:coords_rad.shape[0],:] - coords_rad[node_ordering]))
+    assert diff == 0, "Node coordinates do not match coordinates generated by trimesh."
+
+    # Concatenate all edges from different resolutions and transpose to get shape (2, num_edges)
+    multiscale_edges = np.transpose(np.concatenate(hop_1_edges, axis=0), (1,0))
+
+    # Map the edges to the node ordering
+    inverse_ordering = np.empty_like(node_ordering)
+    inverse_ordering[node_ordering] = np.arange(len(node_ordering))
+    multiscale_edges = inverse_ordering[multiscale_edges]
+
+    LOGGER.debug("multiscale_edges_shape: %s", multiscale_edges.shape)
+    return multiscale_edges
 
 def add_edges_to_nx_graph(
     graph: nx.DiGraph,
@@ -183,6 +208,8 @@ def add_edges_to_nx_graph(
 
         _, vertex_mapping_index = tree.query(r_vertices_rad, k=1)
         neighbour_pairs = create_node_neighbours_list(graph, node_neighbours, vertex_mapping_index)
+        LOGGER.debug("Adding %d edges for resolution %d", len(neighbour_pairs), resolution)
+        #LOGGER.debug("Sample edges: %s", neighbour_pairs)
         graph.add_edges_from(neighbour_pairs)
     return graph