fix(graphs,tests): new test and fix anemoi-graphs tests with gpu (#637)

## Description

New test for `GaussianWeightsAttritbute`. Fix anemoi-graphs tests with
GPU.

##  Additional notes ##
<!-- Include any additional information, caveats, or considerations that
the reviewer should be aware of. -->

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

Author: JPXKQX

graphs/src/anemoi/graphs/edges/attributes.py

@@ -60,7 +60,7 @@ def subset_node_information(self, source_nodes: NodeStorage, target_nodes: NodeS
 
     def forward(self, x: tuple[NodeStorage, NodeStorage], edge_index: Adj, size: Size = None) -> torch.Tensor:
         x = self.subset_node_information(*x)
-        return self.propagate(edge_index, x=x, size=size)
+        return self.propagate(edge_index.to(self.device), x=x, size=size)
 
     @abstractmethod
     def compute(self, x_i: torch.Tensor, x_j: torch.Tensor) -> torch.Tensor: ...

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

@@ -109,8 +109,8 @@ def register_attributes(self, graph: HeteroData, config: DotDict) -> HeteroData:
         """
         for attr_name, attr_config in config.items():
             edge_index = graph[self.name].edge_index
-            edge_builder = instantiate(attr_config)
-            graph[self.name][attr_name] = edge_builder(
+            edge_attribute_builder = instantiate(attr_config)
+            graph[self.name][attr_name] = edge_attribute_builder(
                 x=(graph[self.name[0]], graph[self.name[2]]), edge_index=edge_index
             )
         return graph

graphs/src/anemoi/graphs/processors/post_process.py

@@ -28,6 +28,7 @@
 
 
 class PostProcessor(ABC):
+    """Base PostProcessor class."""
 
     @abstractmethod
     def update_graph(self, graph: HeteroData, **kwargs: Any) -> HeteroData:
@@ -44,6 +45,7 @@ def __init__(
     ) -> None:
         self.nodes_names = (nodes_name,) if isinstance(nodes_name, str) else tuple(nodes_name)
         self.save_mask_indices_to_attr = save_mask_indices_to_attr
+        super().__init__()
 
     def removing_nodes(self, graph: HeteroData, mask: torch.Tensor, nodes_name: str) -> HeteroData:
         """Remove nodes based on the mask passed."""
@@ -113,7 +115,7 @@ def update_graph(self, graph: HeteroData, **kwargs: Any) -> HeteroData:
             The post-processed graph.
         """
         for nodes_name in self.nodes_names:
-            mask = self.compute_mask(graph, nodes_name)
+            mask = self.compute_mask(graph, nodes_name).cpu()
             LOGGER.info(f"Removing {(~mask).sum()} nodes from {nodes_name}.")
             graph = self.removing_nodes(graph, mask, nodes_name)
             graph = self.update_edge_indices(graph, mask, nodes_name)
@@ -248,6 +250,7 @@ class BaseSortEdgeIndex(PostProcessor, ABC):
     def __init__(self, descending: bool = True) -> None:
         assert self.nodes_axis is not None, f"{self.__class__.__name__} must define the nodes_axis class attribute."
         self.descending = descending
+        super().__init__()
 
     def get_sorting_mask(self, edges: dict) -> torch.Tensor:
         sort_indices = torch.sort(edges["edge_index"], descending=self.descending, dim=1)
@@ -311,15 +314,15 @@ def __init__(
         self.source_name = source_name
         self.target_name = target_name
         self.edges_name = (self.source_name, "to", self.target_name)
-        self.mask: torch.Tensor = None
+        super().__init__()
 
-    def removing_edges(self, graph: HeteroData) -> HeteroData:
+    def removing_edges(self, graph: HeteroData, mask: torch.Tensor) -> HeteroData:
         """Remove edges based on the mask passed."""
         for attr_name in graph[self.edges_name].edge_attrs():
             if attr_name == "edge_index":
-                graph[self.edges_name][attr_name] = graph[self.edges_name][attr_name][:, self.mask]
+                graph[self.edges_name][attr_name] = graph[self.edges_name][attr_name].cpu()[:, mask]
             else:
-                graph[self.edges_name][attr_name] = graph[self.edges_name][attr_name][self.mask, :]
+                graph[self.edges_name][attr_name] = graph[self.edges_name][attr_name].cpu()[mask, :]
 
         return graph
 
@@ -351,9 +354,9 @@ def update_graph(self, graph: HeteroData, **kwargs: Any) -> HeteroData:
         HeteroData
             The post-processed graph.
         """
-        self.mask = self.compute_mask(graph)
-        LOGGER.info(f"Removing {(~self.mask).sum()} edges from {self.edges_name}.")
-        graph = self.removing_edges(graph)
+        mask = self.compute_mask(graph).cpu()
+        LOGGER.info(f"Removing {(~mask).sum()} edges from {self.edges_name}.")
+        graph = self.removing_edges(graph, mask)
         graph_config = kwargs.get("graph_config", {})
         graph = self.recompute_attributes(graph, graph_config)
         return graph
@@ -398,14 +401,14 @@ def compute_mask(self, graph: HeteroData) -> torch.Tensor:
         target_nodes = graph[self.target_name]
         edge_index = graph[self.edges_name].edge_index
         lengths = EARTH_RADIUS * EdgeLength()(x=(source_nodes, target_nodes), edge_index=edge_index)
-        mask = torch.where(lengths > self.treshold, False, True).squeeze()
+        mask = torch.where(lengths > self.treshold, False, True).squeeze().cpu()
         cases = [
             (self.source_mask_attr_name, source_nodes, 0),
             (self.target_mask_attr_name, target_nodes, 1),
         ]
         for mask_attr_name, nodes, i in cases:
             if mask_attr_name:
                 attr_mask = nodes[mask_attr_name].squeeze()
-                edge_mask = attr_mask[edge_index[i]]
+                edge_mask = attr_mask[edge_index[i]].cpu()
                 mask = torch.logical_or(mask, ~edge_mask)
         return mask

graphs/tests/edges/test_edge_attributes.py

@@ -16,6 +16,7 @@
 from anemoi.graphs.edges.attributes import AttributeFromTargetNode
 from anemoi.graphs.edges.attributes import EdgeDirection
 from anemoi.graphs.edges.attributes import EdgeLength
+from anemoi.graphs.edges.attributes import GaussianDistanceWeights
 
 TEST_EDGES = ("test_nodes", "to", "test_nodes")
 
@@ -31,10 +32,11 @@ def test_directional_features(graph_nodes_and_edges, norm):
     assert isinstance(edge_attr, torch.Tensor)
 
 
+@pytest.mark.parametrize("edge_attr_cls", [EdgeLength, GaussianDistanceWeights])
 @pytest.mark.parametrize("norm", ["l1", "l2", "unit-max", "unit-std", "unit-range"])
-def test_edge_lengths(graph_nodes_and_edges, norm):
+def test_edge_lengths(edge_attr_cls, graph_nodes_and_edges, norm):
     """Test EdgeLength compute method."""
-    edge_attr_builder = EdgeLength(norm=norm)
+    edge_attr_builder = edge_attr_cls(norm=norm)
     edge_index = graph_nodes_and_edges[TEST_EDGES].edge_index
     source_nodes = graph_nodes_and_edges[TEST_EDGES[0]]
     target_nodes = graph_nodes_and_edges[TEST_EDGES[2]]