From d7052513fed5f61dec567fbecefa743ffd90a65c Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
Date: Fri, 16 Aug 2024 11:52:52 +0000
Subject: [PATCH 1/6] fix: support py3.9

Co-authored-by: Jesper Dramsch <jesper.dramsch@ecmwf.int>
---
 src/anemoi/graphs/create.py               |  8 ++++----
 src/anemoi/graphs/edges/attributes.py     |  9 +++++----
 src/anemoi/graphs/edges/builder.py        |  7 ++++---
 src/anemoi/graphs/edges/directional.py    |  4 ++--
 src/anemoi/graphs/generate/hexagonal.py   | 10 +++++-----
 src/anemoi/graphs/generate/icosahedral.py |  5 +++--
 src/anemoi/graphs/nodes/attributes.py     |  9 ++++-----
 src/anemoi/graphs/nodes/builder.py        | 10 +++++-----
 src/anemoi/graphs/utils.py                |  6 +++---
 9 files changed, 35 insertions(+), 33 deletions(-)

diff --git a/src/anemoi/graphs/create.py b/src/anemoi/graphs/create.py
index d52b37f..9f10c81 100644
--- a/src/anemoi/graphs/create.py
+++ b/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
diff --git a/src/anemoi/graphs/edges/attributes.py b/src/anemoi/graphs/edges/attributes.py
index 9a8d6d8..c65a402 100644
--- a/src/anemoi/graphs/edges/attributes.py
+++ b/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
 
diff --git a/src/anemoi/graphs/edges/builder.py b/src/anemoi/graphs/edges/builder.py
index fdf1738..5fbdb13 100644
--- a/src/anemoi/graphs/edges/builder.py
+++ b/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.
diff --git a/src/anemoi/graphs/edges/directional.py b/src/anemoi/graphs/edges/directional.py
index 9c7cdea..322f37e 100644
--- a/src/anemoi/graphs/edges/directional.py
+++ b/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
diff --git a/src/anemoi/graphs/generate/hexagonal.py b/src/anemoi/graphs/generate/hexagonal.py
index ca35084..9e1bf29 100644
--- a/src/anemoi/graphs/generate/hexagonal.py
+++ b/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.
 
diff --git a/src/anemoi/graphs/generate/icosahedral.py b/src/anemoi/graphs/generate/icosahedral.py
index 47fcec8..332643a 100644
--- a/src/anemoi/graphs/generate/icosahedral.py
+++ b/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.
 
diff --git a/src/anemoi/graphs/nodes/attributes.py b/src/anemoi/graphs/nodes/attributes.py
index 040f134..e0ed1d8 100644
--- a/src/anemoi/graphs/nodes/attributes.py
+++ b/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
diff --git a/src/anemoi/graphs/nodes/builder.py b/src/anemoi/graphs/nodes/builder.py
index a0fa889..0587633 100644
--- a/src/anemoi/graphs/nodes/builder.py
+++ b/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 = 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
diff --git a/src/anemoi/graphs/utils.py b/src/anemoi/graphs/utils.py
index 8999bc6..c895426 100644
--- a/src/anemoi/graphs/utils.py
+++ b/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.

From ae38f875174a26b1bfa6bb6b8d1563fb6893d6f9 Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
Date: Fri, 16 Aug 2024 11:53:26 +0000
Subject: [PATCH 2/6] fix: update & homogeneize changelog

---
 CHANGELOG.md | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e28819a..2623861 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,11 +16,11 @@ 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

From 020f348e34f819425e560c53b28bcf1f0de84988 Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
Date: Fri, 16 Aug 2024 12:59:51 +0000
Subject: [PATCH 3/6] fix: update tests to support py39

---
 tests/edges/test_cutoff.py           | 2 ++
 tests/edges/test_edge_attributes.py  | 2 ++
 tests/edges/test_knn.py              | 2 ++
 tests/edges/test_multiscale_edges.py | 2 ++
 4 files changed, 8 insertions(+)

diff --git a/tests/edges/test_cutoff.py b/tests/edges/test_cutoff.py
index 838134c..0604409 100644
--- a/tests/edges/test_cutoff.py
+++ b/tests/edges/test_cutoff.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 
 from anemoi.graphs.edges import CutOffEdges
diff --git a/tests/edges/test_edge_attributes.py b/tests/edges/test_edge_attributes.py
index 40cba1c..1591014 100644
--- a/tests/edges/test_edge_attributes.py
+++ b/tests/edges/test_edge_attributes.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 import torch
 
diff --git a/tests/edges/test_knn.py b/tests/edges/test_knn.py
index 9f6cae9..01e12ea 100644
--- a/tests/edges/test_knn.py
+++ b/tests/edges/test_knn.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 
 from anemoi.graphs.edges import KNNEdges
diff --git a/tests/edges/test_multiscale_edges.py b/tests/edges/test_multiscale_edges.py
index 91ec04b..3b6e204 100644
--- a/tests/edges/test_multiscale_edges.py
+++ b/tests/edges/test_multiscale_edges.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pytest
 from torch_geometric.data import HeteroData
 

From 5015c8fd1e06f1905ea0374dbaf8a6c77d090dd3 Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
Date: Fri, 16 Aug 2024 14:21:59 +0000
Subject: [PATCH 4/6] fix: style

---
 src/anemoi/graphs/nodes/builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/anemoi/graphs/nodes/builder.py b/src/anemoi/graphs/nodes/builder.py
index 0587633..2afdfc8 100644
--- a/src/anemoi/graphs/nodes/builder.py
+++ b/src/anemoi/graphs/nodes/builder.py
@@ -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: DotDict | None = None) -> HeteroData:
+    def register_attributes(self, graph: HeteroData, config: DotDict = None) -> HeteroData:
         """Register attributes in the nodes of the graph specified.
 
         Parameters

From 84aaa7223fc7f3da1262b2bffb5a9090a4ee651f Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
Date: Fri, 16 Aug 2024 14:44:52 +0000
Subject: [PATCH 5/6] ci: rollback 3.9 tests

---
 .github/workflows/python-publish.yml      | 2 +-
 .github/workflows/python-pull-request.yml | 2 +-
 CHANGELOG.md                              | 1 -
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 7e01780..2cb554a 100644
--- a/.github/workflows/python-publish.yml
+++ b/.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 }}
diff --git a/.github/workflows/python-pull-request.yml b/.github/workflows/python-pull-request.yml
index 08d74fc..0ebecb1 100644
--- a/.github/workflows/python-pull-request.yml
+++ b/.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 }}
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 934703a..2623861 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,7 +27,6 @@ Keep it human-readable, your future self will thank you!
 - ci: fixes and permissions
 
 ### Removed
-ci: testing for python 3.9
 
 ## [0.2.1] - Anemoi-graph Release, bug fix release
 

From 50a5a9aa21fe8076e8e6bdfc02ef227e88a58079 Mon Sep 17 00:00:00 2001
From: Mario Santa Cruz <mariosanta_cruz@hotmail.com>
Date: Mon, 19 Aug 2024 09:11:08 +0000
Subject: [PATCH 6/6] refactor: annotations not used in tests

---
 tests/edges/test_cutoff.py           | 2 --
 tests/edges/test_edge_attributes.py  | 2 --
 tests/edges/test_knn.py              | 2 --
 tests/edges/test_multiscale_edges.py | 2 --
 4 files changed, 8 deletions(-)

diff --git a/tests/edges/test_cutoff.py b/tests/edges/test_cutoff.py
index 0604409..838134c 100644
--- a/tests/edges/test_cutoff.py
+++ b/tests/edges/test_cutoff.py
@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import pytest
 
 from anemoi.graphs.edges import CutOffEdges
diff --git a/tests/edges/test_edge_attributes.py b/tests/edges/test_edge_attributes.py
index 1591014..40cba1c 100644
--- a/tests/edges/test_edge_attributes.py
+++ b/tests/edges/test_edge_attributes.py
@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import pytest
 import torch
 
diff --git a/tests/edges/test_knn.py b/tests/edges/test_knn.py
index 01e12ea..9f6cae9 100644
--- a/tests/edges/test_knn.py
+++ b/tests/edges/test_knn.py
@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import pytest
 
 from anemoi.graphs.edges import KNNEdges
diff --git a/tests/edges/test_multiscale_edges.py b/tests/edges/test_multiscale_edges.py
index 3b6e204..91ec04b 100644
--- a/tests/edges/test_multiscale_edges.py
+++ b/tests/edges/test_multiscale_edges.py
@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import pytest
 from torch_geometric.data import HeteroData