Skip to content

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

Draft
mpvginde wants to merge 12 commits into
refactor/alignmentfrom
feat/ensemble-alignment
Draft

mpvginde wants to merge 12 commits into
refactor/alignmentfrom
feat/ensemble-alignment

Conversation

@mpvginde

Copy link
Copy Markdown
Contributor

Stacked on #44 (and so on #39–#42). Like those, this PR targets refactor/alignment, so until the stack merges the diff includes its commits. To see only this PR's changes, use 68cea40...feat/ensemble-alignment (5 commits). I'll rebase if the stack changes in review.

Why

Time alignment and interpolation already carry a member dimension through. But every run with an ensemble then crashed in broadcast_nans, which the pipeline runs by default. This PR makes the alignment step work for ensembles, on their own or mixed with deterministic forecasts and observations. Verification (chunking of member in Metric, and which metrics apply to which uncertainty type) comes in a follow-up PR.

Changes

  • Uncertainty predicates: ds.mx.is_deterministic(), is_ensemble() and is_quantile(). Like the space and time predicates, they read the trait on every call.
  • broadcast_nans no longer crashes when an ensemble meets a dataset without member. The uncertainty dimension, taken from the trait, is never matched between datasets. If any member is missing at a point, the ensemble counts as missing there, and the mask goes to all its members and to the other datasets. This keeps the ensemble size constant over the verification sample, and ensembles of different sizes work together. Any other difference in dimensions on a shared variable raises a ValueError. Please review the decision in ADR-003 (status: proposed). Its main cost is that one missing member discards the whole ensemble at that point.
  • Delaunay no longer silently drops variables whose last dimension isn't grid_index, such as an ensemble with member after grid_index. point_index takes the place of grid_index, and the other dimensions keep their order.
  • Tests: ds_fcst_ens and ds_grid_ens fixtures, whose member coordinate has standard_name="realization" and units="1". Tests cover time alignment with ensembles (pairwise in all directions, and global), both interpolators with member first or last (numpy and dask), a new tests/test_nans.py, and a mixed deterministic + ensemble Runner.align.

Breaking: broadcast_nans now needs mlwp_uncertainty_trait on every dataset. Datasets from mlwp-data-loaders always have it.

Testing

Not in this PR

🤖 Generated with Claude Code

leifdenby and others added 12 commits October 27, 2026 11:28
accessors/time.py and accessors/space.py held alignment operations
rather than accessor code. Move them into align/time.py and
align/space.py next to the multi-dataset align_time/align_space, and
have the mx accessor import them from there. No behaviour change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…r alignment

Top-level layout becomes `accessor` (the ds.mx accessor) and
`alignment` (time/space/nans alignment plus the interpolators it uses);
the pipeline package follows separately. Public re-exports from
mxalign/__init__.py are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The YAML-driven pipeline (cli, runner, verification, transformations
and the config/dates/save helpers it alone uses) now lives in
mxalign.pipeline. The console script points at mxalign.pipeline.cli:main.
Public re-exports from mxalign/__init__.py are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Replace the untyped Config dict wrapper with pydantic models
(PipelineConfig and one model per section) loaded via load_config().
Unknown or missing keys now fail when the Runner is created instead of
deep inside a run, and the Runner reads typed attributes rather than
popping keys (which also stops align() mutating the config). Keys that
are forwarded to loaders, transformations, metrics and save methods
are kept as pass-through kwargs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…add AGENTS.md

Every module, class, function and method in src/, tests/ and the example
loader now has type annotations on all arguments and return values and a
numpy-style docstring (Parameters/Returns/Raises). This is enforced in
pre-commit (and therefore CI) by ruff's ANN and D rules (numpy convention)
and the numpydoc-validation hook. AGENTS.md describes the package layout,
development tools and coding conventions.

No behaviour changes: apart from docstrings, annotations and import
sorting, the only code edits are sorting mxalign.__all__ and dict literals
in tests.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
is_deterministic(), is_ensemble() and is_quantile() read mlwp_uncertainty_trait on every call, like the space and time predicates (TODO_ENS U4).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add ds_fcst_ens and ds_grid_ens fixtures (member coord with CF realization attrs) and lock in that pairwise and global time alignment and grid-to-point interpolation carry the member dim and the ensemble trait, and never match members between datasets (TODO_ENS U9).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Variables whose last dim wasn't grid_index were silently dropped, e.g. an ensemble with member after grid_index. Move grid_index last for the matmul and put point_index in its place in the output. Variables without grid_index are still skipped, now with a logged warning (TODO_ENS U6).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
broadcast_nans crashed whenever an ensemble met a dataset without member, which the pipeline does by default. The uncertainty dim (from the trait) is now reduced with any() before masks are compared and the mask is applied to all members, so the ensemble size stays constant; members are never matched between datasets. Other dim mismatches raise a ValueError. Decision recorded in ADR-003 (TODO_ENS U1).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants