run pre-commit

Author: frazane

graphs/src/anemoi/graphs/inspect.py

@@ -90,4 +90,6 @@ def inspect(self):
         if self.show_nodes:
             LOGGER.info("Saving interactive plots of nodes ...")
             for nodes_name in self.graph.node_types:
-                plot_interactive_nodes_2d(self.graph, nodes_name, out_file=self.output_path / f"{nodes_name}_nodes.html")
+                plot_interactive_nodes_2d(
+                    self.graph, nodes_name, out_file=self.output_path / f"{nodes_name}_nodes.html"
+                )

graphs/src/anemoi/graphs/plotting/interactive_3d.html.jinja

@@ -953,4 +953,4 @@ nor does it submit to any jurisdiction. #}
     </script>
 </body>
 
-</html>
+</html>

graphs/src/anemoi/graphs/plotting/interactive_3d_html.py

@@ -11,18 +11,20 @@
 
 import numpy as np
 import torch
-from torch_geometric.data import HeteroData
 from jinja2 import Template
+from torch_geometric.data import HeteroData
 
 HTML_TEMPLATE_PATH = Path(__file__).parent / "interactive_3d.html.jinja"
 
 
-def subset_graph(
+def extract_nodes_edges(
     graph: HeteroData,
     nodes: list[str] | None = None,
     edges: list[str] | None = None,
 ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
-    """Load a hetero graph from a file and separate nodes and edges by type.
+    """Extracts nodes and edges from a heterogeneous graph.
+
+    Optionally filters by specified types.
 
     Parameters
     ----------
@@ -98,7 +100,7 @@ def plot_interactive_graph_3d(
     out_file : str | Path, optional
         Name of the file to save the plot. Default is None.
     """
-    nodes, edges = subset_graph(graph)
+    nodes, edges = extract_nodes_edges(graph)
 
     for node_set in nodes:
         node_lats, node_lons = coords_to_latlon(nodes[node_set].numpy())
@@ -140,4 +142,4 @@ def plot_interactive_graph_3d(
     template = Template(HTML_TEMPLATE)
     html_output = template.render(nodes=nodes_embed, edges=edges_embed, max_degree=50, min_degree=1)
     with open(out_file, "w") as f:
-        f.write(html_output)
+        f.write(html_output)