feat: ensemble-aware alignment (uncertainty predicates, NaN broadcasting, Delaunay member dims)

.pre-commit-config.yaml

@@ -24,3 +24,8 @@ repos:
       - id: ruff-check
         args: [--fix]
       - id: ruff-format
+
+  - repo: https://github.com/numpy/numpydoc
+    rev: v1.11.0
+    hooks:
+      - id: numpydoc-validation

AGENTS.md

@@ -0,0 +1,94 @@
+# AGENTS.md
+
+Guidance for AI coding agents and human contributors working on `mxalign`.
+
+## Project
+
+`mxalign` aligns meteorological `xarray` datasets in space and time. Datasets carry
+**traits** (`mlwp_space_trait`, `mlwp_time_trait`, `mlwp_uncertainty_trait` in `ds.attrs`).
+The traits are defined by [mlwp-data-specs](https://github.com/mlwp-tools/mlwp-data-specs) and
+attached by [mlwp-data-loaders](https://github.com/mlwp-tools/mlwp-data-loaders). Loading and
+validating datasets happens in those packages. `mxalign` only reads the traits to decide how to
+align.
+
+Package layout (`src/mxalign/`):
+
+| Package | Contents |
+|---|---|
+| `accessor/` | The `ds.mx` accessor (`MxAccessor`): trait predicates (`is_grid`, `is_forecast`, ...) and thin wrappers around alignment operations |
+| `alignment/` | Time and space alignment (pairwise ops and the multi-dataset `align_time` / `align_space`), NaN broadcasting, and `interpolations/` (interpolator registry, xarray and Delaunay interpolators) |
+| `pipeline/` | The YAML-driven pipeline: `cli` (the `mxalign` command), `runner`, `verification`, `transformations/`, and `config` (pydantic schema for the YAML config) with its `dates` / `save` helpers |
+| `utils/` | Helpers shared by `accessor` and `alignment` (`projections`, `traits`) |
+
+Public functions are re-exported from `mxalign/__init__.py`. Keep those re-exports working when
+moving code.
+
+## Development tools
+
+- Python >= 3.12, managed with [uv](https://docs.astral.sh/uv/).
+- Set up the environment: `uv sync --all-extras`
+- Run the tests: `uv run pytest tests -v`
+- Run the linters: `uvx pre-commit run --all-files`. This runs ruff check (with `--fix`), ruff
+  format, codespell, numpydoc-validation and the standard pre-commit hooks. CI runs the same hooks
+  (`.github/workflows/linting.yaml`), plus the tests (`.github/workflows/tests.yaml`).
+- Add a dependency with `uv add <package>` (or `uv add --optional <extra> <package>`), and commit
+  both `pyproject.toml` and `uv.lock`.
+
+## Coding conventions
+
+- **Type annotations** on every argument (including `*args` / `**kwargs`) and every return value,
+  in `src/` and in `tests/`. Use modern syntax (`X | None`, `list[str]`) and concrete types
+  (`xr.Dataset`, `np.ndarray`, ...) rather than `Any`. Enforced by ruff's `ANN` rules.
+- **Numpy-style docstrings** on every module, class, function and method, in `src/` and in
+  `tests/`. Each one has a one-line imperative summary ending in a period, then `Parameters`,
+  `Returns` (or `Yields`), and `Raises` where applicable. Class docstrings document the
+  `__init__` parameters, so `__init__` itself needs no docstring (it is excluded from
+  numpydoc's checks). pydantic models in `pipeline/config.py` document their fields under
+  `Attributes`. Enforced by ruff's `D` rules (numpy convention) and by `numpydoc lint`.
+  `Examples`, `See Also` and an extended summary are optional. Run `ruff` through pre-commit, which
+  pins the version: newer ruff releases enable more rules by default.
+
+  ```python
+  def align_grid_point(
+      ds1: xr.Dataset, ds2: xr.Dataset, method: str = "xarray", **kwargs: object
+  ) -> xr.Dataset:
+      """Interpolate grid dataset ``ds1`` to the point locations of ``ds2``.
+
+      Parameters
+      ----------
+      ds1 : xr.Dataset
+          Gridded source dataset.
+      ds2 : xr.Dataset
+          Point dataset whose locations are the interpolation targets.
+      method : str, default "xarray"
+          Name of a registered interpolator.
+      **kwargs : object
+          Passed to the interpolator's constructor.
+
+      Returns
+      -------
+      xr.Dataset
+          ``ds1`` interpolated to the points of ``ds2``.
+      """
+  ```
+- **Formatting** with `ruff format`. Don't hand-format against it.
+- **Never mutate inputs**: alignment operations and accessor methods return new datasets (see
+  `MxAccessor`, which reads traits from `ds.attrs` on every call and never modifies `ds`).
+- **Registries**: new interpolators and transformations are registered with
+  `register_interpolator` / `register_transformation` and imported in their package
+  `__init__.py`, so they register on `import mxalign`.
+- **Config**: new pipeline config keys go into the pydantic models in
+  `pipeline/config.py`. Sections reject unknown keys unless they pass kwargs through to a function.
+
+## Tests
+
+- pytest, under `tests/`. Shared fixtures (`ds_fcst`, `ds_obs`, ...) are in `tests/conftest.py`.
+- Test functions are annotated (`-> None`, fixture arguments typed) and have numpy-style docstrings
+  that state the behaviour being checked, the same as package code.
+- Add or update tests together with any behaviour change.
+
+## Changes and decisions
+
+- Add an entry to `CHANGELOG.md` under `[Unreleased]` (Keep a Changelog format) with your GitHub
+  handle. Mark breaking changes with **Breaking:**.
+- Record design decisions as ADRs in `doc/adr/` using `doc/adr/template.md`.

CHANGELOG.md

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Changed
+- All functions, methods and classes (package and tests) have type annotations and numpy-style docstrings, enforced in pre-commit/CI by ruff `ANN`/`D` rules and `numpydoc-validation`. `AGENTS.md` documents the coding conventions and development tools. @leifdenby
+- **Breaking:** the pipeline config is validated against an explicit pydantic schema (`mxalign.pipeline.config.PipelineConfig`, loaded with `load_config`) instead of being read as a free-form dict. Unknown sections or keys, and missing required keys, raise a `ValidationError` when the `Runner` is created; `Runner.config` is now the validated model. Per-dataset `dates` entries override individual fields of the top-level `dates`. @leifdenby
 - **Breaking:** file path patterns (`datasets.<name>.files` with a `dates:` section, and `alignment.save.path`) now use Python `str.format` syntax instead of earthkit-data patterns: `{reference_time:%Y%m%d%H}` instead of `{reference_time:strftime(%Y%m%d%H)}`, `{year:02d}` instead of `{year:int(%02d)}`. `{lead_time}` (integer hours) and `{valid_time}` are now supported in `files` patterns. `earthkit-data` is no longer a dependency. @michielv
 - **Breaking:** temporal alignment always runs in the pipeline; `alignment.time` is optional and only accepts `lead_time` (`reference` | `intersection` | `union`, forecast reference only). Unknown keys such as `return_as` or `method` now raise an error. `align_time` takes `reference` and an explicit `lead_time` instead of `**kwargs`. @michielv
 - Remove dataset loading and dataset validation and outsource it to [mlwp-data-loaders](https://github.com/mlwp-tools/mlwp-data-loaders) and [mlwp-data-specs](https://github.com/mlwp-tools/mlwp-data-specs) @michielv
@@ -17,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added CI action for package build and upload to pypi.org on releases. [\#28](https://github.com/mlwp-tools/mxalign/pull/28) @leifdenby
 - `transform()` and `interpolate()` return the same container type they are given: a one-item list now stays a list instead of being unwrapped to a bare dataset, and a dict with a single dataset no longer raises `KeyError`. @michielv
 - Delaunay interpolation weights are shared across datasets on the same grid, so aligning several models to the same points triangulates the grid once. When a multi-chunk `grid_index` is merged into a single chunk, the leading dimensions are re-chunked automatically so blocks stay near dask's `array.chunk-size`. @michielv
+- `ds.mx` has uncertainty predicates `is_deterministic()`, `is_ensemble()` and `is_quantile()`, which read `mlwp_uncertainty_trait` on every call like the space and time predicates. @michielv
+- **Breaking:** `broadcast_nans` needs the uncertainty trait on every dataset, and never matches the `member` (or `quantile`) dimension between datasets: an ensemble value counts as missing when any of its members is, and is then masked in all members, so the ensemble size stays constant over the verification sample ([ADR-003](doc/adr/ADR-003_ensemble-nan-masking.md)). A shared variable whose other dimensions differ between two datasets raises a `ValueError` that says so. @michielv
 
 ### Fixed
 - `ds.mx.add_xy(crs=...)` and `ds.mx.unstack(crs=...)` no longer replace the dataset held by the cached `ds.mx` accessor, so later `ds.mx` calls work on `ds` itself instead of a copy with the CRS attached. @michielv
@@ -25,6 +29,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - A pipeline config without a `transformations:` section no longer crashes. @michielv
 - `ds.mx.stack()` now stacks regular grids with `latitude`/`longitude` dimensions (it looked for `lat`/`lon`). @michielv
 - Delaunay interpolation could silently reuse the weights of a different grid with the same shape and corner coordinates, giving wrong values; weights are now cached by a hash of the full source and target coordinates. @michielv
+- NaN broadcasting (on by default in the pipeline) no longer crashes when an ensemble is aligned together with an observation or a deterministic forecast. @michielv
+- Delaunay interpolation no longer drops variables whose last dimension is not `grid_index`, such as an ensemble with `member` after `grid_index`. In the output, `point_index` takes the place of `grid_index` and the other dimensions keep their order. Variables without `grid_index` are still skipped, now with a logged warning instead of a `print`. @michielv
 
 ## [0.1.0](https://github.com/mlwp-tools/mxalign/releases/tag/v0.1.0)
 

README.md

@@ -43,7 +43,7 @@ For a more comprehensive interactive example, check out the [introductory notebo
 
 ### Configuration Contents
 
-The configuration file is divided into several main sections:
+The configuration file is divided into several main sections. It is validated on load against the schema in `mxalign.pipeline.config` (`PipelineConfig`): unknown sections or keys, and missing required ones such as `alignment.reference`, raise an error up front. Keys that are passed through to a function, such as loader, transformation, metric and save options, are accepted as-is.
 
 ```yaml
 datasets:
@@ -93,7 +93,7 @@ mxalign slurm path/to/config.yaml --account your_account --queue your_queue --co
 You can also execute the entire configuration-driven pipeline directly from Python using the `Runner` class.
 
 ```python
-from mxalign.runner import Runner
+from mxalign.pipeline.runner import Runner
 
 # Initialize the runner with a YAML config file or a dictionary
 runner = Runner("path/to/config.yaml")

doc/adr/ADR-003_ensemble-nan-masking.md

@@ -0,0 +1,55 @@
+---
+status: proposed
+date: 2026-10-27
+---
+
+# NaN masks for ensembles: reduce over members, never match them
+
+## Context and Problem Statement
+
+`broadcast_nans` propagates missing values across the aligned datasets, so that every model and the reference are verified on the same sample. It matched every dimension by coordinate value. Ensembles break that assumption. Observations and deterministic forecasts have no `member` dimension, and two ensembles have unrelated member labels and often different sizes (e.g. 5 and 31 members). As a result `broadcast_nans` crashed on any run with an ensemble, and the pipeline runs it by default.
+
+What does "missing" mean for an ensemble value, and how are members compared between datasets?
+
+## Decision Drivers
+
+* Every model and the reference should be verified on the same (time, space) sample.
+* Ensemble scores assume a fixed ensemble size M across the sample: fair CRPS and the Brier score are normalised by M, and the spread-skill ratio is compared to `sqrt((M+1)/M)`.
+* Member labels identify members within one ensemble only. Member 3 of model A has nothing to do with member 3 of model B.
+* mxalign should not depend on the position of the `member` dimension, because its order isn't enforced by mlwp-data-specs.
+
+## Considered Options
+
+1. An ensemble value is missing if **any** member is missing. The mask is computed with `member` reduced away and then applied to all members. Members are never matched between datasets.
+2. Keep per-member NaNs, mask the other datasets only where **all** members are missing, and let every metric cope with a varying M.
+3. Match `member` by label between ensembles, like any other dimension, and reduce it only against datasets that don't have it.
+
+## Decision Outcome
+
+Chosen option: 1, because it is the only one that keeps M constant over the verification sample without every metric having to handle NaNs itself. It also handles any number of ensembles of any size the same way.
+
+The uncertainty dimension is taken from the uncertainty trait (`member` for `ensemble`, `quantile` for `quantile`), not from dimension names. The same rule applies to quantile datasets. Apart from the uncertainty dimension, a variable shared by two datasets must have the same dimensions in both. Otherwise `broadcast_nans` raises, because the datasets aren't aligned.
+
+### Consequences
+
+* Good, because M is the same at every verified point, and all models and the reference share one sample.
+* Good, because a dataset tagged `deterministic` that still has a `member` dimension raises an error instead of being silently matched member by member.
+* Bad, because one missing member discards the whole ensemble at that point, e.g. when a single member file is missing. To keep a partial ensemble, subset its members before alignment, or switch `broadcast_nans` off and handle NaNs in the metrics.
+* Neutral, because `broadcast_nans` now requires the uncertainty trait on every dataset, which mlwp-data-loaders always sets.
+
+### Confirmation
+
+`tests/test_nans.py` covers observation/ensemble pairs in both orders, observation, deterministic and ensemble together, ensembles of different sizes and labels, dimension order and dask inputs. `tests/test_runner.py::test_align_mixes_deterministic_and_ensemble` covers the pipeline.
+
+## Pros and Cons of the Options
+
+### 2. Per-member NaNs
+
+* Good, because no member data is discarded.
+* Bad, because M varies over the sample, so fair scores and the spread-skill reference need per-point M, and samples with different M are averaged together.
+* Bad, because every ensemble metric has to mask NaNs jointly itself.
+
+### 3. Match members by label
+
+* Bad, because members of different ensembles are unrelated, so matching labels is meaningless.
+* Bad, because ensembles of different sizes would be masked on the overlapping labels only, which again makes M vary.

examples/advanced.ipynb

@@ -384,7 +384,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from mxalign.runner import Runner\n",
+    "from mxalign.pipeline.runner import Runner\n",
     "\n",
     "runner = Runner(\"./config.yaml\")"
    ]
@@ -455,7 +455,7 @@
     }
    ],
    "source": [
-    "runner.config()"
+    "runner.config"
    ]
   },
   {

examples/myloader.py

@@ -1,7 +1,22 @@
+"""Example custom loader for SYNOP observations, used by the example pipeline config."""
+
 import xarray as xr
 
 
-def load_dataset(file):
+def load_dataset(file: str) -> xr.Dataset:
+    """Load a SYNOP observation netCDF file and attach the mlwp traits.
+
+    Parameters
+    ----------
+    file : str
+        Path to the netCDF file with a ``code`` station dimension.
+
+    Returns
+    -------
+    xr.Dataset
+        Observation point dataset with a ``point_index`` dimension, mlwp
+        traits in its attributes, and CF attributes on its coordinates.
+    """
     ds = xr.open_mfdataset([file], engine="h5netcdf")
 
     ds = ds.rename_dims(code="point_index")

pyproject.toml

@@ -23,10 +23,11 @@ dependencies = [
     "mlwp-data-specs",
     "mlwp-data-loaders",
     "pandas>=2.2",
+    "pydantic>=2",
 ]
 
 [project.scripts]
-mxalign = "mxalign.cli:main"
+mxalign = "mxalign.pipeline.cli:main"
 
 [project.optional-dependencies]
 earthkit = [
@@ -51,6 +52,28 @@ build-backend = "hatchling.build"
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 
+[tool.ruff.lint]
+# ANN: type annotations on all arguments and return values; D: docstrings
+extend-select = ["ANN", "D"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+
+[tool.ruff.lint.per-file-ignores]
+# Orphaned loader that imports removed modules; pending removal or port to mlwp-data-loaders
+"src/mxalign/loaders/ifs_forecast.py" = ["ANN", "D"]
+
+[tool.numpydoc_validation]
+checks = [
+    "all",   # report on all checks, except the below
+    "EX01",  # no "Examples" section required
+    "SA01",  # no "See Also" section required
+    "ES01",  # no "Extended Summary" section required
+]
+# __init__ parameters are documented in the class docstring (numpydoc convention)
+exclude = ['\.__init__$']
+exclude_files = ['src/mxalign/loaders/ifs_forecast\.py']
+
 [dependency-groups]
 dev = [
     "ipykernel>=7.2.0",

src/mxalign/__init__.py

@@ -1,24 +1,29 @@
-from .transformations.transform import transform
-from .transformations.registry import available_transformations, register_transformation
-from .interpolations.interpolate import interpolate
-from .interpolations.registry import available_interpolations, register_interpolator
-from .align.time import align_time
-from .align.space import align_space
+"""Align and verify weather forecast and observation datasets with xarray."""
 
-from . import accessors
-from . import transformations
-from . import interpolations
+from . import accessor, alignment, pipeline
+from .alignment.interpolations.interpolate import interpolate
+from .alignment.interpolations.registry import (
+    available_interpolations,
+    register_interpolator,
+)
+from .alignment.space import align_space
+from .alignment.time import align_time
+from .pipeline.transformations.registry import (
+    available_transformations,
+    register_transformation,
+)
+from .pipeline.transformations.transform import transform
 
 __all__ = [
-    "transform",
+    "accessor",
+    "align_space",
+    "align_time",
+    "alignment",
+    "available_interpolations",
     "available_transformations",
-    "register_transformation",
     "interpolate",
-    "available_interpolations",
+    "pipeline",
     "register_interpolator",
-    "align_time",
-    "align_space",
-    "accessors",
-    "transformations",
-    "interpolations",
+    "register_transformation",
+    "transform",
 ]

src/mxalign/accessor/__init__.py

@@ -0,0 +1,5 @@
+"""Dataset accessors provided by mxalign (registers ``ds.mx`` on import)."""
+
+from . import mx
+
+__all__ = ["mx"]