[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

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

@@ -106,14 +106,15 @@ 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, new_method=self.new_method)
+        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

@@ -8,13 +8,15 @@
 # nor does it submit to any jurisdiction.
 
 
+import logging
 from abc import ABC
 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."""
 
@@ -25,15 +27,17 @@ 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], new_method: bool = False ) -> 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
 
         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"],
+                nodes_coords_rad=nodes["x"],
                 resolutions=scale_resolutions,
                 node_ordering=nodes["_node_ordering"],
             )

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

@@ -8,20 +8,21 @@
 # nor does it submit to any jurisdiction.
 
 
+import logging
 from collections.abc import Iterable
 
 import networkx as nx
 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
 ) -> tuple[nx.DiGraph, np.ndarray, list[int]]:
@@ -135,21 +136,22 @@ 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]))
+    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))
+    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)
@@ -159,6 +161,7 @@ def add_edges_hop_1(nodes_coords_rad, resolutions: list[int], node_ordering: lis
     LOGGER.debug("multiscale_edges_shape: %s", multiscale_edges.shape)
     return multiscale_edges
 
+
 def add_edges_to_nx_graph(
     graph: nx.DiGraph,
     resolutions: list[int],
@@ -209,7 +212,7 @@ 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)
+        # LOGGER.debug("Sample edges: %s", neighbour_pairs)
         graph.add_edges_from(neighbour_pairs)
     return graph