Merge branch 'develop' into 7-local-area-modelling-graphs

.gitattributes

@@ -0,0 +1 @@
+CHANGELOG.md merge=union

.github/CODEOWNERS

@@ -0,0 +1,6 @@
+# CODEOWNERS file
+
+# Protect workflow files
+/.github/ @theissenhelen @jesperdramsch @gmertes
+/.pre-commit-config.yaml @theissenhelen @jesperdramsch @gmertes
+/pyproject.toml @theissenhelen @jesperdramsch @gmertes

.github/ci-hpc-config.yml

@@ -0,0 +1,16 @@
+build:
+  modules:
+    - ninja
+  dependencies:
+    - ecmwf/ecbuild@develop
+    - ecmwf/eccodes@develop
+    - ecmwf/eckit@develop
+    - ecmwf/odc@develop
+  python_dependencies:
+    - ecmwf/anemoi-utils@develop
+    - ecmwf/anemoi-datasets@develop
+  parallel: 64
+
+  pytest_cmd: |
+    python -m pytest -vv -m 'not notebook and not no_cache_init' --cov=. --cov-report=xml
+    python -m coverage report

.github/workflows/changelog-pr-update.yml

@@ -5,6 +5,9 @@ on:
     branches:
       - main
       - develop
+    paths-ignore:
+      - .pre-commit-config.yaml
+      - .readthedocs.yaml
 jobs:
   Check-Changelog:
     name: Check Changelog Action

.github/workflows/ci.yml

@@ -46,7 +46,7 @@ jobs:
   downstream-ci-hpc:
     name: downstream-ci-hpc
     if: ${{ !github.event.pull_request.head.repo.fork && github.event.action != 'labeled' || github.event.label.name == 'approved-for-ci' }}
-    uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci.yml@main
+    uses: ecmwf-actions/downstream-ci/.github/workflows/downstream-ci-hpc.yml@main
     with:
       anemoi-graphs: ecmwf/anemoi-graphs@${{ github.event.pull_request.head.sha || github.sha }}
     secrets: inherit

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

@@ -5,7 +5,7 @@ name: Code Quality checks for PRs
 
 on:
   push:
-  pull_request_target:
+  pull_request:
     types: [opened, synchronize, reopened]
 
 jobs:

.gitignore

@@ -120,6 +120,7 @@ celerybeat.pid
 *.sage.py
 
 # Environments
+.envrc
 .env
 .venv
 env/

.pre-commit-config.yaml

@@ -20,6 +20,12 @@ repos:
   - id: no-commit-to-branch # Prevent committing to main / master
   - id: check-added-large-files # Check for large files added to git
   - id: check-merge-conflict # Check for files that contain merge conflict
+-   repo: https://github.com/pre-commit/pygrep-hooks
+    rev: v1.10.0  # Use the ref you want to point at
+    hooks:
+    -   id: python-use-type-annotations # Check for missing type annotations
+    -   id: python-check-blanket-noqa # Check for # noqa: all
+    -   id: python-no-log-warn # Check for log.warn
 - repo: https://github.com/psf/black-pre-commit-mirror
   rev: 24.8.0
   hooks:
@@ -34,7 +40,7 @@ repos:
     - --force-single-line-imports
     - --profile black
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.6.3
+  rev: v0.6.4
   hooks:
   - id: ruff
     # Next line if for documenation cod snippets
@@ -45,7 +51,7 @@ repos:
     - --exit-non-zero-on-fix
     - --preview
 - repo: https://github.com/sphinx-contrib/sphinx-lint
-  rev: v0.9.1
+  rev: v1.0.0
   hooks:
   - id: sphinx-lint
 # For now, we use it. But it does not support a lot of sphinx features
@@ -59,12 +65,21 @@ repos:
   hooks:
   - id: docconvert
     args: ["numpy"]
-- repo: https://github.com/b8raoult/optional-dependencies-all
-  rev: "0.0.6"
-  hooks:
-  - id: optional-dependencies-all
-    args: ["--inplace", "--exclude-keys=dev,docs,tests", "--group=dev=all,docs,tests"]
 - repo: https://github.com/tox-dev/pyproject-fmt
-  rev: "2.2.1"
+  rev: "2.2.3"
   hooks:
   - id: pyproject-fmt
+-   repo: https://github.com/jshwi/docsig # Check docstrings against function sig
+    rev: v0.60.1
+    hooks:
+    -   id: docsig
+        args:
+        - --ignore-no-params # Allow docstrings without parameters
+        - --check-dunders    # Check dunder methods
+        - --check-overridden # Check overridden methods
+        - --check-protected  # Check protected methods
+        - --check-class      # Check class docstrings
+        - --disable=E113     # Disable empty docstrings
+        - --summary          # Print a summary
+ci:
+  autoupdate_schedule: monthly

CHANGELOG.md

@@ -11,10 +11,14 @@ Keep it human-readable, your future self will thank you!
 ## [Unreleased](https://github.com/ecmwf/anemoi-graphs/compare/0.3.0...HEAD)
 
 ### Added
-- New node builder class, CutOutZarrDatasetNodes, to create nodes from 2 datasets. (#30)
-- New class, KNNAreaMaskBuilder, to specify Area of Interest (AOI) based on a set of nodes. (#30)
-- New node builder classes, LimitedAreaXXXXXNodes, to create nodes within an Area of Interest (AOI). (#30)
-- Expanded MultiScaleEdges to support multi-scale connections in limited area graphs. (#30)
+- ci: hpc-config, CODEOWNERS (#49)
+- feat: New node builder class, CutOutZarrDatasetNodes, to create nodes from 2 datasets. (#30)
+- feat: New class, KNNAreaMaskBuilder, to specify Area of Interest (AOI) based on a set of nodes. (#30)
+- feat: New node builder classes, LimitedAreaXXXXXNodes, to create nodes within an Area of Interest (AOI). (#30)
+- feat: Expanded MultiScaleEdges to support multi-scale connections in limited area graphs. (#30)
+
+### Changed
+- ci: small fixes and updates pre-commit, downsteam-ci (#49)
 
 ## [0.3.0 Anemoi-graphs, minor release](https://github.com/ecmwf/anemoi-graphs/compare/0.2.1...0.3.0) - 2024-09-03
 

pyproject.toml

@@ -10,20 +10,13 @@
 # https://packaging.python.org/en/latest/guides/writing-pyproject-toml/
 
 [build-system]
-requires = [
-  "setuptools>=60",
-  "setuptools-scm>=8",
-]
+requires = [ "setuptools>=60", "setuptools-scm>=8" ]
 
 [project]
 name = "anemoi-graphs"
 
 description = "A package to build graphs for data-driven forecasts."
-keywords = [
-  "ai",
-  "graphs",
-  "tools",
-]
+keywords = [ "ai", "graphs", "tools" ]
 
 license = { file = "LICENSE" }
 authors = [
@@ -46,9 +39,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
 
-dynamic = [
-  "version",
-]
+dynamic = [ "version" ]
 dependencies = [
   "anemoi-datasets[data]>=0.3.3",
   "anemoi-utils>=0.3.6",
@@ -63,20 +54,8 @@ dependencies = [
   "trimesh>=4.1",
 ]
 
-optional-dependencies.all = [
-]
-optional-dependencies.dev = [
-  "nbsphinx",
-  "pandoc",
-  "pytest",
-  "pytest-mock",
-  "requests",
-  "sphinx",
-  "sphinx-argparse",
-  "sphinx-rtd-theme",
-  "termcolor",
-  "tomli",
-]
+optional-dependencies.all = [  ]
+optional-dependencies.dev = [ "anemoi-graphs[docs,tests]" ]
 
 optional-dependencies.docs = [
   "nbsphinx",
@@ -89,10 +68,7 @@ optional-dependencies.docs = [
   "tomli",
 ]
 
-optional-dependencies.tests = [
-  "pytest",
-  "pytest-mock",
-]
+optional-dependencies.tests = [ "pytest", "pytest-mock" ]
 
 urls.Documentation = "https://anemoi-graphs.readthedocs.io/"
 urls.Homepage = "https://github.com/ecmwf/anemoi-graphs/"

src/anemoi/graphs/edges/builder.py

@@ -208,7 +208,7 @@ def __init__(
         assert num_nearest_neighbours > 0, "Number of nearest neighbours must be positive"
         self.num_nearest_neighbours = num_nearest_neighbours
 
-    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage):
+    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage) -> np.ndarray:
         """Compute the adjacency matrix for the KNN method.
 
         Parameters
@@ -217,6 +217,11 @@ def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStor
             The source nodes.
         target_nodes : NodeStorage
             The target nodes.
+
+        Returns
+        -------
+        np.ndarray
+            The adjacency matrix.
         """
         source_coords, target_coords = self.get_node_coordinates(source_nodes, target_nodes)
         assert self.num_nearest_neighbours is not None, "number of neighbors required for knn encoder"
@@ -274,13 +279,13 @@ def __init__(
         cutoff_factor: float,
         source_mask_attr_name: str | None = None,
         target_mask_attr_name: str | None = None,
-    ):
+    ) -> None:
         super().__init__(source_name, target_name, source_mask_attr_name, target_mask_attr_name)
         assert isinstance(cutoff_factor, (int, float)), "Cutoff factor must be a float"
         assert cutoff_factor > 0, "Cutoff factor must be positive"
         self.cutoff_factor = cutoff_factor
 
-    def get_cutoff_radius(self, graph: HeteroData, mask_attr: torch.Tensor | None = None):
+    def get_cutoff_radius(self, graph: HeteroData, mask_attr: torch.Tensor | None = None) -> float:
         """Compute the cut-off radius.
 
         The cut-off radius is computed as the product of the target nodes
@@ -309,7 +314,7 @@ def prepare_node_data(self, graph: HeteroData) -> tuple[NodeStorage, NodeStorage
         self.radius = self.get_cutoff_radius(graph)
         return super().prepare_node_data(graph)
 
-    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage):
+    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage) -> np.ndarray:
         """Get the adjacency matrix for the cut-off method.
 
         Parameters
@@ -318,6 +323,11 @@ def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStor
             The source nodes.
         target_nodes : NodeStorage
             The target nodes.
+
+        Returns
+        -------
+        np.ndarray
+            The adjacency matrix.
         """
         source_coords, target_coords = self.get_node_coordinates(source_nodes, target_nodes)
         LOGGER.info(

src/anemoi/graphs/generate/hex_icosahedron.py

@@ -112,8 +112,6 @@ def add_edges_to_nx_graph(
     depth_children : int
         The number of resolution levels to consider for the connections of children. Defaults to 1, which includes
         connections up to the next resolution level.
-    aoi_mask_builder : KNNAreaMaskBuilder
-        NearestNeighbors with the cloud of points to limit the mesh area, by default None.
 
     Returns
     -------

src/anemoi/graphs/generate/tri_icosahedron.py

@@ -137,8 +137,7 @@ def get_neighbours_within_hops(
     x_hops : int
         Number of hops between 2 nodes to consider them neighbours.
     valid_nodes : list[int], optional
-        List of valid nodes to consider, by default None. It is useful to consider only a subset of the nodes to save
-        computation time.
+        The list of valid nodes to consider, by default None.
 
     Returns
     -------
@@ -176,7 +175,7 @@ def add_neigbours_edges(
         The graph.
     node_idx : int
         The node considered.
-    neighbours : list[int]
+    neighbour_indices : list[int]
         The neighbours of the node.
     self_loops : bool, optional
         Whether is supported to add self-loops, by default False.

src/anemoi/graphs/nodes/attributes.py

@@ -43,6 +43,10 @@ def compute(self, graph: HeteroData, nodes_name: str, *args, **kwargs) -> torch.
             Graph.
         nodes_name : str
             Name of the nodes.
+        args : tuple
+            Additional arguments.
+        kwargs : dict
+            Additional keyword arguments.
 
         Returns
         -------
@@ -70,6 +74,10 @@ def get_raw_values(self, nodes: NodeStorage, *args, **kwargs) -> np.ndarray:
         ----------
         nodes : NodeStorage
             Nodes of the graph.
+        args : tuple
+            Additional arguments.
+        kwargs : dict
+            Additional keyword arguments.
 
         Returns
         -------
@@ -111,6 +119,10 @@ def get_raw_values(self, nodes: NodeStorage, *args, **kwargs) -> np.ndarray:
         ----------
         nodes : NodeStorage
             Nodes of the graph.
+        args : tuple
+            Additional arguments.
+        kwargs : dict
+            Additional keyword arguments.
 
         Returns
         -------

src/anemoi/graphs/nodes/builders/from_refined_icosahedron.py

@@ -25,6 +25,8 @@ class IcosahedralNodes(BaseNodeBuilder, ABC):
     ----------
     resolution : list[int] | int
         Refinement level of the mesh.
+    name : str
+        Name of the nodes.
     """
 
     def __init__(