Fix support for Python 3.9

.github/workflows/python-publish.yml

@@ -16,7 +16,7 @@ jobs:
   checks:
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.9", "3.10"]
     uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2
     with:
       python-version: ${{ matrix.python-version }}

.github/workflows/python-pull-request.yml

@@ -17,7 +17,7 @@ jobs:
   checks:
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.9", "3.10"]
     uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2
     with:
       python-version: ${{ matrix.python-version }}

CHANGELOG.md

@@ -16,18 +16,17 @@ Keep it human-readable, your future self will thank you!
 - Changelog release updater
 
 ### Changed
-- Fix bug in graph cleaning method
-- Fix `anemoi-graphs create`. Config argument is cast to a Path.
-- Fix GraphCreator().clean() to not iterate over a dictionary that may change size during iterations.
-- Fix missing binary dependency
-- **Fix**: Updated `get_raw_values` method in `AreaWeights` to ensure compatibility with `scipy.spatial.SphericalVoronoi` by converting `latitudes` and `longitudes` to NumPy arrays before passing them to the `latlon_rad_to_cartesian` function. This resolves an issue where the function would fail if passed Torch Tensors directly.
+- fix: added support for Python3.9.
+- fix: bug in graph cleaning method
+- fix: `anemoi-graphs create` CLI argument is casted to a Path.
+- ci: fix missing binary dependency in ci-config.yaml
+- fix: Updated `get_raw_values` method in `AreaWeights` to ensure compatibility with `scipy.spatial.SphericalVoronoi` by converting `latitudes` and `longitudes` to NumPy arrays before passing them to the `latlon_rad_to_cartesian` function. This resolves an issue where the function would fail if passed Torch Tensors directly.
 - ci: Reusable workflows for push, PR, and releases
 - ci: ignore docs for downstream ci
 - ci: changed Changelog action to create PR
 - ci: fixes and permissions
 
 ### Removed
-ci: testing for python 3.9
 
 ## [0.2.1] - Anemoi-graph Release, bug fix release
 

src/anemoi/graphs/create.py

@@ -1,8 +1,8 @@
+from __future__ import annotations
+
 import logging
 from itertools import chain
 from pathlib import Path
-from typing import Optional
-from typing import Union
 
 import torch
 from anemoi.utils.config import DotDict
@@ -17,7 +17,7 @@ class GraphCreator:
 
     def __init__(
         self,
-        config: Union[Path, DotDict],
+        config: str | Path | DotDict,
     ):
         if isinstance(config, Path) or isinstance(config, str):
             self.config = DotDict.from_file(config)
@@ -91,7 +91,7 @@ def save(self, graph: HeteroData, save_path: Path, overwrite: bool = False) -> N
         else:
             LOGGER.info("Graph already exists. Use overwrite=True to overwrite.")
 
-    def create(self, save_path: Optional[Path] = None, overwrite: bool = False) -> HeteroData:
+    def create(self, save_path: Path | None = None, overwrite: bool = False) -> HeteroData:
         """Create the graph and save it to the output path.
 
         Parameters

src/anemoi/graphs/edges/attributes.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC
 from abc import abstractmethod
-from typing import Optional
 
 import numpy as np
 import torch
@@ -17,7 +18,7 @@
 class BaseEdgeAttribute(ABC, NormalizerMixin):
     """Base class for edge attributes."""
 
-    def __init__(self, norm: Optional[str] = None) -> None:
+    def __init__(self, norm: str | None = None) -> None:
         self.norm = norm
 
     @abstractmethod
@@ -69,7 +70,7 @@ class EdgeDirection(BaseEdgeAttribute):
         Compute directional attributes.
     """
 
-    def __init__(self, norm: Optional[str] = None, luse_rotated_features: bool = True) -> None:
+    def __init__(self, norm: str | None = None, luse_rotated_features: bool = True) -> None:
         super().__init__(norm)
         self.luse_rotated_features = luse_rotated_features
 
@@ -115,7 +116,7 @@ class EdgeLength(BaseEdgeAttribute):
         Compute edge lengths attributes.
     """
 
-    def __init__(self, norm: Optional[str] = None, invert: bool = False) -> None:
+    def __init__(self, norm: str | None = None, invert: bool = False) -> None:
         super().__init__(norm)
         self.invert = invert
 

src/anemoi/graphs/edges/builder.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC
 from abc import abstractmethod
-from typing import Optional
 
 import networkx as nx
 import numpy as np
@@ -99,7 +100,7 @@ def register_attributes(self, graph: HeteroData, config: DotDict) -> HeteroData:
             graph[self.name][attr_name] = instantiate(attr_config).compute(graph, self.name)
         return graph
 
-    def update_graph(self, graph: HeteroData, attrs_config: Optional[DotDict] = None) -> HeteroData:
+    def update_graph(self, graph: HeteroData, attrs_config: DotDict | None = None) -> HeteroData:
         """Update the graph with the edges.
 
         Parameters
@@ -214,7 +215,7 @@ def __init__(self, source_name: str, target_name: str, cutoff_factor: float):
         assert cutoff_factor > 0, "Cutoff factor must be positive"
         self.cutoff_factor = cutoff_factor
 
-    def get_cutoff_radius(self, graph: HeteroData, mask_attr: Optional[torch.Tensor] = None):
+    def get_cutoff_radius(self, graph: HeteroData, mask_attr: torch.Tensor | None = None):
         """Compute the cut-off radius.
 
         The cut-off radius is computed as the product of the target nodes reference distance and the cut-off factor.

src/anemoi/graphs/edges/directional.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 import numpy as np
 from scipy.spatial.transform import Rotation
@@ -28,7 +28,7 @@ def get_rotation_from_unit_vecs(points: np.ndarray, reference: np.ndarray) -> Ro
     return Rotation.from_rotvec(np.transpose(v_unit * theta))
 
 
-def compute_directions(loc1: np.ndarray, loc2: np.ndarray, pole_vec: Optional[np.ndarray] = None) -> np.ndarray:
+def compute_directions(loc1: np.ndarray, loc2: np.ndarray, pole_vec: np.ndarray | None = None) -> np.ndarray:
     """Compute the direction of the edge joining the nodes considered.
 
     Parameters

src/anemoi/graphs/generate/hexagonal.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from __future__ import annotations
 
 import h3
 import networkx as nx
@@ -8,7 +8,7 @@
 
 def create_hexagonal_nodes(
     resolutions: list[int],
-    area: Optional[dict] = None,
+    area: dict | None = None,
 ) -> tuple[nx.Graph, np.ndarray, list[int]]:
     """Creates a global mesh from a refined icosahedro.
 
@@ -50,7 +50,7 @@ def create_hexagonal_nodes(
 def add_nodes_for_resolution(
     graph: nx.Graph,
     resolution: int,
-    **area_kwargs: Optional[dict],
+    **area_kwargs: dict | None,
 ) -> nx.Graph:
     """Add all nodes at a specified refinement level to a graph.
 
@@ -74,7 +74,7 @@ def add_nodes_for_resolution(
 
 def get_nodes_at_resolution(
     resolution: int,
-    area: Optional[dict] = None,
+    area: dict | None = None,
 ) -> set[str]:
     """Get nodes at a specified refinement level over the entire globe.
 
@@ -160,7 +160,7 @@ def add_neighbour_edges(
 def add_edges_to_children(
     graph: nx.Graph,
     refinement_levels: tuple[int],
-    depth_children: Optional[int] = None,
+    depth_children: int | None = None,
 ) -> nx.Graph:
     """Adds edges to the children of the nodes at the specified resolution levels.
 

src/anemoi/graphs/generate/icosahedral.py

@@ -1,6 +1,7 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Iterable
-from typing import Optional
 
 import networkx as nx
 import numpy as np
@@ -168,7 +169,7 @@ def add_neigbours_edges(
     node_idx: int,
     neighbour_indices: Iterable[int],
     self_loops: bool = False,
-    vertex_mapping_index: Optional[np.ndarray] = None,
+    vertex_mapping_index: np.ndarray | None = None,
 ) -> None:
     """Adds the edges of one node to its neighbours.
 

src/anemoi/graphs/nodes/attributes.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC
 from abc import abstractmethod
-from typing import Optional
 
 import numpy as np
 import torch
@@ -18,7 +19,7 @@
 class BaseWeights(ABC, NormalizerMixin):
     """Base class for the weights of the nodes."""
 
-    def __init__(self, norm: Optional[str] = None) -> None:
+    def __init__(self, norm: str | None = None) -> None:
         self.norm = norm
 
     @abstractmethod
@@ -92,9 +93,7 @@ class AreaWeights(BaseWeights):
         Compute the area attributes for each node.
     """
 
-    def __init__(
-        self, norm: Optional[str] = None, radius: float = 1.0, centre: np.ndarray = np.array([0, 0, 0])
-    ) -> None:
+    def __init__(self, norm: str | None = None, radius: float = 1.0, centre: np.ndarray = np.array([0, 0, 0])) -> None:
         super().__init__(norm)
         self.radius = radius
         self.centre = centre

src/anemoi/graphs/nodes/builder.py

@@ -1,9 +1,9 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC
 from abc import abstractmethod
 from pathlib import Path
-from typing import Optional
-from typing import Union
 
 import numpy as np
 import torch
@@ -44,7 +44,7 @@ def register_nodes(self, graph: HeteroData) -> None:
         graph[self.name].node_type = type(self).__name__
         return graph
 
-    def register_attributes(self, graph: HeteroData, config: Optional[DotDict] = None) -> HeteroData:
+    def register_attributes(self, graph: HeteroData, config: DotDict = None) -> HeteroData:
         """Register attributes in the nodes of the graph specified.
 
         Parameters
@@ -85,7 +85,7 @@ def reshape_coords(self, latitudes: np.ndarray, longitudes: np.ndarray) -> torch
         coords = np.deg2rad(coords)
         return torch.tensor(coords, dtype=torch.float32)
 
-    def update_graph(self, graph: HeteroData, attr_config: Optional[DotDict] = None) -> HeteroData:
+    def update_graph(self, graph: HeteroData, attr_config: DotDict | None = None) -> HeteroData:
         """Update the graph with new nodes.
 
         Parameters
@@ -212,7 +212,7 @@ class IcosahedralNodes(BaseNodeBuilder, ABC):
 
     def __init__(
         self,
-        resolution: Union[int, list[int]],
+        resolution: int | list[int],
         name: str,
     ) -> None:
         self.resolutions = list(range(resolution + 1)) if isinstance(resolution, int) else resolution

src/anemoi/graphs/utils.py

@@ -1,11 +1,11 @@
-from typing import Optional
+from __future__ import annotations
 
 import numpy as np
 import torch
 from sklearn.neighbors import NearestNeighbors
 
 
-def get_nearest_neighbour(coords_rad: torch.Tensor, mask: Optional[torch.Tensor] = None) -> NearestNeighbors:
+def get_nearest_neighbour(coords_rad: torch.Tensor, mask: torch.Tensor | None = None) -> NearestNeighbors:
     """Get NearestNeighbour object fitted to coordinates.
 
     Parameters
@@ -32,7 +32,7 @@ def get_nearest_neighbour(coords_rad: torch.Tensor, mask: Optional[torch.Tensor]
     return nearest_neighbour
 
 
-def get_grid_reference_distance(coords_rad: torch.Tensor, mask: Optional[torch.Tensor] = None) -> float:
+def get_grid_reference_distance(coords_rad: torch.Tensor, mask: torch.Tensor | None = None) -> float:
     """Get the reference distance of the grid.
 
     It is the maximum distance of a node in the mesh with respect to its nearest neighbour.

tests/edges/test_cutoff.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 
 from anemoi.graphs.edges import CutOffEdges

tests/edges/test_edge_attributes.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 import torch
 

tests/edges/test_knn.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 
 from anemoi.graphs.edges import KNNEdges

tests/edges/test_multiscale_edges.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 from torch_geometric.data import HeteroData