Merge branch 'develop'

.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,7 @@
+build:
+  python: '3.10'
+  modules:
+    - ninja
+  python_dependencies:
+    - ecmwf/anemoi-utils@develop
+  parallel: 64

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

@@ -2,6 +2,9 @@ name: Check Changelog Update on PR
 on:
   pull_request:
     types: [assigned, opened, synchronize, reopened, labeled, unlabeled]
+    paths-ignore:
+      - .pre-commit-config.yaml
+      - .readthedocs.yaml
 jobs:
   Check-Changelog:
     name: Check Changelog Action

.github/workflows/ci.yml

@@ -37,7 +37,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-models: ecmwf/anemoi-models@${{ github.event.pull_request.head.sha || github.sha }}
     secrets: inherit

.github/workflows/python-publish.yml

@@ -4,48 +4,22 @@
 name: Upload Python Package
 
 on:
-
-  push: {}
-
   release:
     types: [created]
 
 jobs:
   quality:
-    name: Code QA
-    runs-on: ubuntu-latest
-    steps:
-    - run: sudo apt-get install -y pandoc # Needed by sphinx for notebooks
-    - uses: actions/checkout@v4
-    - uses: actions/setup-python@v5
-      with:
-        python-version: 3.x
-    - uses: pre-commit/action@v3.0.1
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-precommit-run.yml@v2
+    with:
+      skip-hooks: "no-commit-to-branch"
 
   checks:
     strategy:
-      fail-fast: false
       matrix:
-        platform: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.10"]
-
-    name: Python  ${{ matrix.python-version }} on ${{ matrix.platform }}
-    runs-on: ${{ matrix.platform }}
-
-    steps:
-    - uses: actions/checkout@v4
-
-    - uses: actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.python-version }}
-
-    - name: Install
-      run: |
-        pip install -e .[all,tests]
-        pip freeze
-
-    - name: Tests
-      run: pytest
+        python-version: ["3.9", "3.10"]
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2
+    with:
+      python-version: ${{ matrix.python-version }}
 
   deploy:
     needs: [checks, quality]

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

@@ -0,0 +1,19 @@
+---
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+name: Code Quality checks for PRs
+
+on:
+  push:
+  pull_request:
+    types: [opened, synchronize, reopened]
+
+jobs:
+  quality:
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-precommit-run.yml@v2
+    with:
+      skip-hooks: "no-commit-to-branch"
+
+  checks:
+    uses: ecmwf-actions/reusable-workflows/.github/workflows/qa-pytest-pyproject.yml@v2

.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

@@ -10,6 +10,19 @@ Keep it human-readable, your future self will thank you!
 
 ## [Unreleased](https://github.com/ecmwf/anemoi-models/compare/0.3.0...HEAD)
 
+### Added
+- Codeowners file
+- Pygrep precommit hooks
+- Docsig precommit hooks
+- Changelog merge strategy
+- configurabilty of the dropout probability in the the MultiHeadSelfAttention module
+- Variable Bounding as configurable model layers [#13](https://github.com/ecmwf/anemoi-models/issues/13)
+
+### Changed
+- Bugfixes for CI
+
+### Removed
+
 ## [0.3.0](https://github.com/ecmwf/anemoi-models/compare/0.2.1...0.3.0) - Remapping of (meteorological) Variables
 
 ### Added
@@ -22,6 +35,7 @@ Keep it human-readable, your future self will thank you!
 - 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.
+- Feature: Change model to be instantiatable in the interface, addressing [#28](https://github.com/ecmwf/anemoi-models/issues/28) through [#45](https://github.com/ecmwf/anemoi-models/pulls/45)
 
 ### Removed
 

pyproject.toml

@@ -11,20 +11,13 @@
 [build-system]
 build-backend = "setuptools.build_meta"
 
-requires = [
-  "setuptools>=61",
-  "setuptools-scm>=8",
-]
+requires = [ "setuptools>=61", "setuptools-scm>=8" ]
 
 [project]
 name = "anemoi-models"
 description = "A package to hold various functions to support training of ML models."
 readme = "README.md"
-keywords = [
-  "ai",
-  "models",
-  "tools",
-]
+keywords = [ "ai", "models", "tools" ]
 
 license = { file = "LICENSE" }
 authors = [
@@ -47,29 +40,17 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
 
-dynamic = [
-  "version",
-]
+dynamic = [ "version" ]
 dependencies = [
   "anemoi-utils>=0.1.9",
   "einops>=0.6.1",
   "hydra-core>=1.3",
   "torch>=2.2",
   "torch-geometric>=2.3,<2.5",
 ]
-optional-dependencies.all = [
-]
+optional-dependencies.all = [  ]
 
-optional-dependencies.dev = [
-  "hypothesis",
-  "nbsphinx",
-  "pandoc",
-  "pytest",
-  "rstfmt",
-  "sphinx",
-  "sphinx-argparse<0.5",
-  "sphinx-rtd-theme",
-]
+optional-dependencies.dev = [ "anemoi-models[all,docs,tests]" ]
 
 optional-dependencies.docs = [
   "nbsphinx",
@@ -80,10 +61,7 @@ optional-dependencies.docs = [
   "sphinx-rtd-theme",
 ]
 
-optional-dependencies.tests = [
-  "hypothesis",
-  "pytest",
-]
+optional-dependencies.tests = [ "hypothesis", "pytest" ]
 
 urls.Documentation = "https://anemoi-models.readthedocs.io/"
 urls.Homepage = "https://github.com/ecmwf/anemoi-models/"

src/anemoi/models/interface/__init__.py

@@ -14,7 +14,6 @@
 from hydra.utils import instantiate
 from torch_geometric.data import HeteroData
 
-from anemoi.models.models.encoder_processor_decoder import AnemoiModelEncProcDec
 from anemoi.models.preprocessing import Processors
 
 
@@ -73,9 +72,13 @@ def _build_model(self) -> None:
         self.pre_processors = Processors(processors)
         self.post_processors = Processors(processors, inverse=True)
 
-        # Instantiate the model (Can be generalised to other models in the future, here we use AnemoiModelEncProcDec)
-        self.model = AnemoiModelEncProcDec(
-            config=self.config, data_indices=self.data_indices, graph_data=self.graph_data
+        # Instantiate the model
+        self.model = instantiate(
+            self.config.model.model,
+            model_config=self.config,
+            data_indices=self.data_indices,
+            graph_data=self.graph_data,
+            _recursive_=False,  # Disables recursive instantiation by Hydra
         )
 
         # Use the forward method of the model directly

src/anemoi/models/layers/attention.py

@@ -40,19 +40,19 @@ def __init__(
         bias: bool = False,
         is_causal: bool = False,
         window_size: Optional[int] = None,
-        dropout: float = 0.0,
+        dropout_p: float = 0.0,
     ):
         super().__init__()
 
         assert (
             embed_dim % num_heads == 0
         ), f"Embedding dimension ({embed_dim}) must be divisible by number of heads ({num_heads})"
 
-        self.dropout = dropout
         self.num_heads = num_heads
         self.embed_dim = embed_dim
         self.head_dim = embed_dim // num_heads  # q k v
         self.window_size = (window_size, window_size)  # flash attention
+        self.dropout_p = dropout_p
         self.is_causal = is_causal
 
         self.lin_qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=bias)
@@ -86,15 +86,22 @@ def forward(
         query = shard_heads(query, shapes=shapes, mgroup=model_comm_group)
         key = shard_heads(key, shapes=shapes, mgroup=model_comm_group)
         value = shard_heads(value, shapes=shapes, mgroup=model_comm_group)
+        dropout_p = self.dropout_p if self.training else 0.0
 
         if _FLASH_ATTENTION_AVAILABLE:
             query, key, value = (
                 einops.rearrange(t, "batch heads grid vars -> batch grid heads vars") for t in (query, key, value)
             )
-            out = self.attention(query, key, value, causal=False, window_size=self.window_size)
+            out = self.attention(query, key, value, causal=False, window_size=self.window_size, dropout_p=dropout_p)
             out = einops.rearrange(out, "batch grid heads vars -> batch heads grid vars")
         else:
-            out = self.attention(query, key, value, is_causal=False)  # expects (batch heads grid variable) format
+            out = self.attention(
+                query,
+                key,
+                value,
+                is_causal=False,
+                dropout_p=dropout_p,
+            )  # expects (batch heads grid variable) format
 
         out = shard_sequence(out, shapes=shapes, mgroup=model_comm_group)
         out = einops.rearrange(out, "batch heads grid vars -> (batch grid) (heads vars)")

src/anemoi/models/layers/block.py

@@ -55,7 +55,15 @@ def forward(
 class TransformerProcessorBlock(BaseBlock):
     """Transformer block with MultiHeadSelfAttention and MLPs."""
 
-    def __init__(self, num_channels, hidden_dim, num_heads, activation, window_size: int):
+    def __init__(
+        self,
+        num_channels: int,
+        hidden_dim: int,
+        num_heads: int,
+        activation: str,
+        window_size: int,
+        dropout_p: float = 0.0,
+    ):
         super().__init__()
 
         try:
@@ -72,7 +80,7 @@ def __init__(self, num_channels, hidden_dim, num_heads, activation, window_size:
             window_size=window_size,
             bias=False,
             is_causal=False,
-            dropout=0.0,
+            dropout_p=dropout_p,
         )
 
         self.mlp = nn.Sequential(