Merge remote-tracking branch 'origin/develop' into dynamic-graphs

.github/ci-config.yml

@@ -0,0 +1,3 @@
+dependency_branch: develop
+parallelism_factor: 8
+self_build: false # Only for python packages

.github/workflows/ci.yml

@@ -0,0 +1,43 @@
+name: ci
+
+on:
+  # Trigger the workflow on push to master or develop, except tag creation
+  push:
+    branches:
+      - 'main'
+      - 'develop'
+    tags-ignore:
+      - '**'
+    paths:
+      - "src/**"
+      - "tests/**"
+
+  # Trigger the workflow on pull request
+  pull_request: ~
+
+  # Trigger the workflow manuallyp instals
+  workflow_dispatch: ~
+
+  # Trigger after public PR approved for CI
+  pull_request_target:
+    types: [labeled]
+
+jobs:
+  # Run CI including downstream packages on self-hosted runners
+  downstream-ci:
+    name: downstream-ci
+    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
+    with:
+      anemoi-models: ecmwf/anemoi-models@${{ github.event.pull_request.head.sha || github.sha }}
+      codecov_upload: true
+    secrets: inherit
+
+   # Build downstream packages on HPC
+  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
+    with:
+      anemoi-models: ecmwf/anemoi-models@${{ github.event.pull_request.head.sha || github.sha }}
+    secrets: inherit

.github/workflows/label-public-pr.yml

@@ -0,0 +1,10 @@
+# Manage labels of pull requests that originate from forks
+name: label-public-pr
+
+on:
+  pull_request_target:
+    types: [opened, synchronize]
+
+jobs:
+  label:
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/label-pr.yml@v2

.github/workflows/python-publish.yml

@@ -35,7 +35,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
-    - uses: actions/setup-python@v2
+    - uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
 
@@ -48,27 +48,6 @@ jobs:
       run: pytest
 
   deploy:
-
-    if: ${{ github.event_name == 'release' }}
-    runs-on: ubuntu-latest
     needs: [checks, quality]
-
-    steps:
-    - uses: actions/checkout@v4
-
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: 3.x
-
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install build wheel twine
-    - name: Build and publish
-      env:
-        TWINE_USERNAME: __token__
-        TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
-      run: |
-        python -m build
-        twine upload dist/*
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/cd-pypi.yml@v2
+    secrets: inherit

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

@@ -0,0 +1,22 @@
+name: Read the Docs PR Preview
+on:
+  pull_request_target:
+    types:
+      - opened
+      - synchronize
+      - reopened
+    # Execute this action only on PRs that touch
+    # documentation files.
+    paths:
+      - "docs/**"
+
+permissions:
+  pull-requests: write
+
+jobs:
+  documentation-links:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: readthedocs/actions/preview@v1
+        with:
+          project-slug: "anemoi-models"

CHANGELOG.md

@@ -14,18 +14,42 @@ Keep it human-readable, your future self will thank you!
 
 ### Changed
 
+ - Update CI to inherit from common infrastructue reusable workflows
+ - run downstream-ci only when src and tests folders have changed
+ - New error messages for wrongs graphs.
+
 ### Removed
 
-## 0.1.0 Initial Release
+## [0.2.1] - Dependency update
 
 ### Added
-- Documentation
-- Initial code release with models, layers, distributed, preprocessing, and data_indices
-- Added Changelog
 
-### Changed
+- downstream-ci pipeline
+- readthedocs PR update check action
 
 ### Removed
 
-## Git Diffs:
+- anemoi-datasets dependency
+
+## [0.2.0] - Support Heterodata
+
+### Added
+
+- Option to choose the edge attributes
+
+### Changed
+
+- Updated to support new PyTorch Geometric HeteroData structure (defined by `anemoi-graphs` package).
+
+## [0.1.0] - Initial Release
+
+### Added
+- Documentation
+- Initial code release with models, layers, distributed, preprocessing, and data_indices
+- Added Changelog
+
+<!-- Add Git Diffs for Links above -->
+[unreleased]: https://github.com/ecmwf/anemoi-models/compare/0.2.1...HEAD
+[0.2.1]: https://github.com/ecmwf/anemoi-models/compare/0.2.0...0.2.1
+[0.2.0]: https://github.com/ecmwf/anemoi-models/compare/0.1.0...0.2.0
 [0.1.0]: https://github.com/ecmwf/anemoi-models/releases/tag/0.1.0

pyproject.toml

@@ -51,7 +51,6 @@ dynamic = [
   "version",
 ]
 dependencies = [
-  "anemoi-datasets>=0.2.1",
   "anemoi-utils>=0.1.9",
   "einops>=0.6.1",
   "hydra-core>=1.3",

src/anemoi/models/layers/mapper.py

@@ -7,6 +7,7 @@
 # nor does it submit to any jurisdiction.
 #
 
+import logging
 from abc import ABC
 from typing import Optional
 
@@ -30,6 +31,8 @@
 from anemoi.models.layers.graph import TrainableTensor
 from anemoi.models.layers.mlp import MLP
 
+LOGGER = logging.getLogger(__name__)
+
 
 class BaseMapper(nn.Module, ABC):
     """Base Mapper from souce dimension to destination dimension."""
@@ -113,22 +116,31 @@ def pre_process(self, x, shard_shapes, model_comm_group=None):
 
 
 class GraphEdgeMixin:
-    def _register_edges(self, sub_graph: HeteroData, src_size: int, dst_size: int, trainable_size: int) -> None:
+    def _register_edges(
+        self, sub_graph: HeteroData, edge_attributes: list[str], src_size: int, dst_size: int, trainable_size: int
+    ) -> None:
         """Register edge dim, attr, index_base, and increment.
 
         Parameters
         ----------
         sub_graph : HeteroData
             Sub graph of the full structure
+        edge_attributes : list[str]
+            Edge attributes to use.
         src_size : int
             Source size
         dst_size : int
             Target size
         trainable_size : int
             Trainable tensor size
         """
-        self.edge_dim = sub_graph.edge_attr.shape[1] + trainable_size
-        self.register_buffer("edge_attr", sub_graph.edge_attr, persistent=False)
+        assert sub_graph, f"{self.__class__.__name__} needs a valid sub_graph to register edges."
+        assert edge_attributes is not None, "Edge attributes must be provided"
+
+        edge_attr_tensor = torch.cat([sub_graph[attr] for attr in edge_attributes], axis=1)
+
+        self.edge_dim = edge_attr_tensor.shape[1] + trainable_size
+        self.register_buffer("edge_attr", edge_attr_tensor, persistent=False)
         self.register_buffer("edge_index_base", sub_graph.edge_index, persistent=False)
         self.register_buffer(
             "edge_inc", torch.from_numpy(np.asarray([[src_size], [dst_size]], dtype=np.int64)), persistent=True
@@ -174,6 +186,7 @@ def __init__(
         num_heads: int = 16,
         mlp_hidden_ratio: int = 4,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
     ) -> None:
@@ -210,7 +223,7 @@ def __init__(
             activation=activation,
         )
 
-        self._register_edges(sub_graph, src_grid_size, dst_grid_size, trainable_size)
+        self._register_edges(sub_graph, sub_graph_edge_attributes, src_grid_size, dst_grid_size, trainable_size)
 
         self.trainable = TrainableTensor(trainable_size=trainable_size, tensor_size=self.edge_attr.shape[0])
 
@@ -274,6 +287,7 @@ def __init__(
         num_heads: int = 16,
         mlp_hidden_ratio: int = 4,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
     ) -> None:
@@ -312,6 +326,7 @@ def __init__(
             num_heads=num_heads,
             mlp_hidden_ratio=mlp_hidden_ratio,
             sub_graph=sub_graph,
+            sub_graph_edge_attributes=sub_graph_edge_attributes,
             src_grid_size=src_grid_size,
             dst_grid_size=dst_grid_size,
         )
@@ -345,6 +360,7 @@ def __init__(
         num_heads: int = 16,
         mlp_hidden_ratio: int = 4,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
     ) -> None:
@@ -383,6 +399,7 @@ def __init__(
             num_heads=num_heads,
             mlp_hidden_ratio=mlp_hidden_ratio,
             sub_graph=sub_graph,
+            sub_graph_edge_attributes=sub_graph_edge_attributes,
             src_grid_size=src_grid_size,
             dst_grid_size=dst_grid_size,
         )
@@ -415,6 +432,7 @@ def __init__(
         activation: str = "SiLU",
         mlp_extra_layers: int = 0,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
     ) -> None:
@@ -451,7 +469,7 @@ def __init__(
             activation=activation,
         )
 
-        self._register_edges(sub_graph, src_grid_size, dst_grid_size, trainable_size)
+        self._register_edges(sub_graph, sub_graph_edge_attributes, src_grid_size, dst_grid_size, trainable_size)
 
         self.emb_edges = MLP(
             in_features=self.edge_dim,
@@ -518,6 +536,7 @@ def __init__(
         activation: str = "SiLU",
         mlp_extra_layers: int = 0,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
     ) -> None:
@@ -555,6 +574,7 @@ def __init__(
             activation,
             mlp_extra_layers,
             sub_graph=sub_graph,
+            sub_graph_edge_attributes=sub_graph_edge_attributes,
             src_grid_size=src_grid_size,
             dst_grid_size=dst_grid_size,
         )
@@ -602,6 +622,7 @@ def __init__(
         activation: str = "SiLU",
         mlp_extra_layers: int = 0,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
     ) -> None:
@@ -639,6 +660,7 @@ def __init__(
             activation=activation,
             mlp_extra_layers=mlp_extra_layers,
             sub_graph=sub_graph,
+            sub_graph_edge_attributes=sub_graph_edge_attributes,
             src_grid_size=src_grid_size,
             dst_grid_size=dst_grid_size,
         )

src/anemoi/models/layers/processor.py

@@ -171,6 +171,7 @@ def __init__(
         activation: str = "SiLU",
         cpu_offload: bool = False,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
         **kwargs,
@@ -201,7 +202,7 @@ def __init__(
             mlp_extra_layers=mlp_extra_layers,
         )
 
-        self._register_edges(sub_graph, src_grid_size, dst_grid_size, trainable_size)
+        self._register_edges(sub_graph, sub_graph_edge_attributes, src_grid_size, dst_grid_size, trainable_size)
 
         self.trainable = TrainableTensor(trainable_size=trainable_size, tensor_size=self.edge_attr.shape[0])
 
@@ -258,6 +259,7 @@ def __init__(
         activation: str = "GELU",
         cpu_offload: bool = False,
         sub_graph: Optional[HeteroData] = None,
+        sub_graph_edge_attributes: Optional[list[str]] = None,
         src_grid_size: int = 0,
         dst_grid_size: int = 0,
         **kwargs,
@@ -291,7 +293,7 @@ def __init__(
             mlp_hidden_ratio=mlp_hidden_ratio,
         )
 
-        self._register_edges(sub_graph, src_grid_size, dst_grid_size, trainable_size)
+        self._register_edges(sub_graph, sub_graph_edge_attributes, src_grid_size, dst_grid_size, trainable_size)
 
         self.trainable = TrainableTensor(trainable_size=trainable_size, tensor_size=self.edge_attr.shape[0])