From 6b0eaacf9d19e58eb1c11324263e154c428e6b65 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 11:28:28 +0100 Subject: [PATCH 01/12] Move pairwise alignment ops out of accessors/ into align/ 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 --- src/mxalign/accessors/mx.py | 4 +- src/mxalign/accessors/space.py | 35 --------- src/mxalign/accessors/time.py | 128 -------------------------------- src/mxalign/align/space.py | 37 ++++++++++ src/mxalign/align/time.py | 130 ++++++++++++++++++++++++++++++++- 5 files changed, 168 insertions(+), 166 deletions(-) delete mode 100644 src/mxalign/accessors/space.py delete mode 100644 src/mxalign/accessors/time.py diff --git a/src/mxalign/accessors/mx.py b/src/mxalign/accessors/mx.py index 390fc44..b6230ef 100644 --- a/src/mxalign/accessors/mx.py +++ b/src/mxalign/accessors/mx.py @@ -7,8 +7,8 @@ from mlwp_data_specs.specs.traits.time_coordinate import Time from ..utils.projections import create_cartopy_crs, BUILTIN -from . import time as _time -from . import space as _space +from ..align import time as _time +from ..align import space as _space @xr.register_dataset_accessor("mx") diff --git a/src/mxalign/accessors/space.py b/src/mxalign/accessors/space.py deleted file mode 100644 index b8a6ec0..0000000 --- a/src/mxalign/accessors/space.py +++ /dev/null @@ -1,35 +0,0 @@ -import numpy as np -import xarray as xr - -# Tolerance in degrees that the coordinates of two grids can differ while still -# being interpreted as the same grid. 0.0001 degrees ~ 10m at 45 deg latitude. -COORD_TOLERANCE = 0.0001 - - -def align_grid_grid(ds1: xr.Dataset, ds2: xr.Dataset, **kwargs) -> xr.Dataset: - """Return ``ds1`` unchanged if its grid matches ``ds2``'s within ``COORD_TOLERANCE``; raises otherwise.""" - if np.array_equal( - ds1["longitude"].values, ds2["longitude"].values - ) and np.array_equal(ds1["latitude"].values, ds2["latitude"].values): - return ds1 - elif np.allclose( - ds1["longitude"].values, ds2["longitude"].values, atol=COORD_TOLERANCE - ) and np.allclose( - ds1["latitude"].values, ds2["latitude"].values, atol=COORD_TOLERANCE - ): - print( - f"Some lat-lon coordinates differ but within {COORD_TOLERANCE}°, treating as equal" - ) - return ds1 - else: - raise NotImplementedError("Regridding not implemented") - - -def align_grid_point( - ds1: xr.Dataset, ds2: xr.Dataset, method: str = "xarray", **kwargs -) -> xr.Dataset: - """Interpolate grid ``ds1`` to the point locations of ``ds2`` using the named interpolator.""" - from ..interpolations.registry import get_interpolation - - interp_cls = get_interpolation(method) - return interp_cls(ds2, **kwargs).interpolate(ds1.copy()) diff --git a/src/mxalign/accessors/time.py b/src/mxalign/accessors/time.py deleted file mode 100644 index 34e3ed0..0000000 --- a/src/mxalign/accessors/time.py +++ /dev/null @@ -1,128 +0,0 @@ -import numpy as np -import pandas as pd -import xarray as xr - -from mlwp_data_specs.api import TIME_TRAIT_ATTR -from mlwp_data_specs.specs.traits.time_coordinate import Time - - -def _add_valid_time(ds_fcst: xr.Dataset) -> xr.Dataset: - """Add a 2-D ``valid_time`` coordinate (``reference_time + lead_time``) to a forecast dataset.""" - valid_time = ( - ds_fcst["reference_time"].values[:, np.newaxis] + ds_fcst["lead_time"].values - ) - return ds_fcst.assign_coords( - {"valid_time": (["reference_time", "lead_time"], valid_time)} - ) - - -def align_forecast_to_observation( - ds_fcst: xr.Dataset, - ds_obs: xr.Dataset, - lead_time: str | list | np.timedelta64 = "shortest", -) -> xr.Dataset: - """Stack ``ds_fcst`` to valid_time, deduplicate by ``lead_time`` strategy, and reindex to ``ds_obs.valid_time``.""" - ds_with_vt = _add_valid_time(ds_fcst) - ds_stacked = ds_with_vt.stack(time=["reference_time", "lead_time"]).reset_index( - "time" - ) - - vt_vals = ds_stacked.valid_time.values - lt_vals = ds_stacked.lead_time.values - - if isinstance(lead_time, str) and lead_time in ("shortest", "longest"): - df = pd.DataFrame({"vt": vt_vals, "lt": lt_vals}) - agg = "min" if lead_time == "shortest" else "max" - is_candidate = (df.groupby("vt")["lt"].transform(agg) == df["lt"]).values - else: - requested = _parse_lead_times(lead_time) - missing = np.setdiff1d(requested, lt_vals) - if missing.size: - raise ValueError( - f"lead_time(s) {pd.to_timedelta(missing).tolist()} not in forecast " - f"lead times {pd.to_timedelta(np.unique(lt_vals)).tolist()}" - ) - is_candidate = np.isin(lt_vals, requested) - - # Among candidate entries keep the first per valid_time (handles ties: the - # same valid_time reached via several reference_times) - candidate_positions = np.where(is_candidate)[0] - _, first_in_group = np.unique(vt_vals[candidate_positions], return_index=True) - positions = candidate_positions[first_in_group] - - ds_1d = ds_stacked.isel(time=positions) - ds_1d = ds_1d.swap_dims({"time": "valid_time"}) - ds_1d = ds_1d.drop_vars( - [v for v in ["reference_time", "lead_time", "time"] if v in ds_1d.coords] - ) - ds_1d = ds_1d.transpose("valid_time", ...) - - ds_1d = ds_1d.reindex(valid_time=ds_obs.valid_time) - ds_1d.attrs[TIME_TRAIT_ATTR] = Time.OBSERVATION.value - return ds_1d - - -def _parse_lead_times(lead_time) -> np.ndarray: - """Convert one or more lead times (timedelta-like or strings such as ``"6h"``) to ``timedelta64[ns]``.""" - values = np.atleast_1d(np.asarray(lead_time, dtype=object)) - # np.timedelta64 subclasses np.number, so exclude it explicitly - if any( - isinstance(v, (int, float, np.number)) and not isinstance(v, np.timedelta64) - for v in values - ): - raise ValueError( - f"lead_time {lead_time!r} needs a unit, e.g. '6h' or np.timedelta64(6, 'h')." - ) - try: - return pd.to_timedelta(values).to_numpy().astype("timedelta64[ns]") - except ValueError: - raise ValueError( - f"Unknown lead_time option for F→O alignment: {lead_time!r}. Use " - "'shortest', 'longest', or one or more lead times such as '6h'." - ) from None - - -def align_observation_to_forecast( - ds_obs: xr.Dataset, ds_fcst: xr.Dataset -) -> xr.Dataset: - """Broadcast 1-D observations onto ``(reference_time, lead_time)`` by selecting via forecast valid times.""" - ds_fcst_with_vt = _add_valid_time(ds_fcst) - valid_time_2d = ds_fcst_with_vt["valid_time"] # shape (reference_time, lead_time) - - # Reindex obs onto all unique fcst valid_times (NaN-fills fcst valid_times not in obs) - fcst_vt_flat = np.unique(valid_time_2d.values.ravel()) - obs_reindexed = ds_obs.reindex(valid_time=fcst_vt_flat) - - # sel with a 2D DataArray indexer broadcasts 1D obs → (reference_time, lead_time) - ds_out = obs_reindexed.sel(valid_time=valid_time_2d) - - ds_out.attrs[TIME_TRAIT_ATTR] = Time.FORECAST.value - return ds_out - - -def align_observation_to_observation(ds1: xr.Dataset, ds2: xr.Dataset) -> xr.Dataset: - """Reindex ``ds1`` to ``ds2.valid_time`` with NaN-fill for missing times.""" - return ds1.reindex(valid_time=ds2.valid_time) - - -def align_forecast_to_forecast( - ds1: xr.Dataset, ds2: xr.Dataset, lead_time: str = "reference" -) -> xr.Dataset: - """Reindex ``ds1`` to ``ds2.reference_time`` and align lead times per ``lead_time`` strategy.""" - ds_out = ds1.reindex(reference_time=ds2.reference_time) - - if lead_time == "reference": - ds_out = ds_out.reindex(lead_time=ds2.lead_time) - elif lead_time == "intersection": - common_lt = np.intersect1d(ds_out.lead_time.values, ds2.lead_time.values) - ds_out = ds_out.sel(lead_time=common_lt) - elif lead_time == "union": - all_lt = np.union1d(ds_out.lead_time.values, ds2.lead_time.values) - ds_out = ds_out.reindex(lead_time=all_lt) - else: - raise ValueError(f"Unknown lead_time option for F→F alignment: {lead_time!r}") - - # Refresh valid_time - if "valid_time" in ds_out.coords: - ds_out = ds_out.drop_vars("valid_time") - return _add_valid_time(ds_out) diff --git a/src/mxalign/align/space.py b/src/mxalign/align/space.py index 5d4a2a2..8ef60bc 100644 --- a/src/mxalign/align/space.py +++ b/src/mxalign/align/space.py @@ -1,5 +1,10 @@ +import numpy as np import xarray as xr +# Tolerance in degrees that the coordinates of two grids can differ while still +# being interpreted as the same grid. 0.0001 degrees ~ 10m at 45 deg latitude. +COORD_TOLERANCE = 0.0001 + def align_space( datasets: xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset], @@ -39,3 +44,35 @@ def align_space( if keys is not None: return dict(zip(keys, aligned)) return aligned[0] if len(aligned) == 1 else aligned + + +# --- Pairwise alignment (used by ``ds.mx.align_space_with``) --- + + +def align_grid_grid(ds1: xr.Dataset, ds2: xr.Dataset, **kwargs) -> xr.Dataset: + """Return ``ds1`` unchanged if its grid matches ``ds2``'s within ``COORD_TOLERANCE``; raises otherwise.""" + if np.array_equal( + ds1["longitude"].values, ds2["longitude"].values + ) and np.array_equal(ds1["latitude"].values, ds2["latitude"].values): + return ds1 + elif np.allclose( + ds1["longitude"].values, ds2["longitude"].values, atol=COORD_TOLERANCE + ) and np.allclose( + ds1["latitude"].values, ds2["latitude"].values, atol=COORD_TOLERANCE + ): + print( + f"Some lat-lon coordinates differ but within {COORD_TOLERANCE}°, treating as equal" + ) + return ds1 + else: + raise NotImplementedError("Regridding not implemented") + + +def align_grid_point( + ds1: xr.Dataset, ds2: xr.Dataset, method: str = "xarray", **kwargs +) -> xr.Dataset: + """Interpolate grid ``ds1`` to the point locations of ``ds2`` using the named interpolator.""" + from ..interpolations.registry import get_interpolation + + interp_cls = get_interpolation(method) + return interp_cls(ds2, **kwargs).interpolate(ds1.copy()) diff --git a/src/mxalign/align/time.py b/src/mxalign/align/time.py index 8da8cf5..edc191a 100644 --- a/src/mxalign/align/time.py +++ b/src/mxalign/align/time.py @@ -1,8 +1,11 @@ import logging +import numpy as np +import pandas as pd import xarray as xr -from mxalign.accessors.time import align_observation_to_forecast +from mlwp_data_specs.api import TIME_TRAIT_ATTR +from mlwp_data_specs.specs.traits.time_coordinate import Time logger = logging.getLogger(__name__) @@ -141,3 +144,128 @@ def _constrain_to_reference( "within the reference observation's valid_time range." ) return r_star.where(rows_ok, drop=True) + + +# --- Pairwise alignment (used by ``ds.mx.align_time_with``) --- + + +def _add_valid_time(ds_fcst: xr.Dataset) -> xr.Dataset: + """Add a 2-D ``valid_time`` coordinate (``reference_time + lead_time``) to a forecast dataset.""" + valid_time = ( + ds_fcst["reference_time"].values[:, np.newaxis] + ds_fcst["lead_time"].values + ) + return ds_fcst.assign_coords( + {"valid_time": (["reference_time", "lead_time"], valid_time)} + ) + + +def align_forecast_to_observation( + ds_fcst: xr.Dataset, + ds_obs: xr.Dataset, + lead_time: str | list | np.timedelta64 = "shortest", +) -> xr.Dataset: + """Stack ``ds_fcst`` to valid_time, deduplicate by ``lead_time`` strategy, and reindex to ``ds_obs.valid_time``.""" + ds_with_vt = _add_valid_time(ds_fcst) + ds_stacked = ds_with_vt.stack(time=["reference_time", "lead_time"]).reset_index( + "time" + ) + + vt_vals = ds_stacked.valid_time.values + lt_vals = ds_stacked.lead_time.values + + if isinstance(lead_time, str) and lead_time in ("shortest", "longest"): + df = pd.DataFrame({"vt": vt_vals, "lt": lt_vals}) + agg = "min" if lead_time == "shortest" else "max" + is_candidate = (df.groupby("vt")["lt"].transform(agg) == df["lt"]).values + else: + requested = _parse_lead_times(lead_time) + missing = np.setdiff1d(requested, lt_vals) + if missing.size: + raise ValueError( + f"lead_time(s) {pd.to_timedelta(missing).tolist()} not in forecast " + f"lead times {pd.to_timedelta(np.unique(lt_vals)).tolist()}" + ) + is_candidate = np.isin(lt_vals, requested) + + # Among candidate entries keep the first per valid_time (handles ties: the + # same valid_time reached via several reference_times) + candidate_positions = np.where(is_candidate)[0] + _, first_in_group = np.unique(vt_vals[candidate_positions], return_index=True) + positions = candidate_positions[first_in_group] + + ds_1d = ds_stacked.isel(time=positions) + ds_1d = ds_1d.swap_dims({"time": "valid_time"}) + ds_1d = ds_1d.drop_vars( + [v for v in ["reference_time", "lead_time", "time"] if v in ds_1d.coords] + ) + ds_1d = ds_1d.transpose("valid_time", ...) + + ds_1d = ds_1d.reindex(valid_time=ds_obs.valid_time) + ds_1d.attrs[TIME_TRAIT_ATTR] = Time.OBSERVATION.value + return ds_1d + + +def _parse_lead_times(lead_time) -> np.ndarray: + """Convert one or more lead times (timedelta-like or strings such as ``"6h"``) to ``timedelta64[ns]``.""" + values = np.atleast_1d(np.asarray(lead_time, dtype=object)) + # np.timedelta64 subclasses np.number, so exclude it explicitly + if any( + isinstance(v, (int, float, np.number)) and not isinstance(v, np.timedelta64) + for v in values + ): + raise ValueError( + f"lead_time {lead_time!r} needs a unit, e.g. '6h' or np.timedelta64(6, 'h')." + ) + try: + return pd.to_timedelta(values).to_numpy().astype("timedelta64[ns]") + except ValueError: + raise ValueError( + f"Unknown lead_time option for F→O alignment: {lead_time!r}. Use " + "'shortest', 'longest', or one or more lead times such as '6h'." + ) from None + + +def align_observation_to_forecast( + ds_obs: xr.Dataset, ds_fcst: xr.Dataset +) -> xr.Dataset: + """Broadcast 1-D observations onto ``(reference_time, lead_time)`` by selecting via forecast valid times.""" + ds_fcst_with_vt = _add_valid_time(ds_fcst) + valid_time_2d = ds_fcst_with_vt["valid_time"] # shape (reference_time, lead_time) + + # Reindex obs onto all unique fcst valid_times (NaN-fills fcst valid_times not in obs) + fcst_vt_flat = np.unique(valid_time_2d.values.ravel()) + obs_reindexed = ds_obs.reindex(valid_time=fcst_vt_flat) + + # sel with a 2D DataArray indexer broadcasts 1D obs → (reference_time, lead_time) + ds_out = obs_reindexed.sel(valid_time=valid_time_2d) + + ds_out.attrs[TIME_TRAIT_ATTR] = Time.FORECAST.value + return ds_out + + +def align_observation_to_observation(ds1: xr.Dataset, ds2: xr.Dataset) -> xr.Dataset: + """Reindex ``ds1`` to ``ds2.valid_time`` with NaN-fill for missing times.""" + return ds1.reindex(valid_time=ds2.valid_time) + + +def align_forecast_to_forecast( + ds1: xr.Dataset, ds2: xr.Dataset, lead_time: str = "reference" +) -> xr.Dataset: + """Reindex ``ds1`` to ``ds2.reference_time`` and align lead times per ``lead_time`` strategy.""" + ds_out = ds1.reindex(reference_time=ds2.reference_time) + + if lead_time == "reference": + ds_out = ds_out.reindex(lead_time=ds2.lead_time) + elif lead_time == "intersection": + common_lt = np.intersect1d(ds_out.lead_time.values, ds2.lead_time.values) + ds_out = ds_out.sel(lead_time=common_lt) + elif lead_time == "union": + all_lt = np.union1d(ds_out.lead_time.values, ds2.lead_time.values) + ds_out = ds_out.reindex(lead_time=all_lt) + else: + raise ValueError(f"Unknown lead_time option for F→F alignment: {lead_time!r}") + + # Refresh valid_time + if "valid_time" in ds_out.coords: + ds_out = ds_out.drop_vars("valid_time") + return _add_valid_time(ds_out) From 9f7f2df1176636adcffd84f0e4369f872ef6e3c1 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 11:46:21 +0100 Subject: [PATCH 02/12] Rename packages to accessor/ and alignment/, move interpolations under 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 --- src/mxalign/__init__.py | 19 +++++++++++-------- .../{accessors => accessor}/__init__.py | 0 src/mxalign/{accessors => accessor}/mx.py | 4 ++-- src/mxalign/{align => alignment}/__init__.py | 2 ++ .../interpolations/__init__.py | 0 .../{ => alignment}/interpolations/base.py | 2 +- .../interpolations/delaunay.py | 0 .../interpolations/interpolate.py | 0 .../interpolations/registry.py | 0 .../{ => alignment}/interpolations/xarray.py | 0 src/mxalign/{align => alignment}/nans.py | 0 src/mxalign/{align => alignment}/space.py | 2 +- src/mxalign/{align => alignment}/time.py | 0 src/mxalign/runner.py | 4 ++-- tests/test_align_time_global.py | 4 ++-- tests/test_interpolations.py | 8 ++++---- 16 files changed, 25 insertions(+), 20 deletions(-) rename src/mxalign/{accessors => accessor}/__init__.py (100%) rename src/mxalign/{accessors => accessor}/mx.py (99%) rename src/mxalign/{align => alignment}/__init__.py (60%) rename src/mxalign/{ => alignment}/interpolations/__init__.py (100%) rename src/mxalign/{ => alignment}/interpolations/base.py (97%) rename src/mxalign/{ => alignment}/interpolations/delaunay.py (100%) rename src/mxalign/{ => alignment}/interpolations/interpolate.py (100%) rename src/mxalign/{ => alignment}/interpolations/registry.py (100%) rename src/mxalign/{ => alignment}/interpolations/xarray.py (100%) rename src/mxalign/{align => alignment}/nans.py (100%) rename src/mxalign/{align => alignment}/space.py (97%) rename src/mxalign/{align => alignment}/time.py (100%) diff --git a/src/mxalign/__init__.py b/src/mxalign/__init__.py index 1ce5e4f..d104858 100644 --- a/src/mxalign/__init__.py +++ b/src/mxalign/__init__.py @@ -1,13 +1,16 @@ 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 +from .alignment.interpolations.interpolate import interpolate +from .alignment.interpolations.registry import ( + available_interpolations, + register_interpolator, +) +from .alignment.time import align_time +from .alignment.space import align_space -from . import accessors +from . import accessor +from . import alignment from . import transformations -from . import interpolations __all__ = [ "transform", @@ -18,7 +21,7 @@ "register_interpolator", "align_time", "align_space", - "accessors", + "accessor", + "alignment", "transformations", - "interpolations", ] diff --git a/src/mxalign/accessors/__init__.py b/src/mxalign/accessor/__init__.py similarity index 100% rename from src/mxalign/accessors/__init__.py rename to src/mxalign/accessor/__init__.py diff --git a/src/mxalign/accessors/mx.py b/src/mxalign/accessor/mx.py similarity index 99% rename from src/mxalign/accessors/mx.py rename to src/mxalign/accessor/mx.py index b6230ef..d6f39e4 100644 --- a/src/mxalign/accessors/mx.py +++ b/src/mxalign/accessor/mx.py @@ -7,8 +7,8 @@ from mlwp_data_specs.specs.traits.time_coordinate import Time from ..utils.projections import create_cartopy_crs, BUILTIN -from ..align import time as _time -from ..align import space as _space +from ..alignment import time as _time +from ..alignment import space as _space @xr.register_dataset_accessor("mx") diff --git a/src/mxalign/align/__init__.py b/src/mxalign/alignment/__init__.py similarity index 60% rename from src/mxalign/align/__init__.py rename to src/mxalign/alignment/__init__.py index ed4fec2..c341d2d 100644 --- a/src/mxalign/align/__init__.py +++ b/src/mxalign/alignment/__init__.py @@ -1,7 +1,9 @@ from . import time from . import space +from . import interpolations __all__ = [ "time", "space", + "interpolations", ] diff --git a/src/mxalign/interpolations/__init__.py b/src/mxalign/alignment/interpolations/__init__.py similarity index 100% rename from src/mxalign/interpolations/__init__.py rename to src/mxalign/alignment/interpolations/__init__.py diff --git a/src/mxalign/interpolations/base.py b/src/mxalign/alignment/interpolations/base.py similarity index 97% rename from src/mxalign/interpolations/base.py rename to src/mxalign/alignment/interpolations/base.py index c502569..402d9a1 100644 --- a/src/mxalign/interpolations/base.py +++ b/src/mxalign/alignment/interpolations/base.py @@ -1,6 +1,6 @@ import xarray as xr from mlwp_data_specs.specs.traits.spatial_coordinate import Space -from ..utils.traits import update_space_trait +from ...utils.traits import update_space_trait # Attributes set by ds.mx.add_crs / add_grid_mapping that describe a source grid GRID_ONLY_ATTRS = ("crs", "grid_mapping") diff --git a/src/mxalign/interpolations/delaunay.py b/src/mxalign/alignment/interpolations/delaunay.py similarity index 100% rename from src/mxalign/interpolations/delaunay.py rename to src/mxalign/alignment/interpolations/delaunay.py diff --git a/src/mxalign/interpolations/interpolate.py b/src/mxalign/alignment/interpolations/interpolate.py similarity index 100% rename from src/mxalign/interpolations/interpolate.py rename to src/mxalign/alignment/interpolations/interpolate.py diff --git a/src/mxalign/interpolations/registry.py b/src/mxalign/alignment/interpolations/registry.py similarity index 100% rename from src/mxalign/interpolations/registry.py rename to src/mxalign/alignment/interpolations/registry.py diff --git a/src/mxalign/interpolations/xarray.py b/src/mxalign/alignment/interpolations/xarray.py similarity index 100% rename from src/mxalign/interpolations/xarray.py rename to src/mxalign/alignment/interpolations/xarray.py diff --git a/src/mxalign/align/nans.py b/src/mxalign/alignment/nans.py similarity index 100% rename from src/mxalign/align/nans.py rename to src/mxalign/alignment/nans.py diff --git a/src/mxalign/align/space.py b/src/mxalign/alignment/space.py similarity index 97% rename from src/mxalign/align/space.py rename to src/mxalign/alignment/space.py index 8ef60bc..ed23ab5 100644 --- a/src/mxalign/align/space.py +++ b/src/mxalign/alignment/space.py @@ -72,7 +72,7 @@ def align_grid_point( ds1: xr.Dataset, ds2: xr.Dataset, method: str = "xarray", **kwargs ) -> xr.Dataset: """Interpolate grid ``ds1`` to the point locations of ``ds2`` using the named interpolator.""" - from ..interpolations.registry import get_interpolation + from .interpolations.registry import get_interpolation interp_cls = get_interpolation(method) return interp_cls(ds2, **kwargs).interpolate(ds1.copy()) diff --git a/src/mxalign/align/time.py b/src/mxalign/alignment/time.py similarity index 100% rename from src/mxalign/align/time.py rename to src/mxalign/alignment/time.py diff --git a/src/mxalign/runner.py b/src/mxalign/runner.py index 07ba516..3b21117 100644 --- a/src/mxalign/runner.py +++ b/src/mxalign/runner.py @@ -5,8 +5,8 @@ from .utils.config import Config from .transformations.transform import transform -from .align.time import align_time -from .align.nans import broadcast_nans +from .alignment.time import align_time +from .alignment.nans import broadcast_nans from .utils.save import save_dataset, save_metrics from .verification import Metric diff --git a/tests/test_align_time_global.py b/tests/test_align_time_global.py index c725649..9631aca 100644 --- a/tests/test_align_time_global.py +++ b/tests/test_align_time_global.py @@ -19,7 +19,7 @@ import pytest import xarray as xr -from mxalign.align.time import align_time +from mxalign.alignment.time import align_time T0 = np.datetime64("2020-01-01T00:00", "ns") H6 = np.timedelta64(6, "h") @@ -143,7 +143,7 @@ def test_prunes_ref_times_before_reference_window(self, caplog): f2 = fcst(["2020-01-01", "2020-01-02", "2020-01-03"], ["0h", "6h", "12h"]) ref = obs(pd.date_range("2020-01-03", "2020-01-05", freq="3h")) - with caplog.at_level(logging.WARNING, logger="mxalign.align.time"): + with caplog.at_level(logging.WARNING, logger="mxalign.alignment.time"): out = align_time({"f1": f1, "f2": f2, "ref": ref}, reference="ref") np.testing.assert_array_equal( diff --git a/tests/test_interpolations.py b/tests/test_interpolations.py index 0fd8650..12d3504 100644 --- a/tests/test_interpolations.py +++ b/tests/test_interpolations.py @@ -15,10 +15,10 @@ import pytest import xarray as xr -from mxalign.interpolations import delaunay -from mxalign.interpolations.delaunay import DelaunayInterpolator -from mxalign.interpolations.interpolate import interpolate -from mxalign.interpolations.xarray import XarrayInterpolator +from mxalign.alignment.interpolations import delaunay +from mxalign.alignment.interpolations.delaunay import DelaunayInterpolator +from mxalign.alignment.interpolations.interpolate import interpolate +from mxalign.alignment.interpolations.xarray import XarrayInterpolator # --------------------------------------------------------------------------- From 8700ff58cbedf354c57c2f0a9c5ce6d22fae0ee8 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 11:48:21 +0100 Subject: [PATCH 03/12] Move cli, runner, verification and their helpers into pipeline/ 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 --- README.md | 2 +- pyproject.toml | 2 +- src/mxalign/__init__.py | 11 +++++++---- src/mxalign/pipeline/__init__.py | 5 +++++ src/mxalign/{ => pipeline}/cli.py | 0 src/mxalign/{utils => pipeline}/config.py | 0 src/mxalign/{utils => pipeline}/dates.py | 0 src/mxalign/{ => pipeline}/runner.py | 8 ++++---- src/mxalign/{utils => pipeline}/save.py | 0 .../{ => pipeline}/transformations/__init__.py | 0 src/mxalign/{ => pipeline}/transformations/base.py | 0 .../{ => pipeline}/transformations/external.py | 0 .../{ => pipeline}/transformations/registry.py | 0 .../{ => pipeline}/transformations/transform.py | 0 src/mxalign/{ => pipeline}/verification.py | 0 tests/test_dates.py | 6 +++--- tests/test_runner.py | 2 +- tests/test_transformations.py | 4 ++-- 18 files changed, 24 insertions(+), 16 deletions(-) create mode 100644 src/mxalign/pipeline/__init__.py rename src/mxalign/{ => pipeline}/cli.py (100%) rename src/mxalign/{utils => pipeline}/config.py (100%) rename src/mxalign/{utils => pipeline}/dates.py (100%) rename src/mxalign/{ => pipeline}/runner.py (97%) rename src/mxalign/{utils => pipeline}/save.py (100%) rename src/mxalign/{ => pipeline}/transformations/__init__.py (100%) rename src/mxalign/{ => pipeline}/transformations/base.py (100%) rename src/mxalign/{ => pipeline}/transformations/external.py (100%) rename src/mxalign/{ => pipeline}/transformations/registry.py (100%) rename src/mxalign/{ => pipeline}/transformations/transform.py (100%) rename src/mxalign/{ => pipeline}/verification.py (100%) diff --git a/README.md b/README.md index 6f4da9b..b5c57ff 100644 --- a/README.md +++ b/README.md @@ -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") diff --git a/pyproject.toml b/pyproject.toml index 81cd43a..f1a9cee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ ] [project.scripts] -mxalign = "mxalign.cli:main" +mxalign = "mxalign.pipeline.cli:main" [project.optional-dependencies] earthkit = [ diff --git a/src/mxalign/__init__.py b/src/mxalign/__init__.py index d104858..c4549b8 100644 --- a/src/mxalign/__init__.py +++ b/src/mxalign/__init__.py @@ -1,5 +1,8 @@ -from .transformations.transform import transform -from .transformations.registry import available_transformations, register_transformation +from .pipeline.transformations.transform import transform +from .pipeline.transformations.registry import ( + available_transformations, + register_transformation, +) from .alignment.interpolations.interpolate import interpolate from .alignment.interpolations.registry import ( available_interpolations, @@ -10,7 +13,7 @@ from . import accessor from . import alignment -from . import transformations +from . import pipeline __all__ = [ "transform", @@ -23,5 +26,5 @@ "align_space", "accessor", "alignment", - "transformations", + "pipeline", ] diff --git a/src/mxalign/pipeline/__init__.py b/src/mxalign/pipeline/__init__.py new file mode 100644 index 0000000..50dece2 --- /dev/null +++ b/src/mxalign/pipeline/__init__.py @@ -0,0 +1,5 @@ +from . import transformations + +__all__ = [ + "transformations", +] diff --git a/src/mxalign/cli.py b/src/mxalign/pipeline/cli.py similarity index 100% rename from src/mxalign/cli.py rename to src/mxalign/pipeline/cli.py diff --git a/src/mxalign/utils/config.py b/src/mxalign/pipeline/config.py similarity index 100% rename from src/mxalign/utils/config.py rename to src/mxalign/pipeline/config.py diff --git a/src/mxalign/utils/dates.py b/src/mxalign/pipeline/dates.py similarity index 100% rename from src/mxalign/utils/dates.py rename to src/mxalign/pipeline/dates.py diff --git a/src/mxalign/runner.py b/src/mxalign/pipeline/runner.py similarity index 97% rename from src/mxalign/runner.py rename to src/mxalign/pipeline/runner.py index 3b21117..57014b2 100644 --- a/src/mxalign/runner.py +++ b/src/mxalign/pipeline/runner.py @@ -3,11 +3,11 @@ from mlwp_data_loaders import load_and_validate_dataset -from .utils.config import Config +from .config import Config from .transformations.transform import transform -from .alignment.time import align_time -from .alignment.nans import broadcast_nans -from .utils.save import save_dataset, save_metrics +from ..alignment.time import align_time +from ..alignment.nans import broadcast_nans +from .save import save_dataset, save_metrics from .verification import Metric diff --git a/src/mxalign/utils/save.py b/src/mxalign/pipeline/save.py similarity index 100% rename from src/mxalign/utils/save.py rename to src/mxalign/pipeline/save.py diff --git a/src/mxalign/transformations/__init__.py b/src/mxalign/pipeline/transformations/__init__.py similarity index 100% rename from src/mxalign/transformations/__init__.py rename to src/mxalign/pipeline/transformations/__init__.py diff --git a/src/mxalign/transformations/base.py b/src/mxalign/pipeline/transformations/base.py similarity index 100% rename from src/mxalign/transformations/base.py rename to src/mxalign/pipeline/transformations/base.py diff --git a/src/mxalign/transformations/external.py b/src/mxalign/pipeline/transformations/external.py similarity index 100% rename from src/mxalign/transformations/external.py rename to src/mxalign/pipeline/transformations/external.py diff --git a/src/mxalign/transformations/registry.py b/src/mxalign/pipeline/transformations/registry.py similarity index 100% rename from src/mxalign/transformations/registry.py rename to src/mxalign/pipeline/transformations/registry.py diff --git a/src/mxalign/transformations/transform.py b/src/mxalign/pipeline/transformations/transform.py similarity index 100% rename from src/mxalign/transformations/transform.py rename to src/mxalign/pipeline/transformations/transform.py diff --git a/src/mxalign/verification.py b/src/mxalign/pipeline/verification.py similarity index 100% rename from src/mxalign/verification.py rename to src/mxalign/pipeline/verification.py diff --git a/tests/test_dates.py b/tests/test_dates.py index c8e6380..c37a6bb 100644 --- a/tests/test_dates.py +++ b/tests/test_dates.py @@ -2,9 +2,9 @@ import pytest -from mxalign.utils.config import Config -from mxalign.utils.dates import Dates -from mxalign.utils.save import DatasetPath +from mxalign.pipeline.config import Config +from mxalign.pipeline.dates import Dates +from mxalign.pipeline.save import DatasetPath # 3 reference times (daily) × 3 lead times (0, 12, 24h) DATES = dict(start="2020-01-01", end="2020-01-03", period="1D", range="24h", step="12h") diff --git a/tests/test_runner.py b/tests/test_runner.py index a4c52be..b3aa11d 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1,6 +1,6 @@ import pytest -from mxalign.runner import Runner +from mxalign.pipeline.runner import Runner def _runner(datasets: dict, alignment: dict) -> Runner: diff --git a/tests/test_transformations.py b/tests/test_transformations.py index fd737c3..ec85306 100644 --- a/tests/test_transformations.py +++ b/tests/test_transformations.py @@ -1,11 +1,11 @@ import pytest import xarray as xr -from mxalign.transformations.registry import ( +from mxalign.pipeline.transformations.registry import ( available_transformations, register_transformation, ) -from mxalign.transformations.transform import transform +from mxalign.pipeline.transformations.transform import transform @register_transformation("_test_double") From e3e16022db0d66413cd221704bd9bb298fd28e08 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 12:25:09 +0100 Subject: [PATCH 04/12] Update Runner import in advanced example notebook Co-Authored-By: Claude Opus 5.5 --- examples/advanced.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/advanced.ipynb b/examples/advanced.ipynb index 2c15b56..c577565 100644 --- a/examples/advanced.ipynb +++ b/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\")" ] From ba16fc6704f28840325603672be41b6103917cf4 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 12:26:26 +0100 Subject: [PATCH 05/12] Validate pipeline config against an explicit pydantic schema 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 --- CHANGELOG.md | 1 + README.md | 2 +- examples/advanced.ipynb | 2 +- pyproject.toml | 1 + src/mxalign/pipeline/config.py | 178 +++++++++++++++++++++++---------- src/mxalign/pipeline/runner.py | 123 ++++++++++------------- tests/test_config.py | 75 ++++++++++++++ tests/test_dates.py | 7 +- tests/test_runner.py | 13 ++- uv.lock | 148 ++++++++++++++++++++++++++- 10 files changed, 418 insertions(+), 132 deletions(-) create mode 100644 tests/test_config.py diff --git a/CHANGELOG.md b/CHANGELOG.md index cbc0c92..3e3cdbe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Changed +- **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..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 diff --git a/README.md b/README.md index b5c57ff..d359155 100644 --- a/README.md +++ b/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: diff --git a/examples/advanced.ipynb b/examples/advanced.ipynb index c577565..b234aa2 100644 --- a/examples/advanced.ipynb +++ b/examples/advanced.ipynb @@ -455,7 +455,7 @@ } ], "source": [ - "runner.config()" + "runner.config" ] }, { diff --git a/pyproject.toml b/pyproject.toml index f1a9cee..8a70481 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "mlwp-data-specs", "mlwp-data-loaders", "pandas>=2.2", + "pydantic>=2", ] [project.scripts] diff --git a/src/mxalign/pipeline/config.py b/src/mxalign/pipeline/config.py index a4e0682..39b4a6d 100644 --- a/src/mxalign/pipeline/config.py +++ b/src/mxalign/pipeline/config.py @@ -1,62 +1,136 @@ +"""Schema for the pipeline YAML config. + +Sections and keys not listed here are rejected, except where a model allows +extra keys (``extra="allow"``); those are passed on as keyword arguments to +the function the section configures (loader, transformation, metric, save). +""" + +from datetime import datetime +from typing import Literal + import yaml +from pydantic import BaseModel, ConfigDict, model_validator from .dates import Dates -def load_yaml(fn: str) -> dict: - """Load a YAML file and return its contents as a dict.""" - with open(fn, "r") as f: - return yaml.safe_load(f) +class _Strict(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class _PassThrough(BaseModel): + """Model whose unknown keys are kept in ``model_extra`` as function kwargs.""" + model_config = ConfigDict(extra="allow") -class Config: - """Pipeline configuration loaded from a YAML file or a plain dict. + @property + def kwargs(self) -> dict: + return dict(self.model_extra or {}) - Top-level ``dates`` entries are merged into each dataset's loader config - so that file-path patterns can be expanded to concrete paths. + +class DatesConfig(_Strict): + """Date range used to expand ``{reference_time}``/``{lead_time}``/``{valid_time}`` in ``files``. + + Every field is optional so a dataset can override part of the top-level + ``dates`` section; after merging, all five must be set. """ - def __init__(self, config: str | dict) -> None: - self.config = load_yaml(config) if isinstance(config, str) else config - if not isinstance(self.config, dict): - raise TypeError("config should be a dictionary.") - self.dates = self.config.pop("dates", None) - self._init_datasets() - print("Config initialized") - - def __getitem__(self, key: str) -> dict | None: - config = self.config.get(key, None) - if config: - return config.copy() - else: - return config - - def __call__(self) -> dict: - return self.config - - def _init_datasets(self) -> None: - for key, loader in self.config["datasets"].items(): - dates_loader = loader.pop("dates", None) - if self.dates: - if dates_loader: - keys_all = list(set(self.dates.keys()).union(dates_loader.keys())) - dates = { - key: ( - dates_loader[key] - if key in dates_loader.keys() - else self.dates[key] - ) - for key in keys_all - } - else: - dates = self.dates.copy() - else: - if dates_loader: - dates = dates_loader.copy() - else: - dates = None - - if dates: - dates = Dates(**dates) - loader["files"] = dates.substitute(loader["files"]) - self.config["datasets"][key] = loader + start: datetime | str | None = None + end: datetime | str | None = None + period: str | None = None + range: str | None = None + step: str | None = None + + +class DatasetConfig(_PassThrough): + """One entry of ``datasets``; extra keys go to ``load_and_validate_dataset``.""" + + loader: str + files: str | list[str] + variables: str | list[str] | None = None + grid_mapping: str | dict | None = None + dates: DatesConfig | None = None + + +class TransformationConfig(_PassThrough): + """One entry of ``transformations`` (keyed by transformation name); extra keys go to the transformation.""" + + datasets: list[str] | None = None + + +class TimeAlignmentConfig(_Strict): + lead_time: str | list[str] | None = None + + +class SpaceAlignmentConfig(_Strict): + """Interpolator/regridder kwargs, chosen by the kind of space alignment needed.""" + + interpolation: dict | None = None + regrid: dict | None = None + + +class SaveConfig(_PassThrough): + """Output for ``ds.(path, **kwargs)``; extra keys go to the save method.""" + + method: str + path: str + + +class AlignmentSaveConfig(SaveConfig): + datasets: Literal["all", "merge"] = "all" + + +class AlignmentConfig(_Strict): + reference: str + broadcast_nans: bool = True + time: TimeAlignmentConfig | None = None + space: SpaceAlignmentConfig | None = None + save: AlignmentSaveConfig | None = None + + +class MetricConfig(_PassThrough): + """One entry of ``verification.metrics``; extra keys go to the metric function.""" + + function: str + inputs: dict[str, Literal["forecast", "reference"]] + + +class VerificationConfig(_Strict): + reference: str + metrics: dict[str, MetricConfig] | None = None + save: SaveConfig | None = None + + +class PipelineConfig(_Strict): + dates: DatesConfig | None = None + datasets: dict[str, DatasetConfig] + transformations: dict[str, TransformationConfig] | None = None + alignment: AlignmentConfig + verification: VerificationConfig | None = None + + @model_validator(mode="after") + def _expand_files(self) -> "PipelineConfig": + """Expand each dataset's ``files`` pattern(s) over its (merged) ``dates``.""" + top = self.dates.model_dump(exclude_none=True) if self.dates else {} + for name, ds in self.datasets.items(): + own = ds.dates.model_dump(exclude_none=True) if ds.dates else {} + dates = {**top, **own} + if not dates: + continue + missing = set(DatesConfig.model_fields) - dates.keys() + if missing: + raise ValueError( + f"datasets.{name}: dates is missing {', '.join(sorted(missing))}" + ) + patterns = [ds.files] if isinstance(ds.files, str) else ds.files + dates = Dates(**dates) + ds.files = [path for p in patterns for path in dates.substitute(p)] + return self + + +def load_config(config: str | dict) -> PipelineConfig: + """Validate a pipeline config given as a YAML file path or a dict.""" + if isinstance(config, str): + with open(config, "r") as f: + config = yaml.safe_load(f) + return PipelineConfig.model_validate(config) diff --git a/src/mxalign/pipeline/runner.py b/src/mxalign/pipeline/runner.py index 57014b2..080201d 100644 --- a/src/mxalign/pipeline/runner.py +++ b/src/mxalign/pipeline/runner.py @@ -3,7 +3,12 @@ from mlwp_data_loaders import load_and_validate_dataset -from .config import Config +from .config import ( + PipelineConfig, + SpaceAlignmentConfig, + TimeAlignmentConfig, + load_config, +) from .transformations.transform import transform from ..alignment.time import align_time from ..alignment.nans import broadcast_nans @@ -13,7 +18,7 @@ class Runner: def __init__(self, config: str | dict): - self.config = Config(config) + self.config: PipelineConfig = load_config(config) self.datasets = {} def run(self): @@ -26,16 +31,11 @@ def run(self): self.verify() def load_datasets(self): - config = self.config["datasets"] - if config is None: - raise ValueError("No datasets section in the config.") - for name, config_ds in config.items(): - config_ds = config_ds.copy() - loader = config_ds.pop("loader") - variables = config_ds.pop("variables", None) - grid_mapping = config_ds.pop("grid_mapping", None) - - raw_files = config_ds.pop("files") + for name, config_ds in self.config.datasets.items(): + variables = config_ds.variables + grid_mapping = config_ds.grid_mapping + + raw_files = config_ds.files if isinstance(raw_files, str): raw_files = [raw_files] files = [] @@ -48,8 +48,8 @@ def load_datasets(self): ds = load_and_validate_dataset( dataset_path, - loader=loader, - **config_ds, + loader=config_ds.loader, + **config_ds.kwargs, ) if variables is not None: @@ -61,85 +61,74 @@ def load_datasets(self): self.datasets[name] = ds def transform_datasets(self): - config = self.config["transformations"] - if config is None: + if self.config.transformations is None: return - for transformation, config_trans in config.items(): - config_trans = config_trans.copy() + for transformation, config_trans in self.config.transformations.items(): # if no datasets specified, apply to all datasets - names_ds = config_trans.pop("datasets", self.datasets.keys()) + names_ds = config_trans.datasets + if names_ds is None: + names_ds = self.datasets.keys() for name in names_ds: ds = self.datasets[name] self.datasets[name] = transform( - name=transformation, datasets=ds, **config_trans + name=transformation, datasets=ds, **config_trans.kwargs ) def align(self): - config = self.config["alignment"] - reference = config.pop("reference") - brdcst_nans = config.pop("broadcast_nans", True) - config_align_time = config.get("time", None) - config_align_space = config.get("space", None) - config_align_save = config.get("save", None) + config = self.config.alignment + reference = config.reference # align in time (always; the time section only holds options) - self.align_time(config_align_time, reference_name=reference) + self.align_time(config.time, reference_name=reference) # align in space - if config_align_space: - self.align_space(reference=reference, config=config_align_space) + if config.space is not None: + self.align_space(reference=reference, config=config.space) else: print("Skipping spatial alignment") # broadcast NaNs - if brdcst_nans: + if config.broadcast_nans: self.datasets = broadcast_nans(self.datasets) # Save aligned datasets - if config_align_save: - config = config_align_save.copy() - method = config.pop("method") - datasets = config.pop("datasets", "all") - if datasets == "all": + if config.save is not None: + save = config.save + kwargs = {"path": save.path, **save.kwargs} + if save.datasets == "all": for name, ds in self.datasets.items(): - save_dataset(method, name, ds, **config) - elif datasets == "merge": + save_dataset(save.method, name, ds, **kwargs) + elif save.datasets == "merge": ds = xr.concat( self.datasets.values(), dim=xr.Variable("model", list(self.datasets.keys())), ) - save_dataset(method, "merged", ds, **config) - else: - raise ValueError("Unknown option for dataset saving.") + save_dataset(save.method, "merged", ds, **kwargs) def verify(self): - config = self.config["verification"] - reference = self.datasets[config["reference"]] - config_metrics = config.get("metrics", None) - config_save_metrics = config.get("save", None) + config = self.config.verification + if config is None: + return + reference = self.datasets[config.reference] common_vars = set(reference.data_vars) for ds in self.datasets.values(): common_vars.intersection_update(set(ds.data_vars)) common_vars = list(common_vars) - if config_metrics: + if config.metrics: metrics = {} - for metric_name, config_metric in config["metrics"].items(): - config_metric = config_metric.copy() - func_path = config_metric.pop("function") - inputs = config_metric.pop("inputs") - + for metric_name, config_metric in config.metrics.items(): metric = Metric( name=metric_name, - func_path=func_path, + func_path=config_metric.function, ds_ref=reference[common_vars], - inputs=inputs, - **config_metric, + inputs=config_metric.inputs, + **config_metric.kwargs, ) models = {} for ds_name, ds in self.datasets.items(): - if ds_name != config["reference"]: + if ds_name != config.reference: models[ds_name] = metric.compute(ds[common_vars]) models = xr.concat( models.values(), dim=xr.Variable("model", list(models.keys())) @@ -150,28 +139,24 @@ def verify(self): ) self.metrics = metrics.transpose("model", "metric", ...).compute() - if config_save_metrics: - config = config_save_metrics.copy() - method = config.pop("method") - save_metrics(method, self.metrics, **config) - - def align_time(self, config, reference_name=None): - config = dict(config or {}) - lead_time = config.pop("lead_time", None) - if config: - raise ValueError( - f"Unknown option(s) in alignment.time: {', '.join(config)}. " - "Only lead_time is supported." - ) + if config.save is not None: + save = config.save + save_metrics(save.method, self.metrics, path=save.path, **save.kwargs) + + def align_time( + self, config: TimeAlignmentConfig | None, reference_name: str | None = None + ): + lead_time = config.lead_time if config is not None else None self.datasets = align_time( self.datasets, reference=reference_name, lead_time=lead_time ) - def align_space(self, reference, config): + def align_space(self, reference: str, config: SpaceAlignmentConfig): ds_ref = self.datasets[reference] for name, ds in self.datasets.items(): if name != reference: - options = config.get(get_spatial_alignment(ds, ds_ref), {}) + kind = get_spatial_alignment(ds, ds_ref) + options = getattr(config, kind, None) or {} self.datasets[name] = ds.mx.align_space_with(ds_ref, **options) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..64f2b5d --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,75 @@ +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from mxalign.pipeline.config import load_config + +EXAMPLE_CONFIG = Path(__file__).parent.parent / "examples" / "config.yaml" + +DATES = dict(start="2020-01-01", end="2020-01-02", period="1D", range="12h", step="12h") + + +def _config(**sections) -> dict: + return { + "datasets": {"fc": {"loader": "x", "files": "/d/fc.nc"}}, + "alignment": {"reference": "fc"}, + **sections, + } + + +def test_example_config_validates(): + config = load_config(str(EXAMPLE_CONFIG)) + assert config.alignment.reference == "cerra" + assert config.alignment.space.interpolation == {"method": "delaunay"} + assert config.verification.metrics["rmse"].kwargs == { + "dim": ["reference_time", "grid_index"], + "skipna": True, + } + # files patterns are expanded over the dates section + assert len(config.datasets["la-1024-ea-00"].files) > 1 + + +def test_unknown_top_level_key_is_rejected(): + with pytest.raises(ValidationError, match="alignmnet"): + load_config(_config(alignmnet={"reference": "fc"})) + + +def test_missing_alignment_reference_is_rejected(): + config = _config() + config["alignment"] = {"broadcast_nans": False} + with pytest.raises(ValidationError, match="alignment.reference"): + load_config(config) + + +def test_unknown_save_datasets_option_is_rejected(): + save = {"method": "to_netcdf", "path": "out.nc", "datasets": "some"} + with pytest.raises(ValidationError, match="alignment.save.datasets"): + load_config(_config(alignment={"reference": "fc", "save": save})) + + +def test_extra_dataset_keys_are_passed_to_loader(): + config = _config() + config["datasets"]["fc"]["chunks"] = {"time": 1} + assert load_config(config).datasets["fc"].kwargs == {"chunks": {"time": 1}} + + +def test_dataset_dates_override_top_level_dates(): + config = _config(dates=DATES) + config["datasets"]["fc"] = { + "loader": "x", + "files": "/d/{reference_time:%Y%m%d}.nc", + "dates": {"end": "2020-01-03"}, + } + assert load_config(config).datasets["fc"].files == [ + "/d/20200101.nc", + "/d/20200102.nc", + "/d/20200103.nc", + ] + + +def test_incomplete_dates_are_rejected(): + config = _config() + config["datasets"]["fc"]["dates"] = {"start": "2020-01-01"} + with pytest.raises(ValidationError, match="dates is missing end, period"): + load_config(config) diff --git a/tests/test_dates.py b/tests/test_dates.py index c37a6bb..441eac5 100644 --- a/tests/test_dates.py +++ b/tests/test_dates.py @@ -2,7 +2,7 @@ import pytest -from mxalign.pipeline.config import Config +from mxalign.pipeline.config import load_config from mxalign.pipeline.dates import Dates from mxalign.pipeline.save import DatasetPath @@ -45,15 +45,16 @@ def test_substitute_lead_time_requires_whole_hours(): def test_config_expands_files_from_dates_section(): # YAML parses unquoted timestamps to datetime objects - config = Config( + config = load_config( { "dates": {**DATES, "start": datetime.datetime(2020, 1, 1)}, "datasets": { "fc": {"loader": "x", "files": "/d/{reference_time:%Y%m%d}.nc"} }, + "alignment": {"reference": "fc"}, } ) - assert config["datasets"]["fc"]["files"] == [ + assert config.datasets["fc"].files == [ "/d/20200101.nc", "/d/20200102.nc", "/d/20200103.nc", diff --git a/tests/test_runner.py b/tests/test_runner.py index b3aa11d..d84eded 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1,4 +1,5 @@ import pytest +from pydantic import ValidationError from mxalign.pipeline.runner import Runner @@ -19,10 +20,14 @@ def test_align_runs_time_alignment_without_time_section(ds_fcst, ds_obs): assert runner.datasets["obs"].mx.is_forecast() -def test_align_time_rejects_unknown_options(ds_fcst, ds_obs): - runner = _runner({"fcst": ds_fcst, "obs": ds_obs}, {"reference": "obs"}) - with pytest.raises(ValueError, match="Unknown option.*return_as"): - runner.align_time({"return_as": "forecast"}, reference_name="obs") +def test_align_time_rejects_unknown_options(): + with pytest.raises(ValidationError, match="alignment.time.return_as"): + Runner( + { + "datasets": {}, + "alignment": {"reference": "obs", "time": {"return_as": "forecast"}}, + } + ) def test_transform_datasets_without_transformations_section(ds_obs): diff --git a/uv.lock b/uv.lock index 6ce6399..1e99d4f 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 5 +revision = 3 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -204,6 +204,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/76/641ae371508676492379f16e2fa48f4e2c11741bd63c48be4b12a6b09cba/aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e", size = 7490, upload-time = "2025-07-03T22:54:42.156Z" }, ] +[[package]] +name = "annotated-types" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" }, +] + [[package]] name = "appnope" version = "1.0.0" @@ -1808,6 +1817,7 @@ dependencies = [ { name = "mlwp-data-specs" }, { name = "netcdf4" }, { name = "pandas" }, + { name = "pydantic" }, { name = "pyyaml" }, { name = "scipy" }, { name = "xarray" }, @@ -1853,6 +1863,7 @@ requires-dist = [ { name = "mlwp-data-specs", git = "https://github.com/mlwp-tools/mlwp-data-specs?rev=059f382" }, { name = "netcdf4", specifier = ">=1.7.4" }, { name = "pandas", specifier = ">=2.2" }, + { name = "pydantic", specifier = ">=2" }, { name = "pyyaml", specifier = ">=6.0.3" }, { name = "scipy", specifier = ">=1.17.0" }, { name = "xarray", specifier = ">=2026.1.0" }, @@ -2125,7 +2136,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -2432,6 +2443,127 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b6/ae/315bbeee126a769cbde43cf4908acb4b40408d8bf8bbf334dd7368d8dadf/pycparser-3.12-py3-none-any.whl", hash = "sha256:a076f3f489b39b43ed13563a1dafe42447445bec2e294154ec5c57fcdda4c85e", size = 62600, upload-time = "2026-10-22T13:40:02.021Z" }, ] +[[package]] +name = "pydantic" +version = "2.14.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7c/0b/8e10b2e693af8ec54346a14caf36221334775975a747c9ceaa3f8371d96d/pydantic-2.14.1.tar.gz", hash = "sha256:94f478203dd03404682a1ada216965651dd74b1d2d5ffd62e00e0837caab5c26", size = 900362, upload-time = "2026-10-11T18:37:55.396Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/35/ea/a56b9fe5066f3537b7882f77e9c5dfb26d8c2eefdaed9b5fc73d57b4dc22/pydantic-2.14.1-py3-none-any.whl", hash = "sha256:9195d967ec791692a04438115466764fb8b9a27b31f14a760437694f40d6b454", size = 488106, upload-time = "2026-10-11T18:37:53.437Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.50.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a6/24/af4ec4be49fbc810f35b0bdcacc3d433b56bbfa469fd3bfd4ae116cd9bf1/pydantic_core-2.50.1.tar.gz", hash = "sha256:e50d7b94baac6c7d09927fa5ca5800a0c7ee5015c7fcff65beb3a1931b5a6e09", size = 513515, upload-time = "2026-10-11T18:35:44.82Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/94/101b49a6c63f99586755690696eae14532772587d9e39e93f7b8d28a49da/pydantic_core-2.50.1-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:704075d10b74f2f3c6e15407c696d88701df35fc8953f434a431add0d0074db0", size = 2222428, upload-time = "2026-10-11T18:32:20.228Z" }, + { url = "https://files.pythonhosted.org/packages/db/a3/f8b1b39349c33037c73d7a1b794c9c29d8f07d5bd2b74399c1375a529cc6/pydantic_core-2.50.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e8e1d6ce820aa23317e8209a86bd65a540973c12dc7552b48a4f6c8e9926815e", size = 2025779, upload-time = "2026-10-11T18:32:21.723Z" }, + { url = "https://files.pythonhosted.org/packages/d2/2d/862a7d115305cc621358e51ba13f296e0ffdf4fbd984be9f63bc7862730b/pydantic_core-2.50.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c18db21573bd2c6489f9a544b7499f0df2853958c568e5e783536ee1f690af41", size = 2048429, upload-time = "2026-10-11T18:32:23.3Z" }, + { url = "https://files.pythonhosted.org/packages/c0/7d/fed95d18bf2b5abe2d5ab2bbec4184bfc22b6805994a33ccdae8edffd25a/pydantic_core-2.50.1-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:cb57f304525a5e3c13333b772bf9a473f36326e9c821b2e8e1b2fd36f80ae2c3", size = 2173274, upload-time = "2026-10-11T18:32:24.697Z" }, + { url = "https://files.pythonhosted.org/packages/93/1e/c6f22ceac65dc3d291ff79bea0eb74cbd1de8b128f01fee48af03bcea1e1/pydantic_core-2.50.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a27c09d86600f1bf2fe3f37e1ae697faf3143931c09322cd799da94deee923b5", size = 2360403, upload-time = "2026-10-11T18:32:26.281Z" }, + { url = "https://files.pythonhosted.org/packages/2c/97/190b5b6bfe4c72f265feb657913eb30749193927a5ad7b83cfa7866a25a2/pydantic_core-2.50.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:46b3301d3b5c886f77de7546e47274a5842c622ea2020b8c6524c6b66913b4a6", size = 2439799, upload-time = "2026-10-11T18:32:27.844Z" }, + { url = "https://files.pythonhosted.org/packages/90/24/48cd98388e4ac2c08d0800af36d15413db15e20387f1038123a2c86cb2f6/pydantic_core-2.50.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:93ba4e9d8210d941c200431a56b2c0400b131865947903937ed3ec5404307d2e", size = 2160157, upload-time = "2026-10-11T18:32:29.553Z" }, + { url = "https://files.pythonhosted.org/packages/00/86/d909133d0b9979ed14c20f110f21124a0e346caf071bd66ac5af3009062b/pydantic_core-2.50.1-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:e5faeaee74a57d32b3ab3aebad2e348f06d3ba946fc5d28c1728455f00a3d13a", size = 2229523, upload-time = "2026-10-11T18:32:31.027Z" }, + { url = "https://files.pythonhosted.org/packages/90/31/4b43f08131b2098e73a852fd932063385f3103b33839fd9d1abdb89e673d/pydantic_core-2.50.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a3cda0e538208e5d722bbf3698b24f19c0a7d05bc8d5f8a7f9b121ea7fa243d9", size = 2242700, upload-time = "2026-10-11T18:32:32.622Z" }, + { url = "https://files.pythonhosted.org/packages/88/77/de39658569dc7a4c293d110a89629f61145e7538f00d1e4280c28debefd3/pydantic_core-2.50.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:57f51b31ff826e2859120cf4737c5a758a48d96f3e97da40ccee1796d58078ff", size = 2278452, upload-time = "2026-10-11T18:32:34.165Z" }, + { url = "https://files.pythonhosted.org/packages/46/bd/1eeb9e4e791e81210ef9096f58080b32ba15b5b13c108d50aaa3f8ce9f73/pydantic_core-2.50.1-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8daa7ee75245d43ad7d747e5c9ecc1b1d06552f72b14887e9276f787d57375f4", size = 2450098, upload-time = "2026-10-11T18:32:35.547Z" }, + { url = "https://files.pythonhosted.org/packages/1c/e6/b2accfdff17e5a88d6358536702146de00f9339adaa4d5bbe1ee93cb4832/pydantic_core-2.50.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:acbf31f37c53a5ac0c34706c80b4f5107ba20b05fdd3816124bf236ef0c57dd2", size = 2476839, upload-time = "2026-10-11T18:32:37.119Z" }, + { url = "https://files.pythonhosted.org/packages/3b/7e/c8a83eb5d2cd42c2e13507150fb777b1edea5a018adc9cbf2858a1d08285/pydantic_core-2.50.1-cp312-cp312-win32.whl", hash = "sha256:45b11cac094aa25725581d9304eee93c9028516b9ea80dd9e175e13a5a2c840e", size = 2101425, upload-time = "2026-10-11T18:32:38.593Z" }, + { url = "https://files.pythonhosted.org/packages/c5/17/8bbd530b8659e9d963f5b16f6cb8e159f4cd174b4d6304e48969ab12dba8/pydantic_core-2.50.1-cp312-cp312-win_amd64.whl", hash = "sha256:132529c83901437ff642f585216831bf5fd7a91df66829907e155192ead62498", size = 2172522, upload-time = "2026-10-11T18:32:40.112Z" }, + { url = "https://files.pythonhosted.org/packages/af/9f/a79d690b127e4e3cff143d031538ab9aa715f6b2fbeecbd07cd6f2bf9fb7/pydantic_core-2.50.1-cp312-cp312-win_arm64.whl", hash = "sha256:4e834f6a8e4ff772dcc34f58ef5504147a3ea5b0f4eeb13b0f8eb2ca75ac57f1", size = 2161366, upload-time = "2026-10-11T18:32:41.606Z" }, + { url = "https://files.pythonhosted.org/packages/ae/57/0e237d7091d2cd44a35d243b7344227f90440166860469cd036afc23a04d/pydantic_core-2.50.1-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:d5e062c01286d861fd6a1c4ff6e063547b3e713067f2df033c0ff97ac2ca006b", size = 2222774, upload-time = "2026-10-11T18:32:43.103Z" }, + { url = "https://files.pythonhosted.org/packages/f7/b9/c720e56858d4e1539503297ed37063e0c08e0f3541c41b777f6a800f75fa/pydantic_core-2.50.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0c003c3b7f49debb893d2d85ae099ac5959c9839e2f330fadb1fcdf7a6594482", size = 2024738, upload-time = "2026-10-11T18:32:44.587Z" }, + { url = "https://files.pythonhosted.org/packages/4d/b8/fbfc25875219cc060e613170ff10e850c6d8924beb4910da57c2ee3ba1d2/pydantic_core-2.50.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:409e0ea40ec30d9158f33574fd758e689f6045a0f2596701828c27816ca9687d", size = 2047049, upload-time = "2026-10-11T18:32:46.164Z" }, + { url = "https://files.pythonhosted.org/packages/d7/43/34210124d504c553688f2f04b48500b131237528ac545b445e0d6d30e0d5/pydantic_core-2.50.1-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:131059670f1d2444269b8585cb888963994871932447c08b39ac6a51fcfef658", size = 2173415, upload-time = "2026-10-11T18:32:47.722Z" }, + { url = "https://files.pythonhosted.org/packages/65/cf/6e178e8fdc11da5965bef980983bf46326a42f436871dd51ba0a57f39df1/pydantic_core-2.50.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6dbcbee53bf17196a7f745aa9bf5a9603953a1e365b1f020be3207c676a3e7c4", size = 2361719, upload-time = "2026-10-11T18:32:49.217Z" }, + { url = "https://files.pythonhosted.org/packages/11/14/bd5169d356aa91bf777e28ba0b281c192d286d03c2812c9c9db023a2f00e/pydantic_core-2.50.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:325c23f3e35cfbf0fe3486fa5f7260d1e45885173002d30a28ca019994124255", size = 2439803, upload-time = "2026-10-11T18:32:51.025Z" }, + { url = "https://files.pythonhosted.org/packages/1c/bc/d79d000e5203ebef39af839f2ce77a777fcad6e08ad26af9a2fcd114ffc8/pydantic_core-2.50.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:17e722e156d0444ecaefbe640bdb60928752bf2013e2b7a11cdb099aaae19bec", size = 2158402, upload-time = "2026-10-11T18:32:52.714Z" }, + { url = "https://files.pythonhosted.org/packages/1b/a5/4902cb5fd599422c130bd3124ab31ee8b771199bd56662702eed07d3fec7/pydantic_core-2.50.1-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:aa8224f10880d9bf1b5993988ba153d42a8b4f3f4f511f93b1f09c93ff613c72", size = 2229755, upload-time = "2026-10-11T18:32:54.126Z" }, + { url = "https://files.pythonhosted.org/packages/1a/f5/c1481f8669f6060d89110c9b1374173fc8767ca276b84f4870bd8acd5e3f/pydantic_core-2.50.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:41bc8237121bd8dc8d888dfd6279fc166ffc88c1f1bf3a8bf00869680533ca4c", size = 2242303, upload-time = "2026-10-11T18:32:55.641Z" }, + { url = "https://files.pythonhosted.org/packages/04/f9/77fc3c7653ba9b6e42049e17e96b25388187d4a7c274ab1cb07f58ac8419/pydantic_core-2.50.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:45c6266d071c241f2a168d45bf8c54344f0effce35e7e6b73afdec11f3687568", size = 2278345, upload-time = "2026-10-11T18:32:57.38Z" }, + { url = "https://files.pythonhosted.org/packages/95/9b/0579c5d12e7f2b16b27e6782427987341fcce07b0725e02ddb0b74add0c4/pydantic_core-2.50.1-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:1deeacb112d14d3f4fcb16b165f7dbaf76c70ba6e82f37ba042bdab51970a0b8", size = 2449947, upload-time = "2026-10-11T18:32:58.896Z" }, + { url = "https://files.pythonhosted.org/packages/a8/ac/1b677db91eba54cc5922f4de6edc46ba45c2bd712dfdc0130382d158e214/pydantic_core-2.50.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:1c96fd793b73d1b92e65570132505498fe7b21eaef73cdf74e67e5dfba7ac9e4", size = 2476889, upload-time = "2026-10-11T18:33:00.665Z" }, + { url = "https://files.pythonhosted.org/packages/4d/2a/3a9f6624ee3ea9ccba5249dde11418bb4c35780a7a92608f8768bd3fea39/pydantic_core-2.50.1-cp313-cp313-win32.whl", hash = "sha256:06ead20d39ffd6f2f6f2a8f8a6de67ff8bb1b4f14a8a30e058502514ee2ac685", size = 2101118, upload-time = "2026-10-11T18:33:02.329Z" }, + { url = "https://files.pythonhosted.org/packages/2d/1f/323f78ddd9d9938aac420c1abb4e8ba799fc8bdab0acc67c0593b837c979/pydantic_core-2.50.1-cp313-cp313-win_amd64.whl", hash = "sha256:7816e98acc08119dc0f340ab167048ecc54126316330c1f0caf7c6756c88e28f", size = 2170840, upload-time = "2026-10-11T18:33:03.919Z" }, + { url = "https://files.pythonhosted.org/packages/cb/09/8497c52a739ae425c3ac2f7f56414cbc711c67d374346174f40fe2062644/pydantic_core-2.50.1-cp313-cp313-win_arm64.whl", hash = "sha256:c17799a62c142d61b8a3c51752a7cbc87fe2ad4ccfab10e628a77b405075c662", size = 2161633, upload-time = "2026-10-11T18:33:05.518Z" }, + { url = "https://files.pythonhosted.org/packages/49/33/28b96e81677153715e3eafb9f26663a80841e859fde282a380359b0d3fa1/pydantic_core-2.50.1-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:1cf41f1ae3fa155cf167a72689ad044bcc1e3c97e064123677149bdfb5dafc4a", size = 2224428, upload-time = "2026-10-11T18:33:07.153Z" }, + { url = "https://files.pythonhosted.org/packages/94/40/15c06410c9b7b8da5805d27b64e09bd3f900e986728106db2689dbc51513/pydantic_core-2.50.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4df197990c15b5a37c5a277d131d9f2c67de6133f2e5dafd80d9bba4b99f46f9", size = 2023740, upload-time = "2026-10-11T18:33:08.763Z" }, + { url = "https://files.pythonhosted.org/packages/a1/4e/5eb629f6efc2a27e421d789dd4bfacbb5f6d09c80c28b13d1e87d73163d5/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0036473f5583e6a60e50b8b21651511564277a3f05cc5dab8cf579f552cd5f6c", size = 2049435, upload-time = "2026-10-11T18:33:10.366Z" }, + { url = "https://files.pythonhosted.org/packages/ba/8e/f195aebec49ad12318876ac2368c197a7939f5d52f8e156d2238ef8a4588/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:992c3514ec891fa7858099183e4d64e6bd5a5d4ff452fae29df22faa77a006bb", size = 2172975, upload-time = "2026-10-11T18:33:12.368Z" }, + { url = "https://files.pythonhosted.org/packages/e0/f5/7ad9fb83010cd5ea0948409db105676ed779c4e709e2d3d89b9fe2558794/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:739dc730e6be3bd5ec2f4ab5cfc7eb047cc45fc1497b3bafec74ff2ed07df597", size = 2364197, upload-time = "2026-10-11T18:33:14.244Z" }, + { url = "https://files.pythonhosted.org/packages/ca/fb/bf0aab3e78301d202b82a0322ec968cd703b11fc0623fde9a28708936f62/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:32fad3a91e51b6d2039c572db04a5a873260b399f6bd62c3552671fa7a4a2899", size = 2444259, upload-time = "2026-10-11T18:33:15.79Z" }, + { url = "https://files.pythonhosted.org/packages/45/35/38f6d6564fae57d9b12e5676dfa947e0e7d5c46dbeaa7eb3352261d6d299/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:42b54c2c90ad348b5e3a85e03e715d572c1fde357ef104cdfe3b03b697a404ea", size = 2159857, upload-time = "2026-10-11T18:33:17.51Z" }, + { url = "https://files.pythonhosted.org/packages/65/a0/fb0a3ca10f139dcf765b2d312d10cf0f65c57c59993229d00ad12b14ceff/pydantic_core-2.50.1-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:2df1ff41884de2bc4b307bafd7c40a691094fad2ff8e767e5b45a319257bcf4e", size = 2233954, upload-time = "2026-10-11T18:33:19.591Z" }, + { url = "https://files.pythonhosted.org/packages/8c/6a/e63842252702aa4ec6e6b7178ca85e541a77459076592c23ebe2d9840b33/pydantic_core-2.50.1-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fe90228920fd8ff2be62622b6bb8a2b11acd65046d50c6b130614b5879605a20", size = 2243876, upload-time = "2026-10-11T18:33:21.393Z" }, + { url = "https://files.pythonhosted.org/packages/39/25/5991cf8318b37e0dfab47b87541619a1df8501a793cba0d978846cba37a7/pydantic_core-2.50.1-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:844b869f118e22a41a091bdcedda8a71bc1b0f62c38d1a0c3211cece47e1d8fc", size = 2277908, upload-time = "2026-10-11T18:33:23.324Z" }, + { url = "https://files.pythonhosted.org/packages/a7/3e/3ee8baaa6cc25a6961c69168cf9ff0f002d56f4e0d541ad6724c18fb61f3/pydantic_core-2.50.1-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:2eb75304506894a281d346220a4f7481a1b8729577c5ed2a05395991966a8396", size = 2449285, upload-time = "2026-10-11T18:33:24.942Z" }, + { url = "https://files.pythonhosted.org/packages/c0/c7/acbec6deac13fe697a80c275a9b6661db62c4d323bac8a47347b1f39c7cd/pydantic_core-2.50.1-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:6b20a4bffabdad0db2927ac034ae3b8a681b1f7a0182f3e60b479ad2fde21ebb", size = 2479013, upload-time = "2026-10-11T18:33:26.925Z" }, + { url = "https://files.pythonhosted.org/packages/10/08/21a3f237b264389f6219d053ee78fc1b5c2fd402c1b1cb2b6cb8b85f9834/pydantic_core-2.50.1-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:99ba9bc2b8062ea0c326a990f7f00e6530c23579de66dd246e72c4cafef950a5", size = 1484400, upload-time = "2026-10-11T18:33:28.693Z" }, + { url = "https://files.pythonhosted.org/packages/09/ce/077a6d262d12ef09108377ac0717f029f420d773ace63cafd6143af75568/pydantic_core-2.50.1-cp314-cp314-win32.whl", hash = "sha256:cf356f70551d40374eaffb1aa63f1eb6d2006681cbd7a9faea173ce0f4dd7cd2", size = 2103222, upload-time = "2026-10-11T18:33:30.326Z" }, + { url = "https://files.pythonhosted.org/packages/14/4c/350a2415209c43d670eb71d3c040f332c04a30583d39e311b05c7ac15762/pydantic_core-2.50.1-cp314-cp314-win_amd64.whl", hash = "sha256:d32f3acc081cc3923386d88f422cde8892335e95f034e0104bb4cf9310d9915f", size = 2171586, upload-time = "2026-10-11T18:33:32.139Z" }, + { url = "https://files.pythonhosted.org/packages/bf/92/9bea6ca96580a0902fed366f064f0829e404f41889b34543279a1162b888/pydantic_core-2.50.1-cp314-cp314-win_arm64.whl", hash = "sha256:bed5163e03b98bc1fa2eb05d74c63d9c5c95d8ed6254985481640fbf5e237dea", size = 2155773, upload-time = "2026-10-11T18:33:33.896Z" }, + { url = "https://files.pythonhosted.org/packages/86/8c/f121f073cf32bdd5cba7e6230b1ce0ac845b0cf43e44dd597d95af272db2/pydantic_core-2.50.1-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:9572c1369e9c9da2d64a7b7992c786d90ff295abc93964cfe3125e4290768070", size = 2234017, upload-time = "2026-10-11T18:33:35.588Z" }, + { url = "https://files.pythonhosted.org/packages/a4/69/2bb2bcd6146cffe98d760a46c42ae71efd5151d9b2f9c9bf6619a3b32083/pydantic_core-2.50.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2005207aafe1231315718bf6ed5d064a7300fb4772754af35ee72fc68159492e", size = 2001016, upload-time = "2026-10-11T18:33:37.57Z" }, + { url = "https://files.pythonhosted.org/packages/20/b3/fbf854c7d07ec114260c26e9e2071a4381740f9ae09641dbfcbdf2a18c45/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:64f6047f62a6c5ae08d0a6afb035667aa2d97c3d20d69762e034c5ea144d92a5", size = 2023015, upload-time = "2026-10-11T18:33:39.433Z" }, + { url = "https://files.pythonhosted.org/packages/65/20/6de55b2f92cdb614b745c6e9ced639fc4fd7e1e77604825a88bdece6fcbd/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1ef800dd7d85bcdadf4c3076e4c94e43939493558a3b69a1ea830c706d4617bb", size = 2192543, upload-time = "2026-10-11T18:33:41.381Z" }, + { url = "https://files.pythonhosted.org/packages/0c/d9/19e91c94bd5c405945ce2f15526808aea37e162c253160de8ed7bf70b406/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b0135bcdcaa0f23573f286e4cb5e0fd2962700964ed13df085b85f2b97aeab9e", size = 2388704, upload-time = "2026-10-11T18:33:43.167Z" }, + { url = "https://files.pythonhosted.org/packages/b6/bc/2e24c8415eae1a25ee5a5484946a917123bad2e9d01689a759236928175a/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0b3a6f334c6a2345ca15318ff894502a90012536404b37c844a976c76c846e0b", size = 2464118, upload-time = "2026-10-11T18:33:44.856Z" }, + { url = "https://files.pythonhosted.org/packages/bd/1f/945b8053cb061c64e102bcaf7bfb9ed740c0bd4349f9bb978a7d4ddff4ab/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:06e01fbbfdb9be777b316a71b6c49efaf4a08b615d0a98d678cda3023f79d019", size = 2150142, upload-time = "2026-10-11T18:33:46.661Z" }, + { url = "https://files.pythonhosted.org/packages/2b/78/96a3e50bf0d64aaae781107eb9335b7529d05a85793ed6e7241c4cd1d931/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:a29a061fec0b4e2d714f277e70a3a18125ecff803f2fea6eade2f2e53711d112", size = 2245577, upload-time = "2026-10-11T18:33:48.574Z" }, + { url = "https://files.pythonhosted.org/packages/7a/5d/a6038a0322232758a6ebfa709f4ca14a60bf5b93940cd0b345056a557da4/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f5187624823423e1d1b82b1072ac41dc837389e18d3d0572cc19bbee46cd550a", size = 2263728, upload-time = "2026-10-11T18:33:50.527Z" }, + { url = "https://files.pythonhosted.org/packages/0d/4b/76ded3333a457a9344c4f2d63f2d82d0101654b05178fa4ddfe7d3627674/pydantic_core-2.50.1-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:3e46a9eb0a0901dd6275e6b06ac3a464885ef350ec4121fe486869de8053e4bb", size = 2293630, upload-time = "2026-10-11T18:33:52.362Z" }, + { url = "https://files.pythonhosted.org/packages/a9/00/9eca378335c9c1b72bc779bf6e4c2a82ce784f14ec48a0e87102c9870e05/pydantic_core-2.50.1-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:756d669f04e62ec4148ecfe22be6a4484d9b1181a6ef32e205ebfd200540858b", size = 2468048, upload-time = "2026-10-11T18:33:54.118Z" }, + { url = "https://files.pythonhosted.org/packages/35/ca/e3832e9cf93651251de43c8c5c029ed680a3c1a0a1b17ee26c81bed08cbd/pydantic_core-2.50.1-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:c516cc5367ca3448995d42cb994bf3f4c9002d2a7c22eac9622551269ad1b807", size = 2512159, upload-time = "2026-10-11T18:33:55.99Z" }, + { url = "https://files.pythonhosted.org/packages/4b/f2/773469b5a10a39116a2cb17edaf6d722f017dd8da07161b371465311c542/pydantic_core-2.50.1-cp314-cp314t-win32.whl", hash = "sha256:9d1bed94af6a63835461f3cf7502058eb166c58c4778e11d0f433cfb1bd69e19", size = 2098406, upload-time = "2026-10-11T18:33:58.043Z" }, + { url = "https://files.pythonhosted.org/packages/ee/42/0bb74f8f25204b259b11ab7c12dc7f180893b6bd706118b385147fb6efd5/pydantic_core-2.50.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c8dce1f1e0e5358b682a6ad3fa5e31b31d4560997b8e61417e9217c8d60f8a0c", size = 2148481, upload-time = "2026-10-11T18:33:59.913Z" }, + { url = "https://files.pythonhosted.org/packages/47/0d/d801646c9679a4e630e15cf521d4108b93854b42a6e6e02391bd4c6b1095/pydantic_core-2.50.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ceff0acc940be2715bd6ad17b24c0e5304abf44f6efd0f81ee8499e640f9dc86", size = 2174972, upload-time = "2026-10-11T18:34:01.875Z" }, + { url = "https://files.pythonhosted.org/packages/b2/84/23984b763d8862a02a13d27a44b6e8169428fd85ecfed88f54c29108604d/pydantic_core-2.50.1-cp315-cp315-macosx_10_12_x86_64.whl", hash = "sha256:8a6791afa2245e6c6b180122d105941644f5bd410bb18623b408808cc41a3102", size = 2224599, upload-time = "2026-10-11T18:34:04.016Z" }, + { url = "https://files.pythonhosted.org/packages/4d/90/a63cf8586abc1d1a3f6d9b18f0224789ebf003851f092ebda3c7c863fd32/pydantic_core-2.50.1-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:84f34323a61a365b4e9295de6028474754829aaddd59c7bf1a040e7487ef8f3c", size = 2024191, upload-time = "2026-10-11T18:34:05.901Z" }, + { url = "https://files.pythonhosted.org/packages/a6/6a/f34bff9808ffb4907fbf5f5040457d253cacf2da7fce9efbc99ca1b1a44d/pydantic_core-2.50.1-cp315-cp315-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:23edad659e8dbd8ca7e4e877fe6c81573abbdf215bd25a68b53e1272f58b80c7", size = 2048252, upload-time = "2026-10-11T18:34:07.757Z" }, + { url = "https://files.pythonhosted.org/packages/b2/54/13f419bf1eb59852003818e25935aaf175687d978f4c4bca70e08fe40a3b/pydantic_core-2.50.1-cp315-cp315-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a5fce22f1e87d181e924e12da7d81cfe031fb3881a5ddf26ad28f141756ca43", size = 2173376, upload-time = "2026-10-11T18:34:09.592Z" }, + { url = "https://files.pythonhosted.org/packages/6d/6c/b5a34d24cd0c81669d8f8339d74e6815abcf2f8fb48ab4b49b84c09be1d5/pydantic_core-2.50.1-cp315-cp315-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c73622ef819328873b53109ee4f77ceb598bffedd02daf916102be3228866b78", size = 2364559, upload-time = "2026-10-11T18:34:11.534Z" }, + { url = "https://files.pythonhosted.org/packages/eb/8d/d64d6216a8df365082665927ff923f183056f9049fee08e9777c9ac05296/pydantic_core-2.50.1-cp315-cp315-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ce8c25ca38cc0e3d7753ba180808de2c0c8cb24eae0df64491e40921454e9831", size = 2444697, upload-time = "2026-10-11T18:34:13.535Z" }, + { url = "https://files.pythonhosted.org/packages/b8/0d/1b1149f60a00ea21ba5f70e28acbd40feb4598af414f80f53c921fac07c9/pydantic_core-2.50.1-cp315-cp315-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7689580e72a642ab5ec64d5f55b2e33636fa43b4ebe63c0c2c965ef307c7d1aa", size = 2159961, upload-time = "2026-10-11T18:34:15.56Z" }, + { url = "https://files.pythonhosted.org/packages/1e/35/f236549299dcc78e71e945495d7ec67e78d20844d0999c60601685003031/pydantic_core-2.50.1-cp315-cp315-manylinux_2_31_riscv64.whl", hash = "sha256:d5c0e32fdbce7f1e8ef4d11f655694bf5f4175c757a9f1dc2be09b8864e5bcf5", size = 2234581, upload-time = "2026-10-11T18:34:17.502Z" }, + { url = "https://files.pythonhosted.org/packages/6a/85/26901a490522b7f75ef9bb9a7afb73e5bb550f0e1cb5b99a0069183e2eb9/pydantic_core-2.50.1-cp315-cp315-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:40f523349960fa30f3ea51404308ff50f9997a90df639590f47a057c1f32b415", size = 2244090, upload-time = "2026-10-11T18:34:19.402Z" }, + { url = "https://files.pythonhosted.org/packages/ca/2c/481bcfc70ceeb77a778ca6e5b705fe592cb64735c108157f81b7dd9c280e/pydantic_core-2.50.1-cp315-cp315-musllinux_1_1_aarch64.whl", hash = "sha256:d4193206b6587047437f6f11d7e776df23e1c1e23af2a54d9347275614791e10", size = 2277948, upload-time = "2026-10-11T18:34:21.317Z" }, + { url = "https://files.pythonhosted.org/packages/24/eb/f1e09333faa7ba447cde967310f758430cc7c5e987bdab5228816c70030d/pydantic_core-2.50.1-cp315-cp315-musllinux_1_1_armv7l.whl", hash = "sha256:84bc765b282a9d5b7fe0348b8648904f25a6a04b2139da52b1dd30c8ac3a2c8f", size = 2449688, upload-time = "2026-10-11T18:34:23.321Z" }, + { url = "https://files.pythonhosted.org/packages/d1/b3/036bde636db8f76d92996e81aefc75678ab5cec4a07eea1ad0c72a893fc3/pydantic_core-2.50.1-cp315-cp315-musllinux_1_1_x86_64.whl", hash = "sha256:ed1e728b39a383c81035b2459cfcb35d99dfb01f7d6ebe3a913bc1cc5b81e459", size = 2479286, upload-time = "2026-10-11T18:34:25.214Z" }, + { url = "https://files.pythonhosted.org/packages/d0/a3/07f018294ee18d144afeb6df1a47d9e92960f014be45c5ed13519fa5af95/pydantic_core-2.50.1-cp315-cp315-win32.whl", hash = "sha256:bc94f474417604bd383d2cd445d071b07dd55fedceed3ce33407bf1fcc107290", size = 2103683, upload-time = "2026-10-11T18:34:27.42Z" }, + { url = "https://files.pythonhosted.org/packages/5f/98/f9bd7e1f9b6709f155acb9ef826d9c3884fe925f811fd8e55b9b52280bad/pydantic_core-2.50.1-cp315-cp315-win_amd64.whl", hash = "sha256:983a662de2571cb2502fc8ff47b6770b03d025d2eb314c92f77b3f07c74720ed", size = 2171681, upload-time = "2026-10-11T18:34:29.508Z" }, + { url = "https://files.pythonhosted.org/packages/10/87/4bb3e1e7f385c076ab5af4d6dd0571b22042cd8207eb810d5b9fef15ae31/pydantic_core-2.50.1-cp315-cp315-win_arm64.whl", hash = "sha256:94845ff54dc5193f228cab81b2662a04bfbb892e95bdc15edf7399000ce57d54", size = 2156251, upload-time = "2026-10-11T18:34:31.455Z" }, + { url = "https://files.pythonhosted.org/packages/47/47/83643225b08f2aef6c8cc4bbe6e3f79c5e139c4364a6e450d0f399d87774/pydantic_core-2.50.1-cp315-cp315t-macosx_10_12_x86_64.whl", hash = "sha256:4a53d13cdfbedbfa87f08b83c1a0a5efcc767d785a4b41934fa9cb672670493a", size = 2234363, upload-time = "2026-10-11T18:34:33.65Z" }, + { url = "https://files.pythonhosted.org/packages/5e/66/127ca649ba2f2462039dc1e694c6c01e00a3689f023a617364d3507e6d97/pydantic_core-2.50.1-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:efbecf43d321f7b9281441f1f213f7c21c66988b0e06c2730ba13ed47a46bb08", size = 2000972, upload-time = "2026-10-11T18:34:36.037Z" }, + { url = "https://files.pythonhosted.org/packages/5e/4f/e421e0a5d653b1203b090b2e48745976988d7338a647940704b5b9c2b399/pydantic_core-2.50.1-cp315-cp315t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bc1f08f68dac9f9e83845a8039880aba2ab553eb9b2259c3243a313182c253fe", size = 2022925, upload-time = "2026-10-11T18:34:37.972Z" }, + { url = "https://files.pythonhosted.org/packages/d5/4e/ea5568e2491e1a71100f15ae8c2d01ef52db42a184a2d4e716bc79e5eb8f/pydantic_core-2.50.1-cp315-cp315t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:5dfe41f232befddb9c4377f6cfc702b51595e2d78ed082672adf8758d2c4619f", size = 2192867, upload-time = "2026-10-11T18:34:39.921Z" }, + { url = "https://files.pythonhosted.org/packages/ae/5e/3b8c3a35acbe219909ada5defad5d7d9fed845fb2b37bd5ec518f453c119/pydantic_core-2.50.1-cp315-cp315t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:adc06d218a1cadfd2ec4628424d7d79ce4eba69c2965e7e7b55106f0da5208c8", size = 2388687, upload-time = "2026-10-11T18:34:42.184Z" }, + { url = "https://files.pythonhosted.org/packages/b6/aa/7889b4e515f91a2e8c0ae6b5081fec0feb30c4398d1434e14793c60f173a/pydantic_core-2.50.1-cp315-cp315t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2cf91809d0721ab81592ba67bea7694821679c10b1a2e3c3460082b286c1918a", size = 2464415, upload-time = "2026-10-11T18:34:44.384Z" }, + { url = "https://files.pythonhosted.org/packages/08/78/93449e628eb8a6fdcce3eff9043081179d1bc7ce6f1bff32dc5006f41e00/pydantic_core-2.50.1-cp315-cp315t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:23923ab9292c40da026330b1ecf4dc2618c8e86e0422e5d1fbf50d94d64ca4f8", size = 2150823, upload-time = "2026-10-11T18:34:46.392Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f1/72c5bc129fceb0d00f05dc1e67f518c1728de1928c55f81fc13d7690de39/pydantic_core-2.50.1-cp315-cp315t-manylinux_2_31_riscv64.whl", hash = "sha256:f3377c8c2b3ce898423c5e5dd94c7982e30aa7717a7e6ab2470b9de364963709", size = 2245806, upload-time = "2026-10-11T18:34:48.805Z" }, + { url = "https://files.pythonhosted.org/packages/28/2a/922a0e78f3aa6ab837b59f88190233fb8dad546c17995fde9bb3ed3b9b49/pydantic_core-2.50.1-cp315-cp315t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:455a773617b5913bf5c20d0692e5787b119e52c4d40ea644ca31f5758fd31be2", size = 2264160, upload-time = "2026-10-11T18:34:50.876Z" }, + { url = "https://files.pythonhosted.org/packages/52/a8/0f1449e3e1b20941c9372faa18e7b6a092e30cdfa841902473f7989532ac/pydantic_core-2.50.1-cp315-cp315t-musllinux_1_1_aarch64.whl", hash = "sha256:1a9006395dece0e32e704c315eff8a00bede494f6108546cfc5539c89fef4f9a", size = 2294031, upload-time = "2026-10-11T18:34:52.876Z" }, + { url = "https://files.pythonhosted.org/packages/d1/d0/1031f492857de70355fb16524bbb03efce5fd34c93ed4a1ec60be07e0d4b/pydantic_core-2.50.1-cp315-cp315t-musllinux_1_1_armv7l.whl", hash = "sha256:d2d82aa62521c55ddfb000ae70f88cdd8de974078f6024e821dfe5addd0c818f", size = 2468423, upload-time = "2026-10-11T18:34:54.995Z" }, + { url = "https://files.pythonhosted.org/packages/46/52/269ffffa645b8e47906395a39cf9db1151ae9fa4bcd7f47b960bc31baf37/pydantic_core-2.50.1-cp315-cp315t-musllinux_1_1_x86_64.whl", hash = "sha256:009634b83993777ddcd69cad0ffcace43dabde692109528e35f0fde91e386a8b", size = 2512598, upload-time = "2026-10-11T18:34:57.149Z" }, + { url = "https://files.pythonhosted.org/packages/31/5c/e47e28281f20326ff6f3c31d626f0a83e615d2d94ba41bb0ad6ec237184a/pydantic_core-2.50.1-cp315-cp315t-win32.whl", hash = "sha256:3fde4fdc6487a58d944ca87cf5adc95d5f266e872c19599f5f4c0a8a1b1f9f9f", size = 2098535, upload-time = "2026-10-11T18:34:59.313Z" }, + { url = "https://files.pythonhosted.org/packages/03/ad/759e181e69c1b472c60b2049e5f61d5da1deefdd5ce1df4bb2ebcf771f25/pydantic_core-2.50.1-cp315-cp315t-win_amd64.whl", hash = "sha256:1c8632d4ac04e6f91128fca584b3a8a507d81604c24eeaaad00d4be42765c32b", size = 2148645, upload-time = "2026-10-11T18:35:01.571Z" }, + { url = "https://files.pythonhosted.org/packages/6d/56/8a702c27e5be9f47e5f19d8669227424290e4c024e7c370279cbaf244b4e/pydantic_core-2.50.1-cp315-cp315t-win_arm64.whl", hash = "sha256:c3ede305158e75510be50869b319550ab072008c13d64d4ab1e094fb286b6f44", size = 2175297, upload-time = "2026-10-11T18:35:04.079Z" }, + { url = "https://files.pythonhosted.org/packages/a5/09/1bcf160f3cd6333e2a76982fae5963e5032de0bcaaba38464223dc266bfe/pydantic_core-2.50.1-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:76e2e83fa6ec8cdc972d438dafc2522b3a47bee4ec0ae668b29cfb1977ab5242", size = 2217070, upload-time = "2026-10-11T18:35:15.688Z" }, + { url = "https://files.pythonhosted.org/packages/3d/81/f02de95eec7794dfc7a5bc69a834f0e834b31b3009014ce6de903855f775/pydantic_core-2.50.1-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:a51eee75939cf811ac09b278745a6cee7dc873ccfbc8b9af3cc88fe4b7ce25b5", size = 2039027, upload-time = "2026-10-11T18:35:18.07Z" }, + { url = "https://files.pythonhosted.org/packages/4d/16/227746ed3771d9bf2304568b55134849ece34e13b38b1d76de763015c77f/pydantic_core-2.50.1-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ae28183297fb0d2b8dc46a1f01d51f5e45825fc5afe76a835a6cb7fb34821295", size = 2109009, upload-time = "2026-10-11T18:35:20.218Z" }, + { url = "https://files.pythonhosted.org/packages/e5/b9/7664d1592a0e5a74903f357cd5b28ceb7336f3473256c5a4e8c8432f1881/pydantic_core-2.50.1-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:88e492e8b9d0312e7dc13667c30222abf284dc3b79b5302b3607b41a5784ce61", size = 2269883, upload-time = "2026-10-11T18:35:22.48Z" }, +] + [[package]] name = "pygments" version = "2.21.0" @@ -2985,6 +3117,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, ] +[[package]] +name = "typing-inspection" +version = "0.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/26/b09b8010994eccc3c09092e6b34058f36a460eea2d4c3e8b910c695975a0/typing_inspection-0.4.4.tar.gz", hash = "sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47", size = 76928, upload-time = "2026-08-12T12:37:25.997Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/81/4add07e5172b7ac40d8ed5ff580409a7801a4fe26d529bdd915401dabfbe/typing_inspection-0.4.4-py3-none-any.whl", hash = "sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147", size = 14750, upload-time = "2026-08-12T12:37:24.648Z" }, +] + [[package]] name = "tzdata" version = "2026.5" From 5aab7ac8307044571e53e6fb16ef837e4ebecde7 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 13:43:35 +0100 Subject: [PATCH 06/12] Use a non-typo unknown key in config test to satisfy codespell Co-Authored-By: Claude Opus 5.5 --- tests/test_config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 64f2b5d..dd5fbda 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -31,8 +31,8 @@ def test_example_config_validates(): def test_unknown_top_level_key_is_rejected(): - with pytest.raises(ValidationError, match="alignmnet"): - load_config(_config(alignmnet={"reference": "fc"})) + with pytest.raises(ValidationError, match="not_a_section"): + load_config(_config(not_a_section={"reference": "fc"})) def test_missing_alignment_reference_is_rejected(): From 68cea403b2286cccade49f516bfd19158d1aaa5e Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 13:38:57 +0100 Subject: [PATCH 07/12] Annotate and document all functions; enforce with ruff and numpydoc; 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 --- .pre-commit-config.yaml | 5 + AGENTS.md | 94 +++++ CHANGELOG.md | 1 + examples/myloader.py | 17 +- pyproject.toml | 22 + src/mxalign/__init__.py | 35 +- src/mxalign/accessor/__init__.py | 2 + src/mxalign/accessor/mx.py | 336 ++++++++++++++- src/mxalign/alignment/__init__.py | 2 + .../alignment/interpolations/__init__.py | 2 + src/mxalign/alignment/interpolations/base.py | 40 +- .../alignment/interpolations/delaunay.py | 135 +++++- .../alignment/interpolations/interpolate.py | 29 +- .../alignment/interpolations/registry.py | 43 +- .../alignment/interpolations/xarray.py | 68 ++- src/mxalign/alignment/nans.py | 16 +- src/mxalign/alignment/space.py | 68 ++- src/mxalign/alignment/time.py | 197 ++++++++- src/mxalign/pipeline/__init__.py | 2 + src/mxalign/pipeline/cli.py | 45 +- src/mxalign/pipeline/config.py | 187 ++++++++- src/mxalign/pipeline/dates.py | 53 ++- src/mxalign/pipeline/runner.py | 93 ++++- src/mxalign/pipeline/save.py | 57 ++- .../pipeline/transformations/__init__.py | 5 +- src/mxalign/pipeline/transformations/base.py | 54 ++- .../pipeline/transformations/external.py | 49 ++- .../pipeline/transformations/registry.py | 46 ++- .../pipeline/transformations/transform.py | 35 +- src/mxalign/pipeline/verification.py | 114 ++++- src/mxalign/utils/projections.py | 16 +- src/mxalign/utils/traits.py | 56 ++- tests/conftest.py | 29 ++ tests/test_accessor.py | 51 ++- tests/test_align_space.py | 164 +++++++- tests/test_align_time.py | 388 ++++++++++++++++-- tests/test_align_time_global.py | 181 ++++++-- tests/test_config.py | 45 +- tests/test_dates.py | 38 +- tests/test_interpolations.py | 292 +++++++++++-- tests/test_runner.py | 57 ++- tests/test_transformations.py | 38 +- 42 files changed, 2929 insertions(+), 278 deletions(-) create mode 100644 AGENTS.md diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6a0afa9..e6ed3ba 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..a1e9827 --- /dev/null +++ b/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 ` (or `uv add --optional `), 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`. diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e3cdbe..895c682 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ 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..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 diff --git a/examples/myloader.py b/examples/myloader.py index f47c8b1..7989263 100644 --- a/examples/myloader.py +++ b/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") diff --git a/pyproject.toml b/pyproject.toml index 8a70481..13be0c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,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", diff --git a/src/mxalign/__init__.py b/src/mxalign/__init__.py index c4549b8..89ac61e 100644 --- a/src/mxalign/__init__.py +++ b/src/mxalign/__init__.py @@ -1,30 +1,29 @@ -from .pipeline.transformations.transform import transform -from .pipeline.transformations.registry import ( - available_transformations, - register_transformation, -) +"""Align and verify weather forecast and observation datasets with xarray.""" + +from . import accessor, alignment, pipeline from .alignment.interpolations.interpolate import interpolate from .alignment.interpolations.registry import ( available_interpolations, register_interpolator, ) -from .alignment.time import align_time from .alignment.space import align_space - -from . import accessor -from . import alignment -from . import pipeline +from .alignment.time import align_time +from .pipeline.transformations.registry import ( + available_transformations, + register_transformation, +) +from .pipeline.transformations.transform import transform __all__ = [ - "transform", - "available_transformations", - "register_transformation", - "interpolate", - "available_interpolations", - "register_interpolator", - "align_time", - "align_space", "accessor", + "align_space", + "align_time", "alignment", + "available_interpolations", + "available_transformations", + "interpolate", "pipeline", + "register_interpolator", + "register_transformation", + "transform", ] diff --git a/src/mxalign/accessor/__init__.py b/src/mxalign/accessor/__init__.py index 30d84d3..48c51a9 100644 --- a/src/mxalign/accessor/__init__.py +++ b/src/mxalign/accessor/__init__.py @@ -1,3 +1,5 @@ +"""Dataset accessors provided by mxalign (registers ``ds.mx`` on import).""" + from . import mx __all__ = ["mx"] diff --git a/src/mxalign/accessor/mx.py b/src/mxalign/accessor/mx.py index d6f39e4..aa40dd2 100644 --- a/src/mxalign/accessor/mx.py +++ b/src/mxalign/accessor/mx.py @@ -1,4 +1,7 @@ +"""The ``ds.mx`` xarray Dataset accessor for space/time predicates and alignment.""" + import numpy as np +import pandas as pd import xarray as xr import cartopy.crs as ccrs @@ -19,12 +22,34 @@ class MxAccessor: every use (xarray caches the accessor per object, so reading them once would go stale when attrs change). Raises ``KeyError`` when a trait is needed but missing. Never modifies ``ds``; operations return new datasets. + + Parameters + ---------- + ds : xr.Dataset + The dataset the accessor is attached to. """ def __init__(self, ds: xr.Dataset) -> None: self._ds = ds def _trait(self, attr: str) -> str: + """Return the value of the trait attribute ``attr`` from ``ds.attrs``. + + Parameters + ---------- + attr : str + Name of the trait attribute to read. + + Returns + ------- + str + The trait value stored in ``ds.attrs[attr]``. + + Raises + ------ + KeyError + If the dataset has no ``attr`` attribute. + """ try: return self._ds.attrs[attr] except KeyError: @@ -35,32 +60,129 @@ def _trait(self, attr: str) -> str: @property def _space(self) -> Space: + """The space trait of the dataset. + + Returns + ------- + Space + The space trait read from ``ds.attrs``. + + Raises + ------ + KeyError + If the dataset has no space trait attribute. + """ return Space(self._trait(SPACE_TRAIT_ATTR)) @property def _time(self) -> Time: + """The time trait of the dataset. + + Returns + ------- + Time + The time trait read from ``ds.attrs``. + + Raises + ------ + KeyError + If the dataset has no time trait attribute. + """ return Time(self._trait(TIME_TRAIT_ATTR)) # --- Space predicates --- def is_grid(self) -> bool: + """Return whether the dataset has the GRID space trait. + + Returns + ------- + bool + ``True`` if the space trait is ``Space.GRID``. + + Raises + ------ + KeyError + If the dataset has no space trait attribute. + """ return self._space == Space.GRID def is_point(self) -> bool: + """Return whether the dataset has the POINT space trait. + + Returns + ------- + bool + ``True`` if the space trait is ``Space.POINT``. + + Raises + ------ + KeyError + If the dataset has no space trait attribute. + """ return self._space == Space.POINT # --- Time predicates --- def is_forecast(self) -> bool: + """Return whether the dataset has the FORECAST time trait. + + Returns + ------- + bool + ``True`` if the time trait is ``Time.FORECAST``. + + Raises + ------ + KeyError + If the dataset has no time trait attribute. + """ return self._time == Time.FORECAST def is_observation(self) -> bool: + """Return whether the dataset has the OBSERVATION time trait. + + Returns + ------- + bool + ``True`` if the time trait is ``Time.OBSERVATION``. + + Raises + ------ + KeyError + If the dataset has no time trait attribute. + """ return self._time == Time.OBSERVATION # --- Space operations --- def add_crs(self, crs: str | dict | ccrs.Projection) -> xr.Dataset: - """Attach a Cartopy CRS to ``ds.attrs["crs"]``; accepts a name, config dict, or CRS object.""" + """Attach a Cartopy CRS to ``ds.attrs["crs"]``. + + Parameters + ---------- + crs : str or dict or ccrs.Projection + Name of a builtin projection (key into + :data:`mxalign.utils.projections.BUILTIN`, case-insensitive), a + config dict with ``projection``, ``kws_projection`` and optionally + ``kws_globe`` keys (see + :func:`mxalign.utils.projections.create_cartopy_crs`), or a Cartopy + CRS object. + + Returns + ------- + xr.Dataset + A copy of the dataset with ``attrs["crs"]`` set. + + Raises + ------ + ValueError + If the dataset is a point dataset, or ``crs`` is a name that is + not a builtin projection. + KeyError + If the dataset has no space trait attribute, or a config dict + lacks the ``projection`` or ``kws_projection`` key. + """ if self.is_point(): raise ValueError("Cannot add CRS to a point dataset") if isinstance(crs, str): @@ -77,7 +199,29 @@ def add_crs(self, crs: str | dict | ccrs.Projection) -> xr.Dataset: return self._ds.assign_attrs({"crs": crs}) def add_grid_mapping(self, grid_mapping: str | dict) -> xr.Dataset: - """Attach a grid-mapping dict to ``ds.attrs["grid_mapping"]``; accepts a builtin name or dict.""" + """Attach a grid-mapping dict to ``ds.attrs["grid_mapping"]``. + + Parameters + ---------- + grid_mapping : str or dict + Name of a builtin projection (key into + :data:`mxalign.utils.projections.BUILTIN`, case-insensitive) whose + ``kws_grid`` is used, or a grid-mapping dict (``nx``, ``ny``, + ``lon_ll``, ``lat_ll``, ``dx``, ``dy``, ...). + + Returns + ------- + xr.Dataset + A copy of the dataset with ``attrs["grid_mapping"]`` set. + + Raises + ------ + ValueError + If the dataset is a point dataset, or ``grid_mapping`` is a name + that is not a builtin mapping. + KeyError + If the dataset has no space trait attribute. + """ if self.is_point(): raise ValueError("Cannot add grid mapping to a point dataset") if isinstance(grid_mapping, str): @@ -90,7 +234,30 @@ def add_grid_mapping(self, grid_mapping: str | dict) -> xr.Dataset: return self._ds.assign_attrs({"grid_mapping": grid_mapping}) def add_xy(self, crs: str | dict | ccrs.Projection | None = None) -> xr.Dataset: - """Project ``longitude``/``latitude`` to ``xc``/``yc`` coordinates using the dataset CRS.""" + """Project ``longitude``/``latitude`` to ``xc``/``yc`` coordinates. + + Parameters + ---------- + crs : str or dict or ccrs.Projection, optional + CRS to attach first (see :meth:`add_crs`). If ``None``, the CRS + already stored in ``ds.attrs["crs"]`` is used. + + Returns + ------- + xr.Dataset + The dataset with ``xc``/``yc`` coordinates along ``grid_index`` + (grid) or ``point_index`` (point). Returned unchanged if it already + has ``xc`` and ``yc`` coordinates. + + Raises + ------ + ValueError + If no CRS is given or stored, if ``longitude``/``latitude`` are + dimensions of the dataset, or if the dataset is neither grid nor + point. + KeyError + If the dataset has no space trait attribute. + """ ds = self.add_crs(crs) if crs is not None else self._ds crs_obj = ds.attrs.get("crs", None) @@ -122,7 +289,20 @@ def add_xy(self, crs: str | dict | ccrs.Projection | None = None) -> xr.Dataset: raise ValueError("Dataset does not have expected spatial properties") def is_stacked(self) -> bool: - """Return ``True`` if the spatial dims are already collapsed into ``grid_index``.""" + """Return whether the spatial dims are already collapsed into ``grid_index``. + + Returns + ------- + bool + ``False`` if the dataset has ``xc``/``yc`` or + ``longitude``/``latitude`` dimensions, ``True`` if it has a + ``grid_index`` dimension. + + Raises + ------ + ValueError + If the dataset has none of these dimensions. + """ if {"xc", "yc"}.issubset(self._ds.dims) or {"longitude", "latitude"}.issubset( self._ds.dims ): @@ -133,7 +313,21 @@ def is_stacked(self) -> bool: raise ValueError("Dataset does not have expected dimensions for GRID") def stack(self) -> xr.Dataset: - """Collapse 2-D spatial dims into a flat ``grid_index`` dimension.""" + """Collapse 2-D spatial dims into a flat ``grid_index`` dimension. + + Returns + ------- + xr.Dataset + The dataset with ``yc``/``xc`` (or ``latitude``/``longitude``) + stacked into ``grid_index``; returned unchanged if already stacked. + + Raises + ------ + ValueError + If the dataset is a point dataset or has no stackable dimensions. + KeyError + If the dataset has no space trait attribute. + """ if self.is_point(): raise ValueError("POINT datasets cannot be stacked") if self.is_stacked(): @@ -148,9 +342,38 @@ def stack(self) -> xr.Dataset: return self._ds.stack({"grid_index": dims_to_stack}).reset_index("grid_index") def unstack( - self, crs: str | dict | ccrs.Projection | None = None, **kwargs + self, crs: str | dict | ccrs.Projection | None = None, **kwargs: float | None ) -> xr.Dataset: - """Restore a flat ``grid_index`` dim to 2-D ``xc``/``yc`` using grid-mapping metadata.""" + """Restore a flat ``grid_index`` dim to 2-D ``xc``/``yc`` dimensions. + + Grid parameters are taken from ``**kwargs`` where given, and otherwise + from ``ds.attrs["grid_mapping"]``. + + Parameters + ---------- + crs : str or dict or ccrs.Projection, optional + CRS to attach first (see :meth:`add_crs`). If not given, the CRS + already stored in ``ds.attrs["crs"]`` is used. + **kwargs : float or None + Grid parameters ``nx``, ``ny``, ``lon_ll``, ``lat_ll``, ``dx`` and + ``dy``, overriding the values in ``ds.attrs["grid_mapping"]``. + + Returns + ------- + xr.Dataset + The unstacked dataset, with ``attrs["grid_mapping"]`` set to the + grid parameters used; returned unchanged if not stacked. + + Raises + ------ + ValueError + If the dataset is a point dataset, a grid parameter is ``None``, or + ``nx * ny`` does not match the size of ``grid_index``. + KeyError + If a grid parameter is neither given nor found in + ``ds.attrs["grid_mapping"]``, or the dataset has no space trait + or ``crs`` attribute. + """ if self.is_point(): raise ValueError("POINT datasets cannot be unstacked") if not self.is_stacked(): @@ -183,7 +406,19 @@ def unstack( # --- Time operations --- def add_valid_time(self) -> xr.Dataset: - """Add a ``valid_time`` coordinate (``reference_time + lead_time``) to forecast datasets.""" + """Add a ``valid_time`` coordinate (``reference_time + lead_time``) to forecasts. + + Returns + ------- + xr.Dataset + The dataset with a 2-D ``valid_time`` coordinate if it is a + forecast; otherwise the dataset unchanged. + + Raises + ------ + KeyError + If the dataset has no time trait attribute. + """ if self.is_forecast(): return _time._add_valid_time(self._ds) return self._ds @@ -191,7 +426,9 @@ def add_valid_time(self) -> xr.Dataset: # --- Alignment --- def align_time_with( - self, ds2: xr.Dataset, lead_time: str | list | np.timedelta64 | None = None + self, + ds2: xr.Dataset, + lead_time: str | list[str | np.timedelta64] | np.timedelta64 | None = None, ) -> xr.Dataset: """Align this dataset's time axis to match ds2. @@ -204,10 +441,23 @@ def align_time_with( ---------- ds2 : xr.Dataset The reference dataset to align to. - lead_time : str or timedelta or list, optional + lead_time : str or np.timedelta64 or list, optional For Forecast→Observation: "shortest" (default) | "longest" | specific value or list. For Forecast→Forecast: "reference" (default) | "intersection" | "union". Ignored for observation→* cases. + + Returns + ------- + xr.Dataset + This dataset aligned to the time coordinates of ``ds2``. + + Raises + ------ + ValueError + If the time traits cannot be matched, or ``lead_time`` is invalid + for the forecast case. + KeyError + If either dataset has no time trait attribute. """ if self.is_forecast() and ds2.mx.is_observation(): return _time.align_forecast_to_observation( @@ -224,7 +474,7 @@ def align_time_with( else: raise ValueError("Cannot align datasets with unknown time properties") - def align_space_with(self, ds2: xr.Dataset, **kwargs) -> xr.Dataset: + def align_space_with(self, ds2: xr.Dataset, **kwargs: object) -> xr.Dataset: """Align this dataset's spatial grid to match ds2. Always uses "reference" semantics: self is interpolated or reindexed to @@ -234,11 +484,26 @@ def align_space_with(self, ds2: xr.Dataset, **kwargs) -> xr.Dataset: ---------- ds2 : xr.Dataset The reference dataset to align to. - method : str - Interpolation method for grid→point alignment. One of "xarray" or - "delaunay" (default "xarray"). Ignored for grid→grid. - **kwargs - Passed through to the interpolator. + **kwargs : object + Passed through to the pairwise alignment function. For grid→point + alignment this includes ``method``, the interpolator name (one of + "xarray" or "delaunay", default "xarray"), with the remaining + keywords passed to the interpolator. Ignored for grid→grid. + + Returns + ------- + xr.Dataset + This dataset aligned to the spatial coordinates of ``ds2``. + + Raises + ------ + NotImplementedError + For point→point and point→grid alignment, and for grid→grid + alignment of grids that differ (regridding). + ValueError + If the space traits are not compatible. + KeyError + If either dataset has no space trait attribute. """ if self.is_grid(): if ds2.mx.is_grid(): @@ -261,9 +526,42 @@ def _create_multiindex( lat_ll: float, dx: float, dy: float, - **kwargs, -): - """Build the ``(yc, xc)`` MultiIndex for ``ds.grid_index`` from grid-mapping parameters.""" + **kwargs: float, +) -> pd.MultiIndex: + """Build the ``(yc, xc)`` MultiIndex for ``ds.grid_index`` from grid parameters. + + Parameters + ---------- + ds : xr.Dataset + Stacked dataset with a ``grid_index`` dimension and a Cartopy CRS in + ``ds.attrs["crs"]``. + nx : int + Number of grid points in the x direction. + ny : int + Number of grid points in the y direction. + lon_ll : float + Longitude of the lower-left grid point. + lat_ll : float + Latitude of the lower-left grid point. + dx : float + Grid spacing in the x direction, in CRS units. + dy : float + Grid spacing in the y direction, in CRS units. + **kwargs : float + Ignored; allows passing a full grid-mapping dict. + + Returns + ------- + pd.MultiIndex + Product index of ``yc`` and ``xc`` projected coordinates. + + Raises + ------ + ValueError + If ``nx * ny`` does not match the size of ``grid_index``. + KeyError + If the dataset has no ``crs`` attribute. + """ from pandas import MultiIndex if ds.sizes["grid_index"] != nx * ny: diff --git a/src/mxalign/alignment/__init__.py b/src/mxalign/alignment/__init__.py index c341d2d..b161cd4 100644 --- a/src/mxalign/alignment/__init__.py +++ b/src/mxalign/alignment/__init__.py @@ -1,3 +1,5 @@ +"""Temporal and spatial alignment of datasets, including interpolation methods.""" + from . import time from . import space from . import interpolations diff --git a/src/mxalign/alignment/interpolations/__init__.py b/src/mxalign/alignment/interpolations/__init__.py index 19e7e44..fe902f3 100644 --- a/src/mxalign/alignment/interpolations/__init__.py +++ b/src/mxalign/alignment/interpolations/__init__.py @@ -1,3 +1,5 @@ +"""Interpolators from grid to point datasets, registered by name.""" + from . import base from . import xarray from . import delaunay diff --git a/src/mxalign/alignment/interpolations/base.py b/src/mxalign/alignment/interpolations/base.py index 402d9a1..b233b92 100644 --- a/src/mxalign/alignment/interpolations/base.py +++ b/src/mxalign/alignment/interpolations/base.py @@ -1,3 +1,5 @@ +"""Base class shared by all interpolators.""" + import xarray as xr from mlwp_data_specs.specs.traits.spatial_coordinate import Space from ...utils.traits import update_space_trait @@ -12,13 +14,20 @@ class BaseInterpolator: Subclasses must set ``name``, ``source_space``, and ``target_space``, and implement ``_interpolate``. Subclasses must also copy ``source_dataset.attrs`` to the output so time/uncertainty traits are preserved. + + Parameters + ---------- + target_dataset : xr.Dataset + Dataset whose spatial locations the source is interpolated to. + **options : object + Interpolator-specific options, stored as ``self.options``. """ name: str = "base" source_space: Space | None = None target_space: Space | None = None - def __init__(self, target_dataset: xr.Dataset, **options) -> None: + def __init__(self, target_dataset: xr.Dataset, **options: object) -> None: self.target_dataset = target_dataset self.options = options # TODO: Check the properties @@ -26,7 +35,21 @@ def __init__(self, target_dataset: xr.Dataset, **options) -> None: def interpolate( self, source_dataset: xr.Dataset | xr.DataArray ) -> xr.Dataset | xr.DataArray: - """Run interpolation and set the output space trait to ``target_space``.""" + """Run interpolation and set the output space trait to ``target_space``. + + For point targets, the grid-only attributes ``crs`` and + ``grid_mapping`` are dropped from the output. + + Parameters + ---------- + source_dataset : xr.Dataset or xr.DataArray + Data to interpolate to the target locations. + + Returns + ------- + xr.Dataset or xr.DataArray + The interpolated data with the space trait set to ``target_space``. + """ ds_out = self._interpolate(source_dataset) if self.target_space == Space.POINT: # Grid-only attrs don't apply to points (and a cartopy CRS can't be saved) @@ -38,5 +61,16 @@ def interpolate( def _interpolate( self, source_dataset: xr.Dataset | xr.DataArray ) -> xr.Dataset | xr.DataArray: - """Perform the interpolation; subclasses must override this.""" + """Perform the interpolation; subclasses must override this. + + Parameters + ---------- + source_dataset : xr.Dataset or xr.DataArray + Data to interpolate to the target locations. + + Returns + ------- + xr.Dataset or xr.DataArray + The interpolated data. + """ pass diff --git a/src/mxalign/alignment/interpolations/delaunay.py b/src/mxalign/alignment/interpolations/delaunay.py index 7522d68..efa48c2 100644 --- a/src/mxalign/alignment/interpolations/delaunay.py +++ b/src/mxalign/alignment/interpolations/delaunay.py @@ -1,3 +1,5 @@ +"""Grid-to-point linear interpolator based on a Delaunay triangulation.""" + import hashlib import logging from functools import partial @@ -25,11 +27,32 @@ @register_interpolator class DelaunayInterpolator(BaseInterpolator): + """Linear grid-to-point interpolation via Delaunay triangulation. + + The source grid points are triangulated in (latitude, longitude) and each + target point gets barycentric weights from its enclosing triangle. The + sparse weight matrix is cached per source/target coordinate pair. Target + points outside the convex hull of the source grid get NaN. + + Parameters + ---------- + target_dataset : xr.Dataset + Point dataset with ``longitude`` and ``latitude`` coordinates. + **options : object + Interpolator options; only ``method="linear"`` (the default) is + supported. + + Raises + ------ + ValueError + If ``method`` is anything other than ``"linear"``. + """ + name = "delaunay" source_space = Space.GRID target_space = Space.POINT - def __init__(self, target_dataset: xr.Dataset, **options) -> None: + def __init__(self, target_dataset: xr.Dataset, **options: object) -> None: super().__init__(target_dataset, **options) method = self.options.get("method", "linear") if method != "linear": @@ -38,6 +61,27 @@ def __init__(self, target_dataset: xr.Dataset, **options) -> None: ) def _interpolate(self, source_dataset: xr.Dataset) -> xr.Dataset: + """Interpolate a stacked grid dataset to the target points. + + Variables whose last dimension is not ``grid_index`` are skipped. + + Parameters + ---------- + source_dataset : xr.Dataset + Grid dataset with a ``grid_index`` dimension and ``latitude``/ + ``longitude`` coordinates. + + Returns + ------- + xr.Dataset + The interpolated variables along ``point_index``, with the target + ``latitude``/``longitude`` coordinates and the source attributes. + + Raises + ------ + NotImplementedError + If the source dataset has no ``grid_index`` dimension. + """ if "grid_index" not in source_dataset.dims: raise NotImplementedError( "Delaunay interpolation currently only supports stacked grids" @@ -83,12 +127,36 @@ def _interpolate(self, source_dataset: xr.Dataset) -> xr.Dataset: def _fingerprint(points: np.ndarray) -> str: - """Hash of a full coordinate array (~20 ms for a 1M-point grid).""" + """Hash a full coordinate array (~20 ms for a 1M-point grid). + + Parameters + ---------- + points : np.ndarray + Coordinate array to hash. + + Returns + ------- + str + Hex digest of the array contents. + """ return hashlib.blake2b(np.ascontiguousarray(points), digest_size=16).hexdigest() def _get_weights(source_points: np.ndarray, target_points: np.ndarray) -> csr_matrix: - """Return the weight matrix for these source/target points, building it on a cache miss.""" + """Return the cached weight matrix for these points, building it on a miss. + + Parameters + ---------- + source_points : np.ndarray + Source coordinates, shape ``(n_source, 2)`` as (latitude, longitude). + target_points : np.ndarray + Target coordinates, shape ``(n_target, 2)`` as (latitude, longitude). + + Returns + ------- + csr_matrix + Sparse ``(n_target, n_source)`` barycentric weight matrix. + """ key = (_fingerprint(source_points), _fingerprint(target_points)) if key not in _WEIGHTS_CACHE: if len(_WEIGHTS_CACHE) >= _WEIGHTS_CACHE_SIZE: @@ -104,8 +172,25 @@ def _build_weight_matrix( source_points: np.ndarray, target_points: np.ndarray, ) -> csr_matrix: - """Build a sparse ``(n_target, n_source)`` barycentric weight matrix; points outside the convex hull get NaN weights.""" - + """Build a sparse ``(n_target, n_source)`` barycentric weight matrix. + + Target points outside the convex hull of the source points get NaN + weights. + + Parameters + ---------- + triangulation : Delaunay + Delaunay triangulation of ``source_points``. + source_points : np.ndarray + Source coordinates, shape ``(n_source, ndim)``. + target_points : np.ndarray + Target coordinates, shape ``(n_target, ndim)``. + + Returns + ------- + csr_matrix + Sparse ``(n_target, n_source)`` weight matrix. + """ print("Calculating interpolation-weight matrix") n_target = len(target_points) @@ -153,7 +238,25 @@ def _build_weight_matrix( def interpolate_da( da: xr.DataArray, W: csr_matrix, target_points: np.ndarray ) -> xr.DataArray: - """Apply weight matrix ``W`` to ``da`` over its ``grid_index`` dim, returning a ``point_index`` DataArray.""" + """Apply weight matrix ``W`` to ``da`` over its ``grid_index`` dimension. + + The computation is lazy (``map_blocks``); dask-backed inputs whose + ``grid_index`` is split over several chunks are rechunked to one chunk. + + Parameters + ---------- + da : xr.DataArray + Source data whose last dimension is ``grid_index``. + W : csr_matrix + Sparse ``(n_target, n_source)`` weight matrix. + target_points : np.ndarray + Target coordinates, shape ``(n_target, 2)``. + + Returns + ------- + xr.DataArray + Interpolated data with ``grid_index`` replaced by ``point_index``. + """ n_target = len(target_points) leading_dims = da.dims[:-1] @@ -214,7 +317,25 @@ def interpolate_block( W: csr_matrix, target_points: np.ndarray, ) -> xr.DataArray: - """Single-block sparse matmul used by ``map_blocks`` inside ``interpolate_da``.""" + """Apply the sparse weight matrix to a single block. + + Used by ``map_blocks`` inside ``interpolate_da``. + + Parameters + ---------- + block : xr.DataArray + Block of source data whose last dimension is ``grid_index``. + W : csr_matrix + Sparse ``(n_target, n_source)`` weight matrix. + target_points : np.ndarray + Target coordinates, shape ``(n_target, 2)``. + + Returns + ------- + xr.DataArray + Interpolated block with the last dimension replaced by + ``point_index``. + """ data = block.values # shape = (.., npoints) original_shape = data.shape[:-1] data_flat = data.reshape( diff --git a/src/mxalign/alignment/interpolations/interpolate.py b/src/mxalign/alignment/interpolations/interpolate.py index ee8467e..ee01453 100644 --- a/src/mxalign/alignment/interpolations/interpolate.py +++ b/src/mxalign/alignment/interpolations/interpolate.py @@ -1,3 +1,5 @@ +"""Convenience function to interpolate datasets with a named interpolator.""" + import xarray as xr from .registry import get_interpolation @@ -7,9 +9,32 @@ def interpolate( source_datasets: xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset], target_dataset: xr.Dataset, method: str, - **kwargs, + **kwargs: object, ) -> xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset]: - """Interpolate one or more source datasets to ``target_dataset``; returns the same container type.""" + """Interpolate one or more source datasets to ``target_dataset``. + + Parameters + ---------- + source_datasets : xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + Dataset(s) to interpolate. Each is copied before interpolation. + target_dataset : xr.Dataset + Dataset whose spatial locations the sources are interpolated to. + method : str + Name of a registered interpolator, e.g. ``"xarray"`` or ``"delaunay"``. + **kwargs : object + Options passed to the interpolator constructor. + + Returns + ------- + xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + The interpolated dataset(s), in the same container type as + ``source_datasets``. + + Raises + ------ + ValueError + If ``method`` is not a registered interpolator. + """ interp_cls = get_interpolation(method) interpolator = interp_cls(target_dataset, **kwargs) diff --git a/src/mxalign/alignment/interpolations/registry.py b/src/mxalign/alignment/interpolations/registry.py index 9148993..213c321 100644 --- a/src/mxalign/alignment/interpolations/registry.py +++ b/src/mxalign/alignment/interpolations/registry.py @@ -1,20 +1,53 @@ -_INTERPOLATORS: dict[str, type] = {} +"""Registry mapping interpolator names to interpolator classes.""" +from .base import BaseInterpolator -def register_interpolator(cls: type) -> type: - """Register an interpolator class under ``cls.name``; returns the class unchanged.""" +_INTERPOLATORS: dict[str, type[BaseInterpolator]] = {} + + +def register_interpolator(cls: type[BaseInterpolator]) -> type[BaseInterpolator]: + """Register an interpolator class under ``cls.name``. + + Intended for use as a class decorator. + + Parameters + ---------- + cls : type of BaseInterpolator + Interpolator class to register; must define a ``name`` attribute. + + Returns + ------- + type of BaseInterpolator + The class, unchanged. + """ _INTERPOLATORS[cls.name] = cls return cls def available_interpolations() -> list[str]: - """Return the names of all registered interpolators.""" + """Return the names of all registered interpolators. + + Returns + ------- + list of str + Names of the registered interpolators. + """ return list(_INTERPOLATORS.keys()) -def get_interpolation(name: str) -> type: +def get_interpolation(name: str) -> type[BaseInterpolator]: """Return the interpolator class registered under ``name``. + Parameters + ---------- + name : str + Name of a registered interpolator, e.g. ``"xarray"`` or ``"delaunay"``. + + Returns + ------- + type of BaseInterpolator + The registered interpolator class. + Raises ------ ValueError diff --git a/src/mxalign/alignment/interpolations/xarray.py b/src/mxalign/alignment/interpolations/xarray.py index c89f755..76cd9b7 100644 --- a/src/mxalign/alignment/interpolations/xarray.py +++ b/src/mxalign/alignment/interpolations/xarray.py @@ -1,3 +1,5 @@ +"""Grid-to-point interpolator based on ``xarray.Dataset.interp``.""" + from .base import BaseInterpolator from .registry import register_interpolator from mlwp_data_specs.specs.traits.spatial_coordinate import Space @@ -7,11 +9,44 @@ @register_interpolator class XarrayInterpolator(BaseInterpolator): + """Interpolate a grid dataset to points with ``xarray.Dataset.interp``. + + Lat/lon-dimension grids are interpolated directly in longitude/latitude; + projected grids are interpolated in ``xc``/``yc`` after projecting the + target points with the source CRS (stacked grids are unstacked first). + + Takes the constructor arguments of :class:`BaseInterpolator`: a point + ``target_dataset`` with ``longitude`` and ``latitude`` coordinates, and + ``**options`` that are passed to ``xarray.Dataset.interp`` (e.g. + ``method``). + """ + name = "xarray" source_space = Space.GRID target_space = Space.POINT def _interpolate(self, source_dataset: xr.Dataset) -> xr.Dataset: + """Interpolate a grid dataset to the target points. + + Parameters + ---------- + source_dataset : xr.Dataset + Grid dataset with ``longitude``/``latitude`` dimensions, ``xc``/``yc`` + dimensions, or a stacked ``grid_index`` dimension. + + Returns + ------- + xr.Dataset + The source interpolated to the target points. + + Raises + ------ + ValueError + If a stacked source cannot be unstacked, or has unexpected + dimensions. + KeyError + If a projected source has no ``crs`` attribute. + """ if "latitude" in source_dataset.dims and "longitude" in source_dataset.dims: ds_out = self._interpolate_from_latlon(source_dataset) @@ -27,7 +62,25 @@ def _interpolate(self, source_dataset: xr.Dataset) -> xr.Dataset: return ds_out def _interpolate_from_xcyc(self, source_dataset: xr.Dataset) -> xr.Dataset: - """Interpolate a projected (xc/yc) source to target lon/lat points.""" + """Interpolate a projected (xc/yc) source to target lon/lat points. + + Parameters + ---------- + source_dataset : xr.Dataset + Grid dataset with ``xc``/``yc`` dimensions and a Cartopy CRS in + ``attrs["crs"]``. + + Returns + ------- + xr.Dataset + The source interpolated along a ``point_index`` dimension, with + the target ``longitude``/``latitude`` coordinates. + + Raises + ------ + KeyError + If the source dataset has no ``crs`` attribute. + """ import cartopy.crs as ccrs try: @@ -54,7 +107,18 @@ def _interpolate_from_xcyc(self, source_dataset: xr.Dataset) -> xr.Dataset: return ds_out def _interpolate_from_latlon(self, source_dataset: xr.Dataset) -> xr.Dataset: - """Interpolate a lat/lon-dim source directly to target lon/lat points.""" + """Interpolate a lat/lon-dim source directly to target lon/lat points. + + Parameters + ---------- + source_dataset : xr.Dataset + Grid dataset with ``longitude`` and ``latitude`` dimensions. + + Returns + ------- + xr.Dataset + The source interpolated to the target points. + """ longitude = self.target_dataset["longitude"] latitude = self.target_dataset["latitude"] ds_out = source_dataset.interp( diff --git a/src/mxalign/alignment/nans.py b/src/mxalign/alignment/nans.py index f193612..090d7c8 100644 --- a/src/mxalign/alignment/nans.py +++ b/src/mxalign/alignment/nans.py @@ -1,3 +1,5 @@ +"""Propagation of NaN masks across a collection of datasets.""" + import xarray as xr import itertools @@ -5,11 +7,21 @@ def broadcast_nans( datasets: xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset], ) -> xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset]: - """Propagate NaN masks across all datasets so a NaN at any coordinate is NaN in all. + """Propagate NaN masks so a NaN at a coordinate in any dataset is NaN in all. Operates pairwise over deep copies; only shared coordinates and variables are considered. - """ + Parameters + ---------- + datasets : xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + Datasets to propagate NaNs across. A single dataset is returned as is. + + Returns + ------- + xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + Deep copies of the input datasets with NaN masks propagated, in the + same container type as ``datasets``. + """ if isinstance(datasets, xr.Dataset): return datasets elif isinstance(datasets, dict): diff --git a/src/mxalign/alignment/space.py b/src/mxalign/alignment/space.py index ed23ab5..5fdf53c 100644 --- a/src/mxalign/alignment/space.py +++ b/src/mxalign/alignment/space.py @@ -1,3 +1,5 @@ +"""Spatial alignment of datasets to a reference dataset.""" + import numpy as np import xarray as xr @@ -9,7 +11,7 @@ def align_space( datasets: xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset], reference: str | xr.Dataset, - **kwargs, + **kwargs: object, ) -> xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset]: """Align all datasets spatially to a reference dataset. @@ -18,10 +20,19 @@ def align_space( Parameters ---------- - datasets: + datasets : xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset Single dataset, list, or dict of datasets to align. - reference: + reference : str or xr.Dataset Key into *datasets* dict, or a dataset to align to. + **kwargs : object + Forwarded to ``ds.mx.align_space_with``. + + Returns + ------- + xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + The aligned datasets: a dict with the same keys if ``datasets`` is a + dict, otherwise a single dataset if there is only one, else a list. + The reference dataset itself is returned unchanged. """ if isinstance(datasets, dict): keys = list(datasets.keys()) @@ -49,8 +60,29 @@ def align_space( # --- Pairwise alignment (used by ``ds.mx.align_space_with``) --- -def align_grid_grid(ds1: xr.Dataset, ds2: xr.Dataset, **kwargs) -> xr.Dataset: - """Return ``ds1`` unchanged if its grid matches ``ds2``'s within ``COORD_TOLERANCE``; raises otherwise.""" +def align_grid_grid(ds1: xr.Dataset, ds2: xr.Dataset, **kwargs: object) -> xr.Dataset: + """Return ``ds1`` unchanged if its grid matches ``ds2``'s within ``COORD_TOLERANCE``. + + Parameters + ---------- + ds1 : xr.Dataset + Grid dataset to align. + ds2 : xr.Dataset + Reference grid dataset. + **kwargs : object + Ignored. + + Returns + ------- + xr.Dataset + ``ds1``, unchanged. + + Raises + ------ + NotImplementedError + If the ``longitude``/``latitude`` coordinates of the two grids differ + by more than ``COORD_TOLERANCE`` (regridding is not implemented). + """ if np.array_equal( ds1["longitude"].values, ds2["longitude"].values ) and np.array_equal(ds1["latitude"].values, ds2["latitude"].values): @@ -69,9 +101,31 @@ def align_grid_grid(ds1: xr.Dataset, ds2: xr.Dataset, **kwargs) -> xr.Dataset: def align_grid_point( - ds1: xr.Dataset, ds2: xr.Dataset, method: str = "xarray", **kwargs + ds1: xr.Dataset, ds2: xr.Dataset, method: str = "xarray", **kwargs: object ) -> xr.Dataset: - """Interpolate grid ``ds1`` to the point locations of ``ds2`` using the named interpolator.""" + """Interpolate grid ``ds1`` to the point locations of ``ds2``. + + Parameters + ---------- + ds1 : xr.Dataset + Source grid dataset. + ds2 : xr.Dataset + Target point dataset. + method : str, default "xarray" + Name of a registered interpolator, e.g. ``"xarray"`` or ``"delaunay"``. + **kwargs : object + Options passed to the interpolator constructor. + + Returns + ------- + xr.Dataset + ``ds1`` interpolated to the points of ``ds2``. + + Raises + ------ + ValueError + If ``method`` is not a registered interpolator. + """ from .interpolations.registry import get_interpolation interp_cls = get_interpolation(method) diff --git a/src/mxalign/alignment/time.py b/src/mxalign/alignment/time.py index edc191a..16076f8 100644 --- a/src/mxalign/alignment/time.py +++ b/src/mxalign/alignment/time.py @@ -1,3 +1,5 @@ +"""Temporal alignment of forecast and observation datasets.""" + import logging import numpy as np @@ -11,10 +13,10 @@ def align_time( - datasets: list[xr.Dataset] | dict[str, xr.Dataset], + datasets: xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset], reference: str | xr.Dataset, lead_time: str | None = None, -): +) -> xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset]: """Align all datasets temporally to a reference dataset. Two regimes: @@ -45,13 +47,30 @@ def align_time( Parameters ---------- - datasets : list or dict of xr.Dataset + datasets : xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + Single dataset, list, or dict of datasets to align. reference : str or xr.Dataset Key into *datasets* dict, or an xr.Dataset to align to. lead_time : str, optional Only allowed when the reference is a forecast: how the lead times of the other forecasts are matched to the reference's — ``"reference"`` (default), ``"intersection"`` or ``"union"``. + + Returns + ------- + xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + The aligned datasets: a dict with the same keys if ``datasets`` is a + dict, otherwise a single dataset if there is only one, else a list. + + Raises + ------ + ValueError + If ``lead_time`` is given but the reference is not a forecast, if no + ``reference_time`` survives global alignment, or if a pairwise + alignment fails (e.g. an unknown ``lead_time`` option). + KeyError + If ``reference`` is not a key of ``datasets``, or a dataset has no + time trait attribute. """ if isinstance(datasets, dict): keys = list(datasets.keys()) @@ -88,6 +107,31 @@ def _align_time_global( forecasts: list[xr.Dataset], ref_ds: xr.Dataset, ) -> list[xr.Dataset]: + """Align all datasets onto a common ``(reference_time, lead_time)`` grid. + + See the "Global" regime in :func:`align_time` for the algorithm. + + Parameters + ---------- + ds_list : list of xr.Dataset + All datasets to align, forecasts and observations. + forecasts : list of xr.Dataset + The forecast datasets among ``ds_list``. + ref_ds : xr.Dataset + Reference observation dataset bounding the valid times. + + Returns + ------- + list of xr.Dataset + The aligned datasets, in the order of ``ds_list``; all carry the + forecast time trait. + + Raises + ------ + ValueError + If no ``reference_time`` keeps all lead times within the reference + observation's ``valid_time`` range. + """ # R* = intersection of reference_time, L* = union of lead_time. # Pull the aligned *coord* (= the merged index) rather than the array's # data — outer-aligning a coord-DataArray NaN-fills the data where the @@ -134,6 +178,25 @@ def _constrain_to_reference( A ref_time ``r`` is kept only when every ``l ∈ L*`` yields a valid_time in ``[ref_vt.min(), ref_vt.max()]``. ``L*`` is never pruned. + + Parameters + ---------- + r_star : xr.DataArray + Candidate reference times ``R*`` along ``reference_time``. + l_star : xr.DataArray + Lead times ``L*`` along ``lead_time``. + ref_vt : xr.DataArray + Valid times of the reference observation. + + Returns + ------- + xr.DataArray + The subset of ``r_star`` whose valid times all fall within range. + + Raises + ------ + ValueError + If no reference time satisfies the constraint. """ vt = r_star + l_star in_range = (vt >= ref_vt.min()) & (vt <= ref_vt.max()) @@ -150,7 +213,19 @@ def _constrain_to_reference( def _add_valid_time(ds_fcst: xr.Dataset) -> xr.Dataset: - """Add a 2-D ``valid_time`` coordinate (``reference_time + lead_time``) to a forecast dataset.""" + """Add a 2-D ``valid_time`` coordinate (``reference_time + lead_time``) to a forecast. + + Parameters + ---------- + ds_fcst : xr.Dataset + Forecast dataset with ``reference_time`` and ``lead_time`` coordinates. + + Returns + ------- + xr.Dataset + ``ds_fcst`` with a ``valid_time`` coordinate on + ``(reference_time, lead_time)``. + """ valid_time = ( ds_fcst["reference_time"].values[:, np.newaxis] + ds_fcst["lead_time"].values ) @@ -162,9 +237,37 @@ def _add_valid_time(ds_fcst: xr.Dataset) -> xr.Dataset: def align_forecast_to_observation( ds_fcst: xr.Dataset, ds_obs: xr.Dataset, - lead_time: str | list | np.timedelta64 = "shortest", + lead_time: str | list[str | np.timedelta64] | np.timedelta64 = "shortest", ) -> xr.Dataset: - """Stack ``ds_fcst`` to valid_time, deduplicate by ``lead_time`` strategy, and reindex to ``ds_obs.valid_time``.""" + """Turn a forecast into an observation-like dataset on ``ds_obs.valid_time``. + + ``ds_fcst`` is stacked to valid time, deduplicated according to the + ``lead_time`` strategy, and reindexed to ``ds_obs.valid_time``. + + Parameters + ---------- + ds_fcst : xr.Dataset + Forecast dataset with ``reference_time`` and ``lead_time``. + ds_obs : xr.Dataset + Observation dataset with a ``valid_time`` coordinate. + lead_time : str or np.timedelta64 or list, default "shortest" + ``"shortest"`` or ``"longest"`` keeps, per valid time, the entry with + the shortest/longest lead time. Otherwise one or more lead times + (timedelta-like or strings such as ``"6h"``) to keep. Ties for the + same valid time are resolved by keeping the first entry. + + Returns + ------- + xr.Dataset + Dataset with a ``valid_time`` dimension matching ``ds_obs`` (NaN where + the forecast has no data) and the observation time trait. + + Raises + ------ + ValueError + If ``lead_time`` is not a valid option, lacks a unit, or contains lead + times not present in the forecast. + """ ds_with_vt = _add_valid_time(ds_fcst) ds_stacked = ds_with_vt.stack(time=["reference_time", "lead_time"]).reset_index( "time" @@ -205,8 +308,27 @@ def align_forecast_to_observation( return ds_1d -def _parse_lead_times(lead_time) -> np.ndarray: - """Convert one or more lead times (timedelta-like or strings such as ``"6h"``) to ``timedelta64[ns]``.""" +def _parse_lead_times( + lead_time: str | list[str | np.timedelta64] | np.timedelta64, +) -> np.ndarray: + """Convert one or more lead times to a ``timedelta64[ns]`` array. + + Parameters + ---------- + lead_time : str or np.timedelta64 or list + One or more lead times, as timedelta-like values or strings such as + ``"6h"``. + + Returns + ------- + np.ndarray + 1-D array of lead times with dtype ``timedelta64[ns]``. + + Raises + ------ + ValueError + If a lead time is a bare number without unit, or cannot be parsed. + """ values = np.atleast_1d(np.asarray(lead_time, dtype=object)) # np.timedelta64 subclasses np.number, so exclude it explicitly if any( @@ -228,7 +350,25 @@ def _parse_lead_times(lead_time) -> np.ndarray: def align_observation_to_forecast( ds_obs: xr.Dataset, ds_fcst: xr.Dataset ) -> xr.Dataset: - """Broadcast 1-D observations onto ``(reference_time, lead_time)`` by selecting via forecast valid times.""" + """Broadcast 1-D observations onto the ``(reference_time, lead_time)`` forecast grid. + + Observations are selected at the forecast valid times + (``reference_time + lead_time``), NaN-filled where missing. + + Parameters + ---------- + ds_obs : xr.Dataset + Observation dataset with a ``valid_time`` dimension. + ds_fcst : xr.Dataset + Forecast dataset (or coordinate template) with ``reference_time`` and + ``lead_time``. + + Returns + ------- + xr.Dataset + Observations on ``(reference_time, lead_time)``, with the forecast + time trait. + """ ds_fcst_with_vt = _add_valid_time(ds_fcst) valid_time_2d = ds_fcst_with_vt["valid_time"] # shape (reference_time, lead_time) @@ -244,14 +384,49 @@ def align_observation_to_forecast( def align_observation_to_observation(ds1: xr.Dataset, ds2: xr.Dataset) -> xr.Dataset: - """Reindex ``ds1`` to ``ds2.valid_time`` with NaN-fill for missing times.""" + """Reindex ``ds1`` to ``ds2.valid_time`` with NaN-fill for missing times. + + Parameters + ---------- + ds1 : xr.Dataset + Observation dataset to align. + ds2 : xr.Dataset + Reference observation dataset. + + Returns + ------- + xr.Dataset + ``ds1`` reindexed to the valid times of ``ds2``. + """ return ds1.reindex(valid_time=ds2.valid_time) def align_forecast_to_forecast( ds1: xr.Dataset, ds2: xr.Dataset, lead_time: str = "reference" ) -> xr.Dataset: - """Reindex ``ds1`` to ``ds2.reference_time`` and align lead times per ``lead_time`` strategy.""" + """Reindex ``ds1`` to ``ds2.reference_time`` and align lead times by a strategy. + + Parameters + ---------- + ds1 : xr.Dataset + Forecast dataset to align. + ds2 : xr.Dataset + Reference forecast dataset. + lead_time : str, default "reference" + ``"reference"`` reindexes to the lead times of ``ds2``, + ``"intersection"`` keeps only lead times shared by both, and + ``"union"`` reindexes to the lead times of either. + + Returns + ------- + xr.Dataset + ``ds1`` aligned to ``ds2`` with a refreshed ``valid_time`` coordinate. + + Raises + ------ + ValueError + If ``lead_time`` is not one of the supported options. + """ ds_out = ds1.reindex(reference_time=ds2.reference_time) if lead_time == "reference": diff --git a/src/mxalign/pipeline/__init__.py b/src/mxalign/pipeline/__init__.py index 50dece2..492ff6a 100644 --- a/src/mxalign/pipeline/__init__.py +++ b/src/mxalign/pipeline/__init__.py @@ -1,3 +1,5 @@ +"""Config-driven pipeline: loading, transforming, aligning and verifying datasets.""" + from . import transformations __all__ = [ diff --git a/src/mxalign/pipeline/cli.py b/src/mxalign/pipeline/cli.py index 73dafe3..6d2dfd4 100644 --- a/src/mxalign/pipeline/cli.py +++ b/src/mxalign/pipeline/cli.py @@ -1,6 +1,8 @@ +"""Command-line interface for running the pipeline on a local or SLURM dask cluster.""" + import argparse -import sys import logging +import sys # Define log format LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s" @@ -8,10 +10,24 @@ LOG = logging.getLogger(__name__) -def run_local(args): +def run_local(args: argparse.Namespace) -> None: + """Run the pipeline on a local dask cluster. + + Parameters + ---------- + args : argparse.Namespace + Parsed arguments with ``CONFIG``, ``n_workers`` and + ``threads_per_worker``. + + Raises + ------ + SystemExit + With exit code 1 if the pipeline fails; the cluster is closed first. + """ # Only import the necessary modules if function is called # to avoid unnecessary slow imports at the top level from dask.distributed import Client, LocalCluster + from .runner import Runner cluster = LocalCluster( @@ -31,11 +47,25 @@ def run_local(args): sys.exit(1) -def run_slurm(args): +def run_slurm(args: argparse.Namespace) -> None: + """Run the pipeline on a dask cluster of SLURM jobs. + + Parameters + ---------- + args : argparse.Namespace + Parsed arguments with ``CONFIG``, ``queue``, ``account``, ``cores``, + ``memory`` and ``interface``. + + Raises + ------ + SystemExit + With exit code 1 if the pipeline fails; the cluster is closed first. + """ # Only import the necessary modules if function is called # to avoid unnecessary slow imports at the top level from dask.distributed import Client from dask_jobqueue import SLURMCluster + from .runner import Runner cluster = SLURMCluster( @@ -69,8 +99,15 @@ def run_slurm(args): sys.exit(1) -def main(): +def main() -> None: + """Parse the command line and run the selected subcommand. + Raises + ------ + SystemExit + If the arguments are invalid, the command is unknown, or the + pipeline fails. + """ parser = argparse.ArgumentParser(description="mxalign CLI") subparsers = parser.add_subparsers( dest="command", required=True, help="Available commands" diff --git a/src/mxalign/pipeline/config.py b/src/mxalign/pipeline/config.py index 39b4a6d..720471c 100644 --- a/src/mxalign/pipeline/config.py +++ b/src/mxalign/pipeline/config.py @@ -15,6 +15,8 @@ class _Strict(BaseModel): + """Base model that rejects unknown keys.""" + model_config = ConfigDict(extra="forbid") @@ -25,14 +27,35 @@ class _PassThrough(BaseModel): @property def kwargs(self) -> dict: + """Return the extra (unknown) keys of the model as a keyword-argument dict. + + Returns + ------- + dict + A copy of ``model_extra``; empty if there are no extra keys. + """ return dict(self.model_extra or {}) class DatesConfig(_Strict): - """Date range used to expand ``{reference_time}``/``{lead_time}``/``{valid_time}`` in ``files``. - - Every field is optional so a dataset can override part of the top-level - ``dates`` section; after merging, all five must be set. + """Date range used to expand date placeholders in ``files`` patterns. + + The placeholders are ``{reference_time}``, ``{lead_time}`` and + ``{valid_time}``. Every field is optional so a dataset can override part + of the top-level ``dates`` section; after merging, all five must be set. + + Attributes + ---------- + start : datetime or str or None + First reference time of the range. + end : datetime or str or None + Last reference time of the range. + period : str or None + Step between successive reference times (e.g. ``"1D"``). + range : str or None + Forecast horizon; lead times run from 0 to ``range`` inclusive. + step : str or None + Increment between lead times (e.g. ``"1h"``). """ start: datetime | str | None = None @@ -43,7 +66,21 @@ class DatesConfig(_Strict): class DatasetConfig(_PassThrough): - """One entry of ``datasets``; extra keys go to ``load_and_validate_dataset``.""" + """One entry of ``datasets``; extra keys go to ``load_and_validate_dataset``. + + Attributes + ---------- + loader : str + Name of the ``mlwp_data_loaders`` loader used to open the files. + files : str or list of str + File path(s) or path pattern(s); patterns are expanded over ``dates``. + variables : str or list of str or None + Variable(s) to keep after loading; all variables if None. + grid_mapping : str or dict or None + Grid mapping (CRS) to attach to the dataset; none if None. + dates : DatesConfig or None + Dataset-specific date range, merged over the top-level ``dates``. + """ loader: str files: str | list[str] @@ -53,34 +90,98 @@ class DatasetConfig(_PassThrough): class TransformationConfig(_PassThrough): - """One entry of ``transformations`` (keyed by transformation name); extra keys go to the transformation.""" + """One entry of ``transformations``; extra keys go to the transformation. + + Entries are keyed by transformation name. + + Attributes + ---------- + datasets : list of str or None + Names of the datasets to transform; all datasets if None. + """ datasets: list[str] | None = None class TimeAlignmentConfig(_Strict): + """Options for the time alignment step. + + Attributes + ---------- + lead_time : str or list of str or None + How lead times of the other forecasts are matched to the reference + forecast (``"reference"``, ``"intersection"`` or ``"union"``); passed + on to :func:`mxalign.align_time`. + """ + lead_time: str | list[str] | None = None class SpaceAlignmentConfig(_Strict): - """Interpolator/regridder kwargs, chosen by the kind of space alignment needed.""" + """Interpolator/regridder kwargs, chosen by the kind of space alignment needed. + + Attributes + ---------- + interpolation : dict or None + Keyword arguments used when a grid is interpolated to points. + regrid : dict or None + Keyword arguments used when a grid is regridded to another grid. + """ interpolation: dict | None = None regrid: dict | None = None class SaveConfig(_PassThrough): - """Output for ``ds.(path, **kwargs)``; extra keys go to the save method.""" + """Output for ``ds.(path, **kwargs)``; extra keys go to the save method. + + Attributes + ---------- + method : str + Name of the dataset method used to save (e.g. ``"to_netcdf"``). + path : str + Output path, possibly containing placeholders. + """ method: str path: str class AlignmentSaveConfig(SaveConfig): + """Output settings for the aligned datasets. + + Attributes + ---------- + method : str + Name of the dataset method used to save (e.g. ``"to_netcdf"``). + path : str + Output path pattern; may contain ``{name}``, ``{year}``, ``{month}`` + and ``{day}`` placeholders. + datasets : {"all", "merge"} + Save each dataset separately (``"all"``) or concatenated along a new + ``model`` dimension (``"merge"``). + """ + datasets: Literal["all", "merge"] = "all" class AlignmentConfig(_Strict): + """The ``alignment`` section of the pipeline config. + + Attributes + ---------- + reference : str + Name of the dataset that the others are aligned to. + broadcast_nans : bool + Whether to broadcast NaNs across all datasets after alignment. + time : TimeAlignmentConfig or None + Options for time alignment (always performed). + space : SpaceAlignmentConfig or None + Options for space alignment; skipped if None. + save : AlignmentSaveConfig or None + Where to save the aligned datasets; not saved if None. + """ + reference: str broadcast_nans: bool = True time: TimeAlignmentConfig | None = None @@ -89,19 +190,55 @@ class AlignmentConfig(_Strict): class MetricConfig(_PassThrough): - """One entry of ``verification.metrics``; extra keys go to the metric function.""" + """One entry of ``verification.metrics``; extra keys go to the metric function. + + Attributes + ---------- + function : str + Dotted import path of the metric function. + inputs : dict of {str : {"forecast", "reference"}} + Mapping of function argument names to the dataset role passed to them. + """ function: str inputs: dict[str, Literal["forecast", "reference"]] class VerificationConfig(_Strict): + """The ``verification`` section of the pipeline config. + + Attributes + ---------- + reference : str + Name of the dataset used as reference for the metrics. + metrics : dict of {str : MetricConfig} or None + Metrics to compute, keyed by metric name. + save : SaveConfig or None + Where to save the computed metrics; not saved if None. + """ + reference: str metrics: dict[str, MetricConfig] | None = None save: SaveConfig | None = None class PipelineConfig(_Strict): + """Top-level pipeline config. + + Attributes + ---------- + dates : DatesConfig or None + Default date range used to expand dataset ``files`` patterns. + datasets : dict of {str : DatasetConfig} + Datasets to load, keyed by name. + transformations : dict of {str : TransformationConfig} or None + Transformations to apply, keyed by transformation name. + alignment : AlignmentConfig + Alignment settings. + verification : VerificationConfig or None + Verification settings; verification is skipped if None. + """ + dates: DatesConfig | None = None datasets: dict[str, DatasetConfig] transformations: dict[str, TransformationConfig] | None = None @@ -110,7 +247,19 @@ class PipelineConfig(_Strict): @model_validator(mode="after") def _expand_files(self) -> "PipelineConfig": - """Expand each dataset's ``files`` pattern(s) over its (merged) ``dates``.""" + """Expand each dataset's ``files`` pattern(s) over its (merged) ``dates``. + + Returns + ------- + PipelineConfig + The config itself, with ``files`` of each dataset that has dates + replaced by the expanded list of paths. + + Raises + ------ + ValueError + If the merged ``dates`` of a dataset are non-empty but incomplete. + """ top = self.dates.model_dump(exclude_none=True) if self.dates else {} for name, ds in self.datasets.items(): own = ds.dates.model_dump(exclude_none=True) if ds.dates else {} @@ -129,7 +278,23 @@ def _expand_files(self) -> "PipelineConfig": def load_config(config: str | dict) -> PipelineConfig: - """Validate a pipeline config given as a YAML file path or a dict.""" + """Validate a pipeline config given as a YAML file path or a dict. + + Parameters + ---------- + config : str or dict + Path to a YAML config file, or the already-parsed config. + + Returns + ------- + PipelineConfig + The validated config. + + Raises + ------ + pydantic.ValidationError + If the config does not match the schema. + """ if isinstance(config, str): with open(config, "r") as f: config = yaml.safe_load(f) diff --git a/src/mxalign/pipeline/dates.py b/src/mxalign/pipeline/dates.py index 56d6f1a..499bb29 100644 --- a/src/mxalign/pipeline/dates.py +++ b/src/mxalign/pipeline/dates.py @@ -1,3 +1,5 @@ +"""Expansion of date-templated file path patterns over reference and lead times.""" + import itertools import re import string @@ -14,13 +16,15 @@ class Dates: Parameters ---------- - start, end: - First and last reference time of the range. - period: + start : str or datetime or np.datetime64 + First reference time of the range. + end : str or datetime or np.datetime64 + Last reference time of the range. + period : str or np.timedelta64 Step between successive reference times (e.g. ``"1D"``). - range: + range : str or np.timedelta64 Forecast horizon — lead times run from 0 to ``range`` inclusive. - step: + step : str or np.timedelta64 Increment between lead times (e.g. ``"1h"``). """ @@ -31,18 +35,33 @@ def __init__( period: str | np.timedelta64, range: str | np.timedelta64, step: str | np.timedelta64, - ): + ) -> None: self.reference_times = pd.date_range(start, end, freq=pd.Timedelta(period)) self.lead_times = pd.timedelta_range( pd.Timedelta(0), pd.Timedelta(range), freq=pd.Timedelta(step) ) def substitute(self, path: str) -> list[str]: - """Expand ``path`` over all (reference_time, lead_time) pairs; return sorted unique paths. + """Expand ``path`` over all (reference_time, lead_time) pairs. + + Parameters + ---------- + path : str + Path pattern in :meth:`str.format` syntax with the fields + ``reference_time`` and ``valid_time`` (datetimes, e.g. + ``{reference_time:%Y%m%d%H}``) and ``lead_time`` (integer hours, + e.g. ``{lead_time:03d}``). - ``path`` uses :meth:`str.format` syntax with the fields ``reference_time`` and - ``valid_time`` (datetimes, e.g. ``{reference_time:%Y%m%d%H}``) and ``lead_time`` - (integer hours, e.g. ``{lead_time:03d}``). + Returns + ------- + list of str + Sorted unique expanded paths. + + Raises + ------ + ValueError + If ``path`` uses ``{lead_time}`` and a lead time is not a whole + number of hours. """ fields = _field_names(path) if "lead_time" in fields and (self.lead_times % HOUR != pd.Timedelta(0)).any(): @@ -65,7 +84,19 @@ def substitute(self, path: str) -> list[str]: def _field_names(path: str) -> set[str]: - """Return the top-level replacement field names used in a :meth:`str.format` pattern.""" + """Return the top-level replacement field names used in a format pattern. + + Parameters + ---------- + path : str + Pattern in :meth:`str.format` syntax. + + Returns + ------- + set of str + Field names, with attribute access and indexing stripped (e.g. + ``{a.b}`` and ``{a[0]}`` both give ``"a"``). + """ return { re.split(r"[.\[]", name)[0] for _, name, _, _ in string.Formatter().parse(path) diff --git a/src/mxalign/pipeline/runner.py b/src/mxalign/pipeline/runner.py index 080201d..c7bae49 100644 --- a/src/mxalign/pipeline/runner.py +++ b/src/mxalign/pipeline/runner.py @@ -1,27 +1,39 @@ +"""Runner that executes the config-driven load, transform, align and verify pipeline.""" + import os -import xarray as xr +import xarray as xr from mlwp_data_loaders import load_and_validate_dataset +from ..alignment.nans import broadcast_nans +from ..alignment.time import align_time from .config import ( PipelineConfig, SpaceAlignmentConfig, TimeAlignmentConfig, load_config, ) -from .transformations.transform import transform -from ..alignment.time import align_time -from ..alignment.nans import broadcast_nans from .save import save_dataset, save_metrics +from .transformations.transform import transform from .verification import Metric class Runner: - def __init__(self, config: str | dict): + """Run the verification pipeline described by a pipeline config. + + Parameters + ---------- + config : str or dict + Path to a YAML config file, or the already-parsed config; validated + with :func:`load_config`. + """ + + def __init__(self, config: str | dict) -> None: self.config: PipelineConfig = load_config(config) self.datasets = {} - def run(self): + def run(self) -> None: + """Load, transform, align and verify the datasets.""" # 1. Load the datasets self.load_datasets() @@ -30,7 +42,12 @@ def run(self): self.align() self.verify() - def load_datasets(self): + def load_datasets(self) -> None: + """Load every configured dataset into ``self.datasets``. + + Missing files are skipped with a message. Optionally selects the + configured variables and attaches the configured grid mapping. + """ for name, config_ds in self.config.datasets.items(): variables = config_ds.variables grid_mapping = config_ds.grid_mapping @@ -60,7 +77,12 @@ def load_datasets(self): self.datasets[name] = ds - def transform_datasets(self): + def transform_datasets(self) -> None: + """Apply the configured transformations to the loaded datasets. + + Each transformation is applied to the datasets it lists, or to all + datasets if it lists none. + """ if self.config.transformations is None: return for transformation, config_trans in self.config.transformations.items(): @@ -74,7 +96,12 @@ def transform_datasets(self): name=transformation, datasets=ds, **config_trans.kwargs ) - def align(self): + def align(self) -> None: + """Align the datasets in time and space to the reference dataset. + + Afterwards NaNs are optionally broadcast across datasets and the + aligned datasets are optionally saved, as configured. + """ config = self.config.alignment reference = config.reference @@ -105,7 +132,13 @@ def align(self): ) save_dataset(save.method, "merged", ds, **kwargs) - def verify(self): + def verify(self) -> None: + """Compute the configured metrics against the reference dataset. + + The metrics are computed for every non-reference dataset on the + variables common to all datasets, stored in ``self.metrics`` with + ``model`` and ``metric`` dimensions, and optionally saved. + """ config = self.config.verification if config is None: return @@ -145,13 +178,32 @@ def verify(self): def align_time( self, config: TimeAlignmentConfig | None, reference_name: str | None = None - ): + ) -> None: + """Align ``self.datasets`` in time to the reference dataset. + + Parameters + ---------- + config : TimeAlignmentConfig or None + Time alignment options; defaults are used if None. + reference_name : str or None + Name of the reference dataset. + """ lead_time = config.lead_time if config is not None else None self.datasets = align_time( self.datasets, reference=reference_name, lead_time=lead_time ) - def align_space(self, reference: str, config: SpaceAlignmentConfig): + def align_space(self, reference: str, config: SpaceAlignmentConfig) -> None: + """Align every non-reference dataset in space to the reference dataset. + + Parameters + ---------- + reference : str + Name of the reference dataset. + config : SpaceAlignmentConfig + Interpolator/regridder options, picked by the kind of alignment + each dataset needs (see :func:`get_spatial_alignment`). + """ ds_ref = self.datasets[reference] for name, ds in self.datasets.items(): if name != reference: @@ -160,7 +212,22 @@ def align_space(self, reference: str, config: SpaceAlignmentConfig): self.datasets[name] = ds.mx.align_space_with(ds_ref, **options) -def get_spatial_alignment(ds, reference): +def get_spatial_alignment(ds: xr.Dataset, reference: xr.Dataset) -> str: + """Return the kind of spatial alignment needed to align ``ds`` to ``reference``. + + Parameters + ---------- + ds : xr.Dataset + Dataset to align. + reference : xr.Dataset + Reference dataset. + + Returns + ------- + str + ``"interpolation"`` for grid to points, ``"regrid"`` for grid to + grid, and ``"null"`` otherwise. + """ if reference.mx.is_point() and ds.mx.is_grid(): return "interpolation" if reference.mx.is_grid() and ds.mx.is_grid(): diff --git a/src/mxalign/pipeline/save.py b/src/mxalign/pipeline/save.py index 95fa5bc..7cbdfc1 100644 --- a/src/mxalign/pipeline/save.py +++ b/src/mxalign/pipeline/save.py @@ -1,12 +1,22 @@ +"""Helpers for saving aligned datasets and computed metrics to disk.""" + import xarray as xr class DatasetPath: - """Resolve a path pattern to a concrete file path using date components from a dataset. + """Resolve a path pattern to a concrete file path using a dataset's dates. The dominant year/month/day (by count) in ``reference_time`` (forecast) or ``valid_time`` (observation) is used to fill ``{year}``, ``{month}``, ``{day}`` placeholders via :meth:`substitute`. + + Parameters + ---------- + name : str + Dataset name, used to fill the ``{name}`` placeholder. + ds : xr.Dataset + Forecast or observation dataset whose dates determine the + ``{year}``, ``{month}`` and ``{day}`` placeholders. """ def __init__(self, name: str, ds: xr.Dataset) -> None: @@ -43,14 +53,38 @@ def __init__(self, name: str, ds: xr.Dataset) -> None: self.day = int(days.isel(day=days.argmax())["day"].values) def substitute(self, path: str) -> str: - """Expand ``path`` (:meth:`str.format` syntax) with ``{name}``, ``{year}``, ``{month}``, ``{day}``.""" + """Fill the ``{name}``, ``{year}``, ``{month}`` and ``{day}`` fields of ``path``. + + Parameters + ---------- + path : str + Path pattern in :meth:`str.format` syntax. + + Returns + ------- + str + The path with the placeholders filled in. + """ return path.format( name=self.name, year=self.year, month=self.month, day=self.day ) -def save_dataset(method: str, name: str, ds: xr.Dataset, **kwargs) -> None: - """Save ``ds`` to a path derived from its dominant date using ``ds.``.""" +def save_dataset(method: str, name: str, ds: xr.Dataset, **kwargs: object) -> None: + """Save ``ds`` to a path derived from its dominant date using ``ds.``. + + Parameters + ---------- + method : str + Name of the dataset method used to save (e.g. ``"to_netcdf"``). + name : str + Dataset name, used to fill the ``{name}`` placeholder of the path. + ds : xr.Dataset + Dataset to save. + **kwargs + Must contain ``path``, a pattern filled via :class:`DatasetPath`; the + remaining keyword arguments are passed on to ``ds.``. + """ save_fn = getattr(ds, method) dataset = DatasetPath(name, ds) path = dataset.substitute(kwargs.pop("path")) @@ -58,8 +92,19 @@ def save_dataset(method: str, name: str, ds: xr.Dataset, **kwargs) -> None: save_fn(path, **kwargs) -def save_metrics(method: str, ds: xr.Dataset, **kwargs) -> None: - """Save ``ds`` to a fixed path using ``ds.``.""" +def save_metrics(method: str, ds: xr.Dataset, **kwargs: object) -> None: + """Save ``ds`` to a fixed path using ``ds.``. + + Parameters + ---------- + method : str + Name of the dataset method used to save (e.g. ``"to_netcdf"``). + ds : xr.Dataset + Dataset (of metrics) to save. + **kwargs + Must contain ``path``, the output path; the remaining keyword + arguments are passed on to ``ds.``. + """ save_fn = getattr(ds, method) path = kwargs.pop("path") print(f"Saving to {path}") diff --git a/src/mxalign/pipeline/transformations/__init__.py b/src/mxalign/pipeline/transformations/__init__.py index c337718..ac68a83 100644 --- a/src/mxalign/pipeline/transformations/__init__.py +++ b/src/mxalign/pipeline/transformations/__init__.py @@ -1,5 +1,6 @@ -from . import base -from . import external +"""Dataset transformations and the registry they are registered in.""" + +from . import base, external __all__ = [ "base", diff --git a/src/mxalign/pipeline/transformations/base.py b/src/mxalign/pipeline/transformations/base.py index 83d740c..dd6bc69 100644 --- a/src/mxalign/pipeline/transformations/base.py +++ b/src/mxalign/pipeline/transformations/base.py @@ -1,3 +1,5 @@ +"""Built-in transformations: renaming, unit conversion and wind speed.""" + import xarray as xr from .registry import register_transformation @@ -5,7 +7,21 @@ @register_transformation("rename") def transform_rename(ds: xr.Dataset, rename_dict: dict[str, list[str]]) -> xr.Dataset: - """Rename variables in ``ds``; ``rename_dict`` maps new names to lists of old names.""" + """Rename variables in ``ds`` according to a mapping of new to old names. + + Parameters + ---------- + ds : xr.Dataset + Dataset whose variables are renamed. + rename_dict : dict of str to list of str + Mapping of new names to lists of old names; any variable in ``ds`` + whose name is in a list is renamed to the corresponding key. + + Returns + ------- + xr.Dataset + Dataset with the variables renamed. + """ new_dict = {} for new_name, old_names in rename_dict.items(): for name in ds.keys(): @@ -20,7 +36,22 @@ def transform_kelvin_to_celcius( variables: str | list[str], inverse: bool = False, ) -> xr.Dataset: - """Convert ``variables`` between Kelvin and Celsius; set ``inverse=True`` for °C → K.""" + """Convert ``variables`` between Kelvin and Celsius. + + Parameters + ---------- + ds : xr.Dataset + Dataset holding the variables; modified in place. + variables : str or list of str + Name(s) of the variables to convert. + inverse : bool, default False + Convert from Celsius to Kelvin instead of Kelvin to Celsius. + + Returns + ------- + xr.Dataset + ``ds`` with the variables converted. + """ T_C2K = 273.15 if isinstance(variables, str): variables = [variables] @@ -32,7 +63,24 @@ def transform_kelvin_to_celcius( @register_transformation("uv_to_speed") def transform(ds: xr.Dataset, u: str, v: str, speed: str) -> xr.Dataset: - """Compute wind speed from ``u`` and ``v`` components and store it as ``speed``.""" + """Compute wind speed from ``u`` and ``v`` components and store it as ``speed``. + + Parameters + ---------- + ds : xr.Dataset + Dataset holding the wind components; modified in place. + u : str + Name of the eastward wind component variable. + v : str + Name of the northward wind component variable. + speed : str + Name of the wind speed variable to write. + + Returns + ------- + xr.Dataset + ``ds`` with the ``speed`` variable added. + """ import numpy as np result = np.sqrt(ds[u] ** 2 + ds[v] ** 2) diff --git a/src/mxalign/pipeline/transformations/external.py b/src/mxalign/pipeline/transformations/external.py index 00ed7c3..9af1e0d 100644 --- a/src/mxalign/pipeline/transformations/external.py +++ b/src/mxalign/pipeline/transformations/external.py @@ -1,3 +1,5 @@ +"""Transformation that calls an external function given by its import path.""" + from collections.abc import Callable import xarray as xr @@ -11,18 +13,35 @@ def transform( func_path: str, inputs: dict[str, str], output: str, - **kwargs, + **kwargs: object, ) -> xr.Dataset: """Call an external function and store its result as ``output`` in ``ds``. Parameters ---------- - func_path: + ds : xr.Dataset + Dataset providing the input variables; modified in place. + func_path : str Dotted path to the function, e.g. ``"mypackage.module.func"``. - inputs: + inputs : dict of str to str Mapping of function argument names to variable names in ``ds``. - output: - Name of the new variable written back into ``ds``. + output : str + Name of the new variable written back into ``ds``; the function's + result must have the shape of ``ds.dims``. + **kwargs + Additional keyword arguments passed on to the function. + + Returns + ------- + xr.Dataset + ``ds`` with the ``output`` variable added. + + Raises + ------ + ImportError + If the module in ``func_path`` cannot be imported. + AttributeError + If the module has no attribute with the function name. """ func = _resolve_function(func_path) input_kwargs = {arg_name: ds[var_name] for arg_name, var_name in inputs.items()} @@ -32,7 +51,25 @@ def transform( def _resolve_function(func_path: str) -> Callable: - """Import and return the callable at ``func_path`` (e.g. ``"pkg.module.func"``).""" + """Import and return the callable at ``func_path`` (e.g. ``"pkg.module.func"``). + + Parameters + ---------- + func_path : str + Dotted path of the form ``"."``. + + Returns + ------- + Callable + The resolved function. + + Raises + ------ + ImportError + If the module cannot be imported. + AttributeError + If the module has no attribute with the function name. + """ import importlib module_path, func_name = func_path.rsplit(".", 1) diff --git a/src/mxalign/pipeline/transformations/registry.py b/src/mxalign/pipeline/transformations/registry.py index d86cc85..7df58c6 100644 --- a/src/mxalign/pipeline/transformations/registry.py +++ b/src/mxalign/pipeline/transformations/registry.py @@ -1,12 +1,38 @@ +"""Registry mapping transformation names to transformation functions.""" + from collections.abc import Callable _TRANSFORMATION_REGISTRY: dict[str, Callable] = {} def register_transformation(name: str) -> Callable: - """Decorator that registers a transformation function under ``name``.""" + """Return a decorator that registers a transformation function under ``name``. + + Parameters + ---------- + name : str + Name under which the decorated function is registered; an existing + registration with the same name is replaced. + + Returns + ------- + Callable + Decorator that registers the function and returns it unchanged. + """ def decorator(func: Callable) -> Callable: + """Register ``func`` under ``name``. + + Parameters + ---------- + func : Callable + Transformation function taking a dataset as first argument. + + Returns + ------- + Callable + ``func`` itself, unchanged. + """ _TRANSFORMATION_REGISTRY[name] = func return func @@ -14,13 +40,29 @@ def decorator(func: Callable) -> Callable: def available_transformations() -> list[str]: - """Return the names of all registered transformations.""" + """Return the names of all registered transformations. + + Returns + ------- + list of str + Registered transformation names, in registration order. + """ return list(_TRANSFORMATION_REGISTRY.keys()) def get_transformation(name: str) -> Callable: """Return the transformation function registered under ``name``. + Parameters + ---------- + name : str + Name of the registered transformation. + + Returns + ------- + Callable + The registered transformation function. + Raises ------ ValueError diff --git a/src/mxalign/pipeline/transformations/transform.py b/src/mxalign/pipeline/transformations/transform.py index 5198145..78743c0 100644 --- a/src/mxalign/pipeline/transformations/transform.py +++ b/src/mxalign/pipeline/transformations/transform.py @@ -1,3 +1,5 @@ +"""Entry point for applying a registered transformation to datasets.""" + import xarray as xr from .registry import get_transformation @@ -6,10 +8,37 @@ def transform( name: str, datasets: xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset], - *args, - **kwargs, + *args: object, + **kwargs: object, ) -> xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset]: - """Apply the named transformation to one or more datasets; returns the same container type.""" + """Apply the named transformation to one or more datasets. + + Each dataset is copied before being passed to the transformation, so the + inputs are not modified. + + Parameters + ---------- + name : str + Name of a registered transformation (see + :func:`available_transformations`). + datasets : xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + Dataset(s) to transform. + *args + Positional arguments passed on to the transformation. + **kwargs + Keyword arguments passed on to the transformation. + + Returns + ------- + xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset + The transformed dataset(s), in the same container type as + ``datasets``. + + Raises + ------ + ValueError + If ``name`` is not a registered transformation. + """ func = get_transformation(name) if isinstance(datasets, dict): return {key: func(ds.copy(), *args, **kwargs) for key, ds in datasets.items()} diff --git a/src/mxalign/pipeline/verification.py b/src/mxalign/pipeline/verification.py index 88ad172..0fa0234 100644 --- a/src/mxalign/pipeline/verification.py +++ b/src/mxalign/pipeline/verification.py @@ -1,9 +1,52 @@ -from .transformations.external import _resolve_function +"""Metric computation for the verification step of the pipeline.""" + from functools import partial +import xarray as xr + +from .transformations.external import _resolve_function + class Metric: - def __init__(self, name, func_path, ds_ref, inputs, **kwargs): + """A verification metric with the reference dataset bound in. + + Parameters + ---------- + name : str + Name of the metric. + func_path : str + Dotted import path of the metric function. + ds_ref : xr.Dataset + Reference dataset passed to the arguments marked ``"reference"``. + inputs : dict of str to str + Mapping of function argument names to the dataset role passed to + them: ``"reference"`` for ``ds_ref``; any other value marks the + argument that receives the dataset given to :meth:`compute`. + **kwargs + Additional keyword arguments passed on to the metric function. + ``dim`` is also used to rechunk the datasets for ``xskillscore`` + functions. + + Raises + ------ + ValueError + If more than one argument in ``inputs`` is not ``"reference"``. + ImportError + If the module in ``func_path`` cannot be imported. + AttributeError + If the module has no attribute with the function name. + IndexError + If no argument in ``inputs`` is a non-reference (forecast) argument. + """ + + def __init__( + self, + name: str, + func_path: str, + ds_ref: xr.Dataset, + inputs: dict[str, str], + **kwargs: object, + ) -> None: self.name = name func = _resolve_function(func_path) self._is_xskillscore = func.__module__.startswith("xskillscore") @@ -26,13 +69,37 @@ def __init__(self, name, func_path, ds_ref, inputs, **kwargs): self._func = partial(func, **partial_kwargs) self._kwarg_ds = kwarg_ds[0] - def compute(self, ds): + def compute(self, ds: xr.Dataset) -> xr.Dataset: + """Compute the metric of ``ds`` against the reference dataset. + + Parameters + ---------- + ds : xr.Dataset + Dataset (e.g. a forecast) to verify. + + Returns + ------- + xr.Dataset + The result of the metric function. + """ if self._is_xskillscore: ds = self._rechunk(ds) kwarg_ds = {self._kwarg_ds: ds} return self._func(**kwarg_ds) - def _rechunk(self, ds): + def _rechunk(self, ds: xr.Dataset) -> xr.Dataset: + """Rechunk ``ds`` to one chunk along ``dim`` and size-1 chunks elsewhere. + + Parameters + ---------- + ds : xr.Dataset + Dataset to rechunk. + + Returns + ------- + xr.Dataset + The rechunked dataset, or ``ds`` unchanged if no ``dim`` was given. + """ if self._dim is None: return ds dim = [self._dim] if isinstance(self._dim, str) else self._dim @@ -43,7 +110,44 @@ def _rechunk(self, ds): return ds.chunk(chunks) -def verify(fcst, obs, func_path, inputs, **kwargs): +def verify( + fcst: xr.Dataset, + obs: xr.Dataset, + func_path: str, + inputs: dict[str, str], + **kwargs: object, +) -> xr.Dataset: + """Compute a metric function of a forecast against an observation. + + Parameters + ---------- + fcst : xr.Dataset + Forecast dataset. + obs : xr.Dataset + Observation dataset. + func_path : str + Dotted import path of the metric function. + inputs : dict of str to str + Mapping of function argument names to ``"forecast"`` or + ``"observation"``, selecting which dataset each argument receives. + **kwargs + Additional keyword arguments passed on to the metric function. + + Returns + ------- + xr.Dataset + The result of the metric function. + + Raises + ------ + KeyError + If a value in ``inputs`` is neither ``"forecast"`` nor + ``"observation"``. + ImportError + If the module in ``func_path`` cannot be imported. + AttributeError + If the module has no attribute with the function name. + """ func = _resolve_function(func_path=func_path) datasets = { "forecast": fcst, diff --git a/src/mxalign/utils/projections.py b/src/mxalign/utils/projections.py index 62f9d3c..6aa589d 100644 --- a/src/mxalign/utils/projections.py +++ b/src/mxalign/utils/projections.py @@ -1,3 +1,5 @@ +"""Cartopy projection helpers and builtin grid definitions (CERRA, UWCW).""" + import cartopy.crs as ccrs @@ -10,19 +12,23 @@ def create_cartopy_crs( Parameters ---------- - projection: + projection : str Key into :data:`PROJECTIONS` (e.g. ``"lcc"``, ``"latlon"``). - kws_projection: + kws_projection : dict Keyword arguments forwarded to the projection constructor. - kws_globe: - Optional keyword arguments forwarded to :class:`cartopy.crs.Globe`. + kws_globe : dict, optional + Keyword arguments forwarded to :class:`cartopy.crs.Globe`. + + Returns + ------- + ccrs.Projection + The constructed Cartopy CRS. Raises ------ ValueError If ``projection`` is not a key in :data:`PROJECTIONS`. """ - # - Get the cartopy projection class try: proj_cls = PROJECTIONS[projection] diff --git a/src/mxalign/utils/traits.py b/src/mxalign/utils/traits.py index acac041..f7b727a 100644 --- a/src/mxalign/utils/traits.py +++ b/src/mxalign/utils/traits.py @@ -1,3 +1,5 @@ +"""Helpers to validate and set the mlwp space, time and uncertainty traits on datasets.""" + import xarray as xr from mlwp_data_specs.api import ( @@ -23,21 +25,69 @@ def update_space_trait(ds: xr.Dataset, new_trait: Space) -> xr.Dataset: - """Validate and set the space trait attribute on ``ds`` in-place.""" + """Validate and set the space trait attribute on ``ds`` in-place. + + The validation report is not inspected, so a failed validation does not + raise. + + Parameters + ---------- + ds : xr.Dataset + Dataset to update. + new_trait : Space + Space trait to validate against and set. + + Returns + ------- + xr.Dataset + The same dataset object, with the space trait attribute set. + """ validate_space_dataset(ds, trait=new_trait) ds.attrs[SPACE_TRAIT_ATTR] = new_trait.value return ds def update_time_trait(ds: xr.Dataset, new_trait: Time) -> xr.Dataset: - """Validate and set the time trait attribute on ``ds`` in-place.""" + """Validate and set the time trait attribute on ``ds`` in-place. + + The validation report is not inspected, so a failed validation does not + raise. + + Parameters + ---------- + ds : xr.Dataset + Dataset to update. + new_trait : Time + Time trait to validate against and set. + + Returns + ------- + xr.Dataset + The same dataset object, with the time trait attribute set. + """ validate_time_dataset(ds, trait=new_trait) ds.attrs[TIME_TRAIT_ATTR] = new_trait.value return ds def update_uncertainty_trait(ds: xr.Dataset, new_trait: Uncertainty) -> xr.Dataset: - """Validate and set the uncertainty trait attribute on ``ds`` in-place.""" + """Validate and set the uncertainty trait attribute on ``ds`` in-place. + + The validation report is not inspected, so a failed validation does not + raise. + + Parameters + ---------- + ds : xr.Dataset + Dataset to update. + new_trait : Uncertainty + Uncertainty trait to validate against and set. + + Returns + ------- + xr.Dataset + The same dataset object, with the uncertainty trait attribute set. + """ validate_uncertainty_dataset(ds, trait=new_trait) ds.attrs[UNCERTAINTY_TRAIT_ATTR] = new_trait.value return ds diff --git a/tests/conftest.py b/tests/conftest.py index 1409430..d62ac11 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,5 @@ +"""Shared fixtures and time coordinates for the ``mxalign`` test suite.""" + import numpy as np import pytest import xarray as xr @@ -26,6 +28,18 @@ def _props(time: str) -> dict: + """Build the ``mlwp_*`` trait attributes for a deterministic point dataset. + + Parameters + ---------- + time : str + Value of the ``mlwp_time_trait`` attribute, e.g. ``"forecast"``. + + Returns + ------- + dict + Trait attributes to attach to a dataset's ``attrs``. + """ return { "mlwp_time_trait": time, "mlwp_space_trait": "point", @@ -35,6 +49,14 @@ def _props(time: str) -> dict: @pytest.fixture def ds_fcst() -> xr.Dataset: + """Provide a point forecast with 3 reference times and 4 lead times. + + Returns + ------- + xr.Dataset + Forecast with ``temp[i, j] = i * 10 + j`` over ``reference_time`` and + ``lead_time``. + """ return xr.Dataset( {"temp": (["reference_time", "lead_time"], FORECAST_VALUES)}, coords={"reference_time": REFERENCE_TIMES, "lead_time": LEAD_TIMES}, @@ -44,6 +66,13 @@ def ds_fcst() -> xr.Dataset: @pytest.fixture def ds_obs() -> xr.Dataset: + """Provide a point observation with 6 valid times from T0-6h to T0+24h. + + Returns + ------- + xr.Dataset + Observation with ``temp`` values 0, 10, ..., 50 along ``valid_time``. + """ return xr.Dataset( {"temp": ("valid_time", OBS_VALUES)}, coords={"valid_time": OBS_TIMES}, diff --git a/tests/test_accessor.py b/tests/test_accessor.py index 75dc9e4..777c8da 100644 --- a/tests/test_accessor.py +++ b/tests/test_accessor.py @@ -10,7 +10,13 @@ @pytest.fixture def ds_grid() -> xr.Dataset: - """2×2 grid stacked to grid_index, lon/lat in 0..1.""" + """Provide a 2×2 grid stacked to ``grid_index``, with lon/lat in 0..1. + + Returns + ------- + xr.Dataset + Stacked grid forecast with ``temp`` values 0..3. + """ return xr.Dataset( {"temp": ("grid_index", np.arange(4.0))}, coords={ @@ -22,13 +28,23 @@ def ds_grid() -> xr.Dataset: class TestTraits: - def test_predicates_follow_attr_changes(self, ds_grid): + """Group tests of reading traits through the accessor.""" + + def test_predicates_follow_attr_changes(self, ds_grid: xr.Dataset) -> None: + """Check that trait predicates reflect changes to the dataset attrs. + + Parameters + ---------- + ds_grid : xr.Dataset + Stacked grid forecast fixture. + """ assert ds_grid.mx.is_forecast() ds_grid.attrs["mlwp_time_trait"] = "observation" assert ds_grid.mx.is_observation() assert not ds_grid.mx.is_forecast() - def test_missing_trait_raises_on_use_not_on_access(self): + def test_missing_trait_raises_on_use_not_on_access(self) -> None: + """Check that a missing trait raises only when a predicate needs it.""" ds = xr.Dataset({"temp": ("grid_index", np.arange(4.0))}) assert ds.mx.is_stacked() # needs no traits with pytest.raises(KeyError, match="mlwp_space_trait"): @@ -36,7 +52,18 @@ def test_missing_trait_raises_on_use_not_on_access(self): class TestNoMutation: - def test_add_xy_with_crs_leaves_dataset_untouched(self, ds_grid): + """Group tests that accessor methods do not mutate the source dataset.""" + + def test_add_xy_with_crs_leaves_dataset_untouched( + self, ds_grid: xr.Dataset + ) -> None: + """Check that ``add_xy`` with an explicit CRS leaves the input unchanged. + + Parameters + ---------- + ds_grid : xr.Dataset + Stacked grid forecast fixture. + """ out = ds_grid.mx.add_xy(crs=ccrs.PlateCarree()) np.testing.assert_allclose(out["xc"].values, ds_grid["longitude"].values) @@ -45,7 +72,16 @@ def test_add_xy_with_crs_leaves_dataset_untouched(self, ds_grid): with pytest.raises(ValueError, match="No CRS"): ds_grid.mx.add_xy() - def test_unstack_with_crs_leaves_dataset_untouched(self, ds_grid): + def test_unstack_with_crs_leaves_dataset_untouched( + self, ds_grid: xr.Dataset + ) -> None: + """Check that ``unstack`` with an explicit CRS leaves the input unchanged. + + Parameters + ---------- + ds_grid : xr.Dataset + Stacked grid forecast fixture. + """ out = ds_grid.mx.unstack( crs=ccrs.PlateCarree(), nx=2, ny=2, lon_ll=0, lat_ll=0, dx=1, dy=1 ) @@ -56,7 +92,10 @@ def test_unstack_with_crs_leaves_dataset_untouched(self, ds_grid): class TestStack: - def test_stack_regular_latlon_grid(self): + """Group tests of stacking grids to ``grid_index``.""" + + def test_stack_regular_latlon_grid(self) -> None: + """Check that a regular lat/lon grid stacks to a 1D ``grid_index``.""" ds = xr.Dataset( {"temp": (["latitude", "longitude"], np.zeros((2, 3)))}, coords={"latitude": [0.0, 1.0], "longitude": [0.0, 1.0, 2.0]}, diff --git a/tests/test_align_space.py b/tests/test_align_space.py index a621299..39e6932 100644 --- a/tests/test_align_space.py +++ b/tests/test_align_space.py @@ -20,13 +20,26 @@ import mxalign # registers ds.mx accessor # noqa: F401 - # --------------------------------------------------------------------------- # Helpers / fixtures # --------------------------------------------------------------------------- -def _props(space, time="forecast"): +def _props(space: str, time: str = "forecast") -> dict: + """Build the ``mlwp_*`` trait attributes for a deterministic dataset. + + Parameters + ---------- + space : str + Value of the ``mlwp_space_trait`` attribute, e.g. ``"grid"``. + time : str, default "forecast" + Value of the ``mlwp_time_trait`` attribute. + + Returns + ------- + dict + Trait attributes to attach to a dataset's ``attrs``. + """ return { "mlwp_space_trait": space, "mlwp_time_trait": time, @@ -35,7 +48,14 @@ def _props(space, time="forecast"): @pytest.fixture -def ds_grid(): +def ds_grid() -> xr.Dataset: + """Provide a 3×3 lat/lon grid forecast with ``temp[i, j] = lat[i] + lon[j]``. + + Returns + ------- + xr.Dataset + Grid forecast on latitudes and longitudes 0, 1, 2. + """ lats = np.array([0.0, 1.0, 2.0]) lons = np.array([0.0, 1.0, 2.0]) temp = lats[:, np.newaxis] + lons[np.newaxis, :] @@ -47,7 +67,14 @@ def ds_grid(): @pytest.fixture -def ds_point(): +def ds_point() -> xr.Dataset: + """Provide 3 observation points, two on grid nodes and one interior. + + Returns + ------- + xr.Dataset + Point observation along ``point_index``. + """ return xr.Dataset( {"temp": ("point_index", np.array([0.0, 1.0, 2.0]))}, coords={ @@ -64,12 +91,28 @@ def ds_point(): class TestGridToGrid: - def test_identical_grids_return_self(self, ds_grid): + """Group tests of aligning a grid dataset to another grid.""" + + def test_identical_grids_return_self(self, ds_grid: xr.Dataset) -> None: + """Check that aligning a grid to itself returns the same object. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + """ result = ds_grid.mx.align_space_with(ds_grid) assert result is ds_grid - def test_within_tolerance_treated_as_equal(self, ds_grid): + def test_within_tolerance_treated_as_equal(self, ds_grid: xr.Dataset) -> None: + """Check that grids differing within tolerance are treated as equal. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + """ ds_grid2 = ds_grid.assign_coords( latitude=ds_grid.latitude + 1e-5, longitude=ds_grid.longitude + 1e-5, @@ -78,7 +121,14 @@ def test_within_tolerance_treated_as_equal(self, ds_grid): assert result is ds_grid - def test_different_grids_raise(self, ds_grid): + def test_different_grids_raise(self, ds_grid: xr.Dataset) -> None: + """Check that aligning to a different grid raises ``NotImplementedError``. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + """ ds_grid2 = ds_grid.assign_coords( latitude=ds_grid.latitude + 10.0, ).assign_attrs(_props("grid")) @@ -86,7 +136,14 @@ def test_different_grids_raise(self, ds_grid): with pytest.raises(NotImplementedError): ds_grid.mx.align_space_with(ds_grid2) - def test_result_stays_grid(self, ds_grid): + def test_result_stays_grid(self, ds_grid: xr.Dataset) -> None: + """Check that grid-to-grid alignment keeps the grid space trait. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + """ result = ds_grid.mx.align_space_with(ds_grid) assert result.mx.is_grid() @@ -99,18 +156,53 @@ def test_result_stays_grid(self, ds_grid): class TestGridToPoint: - def test_result_has_point_trait(self, ds_grid, ds_point): + """Group tests of interpolating a grid dataset to point locations.""" + + def test_result_has_point_trait( + self, ds_grid: xr.Dataset, ds_point: xr.Dataset + ) -> None: + """Check that grid-to-point alignment sets the point space trait. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + ds_point : xr.Dataset + Point observation fixture. + """ result = ds_grid.mx.align_space_with(ds_point, method="xarray") assert result.mx.is_point() assert not result.mx.is_grid() - def test_result_has_point_index_dim(self, ds_grid, ds_point): + def test_result_has_point_index_dim( + self, ds_grid: xr.Dataset, ds_point: xr.Dataset + ) -> None: + """Check that grid-to-point alignment produces a ``point_index`` dim. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + ds_point : xr.Dataset + Point observation fixture. + """ result = ds_grid.mx.align_space_with(ds_point, method="xarray") assert "point_index" in result.dims - def test_result_has_target_latlon_coords(self, ds_grid, ds_point): + def test_result_has_target_latlon_coords( + self, ds_grid: xr.Dataset, ds_point: xr.Dataset + ) -> None: + """Check that grid-to-point alignment copies the target lat/lon coords. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + ds_point : xr.Dataset + Point observation fixture. + """ result = ds_grid.mx.align_space_with(ds_point, method="xarray") np.testing.assert_array_equal( @@ -120,14 +212,36 @@ def test_result_has_target_latlon_coords(self, ds_grid, ds_point): result["longitude"].values, ds_point["longitude"].values ) - def test_interpolated_values_xarray(self, ds_grid, ds_point): + def test_interpolated_values_xarray( + self, ds_grid: xr.Dataset, ds_point: xr.Dataset + ) -> None: + """Check that the xarray method interpolates the expected point values. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + ds_point : xr.Dataset + Point observation fixture. + """ result = ds_grid.mx.align_space_with(ds_point, method="xarray") assert result["temp"].isel(point_index=0).item() == pytest.approx(0.0) assert result["temp"].isel(point_index=1).item() == pytest.approx(1.0) assert result["temp"].isel(point_index=2).item() == pytest.approx(2.0) - def test_interpolated_values_delaunay(self, ds_grid, ds_point): + def test_interpolated_values_delaunay( + self, ds_grid: xr.Dataset, ds_point: xr.Dataset + ) -> None: + """Check that the Delaunay method interpolates the expected point values. + + Parameters + ---------- + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + ds_point : xr.Dataset + Point observation fixture. + """ ds_stacked = ds_grid.stack(grid_index=["latitude", "longitude"]).reset_index( "grid_index" ) @@ -145,10 +259,30 @@ def test_interpolated_values_delaunay(self, ds_grid, ds_point): class TestPointAlignmentNotImplemented: - def test_point_to_grid_raises(self, ds_point, ds_grid): + """Group tests that aligning from point datasets is not implemented.""" + + def test_point_to_grid_raises( + self, ds_point: xr.Dataset, ds_grid: xr.Dataset + ) -> None: + """Check that aligning a point dataset to a grid raises. + + Parameters + ---------- + ds_point : xr.Dataset + Point observation fixture. + ds_grid : xr.Dataset + Lat/lon grid forecast fixture. + """ with pytest.raises(NotImplementedError): ds_point.mx.align_space_with(ds_grid) - def test_point_to_point_raises(self, ds_point): + def test_point_to_point_raises(self, ds_point: xr.Dataset) -> None: + """Check that aligning a point dataset to points raises. + + Parameters + ---------- + ds_point : xr.Dataset + Point observation fixture. + """ with pytest.raises(NotImplementedError): ds_point.mx.align_space_with(ds_point) diff --git a/tests/test_align_time.py b/tests/test_align_time.py index 667a7cb..5d82b57 100644 --- a/tests/test_align_time.py +++ b/tests/test_align_time.py @@ -21,6 +21,7 @@ import numpy as np import pytest import xarray as xr +from numpy.typing import ArrayLike T0 = np.datetime64("2020-01-01T00:00", "ns") H6 = np.timedelta64(6, "h") @@ -32,6 +33,18 @@ def _props(time: str) -> dict: + """Build the ``mlwp_*`` trait attributes for a deterministic point dataset. + + Parameters + ---------- + time : str + Value of the ``mlwp_time_trait`` attribute, e.g. ``"forecast"``. + + Returns + ------- + dict + Trait attributes to attach to a dataset's ``attrs``. + """ return { "mlwp_time_trait": time, "mlwp_space_trait": "point", @@ -39,7 +52,21 @@ def _props(time: str) -> dict: } -def obs(valid_times, values): +def obs(valid_times: ArrayLike, values: ArrayLike) -> xr.Dataset: + """Build a point observation dataset. + + Parameters + ---------- + valid_times : ArrayLike + Values of the ``valid_time`` coordinate. + values : ArrayLike + Values of ``temp`` along ``valid_time``. + + Returns + ------- + xr.Dataset + Observation dataset with a ``temp`` variable. + """ return xr.Dataset( {"temp": ("valid_time", np.asarray(values, dtype=float))}, coords={"valid_time": np.asarray(valid_times)}, @@ -47,7 +74,25 @@ def obs(valid_times, values): ) -def fcst(reference_times, lead_times, values): +def fcst( + reference_times: ArrayLike, lead_times: ArrayLike, values: ArrayLike +) -> xr.Dataset: + """Build a point forecast dataset. + + Parameters + ---------- + reference_times : ArrayLike + Values of the ``reference_time`` coordinate. + lead_times : ArrayLike + Values of the ``lead_time`` coordinate. + values : ArrayLike + Values of ``temp`` with shape ``(reference_time, lead_time)``. + + Returns + ------- + xr.Dataset + Forecast dataset with a ``temp`` variable. + """ return xr.Dataset( {"temp": (["reference_time", "lead_time"], np.asarray(values, dtype=float))}, coords={ @@ -64,7 +109,18 @@ def fcst(reference_times, lead_times, values): class TestForecastToObservation: - def test_shortest_lead_time(self, ds_fcst, ds_obs): + """Group tests of aligning a forecast to observation valid times.""" + + def test_shortest_lead_time(self, ds_fcst: xr.Dataset, ds_obs: xr.Dataset) -> None: + """Check that ``lead_time="shortest"`` picks the shortest lead per valid time. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + """ result = ds_fcst.mx.align_time_with(ds_obs, lead_time="shortest") assert result.mx.is_observation() @@ -88,7 +144,16 @@ def test_shortest_lead_time(self, ds_fcst, ds_obs): result["temp"].sel(valid_time=T0 + 4 * H6).item() == 22.0 ) # (T0+12h, 12h) is shortest - def test_longest_lead_time(self, ds_fcst, ds_obs): + def test_longest_lead_time(self, ds_fcst: xr.Dataset, ds_obs: xr.Dataset) -> None: + """Check that ``lead_time="longest"`` picks the longest lead per valid time. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + """ result = ds_fcst.mx.align_time_with(ds_obs, lead_time="longest") assert result.mx.is_observation() @@ -108,7 +173,16 @@ def test_longest_lead_time(self, ds_fcst, ds_obs): result["temp"].sel(valid_time=T0 + 4 * H6).item() == 13.0 ) # (T0+6h, 18h) beats (T0+12h, 12h) - def test_specific_lead_time(self, ds_fcst, ds_obs): + def test_specific_lead_time(self, ds_fcst: xr.Dataset, ds_obs: xr.Dataset) -> None: + """Check that a specific ``lead_time`` only fills times it produces. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + """ lt = np.timedelta64(6, "h") result = ds_fcst.mx.align_time_with(ds_obs, lead_time=lt) @@ -121,7 +195,16 @@ def test_specific_lead_time(self, ds_fcst, ds_obs): assert result["temp"].sel(valid_time=T0 + 3 * H6).item() == 21.0 # (T0+12h, 6h) assert np.isnan(result["temp"].sel(valid_time=T0 + 4 * H6).item()) - def test_lead_time_as_string(self, ds_fcst, ds_obs): + def test_lead_time_as_string(self, ds_fcst: xr.Dataset, ds_obs: xr.Dataset) -> None: + """Check that a string ``lead_time`` matches the equivalent timedelta. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + """ expected = ds_fcst.mx.align_time_with(ds_obs, lead_time=np.timedelta64(6, "h")) result = ds_fcst.mx.align_time_with(ds_obs, lead_time="6h") xr.testing.assert_identical(result, expected) @@ -131,7 +214,23 @@ def test_lead_time_as_string(self, ds_fcst, ds_obs): [["6h", "12h"], np.array([6, 12], dtype="timedelta64[h]")], ids=["list", "ndarray"], ) - def test_multiple_lead_times(self, ds_fcst, ds_obs, lead_time): + def test_multiple_lead_times( + self, + ds_fcst: xr.Dataset, + ds_obs: xr.Dataset, + lead_time: list[str] | np.ndarray, + ) -> None: + """Check that several lead times are combined, earliest reference first. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + lead_time : list[str] | np.ndarray + Lead times to select, as strings or a timedelta array. + """ result = ds_fcst.mx.align_time_with(ds_obs, lead_time=lead_time) assert np.isnan(result["temp"].sel(valid_time=T0).item()) # only 0h reaches T0 @@ -149,16 +248,57 @@ def test_multiple_lead_times(self, ds_fcst, ds_obs, lead_time): (["6h", "24h"], "not in forecast lead times"), ], ) - def test_invalid_lead_time_raises(self, ds_fcst, ds_obs, lead_time, match): + def test_invalid_lead_time_raises( + self, + ds_fcst: xr.Dataset, + ds_obs: xr.Dataset, + lead_time: str | int | list[str], + match: str, + ) -> None: + """Check that an invalid ``lead_time`` raises ``ValueError``. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + lead_time : str | int | list[str] + Invalid lead time option. + match : str + Pattern expected in the error message. + """ with pytest.raises(ValueError, match=match): ds_fcst.mx.align_time_with(ds_obs, lead_time=lead_time) - def test_nan_filled_for_times_not_in_forecast(self, ds_fcst, ds_obs): + def test_nan_filled_for_times_not_in_forecast( + self, ds_fcst: xr.Dataset, ds_obs: xr.Dataset + ) -> None: + """Check that valid times the forecast never reaches are NaN. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + """ result = ds_fcst.mx.align_time_with(ds_obs, lead_time="shortest") # T0-6h is in obs but never produced by any (ref_time, lead_time) pair assert np.isnan(result["temp"].sel(valid_time=T0 - H6).item()) - def test_result_has_observation_property(self, ds_fcst, ds_obs): + def test_result_has_observation_property( + self, ds_fcst: xr.Dataset, ds_obs: xr.Dataset + ) -> None: + """Check that the result carries the observation time trait. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + """ result = ds_fcst.mx.align_time_with(ds_obs, lead_time="shortest") assert result.mx.is_observation() assert not result.mx.is_forecast() @@ -170,7 +310,20 @@ def test_result_has_observation_property(self, ds_fcst, ds_obs): class TestObservationToForecast: - def test_values_placed_at_correct_positions(self, ds_obs, ds_fcst): + """Group tests of reshaping an observation onto a forecast grid.""" + + def test_values_placed_at_correct_positions( + self, ds_obs: xr.Dataset, ds_fcst: xr.Dataset + ) -> None: + """Check that observations land at the matching (reference, lead) cells. + + Parameters + ---------- + ds_obs : xr.Dataset + Point observation fixture. + ds_fcst : xr.Dataset + Point forecast fixture. + """ result = ds_obs.mx.align_time_with(ds_fcst) assert result.mx.is_forecast() @@ -202,7 +355,18 @@ def test_values_placed_at_correct_positions(self, ds_obs, ds_fcst): == 50.0 ) - def test_nan_where_obs_missing(self, ds_obs, ds_fcst): + def test_nan_where_obs_missing( + self, ds_obs: xr.Dataset, ds_fcst: xr.Dataset + ) -> None: + """Check that cells whose valid time has no observation are NaN. + + Parameters + ---------- + ds_obs : xr.Dataset + Point observation fixture. + ds_fcst : xr.Dataset + Point forecast fixture. + """ result = ds_obs.mx.align_time_with(ds_fcst) # T0+30h is not in obs; it appears at (T0+12h, lead=18h) @@ -215,7 +379,18 @@ def test_nan_where_obs_missing(self, ds_obs, ds_fcst): .item() ) - def test_obs_value_repeated_for_shared_valid_times(self, ds_obs, ds_fcst): + def test_obs_value_repeated_for_shared_valid_times( + self, ds_obs: xr.Dataset, ds_fcst: xr.Dataset + ) -> None: + """Check that one observation fills every cell sharing its valid time. + + Parameters + ---------- + ds_obs : xr.Dataset + Point observation fixture. + ds_fcst : xr.Dataset + Point forecast fixture. + """ result = ds_obs.mx.align_time_with(ds_fcst) # T0+12h appears at (T0,12h), (T0+6h,6h), (T0+12h,0h) — all should equal 30.0 @@ -238,7 +413,18 @@ def test_obs_value_repeated_for_shared_valid_times(self, ds_obs, ds_fcst): == 30.0 ) - def test_result_has_forecast_property(self, ds_obs, ds_fcst): + def test_result_has_forecast_property( + self, ds_obs: xr.Dataset, ds_fcst: xr.Dataset + ) -> None: + """Check that the result carries the forecast time trait. + + Parameters + ---------- + ds_obs : xr.Dataset + Point observation fixture. + ds_fcst : xr.Dataset + Point forecast fixture. + """ result = ds_obs.mx.align_time_with(ds_fcst) assert result.mx.is_forecast() assert not result.mx.is_observation() @@ -250,17 +436,44 @@ def test_result_has_forecast_property(self, ds_obs, ds_fcst): class TestObservationToObservation: + """Group tests of aligning an observation to another observation.""" + @pytest.fixture - def ds_obs1(self): + def ds_obs1(self) -> xr.Dataset: + """Provide an observation at T0 .. T0+18h with values 0..3. + + Returns + ------- + xr.Dataset + Source observation dataset. + """ times = np.array([T0, T0 + H6, T0 + 2 * H6, T0 + 3 * H6]) return obs(times, [0.0, 1.0, 2.0, 3.0]) @pytest.fixture - def ds_obs2(self): + def ds_obs2(self) -> xr.Dataset: + """Provide an observation at T0+6h .. T0+24h with values 10..40. + + Returns + ------- + xr.Dataset + Target observation dataset. + """ times = np.array([T0 + H6, T0 + 2 * H6, T0 + 3 * H6, T0 + 4 * H6]) return obs(times, [10.0, 20.0, 30.0, 40.0]) - def test_reindexes_to_ds2_valid_times(self, ds_obs1, ds_obs2): + def test_reindexes_to_ds2_valid_times( + self, ds_obs1: xr.Dataset, ds_obs2: xr.Dataset + ) -> None: + """Check that the result is reindexed onto the target valid times. + + Parameters + ---------- + ds_obs1 : xr.Dataset + Source observation fixture. + ds_obs2 : xr.Dataset + Target observation fixture. + """ result = ds_obs1.mx.align_time_with(ds_obs2) assert list(result.valid_time.values) == list(ds_obs2.valid_time.values) @@ -271,11 +484,33 @@ def test_reindexes_to_ds2_valid_times(self, ds_obs1, ds_obs2): assert result["temp"].sel(valid_time=T0 + 2 * H6).item() == 2.0 assert result["temp"].sel(valid_time=T0 + 3 * H6).item() == 3.0 - def test_ds1_only_times_are_dropped(self, ds_obs1, ds_obs2): + def test_ds1_only_times_are_dropped( + self, ds_obs1: xr.Dataset, ds_obs2: xr.Dataset + ) -> None: + """Check that valid times only in the source are dropped. + + Parameters + ---------- + ds_obs1 : xr.Dataset + Source observation fixture. + ds_obs2 : xr.Dataset + Target observation fixture. + """ result = ds_obs1.mx.align_time_with(ds_obs2) assert T0 not in result.valid_time.values # only in ds1 - def test_result_stays_observation(self, ds_obs1, ds_obs2): + def test_result_stays_observation( + self, ds_obs1: xr.Dataset, ds_obs2: xr.Dataset + ) -> None: + """Check that observation-to-observation alignment keeps the trait. + + Parameters + ---------- + ds_obs1 : xr.Dataset + Source observation fixture. + ds_obs2 : xr.Dataset + Target observation fixture. + """ result = ds_obs1.mx.align_time_with(ds_obs2) assert result.mx.is_observation() @@ -286,21 +521,52 @@ def test_result_stays_observation(self, ds_obs1, ds_obs2): class TestForecastToForecast: + """Group tests of aligning a forecast to another forecast.""" + @pytest.fixture - def ds_fcst2(self): + def ds_fcst2(self) -> xr.Dataset: + """Provide a zero-valued forecast at T0+6h .. T0+18h with 6h and 12h leads. + + Returns + ------- + xr.Dataset + Target forecast dataset. + """ ref = np.array([T0 + H6, T0 + 2 * H6, T0 + 3 * H6]) lead = np.array([np.timedelta64(6, "h"), np.timedelta64(12, "h")]) values = np.zeros((3, 2)) return fcst(ref, lead, values) - def test_reindexes_to_ds2_reference_times(self, ds_fcst, ds_fcst2): + def test_reindexes_to_ds2_reference_times( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that the result is reindexed onto the target reference times. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="reference") np.testing.assert_array_equal( result.reference_time.values, ds_fcst2.reference_time.values ) - def test_lead_time_reference_drops_ds1_only_leads(self, ds_fcst, ds_fcst2): + def test_lead_time_reference_drops_ds1_only_leads( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that ``lead_time="reference"`` drops leads only in the source. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="reference") np.testing.assert_array_equal( @@ -309,13 +575,35 @@ def test_lead_time_reference_drops_ds1_only_leads(self, ds_fcst, ds_fcst2): assert np.timedelta64(0, "h") not in result.lead_time.values assert np.timedelta64(18, "h") not in result.lead_time.values - def test_nan_for_ref_times_not_in_ds1(self, ds_fcst, ds_fcst2): + def test_nan_for_ref_times_not_in_ds1( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that reference times missing from the source are all NaN. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="reference") # T0+18h is in ds_fcst2 but not ds_fcst → all NaN assert np.isnan(result["temp"].sel(reference_time=T0 + 3 * H6).values).all() - def test_values_preserved_for_common_ref_times(self, ds_fcst, ds_fcst2): + def test_values_preserved_for_common_ref_times( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that source values are kept at shared (reference, lead) cells. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="reference") assert ( @@ -343,7 +631,18 @@ def test_values_preserved_for_common_ref_times(self, ds_fcst, ds_fcst2): == 22.0 ) - def test_lead_time_intersection_keeps_common_leads(self, ds_fcst, ds_fcst2): + def test_lead_time_intersection_keeps_common_leads( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that ``lead_time="intersection"`` keeps only shared leads. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="intersection") expected_lead = np.array([np.timedelta64(6, "h"), np.timedelta64(12, "h")]) @@ -351,16 +650,49 @@ def test_lead_time_intersection_keeps_common_leads(self, ds_fcst, ds_fcst2): assert np.timedelta64(0, "h") not in result.lead_time.values assert np.timedelta64(18, "h") not in result.lead_time.values - def test_lead_time_union_keeps_all_leads(self, ds_fcst, ds_fcst2): + def test_lead_time_union_keeps_all_leads( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that ``lead_time="union"`` keeps the leads of both forecasts. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="union") expected_lead = np.array([np.timedelta64(h, "h") for h in [0, 6, 12, 18]]) np.testing.assert_array_equal(result.lead_time.values, expected_lead) - def test_result_has_valid_time_coord(self, ds_fcst, ds_fcst2): + def test_result_has_valid_time_coord( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that the aligned forecast has a ``valid_time`` coordinate. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="reference") assert "valid_time" in result.coords - def test_result_stays_forecast(self, ds_fcst, ds_fcst2): + def test_result_stays_forecast( + self, ds_fcst: xr.Dataset, ds_fcst2: xr.Dataset + ) -> None: + """Check that forecast-to-forecast alignment keeps the forecast trait. + + Parameters + ---------- + ds_fcst : xr.Dataset + Source point forecast fixture. + ds_fcst2 : xr.Dataset + Target forecast fixture. + """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="reference") assert result.mx.is_forecast() diff --git a/tests/test_align_time_global.py b/tests/test_align_time_global.py index 9631aca..9b28e4f 100644 --- a/tests/test_align_time_global.py +++ b/tests/test_align_time_global.py @@ -18,6 +18,7 @@ import pandas as pd import pytest import xarray as xr +from numpy.typing import ArrayLike from mxalign.alignment.time import align_time @@ -26,6 +27,18 @@ def _props(time: str) -> dict: + """Build the ``mlwp_*`` trait attributes for a deterministic point dataset. + + Parameters + ---------- + time : str + Value of the ``mlwp_time_trait`` attribute, e.g. ``"forecast"``. + + Returns + ------- + dict + Trait attributes to attach to a dataset's ``attrs``. + """ return { "mlwp_time_trait": time, "mlwp_space_trait": "point", @@ -33,7 +46,21 @@ def _props(time: str) -> dict: } -def fcst(reference_times, lead_times): +def fcst(reference_times: ArrayLike, lead_times: ArrayLike) -> xr.Dataset: + """Build a point forecast with ``temp`` counting up from 0 in row-major order. + + Parameters + ---------- + reference_times : ArrayLike + Reference times, parseable by ``pd.to_datetime``. + lead_times : ArrayLike + Lead times, parseable by ``pd.to_timedelta``. + + Returns + ------- + xr.Dataset + Forecast dataset over ``reference_time`` and ``lead_time``. + """ rt = pd.to_datetime(np.asarray(reference_times)) lt = pd.to_timedelta(np.asarray(lead_times)) values = np.arange(len(rt) * len(lt), dtype=float).reshape(len(rt), len(lt)) @@ -44,7 +71,21 @@ def fcst(reference_times, lead_times): ) -def obs(valid_times, offset=0.0): +def obs(valid_times: ArrayLike, offset: float = 0.0) -> xr.Dataset: + """Build a point observation with ``temp`` counting up from ``offset``. + + Parameters + ---------- + valid_times : ArrayLike + Valid times, parseable by ``pd.to_datetime``. + offset : float, default 0.0 + Value of the first observation. + + Returns + ------- + xr.Dataset + Observation dataset along ``valid_time``. + """ vt = pd.to_datetime(np.asarray(valid_times)) return xr.Dataset( {"temp": (["valid_time"], np.arange(len(vt), dtype=float) + offset)}, @@ -59,8 +100,17 @@ def obs(valid_times, offset=0.0): class TestGlobalForecastsObservation: + """Group tests of the global regime with two forecasts and an observation.""" + @pytest.fixture - def datasets(self): + def datasets(self) -> dict[str, xr.Dataset]: + """Provide two partly overlapping forecasts and a reference observation. + + Returns + ------- + dict[str, xr.Dataset] + Datasets ``f1``, ``f2`` and the observation ``ref``. + """ f1 = fcst( ["2020-01-01", "2020-01-02", "2020-01-03"], ["0h", "6h", "12h"], @@ -72,7 +122,16 @@ def datasets(self): ref = obs(pd.date_range("2020-01-01", "2020-01-05", freq="3h"), offset=100.0) return {"f1": f1, "f2": f2, "ref": ref} - def test_r_star_is_intersection_of_forecast_ref_times(self, datasets): + def test_r_star_is_intersection_of_forecast_ref_times( + self, datasets: dict[str, xr.Dataset] + ) -> None: + """Check that R* is the intersection of the forecasts' reference times. + + Parameters + ---------- + datasets : dict[str, xr.Dataset] + Two forecasts and a reference observation. + """ out = align_time(datasets, reference="ref") np.testing.assert_array_equal( out["f1"].reference_time.values.astype("datetime64[D]"), @@ -80,7 +139,16 @@ def test_r_star_is_intersection_of_forecast_ref_times(self, datasets): ) assert out["f2"].sizes["reference_time"] == 2 - def test_l_star_is_union_of_forecast_lead_times(self, datasets): + def test_l_star_is_union_of_forecast_lead_times( + self, datasets: dict[str, xr.Dataset] + ) -> None: + """Check that L* is the union of the forecasts' lead times. + + Parameters + ---------- + datasets : dict[str, xr.Dataset] + Two forecasts and a reference observation. + """ out = align_time(datasets, reference="ref") expected = np.array([0, 6, 12, 18], dtype="int64") np.testing.assert_array_equal( @@ -88,18 +156,43 @@ def test_l_star_is_union_of_forecast_lead_times(self, datasets): expected, ) - def test_all_outputs_share_grid(self, datasets): + def test_all_outputs_share_grid(self, datasets: dict[str, xr.Dataset]) -> None: + """Check that all aligned datasets end up with the same shape. + + Parameters + ---------- + datasets : dict[str, xr.Dataset] + Two forecasts and a reference observation. + """ out = align_time(datasets, reference="ref") shapes = {k: out[k].temp.shape for k in out} assert len(set(shapes.values())) == 1, f"shapes differ: {shapes}" - def test_reference_observation_becomes_forecast_shaped(self, datasets): + def test_reference_observation_becomes_forecast_shaped( + self, datasets: dict[str, xr.Dataset] + ) -> None: + """Check that the reference observation is reshaped to forecast dims. + + Parameters + ---------- + datasets : dict[str, xr.Dataset] + Two forecasts and a reference observation. + """ out = align_time(datasets, reference="ref") assert out["ref"].mx.is_forecast() assert out["ref"].temp.dims == ("reference_time", "lead_time") assert out["ref"].valid_time.dims == ("reference_time", "lead_time") - def test_forecast_nan_filled_for_missing_lead_times(self, datasets): + def test_forecast_nan_filled_for_missing_lead_times( + self, datasets: dict[str, xr.Dataset] + ) -> None: + """Check that forecasts are NaN at lead times they do not have. + + Parameters + ---------- + datasets : dict[str, xr.Dataset] + Two forecasts and a reference observation. + """ out = align_time(datasets, reference="ref") # f1 has no 18h lead → its 18h column is all-NaN across R* assert np.isnan(out["f1"].temp.sel(lead_time="18h").values).all() @@ -115,7 +208,10 @@ def test_forecast_nan_filled_for_missing_lead_times(self, datasets): class TestOverlappingValidTimes: - def test_observation_broadcast_when_runs_overlap(self): + """Group tests of forecast runs whose valid times overlap.""" + + def test_observation_broadcast_when_runs_overlap(self) -> None: + """Check that one observation fills every cell sharing its valid time.""" # Runs 6h apart with leads up to 12h → each valid_time is hit up to 3 times f1 = fcst( ["2020-01-01 00:00", "2020-01-01 06:00", "2020-01-01 12:00"], @@ -138,7 +234,18 @@ def test_observation_broadcast_when_runs_overlap(self): class TestReferenceTimePruning: - def test_prunes_ref_times_before_reference_window(self, caplog): + """Group tests of pruning R* to the reference valid-time window.""" + + def test_prunes_ref_times_before_reference_window( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Check that reference times before the reference window are pruned. + + Parameters + ---------- + caplog : pytest.LogCaptureFixture + Captures the pruning warning. + """ f1 = fcst(["2020-01-01", "2020-01-02", "2020-01-03"], ["0h", "6h", "12h"]) f2 = fcst(["2020-01-01", "2020-01-02", "2020-01-03"], ["0h", "6h", "12h"]) ref = obs(pd.date_range("2020-01-03", "2020-01-05", freq="3h")) @@ -154,8 +261,11 @@ def test_prunes_ref_times_before_reference_window(self, caplog): assert not np.isnan(out["ref"].temp.values).any() assert any("pruned" in r.message for r in caplog.records) - def test_keeps_long_lead_time_when_reference_covers_it(self): - """Forecast F2 has a long lead_time; F1 doesn't. L* must keep it.""" + def test_keeps_long_lead_time_when_reference_covers_it(self) -> None: + """Check that L* keeps a long lead time that only one forecast has. + + Forecast F2 has a long lead_time; F1 doesn't. L* must keep it. + """ f1 = fcst(["2020-01-02", "2020-01-03"], ["0h", "6h", "12h"]) f2 = fcst(["2020-01-02", "2020-01-03"], ["0h", "6h", "12h", "240h"]) ref = obs(pd.date_range("2020-01-01", "2020-01-31", freq="3h")) @@ -168,8 +278,11 @@ def test_keeps_long_lead_time_when_reference_covers_it(self): assert np.isnan(out["f1"].temp.sel(lead_time="240h").values).all() assert not np.isnan(out["f1"].temp.sel(lead_time="0h").values).any() - def test_raises_when_no_ref_time_fits_all_lead_times(self): - """Long lead_time + short reference → R* empties → ValueError.""" + def test_raises_when_no_ref_time_fits_all_lead_times(self) -> None: + """Check that an empty R* after pruning raises ``ValueError``. + + Long lead_time + short reference → R* empties → ValueError. + """ f1 = fcst(["2020-01-02", "2020-01-03"], ["0h"]) f2 = fcst(["2020-01-02", "2020-01-03"], ["0h", "240h"]) ref = obs(pd.date_range("2020-01-01", "2020-01-04", freq="3h")) @@ -184,7 +297,10 @@ def test_raises_when_no_ref_time_fits_all_lead_times(self): class TestMultipleObservations: - def test_both_observations_reshape_to_forecast_grid(self): + """Group tests of global alignment with several observations.""" + + def test_both_observations_reshape_to_forecast_grid(self) -> None: + """Check that every observation is reshaped onto the forecast grid.""" f1 = fcst(["2020-01-02", "2020-01-03"], ["0h", "6h", "12h"]) f2 = fcst(["2020-01-02", "2020-01-03"], ["6h", "12h", "18h"]) o1 = obs(pd.date_range("2020-01-01", "2020-01-05", freq="3h"), offset=100.0) @@ -208,9 +324,14 @@ def test_both_observations_reshape_to_forecast_grid(self): class TestPairwiseFallback: - def test_forecast_reference_uses_pairwise(self): - """When the reference is a forecast, no global pruning happens — - each non-ref dataset is aligned pairwise to the reference.""" + """Group tests of falling back to pairwise alignment.""" + + def test_forecast_reference_uses_pairwise(self) -> None: + """Check that a forecast reference triggers pairwise alignment. + + When the reference is a forecast, no global pruning happens — + each non-ref dataset is aligned pairwise to the reference. + """ f1 = fcst(["2020-01-01", "2020-01-02", "2020-01-03"], ["0h", "6h", "12h"]) f_ref = fcst(["2020-01-02", "2020-01-03"], ["6h", "12h"]) @@ -225,8 +346,11 @@ def test_forecast_reference_uses_pairwise(self): ) assert out["ref"] is f_ref - def test_all_observations_uses_pairwise(self): - """Reference is obs and no forecasts → pairwise reindex.""" + def test_all_observations_uses_pairwise(self) -> None: + """Check that observations without forecasts are aligned pairwise. + + Reference is obs and no forecasts → pairwise reindex. + """ o1 = obs(pd.date_range("2020-01-01", "2020-01-05", freq="6h")) o2 = obs(pd.date_range("2020-01-02", "2020-01-04", freq="6h")) @@ -242,7 +366,10 @@ def test_all_observations_uses_pairwise(self): class TestLeadTimeOption: - def test_forecast_reference_forwards_lead_time(self): + """Group tests of the ``lead_time`` option of ``align_time``.""" + + def test_forecast_reference_forwards_lead_time(self) -> None: + """Check that ``lead_time`` is forwarded to pairwise forecast alignment.""" f1 = fcst(["2020-01-02", "2020-01-03"], ["0h", "6h"]) f_ref = fcst(["2020-01-02", "2020-01-03"], ["6h", "12h"]) @@ -251,7 +378,8 @@ def test_forecast_reference_forwards_lead_time(self): leads = out["f1"].lead_time.values.astype("timedelta64[h]").astype("int64") np.testing.assert_array_equal(leads, [0, 6, 12]) - def test_forecast_reference_rejects_unknown_lead_time(self): + def test_forecast_reference_rejects_unknown_lead_time(self) -> None: + """Check that an invalid forecast-to-forecast ``lead_time`` raises.""" f1 = fcst(["2020-01-02"], ["0h"]) f_ref = fcst(["2020-01-02"], ["0h"]) @@ -259,7 +387,14 @@ def test_forecast_reference_rejects_unknown_lead_time(self): align_time({"f1": f1, "ref": f_ref}, reference="ref", lead_time="shortest") @pytest.mark.parametrize("with_forecast", [True, False]) - def test_observation_reference_rejects_lead_time(self, with_forecast): + def test_observation_reference_rejects_lead_time(self, with_forecast: bool) -> None: + """Check that ``lead_time`` is rejected when the reference is an observation. + + Parameters + ---------- + with_forecast : bool + Whether to include a forecast alongside the reference observation. + """ datasets = {"ref": obs(pd.date_range("2020-01-01", "2020-01-05", freq="6h"))} if with_forecast: datasets["f1"] = fcst(["2020-01-02"], ["0h", "6h"]) diff --git a/tests/test_config.py b/tests/test_config.py index dd5fbda..3c02354 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,3 +1,5 @@ +"""Tests for loading and validating pipeline configuration.""" + from pathlib import Path import pytest @@ -7,10 +9,28 @@ EXAMPLE_CONFIG = Path(__file__).parent.parent / "examples" / "config.yaml" -DATES = dict(start="2020-01-01", end="2020-01-02", period="1D", range="12h", step="12h") +DATES = { + "start": "2020-01-01", + "end": "2020-01-02", + "period": "1D", + "range": "12h", + "step": "12h", +} + + +def _config(**sections: dict) -> dict: + """Build a minimal valid config dict, extended with extra sections. + Parameters + ---------- + **sections : dict + Extra top-level sections to add to (or override in) the config. -def _config(**sections) -> dict: + Returns + ------- + dict + Config dict with a single ``fc`` dataset used as alignment reference. + """ return { "datasets": {"fc": {"loader": "x", "files": "/d/fc.nc"}}, "alignment": {"reference": "fc"}, @@ -18,7 +38,8 @@ def _config(**sections) -> dict: } -def test_example_config_validates(): +def test_example_config_validates() -> None: + """Check that the example config in ``examples/`` loads and validates.""" config = load_config(str(EXAMPLE_CONFIG)) assert config.alignment.reference == "cerra" assert config.alignment.space.interpolation == {"method": "delaunay"} @@ -30,31 +51,36 @@ def test_example_config_validates(): assert len(config.datasets["la-1024-ea-00"].files) > 1 -def test_unknown_top_level_key_is_rejected(): +def test_unknown_top_level_key_is_rejected() -> None: + """Check that an unknown top-level config key raises a validation error.""" with pytest.raises(ValidationError, match="not_a_section"): load_config(_config(not_a_section={"reference": "fc"})) -def test_missing_alignment_reference_is_rejected(): +def test_missing_alignment_reference_is_rejected() -> None: + """Check that a config without ``alignment.reference`` is rejected.""" config = _config() config["alignment"] = {"broadcast_nans": False} with pytest.raises(ValidationError, match="alignment.reference"): load_config(config) -def test_unknown_save_datasets_option_is_rejected(): +def test_unknown_save_datasets_option_is_rejected() -> None: + """Check that an unknown ``alignment.save.datasets`` option is rejected.""" save = {"method": "to_netcdf", "path": "out.nc", "datasets": "some"} with pytest.raises(ValidationError, match="alignment.save.datasets"): load_config(_config(alignment={"reference": "fc", "save": save})) -def test_extra_dataset_keys_are_passed_to_loader(): +def test_extra_dataset_keys_are_passed_to_loader() -> None: + """Check that unknown dataset keys are collected as loader kwargs.""" config = _config() config["datasets"]["fc"]["chunks"] = {"time": 1} assert load_config(config).datasets["fc"].kwargs == {"chunks": {"time": 1}} -def test_dataset_dates_override_top_level_dates(): +def test_dataset_dates_override_top_level_dates() -> None: + """Check that per-dataset dates override the top-level dates section.""" config = _config(dates=DATES) config["datasets"]["fc"] = { "loader": "x", @@ -68,7 +94,8 @@ def test_dataset_dates_override_top_level_dates(): ] -def test_incomplete_dates_are_rejected(): +def test_incomplete_dates_are_rejected() -> None: + """Check that a dates section missing required fields is rejected.""" config = _config() config["datasets"]["fc"]["dates"] = {"start": "2020-01-01"} with pytest.raises(ValidationError, match="dates is missing end, period"): diff --git a/tests/test_dates.py b/tests/test_dates.py index 441eac5..a5e6e24 100644 --- a/tests/test_dates.py +++ b/tests/test_dates.py @@ -1,21 +1,32 @@ +"""Tests for date-based file path substitution.""" + import datetime import pytest +import xarray as xr from mxalign.pipeline.config import load_config from mxalign.pipeline.dates import Dates from mxalign.pipeline.save import DatasetPath # 3 reference times (daily) × 3 lead times (0, 12, 24h) -DATES = dict(start="2020-01-01", end="2020-01-03", period="1D", range="24h", step="12h") +DATES = { + "start": "2020-01-01", + "end": "2020-01-03", + "period": "1D", + "range": "24h", + "step": "12h", +} -def test_substitute_reference_time(): +def test_substitute_reference_time() -> None: + """Check that ``reference_time`` placeholders expand to one path per run.""" paths = Dates(**DATES).substitute("/d/{reference_time:%Y%m%d%H}.nc") assert paths == ["/d/2020010100.nc", "/d/2020010200.nc", "/d/2020010300.nc"] -def test_substitute_reference_and_lead_time(): +def test_substitute_reference_and_lead_time() -> None: + """Check that reference and lead time placeholders expand as a product.""" paths = Dates(**DATES).substitute("/d/{reference_time:%Y%m%d}_+{lead_time:03d}.nc") assert len(paths) == 9 assert paths[:3] == [ @@ -25,7 +36,8 @@ def test_substitute_reference_and_lead_time(): ] -def test_substitute_valid_time_is_deduplicated_and_sorted(): +def test_substitute_valid_time_is_deduplicated_and_sorted() -> None: + """Check that ``valid_time`` placeholders give sorted, unique paths.""" paths = Dates(**DATES).substitute("/d/{valid_time:%Y%m%d%H}.nc") # 2020-01-01T00 … 2020-01-04T00 every 12h; overlaps between references collapse assert len(paths) == 7 @@ -33,17 +45,20 @@ def test_substitute_valid_time_is_deduplicated_and_sorted(): assert paths[-1] == "/d/2020010400.nc" -def test_substitute_without_fields_returns_single_path(): +def test_substitute_without_fields_returns_single_path() -> None: + """Check that a pattern without placeholders yields a single path.""" assert Dates(**DATES).substitute("/d/store.zarr") == ["/d/store.zarr"] -def test_substitute_lead_time_requires_whole_hours(): +def test_substitute_lead_time_requires_whole_hours() -> None: + """Check that a sub-hourly step raises when substituting ``lead_time``.""" dates = Dates(**{**DATES, "step": "30min"}) with pytest.raises(ValueError, match="whole hours"): dates.substitute("/d/{lead_time}.nc") -def test_config_expands_files_from_dates_section(): +def test_config_expands_files_from_dates_section() -> None: + """Check that config loading expands file patterns over the dates section.""" # YAML parses unquoted timestamps to datetime objects config = load_config( { @@ -61,6 +76,13 @@ def test_config_expands_files_from_dates_section(): ] -def test_dataset_path_substitute(ds_fcst): +def test_dataset_path_substitute(ds_fcst: xr.Dataset) -> None: + """Check that ``DatasetPath`` fills name and date fields from the dataset. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + """ path = DatasetPath("model", ds_fcst).substitute("{name}_{year}{month:02d}.nc") assert path == "model_202001.nc" diff --git a/tests/test_interpolations.py b/tests/test_interpolations.py index 12d3504..c198953 100644 --- a/tests/test_interpolations.py +++ b/tests/test_interpolations.py @@ -9,25 +9,40 @@ Trait / accessor-level tests live in test_align_space.py. """ +from pathlib import Path + import cartopy.crs as ccrs import dask import numpy as np import pytest import xarray as xr +from numpy.typing import ArrayLike from mxalign.alignment.interpolations import delaunay from mxalign.alignment.interpolations.delaunay import DelaunayInterpolator from mxalign.alignment.interpolations.interpolate import interpolate from mxalign.alignment.interpolations.xarray import XarrayInterpolator - # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- -def _latlon_grid(lats, lons) -> xr.Dataset: - """Regular lat/lon grid with temp[i, j] = lat[i] + lon[j].""" +def _latlon_grid(lats: ArrayLike, lons: ArrayLike) -> xr.Dataset: + """Build a regular lat/lon grid with ``temp[i, j] = lat[i] + lon[j]``. + + Parameters + ---------- + lats : ArrayLike + Latitudes of the grid. + lons : ArrayLike + Longitudes of the grid. + + Returns + ------- + xr.Dataset + Grid dataset with ``latitude`` and ``longitude`` dims. + """ lats, lons = np.asarray(lats), np.asarray(lons) temp = lats[:, np.newaxis] + lons[np.newaxis, :] return xr.Dataset( @@ -36,7 +51,21 @@ def _latlon_grid(lats, lons) -> xr.Dataset: ) -def _stacked_grid(lats, lons) -> xr.Dataset: +def _stacked_grid(lats: ArrayLike, lons: ArrayLike) -> xr.Dataset: + """Build a regular lat/lon grid stacked to a ``grid_index`` dim. + + Parameters + ---------- + lats : ArrayLike + Latitudes of the grid. + lons : ArrayLike + Longitudes of the grid. + + Returns + ------- + xr.Dataset + Grid dataset with ``temp = lat + lon`` along ``grid_index``. + """ return ( _latlon_grid(lats, lons) .stack(grid_index=["latitude", "longitude"]) @@ -45,22 +74,47 @@ def _stacked_grid(lats, lons) -> xr.Dataset: @pytest.fixture -def grid_latlon(): - """3×3 regular lat/lon grid. temp[i, j] = lat[i] + lon[j].""" +def grid_latlon() -> xr.Dataset: + """Provide a 3×3 regular lat/lon grid with ``temp[i, j] = lat[i] + lon[j]``. + + Returns + ------- + xr.Dataset + Grid dataset on latitudes and longitudes 0, 1, 2. + """ return _latlon_grid([0.0, 1.0, 2.0], [0.0, 1.0, 2.0]) @pytest.fixture -def grid_stacked(grid_latlon): - """grid_latlon stacked to grid_index (required by DelaunayInterpolator).""" +def grid_stacked(grid_latlon: xr.Dataset) -> xr.Dataset: + """Provide ``grid_latlon`` stacked to ``grid_index``. + + Stacking is required by ``DelaunayInterpolator``. + + Parameters + ---------- + grid_latlon : xr.Dataset + Regular lat/lon grid fixture. + + Returns + ------- + xr.Dataset + Grid dataset along ``grid_index``. + """ return grid_latlon.stack(grid_index=["latitude", "longitude"]).reset_index( "grid_index" ) @pytest.fixture -def target_points(): - """3 target points: two at grid nodes, one at an interior location.""" +def target_points() -> xr.Dataset: + """Provide 3 target points: two at grid nodes, one at an interior location. + + Returns + ------- + xr.Dataset + Coordinate-only dataset with lat/lon along ``point_index``. + """ return xr.Dataset( coords={ "latitude": ("point_index", np.array([0.0, 0.5, 1.0])), @@ -75,25 +129,71 @@ def target_points(): class TestXarrayInterpolator: - def test_values_at_grid_nodes(self, grid_latlon, target_points): + """Group tests of ``XarrayInterpolator``.""" + + def test_values_at_grid_nodes( + self, grid_latlon: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that values at grid nodes are reproduced exactly. + + Parameters + ---------- + grid_latlon : xr.Dataset + Regular lat/lon grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = XarrayInterpolator(target_points)._interpolate(grid_latlon) # lat=0, lon=0 → 0+0=0; lat=1, lon=1 → 1+1=2 assert result["temp"].isel(point_index=0).item() == pytest.approx(0.0) assert result["temp"].isel(point_index=2).item() == pytest.approx(2.0) - def test_value_at_interior_point(self, grid_latlon, target_points): + def test_value_at_interior_point( + self, grid_latlon: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that an interior point is interpolated bilinearly. + + Parameters + ---------- + grid_latlon : xr.Dataset + Regular lat/lon grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = XarrayInterpolator(target_points)._interpolate(grid_latlon) # lat=0.5, lon=0.5 → 0.5+0.5=1.0 (exact for bilinear on linear function) assert result["temp"].isel(point_index=1).item() == pytest.approx(1.0) - def test_output_has_point_index_dim(self, grid_latlon, target_points): + def test_output_has_point_index_dim( + self, grid_latlon: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that the output has a ``point_index`` dim. + + Parameters + ---------- + grid_latlon : xr.Dataset + Regular lat/lon grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = XarrayInterpolator(target_points)._interpolate(grid_latlon) assert "point_index" in result.dims - def test_output_has_latlon_coords_from_target(self, grid_latlon, target_points): + def test_output_has_latlon_coords_from_target( + self, grid_latlon: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that the output carries the target lat/lon coords. + + Parameters + ---------- + grid_latlon : xr.Dataset + Regular lat/lon grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = XarrayInterpolator(target_points)._interpolate(grid_latlon) np.testing.assert_array_equal( @@ -103,7 +203,14 @@ def test_output_has_latlon_coords_from_target(self, grid_latlon, target_points): result["longitude"].values, target_points["longitude"].values ) - def test_outside_domain_is_nan(self, grid_latlon): + def test_outside_domain_is_nan(self, grid_latlon: xr.Dataset) -> None: + """Check that a point outside the grid domain is NaN. + + Parameters + ---------- + grid_latlon : xr.Dataset + Regular lat/lon grid fixture. + """ far_point = xr.Dataset( coords={ "latitude": ("point_index", np.array([10.0])), @@ -114,7 +221,16 @@ def test_outside_domain_is_nan(self, grid_latlon): assert np.isnan(result["temp"].isel(point_index=0).item()) - def test_projected_source_keeps_target_point_coords(self, target_points): + def test_projected_source_keeps_target_point_coords( + self, target_points: xr.Dataset + ) -> None: + """Check that a projected source keeps the target's point coordinates. + + Parameters + ---------- + target_points : xr.Dataset + Target point locations fixture. + """ # xc/yc source with an identity projection, so xc == lon and yc == lat xc = yc = np.array([0.0, 1.0, 2.0]) source = xr.Dataset( @@ -143,24 +259,70 @@ def test_projected_source_keeps_target_point_coords(self, target_points): class TestDelaunayInterpolator: - def test_values_at_grid_nodes(self, grid_stacked, target_points): + """Group tests of ``DelaunayInterpolator``.""" + + def test_values_at_grid_nodes( + self, grid_stacked: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that values at grid nodes are reproduced exactly. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = DelaunayInterpolator(target_points)._interpolate(grid_stacked) assert result["temp"].isel(point_index=0).item() == pytest.approx(0.0) assert result["temp"].isel(point_index=2).item() == pytest.approx(2.0) - def test_value_at_interior_point(self, grid_stacked, target_points): + def test_value_at_interior_point( + self, grid_stacked: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that an interior point is interpolated barycentrically. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = DelaunayInterpolator(target_points)._interpolate(grid_stacked) # Barycentric interpolation is exact for any linear function assert result["temp"].isel(point_index=1).item() == pytest.approx(1.0) - def test_output_has_point_index_dim(self, grid_stacked, target_points): + def test_output_has_point_index_dim( + self, grid_stacked: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that the output has a ``point_index`` dim. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = DelaunayInterpolator(target_points)._interpolate(grid_stacked) assert "point_index" in result.dims - def test_output_has_latlon_coords_from_target(self, grid_stacked, target_points): + def test_output_has_latlon_coords_from_target( + self, grid_stacked: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that the output carries the target lat/lon coords. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ result = DelaunayInterpolator(target_points)._interpolate(grid_stacked) np.testing.assert_array_equal( @@ -171,8 +333,22 @@ def test_output_has_latlon_coords_from_target(self, grid_stacked, target_points) ) def test_weights_shared_across_instances( - self, grid_stacked, target_points, monkeypatch - ): + self, + grid_stacked: xr.Dataset, + target_points: xr.Dataset, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Check that interpolation weights are cached across interpolator instances. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + target_points : xr.Dataset + Target point locations fixture. + monkeypatch : pytest.MonkeyPatch + Used to reset the weights cache and count weight builds. + """ monkeypatch.setattr(delaunay, "_WEIGHTS_CACHE", {}) calls = [] build = delaunay._build_weight_matrix @@ -189,8 +365,22 @@ def test_weights_shared_across_instances( assert len(calls) == 1 def test_cache_distinguishes_grids_with_same_extent( - self, grid_stacked, target_points, monkeypatch - ): + self, + grid_stacked: xr.Dataset, + target_points: xr.Dataset, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Check that the weights cache tells apart grids with the same extent. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + target_points : xr.Dataset + Target point locations fixture. + monkeypatch : pytest.MonkeyPatch + Used to reset the weights cache. + """ monkeypatch.setattr(delaunay, "_WEIGHTS_CACHE", {}) # Same shape, first latitude and last longitude, different interior other = _stacked_grid(lats=[0.0, 1.0, 2.0], lons=[0.0, 1.5, 2.0]) @@ -200,7 +390,14 @@ def test_cache_distinguishes_grids_with_same_extent( assert result["temp"].isel(point_index=1).item() == pytest.approx(1.0) - def test_rechunk_keeps_block_size_bounded(self, target_points): + def test_rechunk_keeps_block_size_bounded(self, target_points: xr.Dataset) -> None: + """Check that rechunking ``grid_index`` keeps blocks within the chunk budget. + + Parameters + ---------- + target_points : xr.Dataset + Target point locations fixture. + """ grid = _stacked_grid(lats=np.linspace(0, 2, 20), lons=np.linspace(0, 2, 50)) n_steps, n_source = 64, grid.sizes["grid_index"] source = ( @@ -217,7 +414,14 @@ def test_rechunk_keeps_block_size_bounded(self, target_points): assert rows_per_block * n_source * 8 <= 64 * 1024 np.testing.assert_allclose(result["temp"].isel(step=0).values, [0.0, 1.0, 2.0]) - def test_outside_convex_hull_is_nan(self, grid_stacked): + def test_outside_convex_hull_is_nan(self, grid_stacked: xr.Dataset) -> None: + """Check that a point outside the convex hull is NaN. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + """ far_point = xr.Dataset( coords={ "latitude": ("point_index", np.array([10.0])), @@ -247,8 +451,27 @@ def test_outside_convex_hull_is_nan(self, grid_stacked): [(XarrayInterpolator, "grid_latlon"), (DelaunayInterpolator, "grid_stacked")], ) def test_point_output_drops_grid_attrs( - interpolator, source, target_points, tmp_path, request -): + interpolator: type[XarrayInterpolator] | type[DelaunayInterpolator], + source: str, + target_points: xr.Dataset, + tmp_path: Path, + request: pytest.FixtureRequest, +) -> None: + """Check that point output drops grid-only attrs and can be saved to netCDF. + + Parameters + ---------- + interpolator : type[XarrayInterpolator] | type[DelaunayInterpolator] + Interpolator class under test. + source : str + Name of the source grid fixture. + target_points : xr.Dataset + Target point locations fixture. + tmp_path : Path + Temporary directory to write the output to. + request : pytest.FixtureRequest + Used to look up the source grid fixture by name. + """ source = request.getfixturevalue(source).assign_attrs(GRID_ATTRS) result = interpolator(target_points).interpolate(source) @@ -258,7 +481,18 @@ def test_point_output_drops_grid_attrs( result.to_netcdf(tmp_path / "points.nc") # a cartopy CRS attr would fail here -def test_interpolate_returns_same_container_type(grid_latlon, target_points): +def test_interpolate_returns_same_container_type( + grid_latlon: xr.Dataset, target_points: xr.Dataset +) -> None: + """Check that ``interpolate`` returns the same container type it is given. + + Parameters + ---------- + grid_latlon : xr.Dataset + Regular lat/lon grid fixture. + target_points : xr.Dataset + Target point locations fixture. + """ out = interpolate( {"a": grid_latlon.assign_attrs(GRID_ATTRS)}, target_points, "xarray" ) diff --git a/tests/test_runner.py b/tests/test_runner.py index d84eded..9bba149 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1,16 +1,46 @@ +"""Tests for the pipeline ``Runner``.""" + +from pathlib import Path + import pytest +import xarray as xr from pydantic import ValidationError from mxalign.pipeline.runner import Runner def _runner(datasets: dict, alignment: dict) -> Runner: + """Build a ``Runner`` with pre-loaded datasets and an alignment section. + + Parameters + ---------- + datasets : dict + Datasets by name, assigned directly instead of being loaded. + alignment : dict + The ``alignment`` section of the config. + + Returns + ------- + Runner + Runner ready to transform or align the given datasets. + """ runner = Runner({"datasets": {}, "alignment": alignment}) runner.datasets = datasets return runner -def test_align_runs_time_alignment_without_time_section(ds_fcst, ds_obs): +def test_align_runs_time_alignment_without_time_section( + ds_fcst: xr.Dataset, ds_obs: xr.Dataset +) -> None: + """Check that ``align`` runs time alignment when no time section is given. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_obs : xr.Dataset + Point observation fixture. + """ runner = _runner( {"fcst": ds_fcst, "obs": ds_obs}, {"reference": "obs", "broadcast_nans": False}, @@ -20,7 +50,8 @@ def test_align_runs_time_alignment_without_time_section(ds_fcst, ds_obs): assert runner.datasets["obs"].mx.is_forecast() -def test_align_time_rejects_unknown_options(): +def test_align_time_rejects_unknown_options() -> None: + """Check that unknown ``alignment.time`` options are rejected.""" with pytest.raises(ValidationError, match="alignment.time.return_as"): Runner( { @@ -30,13 +61,31 @@ def test_align_time_rejects_unknown_options(): ) -def test_transform_datasets_without_transformations_section(ds_obs): +def test_transform_datasets_without_transformations_section( + ds_obs: xr.Dataset, +) -> None: + """Check that datasets pass through unchanged without a transformations section. + + Parameters + ---------- + ds_obs : xr.Dataset + Point observation fixture. + """ runner = _runner({"obs": ds_obs}, {"reference": "obs"}) runner.transform_datasets() assert runner.datasets["obs"] is ds_obs -def test_align_saves_merged_datasets(ds_obs, tmp_path): +def test_align_saves_merged_datasets(ds_obs: xr.Dataset, tmp_path: Path) -> None: + """Check that ``align`` saves a single merged file with ``datasets: merge``. + + Parameters + ---------- + ds_obs : xr.Dataset + Point observation fixture. + tmp_path : Path + Temporary directory to save the output in. + """ runner = _runner( {"a": ds_obs, "b": ds_obs.copy()}, { diff --git a/tests/test_transformations.py b/tests/test_transformations.py index ec85306..6021b58 100644 --- a/tests/test_transformations.py +++ b/tests/test_transformations.py @@ -1,3 +1,7 @@ +"""Tests for the transformation registry and ``transform``.""" + +from collections.abc import Callable + import pytest import xarray as xr @@ -9,22 +13,37 @@ @register_transformation("_test_double") -def _double(ds): +def _double(ds: xr.Dataset) -> xr.Dataset: + """Double every value in the dataset. + + Parameters + ---------- + ds : xr.Dataset + Dataset to transform. + + Returns + ------- + xr.Dataset + Dataset with all values multiplied by 2. + """ return ds * 2 -def test_registered_transformation_is_available(): +def test_registered_transformation_is_available() -> None: + """Check that a registered transformation is listed as available.""" assert "_test_double" in available_transformations() -def test_transform_applies_function_without_mutating_input(): +def test_transform_applies_function_without_mutating_input() -> None: + """Check that ``transform`` applies the function and leaves the input as is.""" ds = xr.Dataset({"t2m": ("x", [1.0, 2.0, 3.0])}) out = transform("_test_double", ds) assert out["t2m"].values.tolist() == [2.0, 4.0, 6.0] assert ds["t2m"].values.tolist() == [1.0, 2.0, 3.0] -def test_unknown_transformation_raises(): +def test_unknown_transformation_raises() -> None: + """Check that an unknown transformation name raises ``ValueError``.""" with pytest.raises(ValueError, match="Unknown transformation"): transform("does_not_exist", xr.Dataset()) @@ -32,7 +51,16 @@ def test_unknown_transformation_raises(): @pytest.mark.parametrize( "container", [lambda ds: {"a": ds}, lambda ds: [ds]], ids=["dict", "list"] ) -def test_transform_returns_same_container_type(container): +def test_transform_returns_same_container_type( + container: Callable[[xr.Dataset], dict[str, xr.Dataset] | list[xr.Dataset]], +) -> None: + """Check that ``transform`` returns the same container type it is given. + + Parameters + ---------- + container : Callable[[xr.Dataset], dict[str, xr.Dataset] | list[xr.Dataset]] + Function wrapping a dataset in a dict or a list. + """ datasets = container(xr.Dataset({"t2m": ("x", [1.0])})) out = transform("_test_double", datasets) assert type(out) is type(datasets) From e89897624898d004bdcf7b07884084dd45f3be8d Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 27 Oct 2026 13:27:01 +0000 Subject: [PATCH 08/12] add uncertainty predicates to the mx accessor 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 --- src/mxalign/accessor/mx.py | 78 +++++++++++++++++++++++++++++++++++--- tests/test_accessor.py | 30 +++++++++++++++ 2 files changed, 103 insertions(+), 5 deletions(-) diff --git a/src/mxalign/accessor/mx.py b/src/mxalign/accessor/mx.py index aa40dd2..1a74082 100644 --- a/src/mxalign/accessor/mx.py +++ b/src/mxalign/accessor/mx.py @@ -1,13 +1,18 @@ -"""The ``ds.mx`` xarray Dataset accessor for space/time predicates and alignment.""" +"""The ``ds.mx`` xarray Dataset accessor for trait predicates and alignment.""" import numpy as np import pandas as pd import xarray as xr import cartopy.crs as ccrs -from mlwp_data_specs.api import SPACE_TRAIT_ATTR, TIME_TRAIT_ATTR +from mlwp_data_specs.api import ( + SPACE_TRAIT_ATTR, + TIME_TRAIT_ATTR, + UNCERTAINTY_TRAIT_ATTR, +) from mlwp_data_specs.specs.traits.spatial_coordinate import Space from mlwp_data_specs.specs.traits.time_coordinate import Time +from mlwp_data_specs.specs.traits.uncertainty import Uncertainty from ..utils.projections import create_cartopy_crs, BUILTIN from ..alignment import time as _time @@ -16,10 +21,10 @@ @xr.register_dataset_accessor("mx") class MxAccessor: - """``ds.mx`` accessor — space/time predicates and alignment operations. + """``ds.mx`` accessor — trait predicates and alignment operations. - Reads ``mlwp_space_trait`` and ``mlwp_time_trait`` from ``ds.attrs`` on - every use (xarray caches the accessor per object, so reading them once + Reads ``mlwp_space_trait``, ``mlwp_time_trait`` and + ``mlwp_uncertainty_trait`` from ``ds.attrs`` on every use (xarray caches the accessor per object, so reading them once would go stale when attrs change). Raises ``KeyError`` when a trait is needed but missing. Never modifies ``ds``; operations return new datasets. @@ -90,6 +95,22 @@ def _time(self) -> Time: """ return Time(self._trait(TIME_TRAIT_ATTR)) + @property + def _uncertainty(self) -> Uncertainty: + """The uncertainty trait of the dataset. + + Returns + ------- + Uncertainty + The uncertainty trait read from ``ds.attrs``. + + Raises + ------ + KeyError + If the dataset has no uncertainty trait attribute. + """ + return Uncertainty(self._trait(UNCERTAINTY_TRAIT_ATTR)) + # --- Space predicates --- def is_grid(self) -> bool: @@ -154,6 +175,53 @@ def is_observation(self) -> bool: """ return self._time == Time.OBSERVATION + # --- Uncertainty predicates --- + + def is_deterministic(self) -> bool: + """Return whether the dataset has the DETERMINISTIC uncertainty trait. + + Returns + ------- + bool + ``True`` if the uncertainty trait is ``Uncertainty.DETERMINISTIC``. + + Raises + ------ + KeyError + If the dataset has no uncertainty trait attribute. + """ + return self._uncertainty == Uncertainty.DETERMINISTIC + + def is_ensemble(self) -> bool: + """Return whether the dataset has the ENSEMBLE uncertainty trait. + + Returns + ------- + bool + ``True`` if the uncertainty trait is ``Uncertainty.ENSEMBLE``. + + Raises + ------ + KeyError + If the dataset has no uncertainty trait attribute. + """ + return self._uncertainty == Uncertainty.ENSEMBLE + + def is_quantile(self) -> bool: + """Return whether the dataset has the QUANTILE uncertainty trait. + + Returns + ------- + bool + ``True`` if the uncertainty trait is ``Uncertainty.QUANTILE``. + + Raises + ------ + KeyError + If the dataset has no uncertainty trait attribute. + """ + return self._uncertainty == Uncertainty.QUANTILE + # --- Space operations --- def add_crs(self, crs: str | dict | ccrs.Projection) -> xr.Dataset: diff --git a/tests/test_accessor.py b/tests/test_accessor.py index 777c8da..a5c8cd7 100644 --- a/tests/test_accessor.py +++ b/tests/test_accessor.py @@ -50,6 +50,36 @@ def test_missing_trait_raises_on_use_not_on_access(self) -> None: with pytest.raises(KeyError, match="mlwp_space_trait"): ds.mx.is_grid() + @pytest.mark.parametrize("trait", ["deterministic", "ensemble", "quantile"]) + def test_uncertainty_predicates(self, ds_grid: xr.Dataset, trait: str) -> None: + """Check that exactly the predicate matching the uncertainty trait holds. + + Parameters + ---------- + ds_grid : xr.Dataset + Stacked grid forecast fixture. + trait : str + Value of the ``mlwp_uncertainty_trait`` attribute. + """ + ds_grid.attrs["mlwp_uncertainty_trait"] = trait + predicates = { + "deterministic": ds_grid.mx.is_deterministic(), + "ensemble": ds_grid.mx.is_ensemble(), + "quantile": ds_grid.mx.is_quantile(), + } + assert predicates == {name: name == trait for name in predicates} + + def test_missing_uncertainty_trait_raises(self, ds_grid: xr.Dataset) -> None: + """Check that uncertainty predicates raise when the trait is missing. + + Parameters + ---------- + ds_grid : xr.Dataset + Stacked grid forecast fixture, which has no uncertainty trait. + """ + with pytest.raises(KeyError, match="mlwp_uncertainty_trait"): + ds_grid.mx.is_ensemble() + class TestNoMutation: """Group tests that accessor methods do not mutate the source dataset.""" From fdd9f8385d088bf4ce9ec02b2ac62b33632887f5 Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 27 Oct 2026 13:27:01 +0000 Subject: [PATCH 09/12] test ensemble time and space alignment 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 --- tests/conftest.py | 86 ++++++++++++++++- tests/test_align_space.py | 25 +++++ tests/test_align_time.py | 159 ++++++++++++++++++++++++++++++++ tests/test_align_time_global.py | 38 ++++++++ 4 files changed, 304 insertions(+), 4 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index d62ac11..ac82c77 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,14 +26,27 @@ OBS_TIMES = np.array([T0 + i * H6 for i in range(-1, 5)]) # T0-6h … T0+24h OBS_VALUES = np.arange(len(OBS_TIMES), dtype=float) * 10.0 # 0, 10, 20, 30, 40, 50 +# Ensemble members: CF "realization" coordinate (TODO_ENS.md, convention C2). +# Ensemble value convention: member m adds 100 * m to the deterministic value. +MEMBERS = np.array([0, 1, 2]) +MEMBER_ATTRS = { + "standard_name": "realization", + "units": "1", + "long_name": "ensemble member number", +} -def _props(time: str) -> dict: - """Build the ``mlwp_*`` trait attributes for a deterministic point dataset. + +def _props(time: str, space: str = "point", uncertainty: str = "deterministic") -> dict: + """Build the ``mlwp_*`` trait attributes for a dataset. Parameters ---------- time : str Value of the ``mlwp_time_trait`` attribute, e.g. ``"forecast"``. + space : str, default "point" + Value of the ``mlwp_space_trait`` attribute. + uncertainty : str, default "deterministic" + Value of the ``mlwp_uncertainty_trait`` attribute. Returns ------- @@ -42,11 +55,27 @@ def _props(time: str) -> dict: """ return { "mlwp_time_trait": time, - "mlwp_space_trait": "point", - "mlwp_uncertainty_trait": "deterministic", + "mlwp_space_trait": space, + "mlwp_uncertainty_trait": uncertainty, } +def _member_coord(members: np.ndarray = MEMBERS) -> xr.DataArray: + """Build a ``member`` coordinate with the CF ``realization`` attributes. + + Parameters + ---------- + members : np.ndarray, default MEMBERS + Member numbers. + + Returns + ------- + xr.DataArray + Integer ``member`` coordinate. + """ + return xr.DataArray(members, dims="member", attrs=MEMBER_ATTRS) + + @pytest.fixture def ds_fcst() -> xr.Dataset: """Provide a point forecast with 3 reference times and 4 lead times. @@ -78,3 +107,52 @@ def ds_obs() -> xr.Dataset: coords={"valid_time": OBS_TIMES}, attrs=_props("observation"), ) + + +@pytest.fixture +def ds_fcst_ens() -> xr.Dataset: + """Provide a 3-member point forecast ensemble on the times of ``ds_fcst``. + + Returns + ------- + xr.Dataset + Ensemble with ``temp[m, i, j] = 100 * m + i * 10 + j`` over ``member``, + ``reference_time`` and ``lead_time`` (``member`` leftmost, as CF + recommends). + """ + values = 100.0 * MEMBERS[:, np.newaxis, np.newaxis] + FORECAST_VALUES + return xr.Dataset( + {"temp": (["member", "reference_time", "lead_time"], values)}, + coords={ + "member": _member_coord(), + "reference_time": REFERENCE_TIMES, + "lead_time": LEAD_TIMES, + }, + attrs=_props("forecast", uncertainty="ensemble"), + ) + + +@pytest.fixture +def ds_grid_ens() -> xr.Dataset: + """Provide a 3-member forecast ensemble on a stacked 3×3 lat/lon grid. + + Returns + ------- + xr.Dataset + Ensemble with ``temp[m, g] = 100 * m + lat[g] + lon[g]`` over + ``member`` and ``grid_index``, for latitudes and longitudes 0°, 1°, + 2°. The field is linear in space, so both interpolators reproduce it + exactly. + """ + lat, lon = np.meshgrid([0.0, 1.0, 2.0], [0.0, 1.0, 2.0], indexing="ij") + lat, lon = lat.ravel(), lon.ravel() + values = 100.0 * MEMBERS[:, np.newaxis] + (lat + lon) + return xr.Dataset( + {"temp": (["member", "grid_index"], values)}, + coords={ + "member": _member_coord(), + "latitude": ("grid_index", lat), + "longitude": ("grid_index", lon), + }, + attrs=_props("forecast", space="grid", uncertainty="ensemble"), + ) diff --git a/tests/test_align_space.py b/tests/test_align_space.py index 39e6932..2e3bf4c 100644 --- a/tests/test_align_space.py +++ b/tests/test_align_space.py @@ -252,6 +252,31 @@ def test_interpolated_values_delaunay( assert result["temp"].isel(point_index=1).item() == pytest.approx(1.0) assert result["temp"].isel(point_index=2).item() == pytest.approx(2.0) + @pytest.mark.parametrize("method", ["xarray", "delaunay"]) + def test_ensemble_keeps_members_and_trait( + self, ds_grid_ens: xr.Dataset, ds_point: xr.Dataset, method: str + ) -> None: + """Check that an ensemble stays an ensemble with the same members. + + Parameters + ---------- + ds_grid_ens : xr.Dataset + Stacked grid ensemble fixture. + ds_point : xr.Dataset + Point observation fixture. + method : str + Interpolator to use; the xarray one needs a lat/lon-dim source. + """ + source = ds_grid_ens + if method == "xarray": + source = source.set_index(grid_index=["latitude", "longitude"]).unstack() + + result = source.mx.align_space_with(ds_point, method=method) + + assert result.mx.is_point() + assert result.mx.is_ensemble() + xr.testing.assert_identical(result["member"], ds_grid_ens["member"]) + # --------------------------------------------------------------------------- # Case 3: Point → * (not implemented) diff --git a/tests/test_align_time.py b/tests/test_align_time.py index 5d82b57..b3e3830 100644 --- a/tests/test_align_time.py +++ b/tests/test_align_time.py @@ -21,6 +21,10 @@ import numpy as np import pytest import xarray as xr +from mlwp_data_specs.specs.traits.uncertainty import Uncertainty +from mlwp_data_specs.specs.traits.uncertainty import ( + validate_dataset as validate_uncertainty, +) from numpy.typing import ArrayLike T0 = np.datetime64("2020-01-01T00:00", "ns") @@ -696,3 +700,158 @@ def test_result_stays_forecast( """ result = ds_fcst.mx.align_time_with(ds_fcst2, lead_time="reference") assert result.mx.is_forecast() + + +# --------------------------------------------------------------------------- +# Ensembles: the member dimension rides along and is never matched +# --------------------------------------------------------------------------- + + +def _assert_ensemble(ds: xr.Dataset, members: ArrayLike) -> None: + """Assert that ``ds`` is still a valid ensemble with the given members. + + Parameters + ---------- + ds : xr.Dataset + Dataset to check. + members : ArrayLike + Expected values of the ``member`` coordinate. + """ + assert ds.mx.is_ensemble() + np.testing.assert_array_equal(ds["member"].values, members) + assert ds["member"].attrs["standard_name"] == "realization" + report, _ = validate_uncertainty(ds, trait=Uncertainty.ENSEMBLE) + assert not report.has_fails() + + +def _per_member(ds: xr.Dataset, member: int) -> xr.DataArray: + """Return ``temp`` of one member, minus the 100-per-member offset. + + Parameters + ---------- + ds : xr.Dataset + Ensemble built from the ``ds_fcst_ens`` fixture. + member : int + Member to select. + + Returns + ------- + xr.DataArray + The member's values, comparable to the deterministic ``ds_fcst``. + """ + return ds["temp"].sel(member=member, drop=True) - 100.0 * member + + +class TestEnsemble: + """Group tests of time alignment with a ``member`` dimension.""" + + @pytest.mark.parametrize("lead_time", ["shortest", "longest", "6h"]) + def test_forecast_ensemble_to_observation( + self, + ds_fcst_ens: xr.Dataset, + ds_fcst: xr.Dataset, + ds_obs: xr.Dataset, + lead_time: str, + ) -> None: + """Check that each member is aligned like the deterministic forecast. + + Parameters + ---------- + ds_fcst_ens : xr.Dataset + Point forecast ensemble fixture. + ds_fcst : xr.Dataset + Point forecast fixture, equal to member 0. + ds_obs : xr.Dataset + Point observation fixture. + lead_time : str + Lead time strategy for forecast-to-observation alignment. + """ + result = ds_fcst_ens.mx.align_time_with(ds_obs, lead_time=lead_time) + expected = ds_fcst.mx.align_time_with(ds_obs, lead_time=lead_time) + + assert result.mx.is_observation() + _assert_ensemble(result, ds_fcst_ens["member"]) + for member in result["member"].values: + xr.testing.assert_equal(_per_member(result, member), expected["temp"]) + + def test_observation_to_forecast_ensemble_adds_no_member( + self, ds_obs: xr.Dataset, ds_fcst_ens: xr.Dataset, ds_fcst: xr.Dataset + ) -> None: + """Check that observations reshape onto an ensemble's times only. + + Parameters + ---------- + ds_obs : xr.Dataset + Point observation fixture. + ds_fcst_ens : xr.Dataset + Point forecast ensemble fixture. + ds_fcst : xr.Dataset + Point forecast fixture with the same times as the ensemble. + """ + result = ds_obs.mx.align_time_with(ds_fcst_ens) + + xr.testing.assert_identical(result, ds_obs.mx.align_time_with(ds_fcst)) + assert result.mx.is_deterministic() + + def test_forecast_ensemble_to_forecast( + self, ds_fcst_ens: xr.Dataset, ds_fcst: xr.Dataset + ) -> None: + """Check that each member is aligned like the deterministic forecast. + + Parameters + ---------- + ds_fcst_ens : xr.Dataset + Point forecast ensemble fixture. + ds_fcst : xr.Dataset + Point forecast fixture, equal to member 0. + """ + target = ds_fcst.isel(reference_time=[1, 2], lead_time=[0, 1]) + + result = ds_fcst_ens.mx.align_time_with(target) + expected = ds_fcst.mx.align_time_with(target) + + _assert_ensemble(result, ds_fcst_ens["member"]) + for member in result["member"].values: + xr.testing.assert_equal(_per_member(result, member), expected["temp"]) + + def test_forecast_to_forecast_ensemble_adds_no_member( + self, ds_fcst: xr.Dataset, ds_fcst_ens: xr.Dataset + ) -> None: + """Check that a deterministic forecast aligns to an ensemble's times only. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_fcst_ens : xr.Dataset + Point forecast ensemble fixture. + """ + target = ds_fcst_ens.isel(reference_time=[1, 2], lead_time=[0, 1]) + + result = ds_fcst.mx.align_time_with(target) + + xr.testing.assert_identical( + result, ds_fcst.mx.align_time_with(target.isel(member=0, drop=True)) + ) + + def test_members_are_not_matched_between_ensembles( + self, ds_fcst_ens: xr.Dataset + ) -> None: + """Check that aligning to an ensemble with other members keeps all members. + + Member 1 of one ensemble has nothing to do with member 1 of another, so + the ``member`` dimension must not be reindexed. + + Parameters + ---------- + ds_fcst_ens : xr.Dataset + Point forecast ensemble fixture with members 0, 1, 2. + """ + target = ds_fcst_ens.isel(member=[1, 2]).assign_coords(member=[10, 11]) + + result = ds_fcst_ens.mx.align_time_with(target) + + _assert_ensemble(result, ds_fcst_ens["member"]) + xr.testing.assert_equal( + result["temp"].drop_vars("valid_time"), ds_fcst_ens["temp"] + ) diff --git a/tests/test_align_time_global.py b/tests/test_align_time_global.py index 9b28e4f..115114d 100644 --- a/tests/test_align_time_global.py +++ b/tests/test_align_time_global.py @@ -401,3 +401,41 @@ def test_observation_reference_rejects_lead_time(self, with_forecast: bool) -> N with pytest.raises(ValueError, match="only supported when the reference"): align_time(datasets, reference="ref", lead_time="union") + + +# --------------------------------------------------------------------------- +# Ensembles in the global regime +# --------------------------------------------------------------------------- + + +class TestEnsemble: + """Group tests of global alignment with an ensemble among the forecasts.""" + + def test_ensemble_keeps_members_on_common_grid( + self, ds_fcst: xr.Dataset, ds_fcst_ens: xr.Dataset, ds_obs: xr.Dataset + ) -> None: + """Check that an ensemble is put on the common grid member by member. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture, equal to member 0 of the ensemble. + ds_fcst_ens : xr.Dataset + Point forecast ensemble fixture. + ds_obs : xr.Dataset + Point observation fixture, used as the reference. + """ + out = align_time( + {"det": ds_fcst, "ens": ds_fcst_ens, "obs": ds_obs}, reference="obs" + ) + + ens = out["ens"] + assert ens.mx.is_ensemble() + assert ens["member"].attrs["standard_name"] == "realization" + for member in ens["member"].values: + xr.testing.assert_identical( + ens["temp"].sel(member=member, drop=True) - 100.0 * member, + out["det"]["temp"], + ) + # The observation is reshaped onto the time grid only + assert set(out["obs"]["temp"].dims) == {"reference_time", "lead_time"} From cc2aff25995cef94ddfb870831528a277409cf36 Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 27 Oct 2026 13:27:05 +0000 Subject: [PATCH 10/12] interpolate variables with grid_index in any position (delaunay) 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 --- .../alignment/interpolations/delaunay.py | 19 ++-- tests/test_interpolations.py | 89 +++++++++++++++++++ 2 files changed, 102 insertions(+), 6 deletions(-) diff --git a/src/mxalign/alignment/interpolations/delaunay.py b/src/mxalign/alignment/interpolations/delaunay.py index efa48c2..913b23c 100644 --- a/src/mxalign/alignment/interpolations/delaunay.py +++ b/src/mxalign/alignment/interpolations/delaunay.py @@ -63,7 +63,9 @@ def __init__(self, target_dataset: xr.Dataset, **options: object) -> None: def _interpolate(self, source_dataset: xr.Dataset) -> xr.Dataset: """Interpolate a stacked grid dataset to the target points. - Variables whose last dimension is not ``grid_index`` are skipped. + Variables without a ``grid_index`` dimension are skipped. In the + output, ``point_index`` takes the place of ``grid_index`` and the other + dimensions keep their order. Parameters ---------- @@ -110,13 +112,18 @@ def _interpolate(self, source_dataset: xr.Dataset) -> xr.Dataset: arrays_out = {} for var in source_dataset.data_vars: da = source_dataset[var] - if da.dims[-1] != "grid_index": - print( - f"Skipping variable '{var}' - doesn't end with spatial dimension grid_index" + if "grid_index" not in da.dims: + logger.warning( + "Delaunay interpolation: skipping variable '%s' without a " + "grid_index dimension.", + var, ) continue - else: - arrays_out[var] = interpolate_da(da, W, target_points) + # The weight matmul needs grid_index last; put point_index back in + # its place afterwards so e.g. a trailing member dim stays trailing + out_dims = [("point_index" if d == "grid_index" else d) for d in da.dims] + da = da.transpose(..., "grid_index") + arrays_out[var] = interpolate_da(da, W, target_points).transpose(*out_dims) ds_out = xr.Dataset(arrays_out).assign_coords( latitude=self.target_dataset["latitude"], diff --git a/tests/test_interpolations.py b/tests/test_interpolations.py index c198953..c8351ea 100644 --- a/tests/test_interpolations.py +++ b/tests/test_interpolations.py @@ -9,6 +9,7 @@ Trait / accessor-level tests live in test_align_space.py. """ +import logging from pathlib import Path import cartopy.crs as ccrs @@ -252,6 +253,28 @@ def test_projected_source_keeps_target_point_coords( np.testing.assert_array_equal(result["station"].values, ["a", "b", "c"]) np.testing.assert_allclose(result["temp"].values, [0.0, 1.0, 2.0]) + def test_ensemble_values_per_member( + self, ds_grid_ens: xr.Dataset, target_points: xr.Dataset + ) -> None: + """Check that each member is interpolated on its own. + + Parameters + ---------- + ds_grid_ens : xr.Dataset + Stacked grid ensemble fixture, ``temp = 100 * member + lat + lon``. + target_points : xr.Dataset + Target point locations fixture. + """ + source = ds_grid_ens.set_index(grid_index=["latitude", "longitude"]).unstack() + + result = XarrayInterpolator(target_points)._interpolate(source) + + members = ds_grid_ens["member"].values + np.testing.assert_allclose( + result["temp"].transpose("member", "point_index").values, + 100.0 * members[:, np.newaxis] + np.array([0.0, 1.0, 2.0]), + ) + # --------------------------------------------------------------------------- # DelaunayInterpolator @@ -432,6 +455,72 @@ def test_outside_convex_hull_is_nan(self, grid_stacked: xr.Dataset) -> None: assert np.isnan(result["temp"].isel(point_index=0).item()) + @pytest.mark.parametrize("chunked", [False, True]) + @pytest.mark.parametrize( + "dims", [("member", "grid_index"), ("grid_index", "member")] + ) + def test_ensemble_keeps_dim_order( + self, + ds_grid_ens: xr.Dataset, + target_points: xr.Dataset, + dims: tuple[str, str], + chunked: bool, + ) -> None: + """Check that ``member`` is interpolated wherever it sits in the dims. + + ``point_index`` takes the place of ``grid_index`` in the output, so a + ``member`` dimension after ``grid_index`` stays after ``point_index``. + + Parameters + ---------- + ds_grid_ens : xr.Dataset + Stacked grid ensemble fixture, ``temp = 100 * member + lat + lon``. + target_points : xr.Dataset + Target point locations fixture. + dims : tuple[str, str] + Dimension order of the source variable. + chunked : bool + Whether the source is dask-backed with several ``grid_index`` chunks. + """ + source = ds_grid_ens.transpose(*dims) + if chunked: + source = source.chunk({"grid_index": 4, "member": 1}) + + result = DelaunayInterpolator(target_points)._interpolate(source) + + out_dims = tuple("point_index" if d == "grid_index" else d for d in dims) + assert result["temp"].dims == out_dims + members = ds_grid_ens["member"].values + np.testing.assert_allclose( + result["temp"].transpose("member", "point_index").values, + 100.0 * members[:, np.newaxis] + np.array([0.0, 1.0, 2.0]), + ) + + def test_variable_without_grid_index_is_skipped( + self, + grid_stacked: xr.Dataset, + target_points: xr.Dataset, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Check that a variable without ``grid_index`` is dropped with a warning. + + Parameters + ---------- + grid_stacked : xr.Dataset + Stacked grid fixture. + target_points : xr.Dataset + Target point locations fixture. + caplog : pytest.LogCaptureFixture + Captures the warning. + """ + source = grid_stacked.assign(offset=("step", [1.0, 2.0])) + + with caplog.at_level(logging.WARNING, logger=delaunay.__name__): + result = DelaunayInterpolator(target_points)._interpolate(source) + + assert list(result.data_vars) == ["temp"] + assert any("'offset'" in r.getMessage() for r in caplog.records) + # --------------------------------------------------------------------------- # Shared interpolator behaviour From 5e7163bf6b65a1ad82860c3a5072016c64a59cec Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 27 Oct 2026 13:27:09 +0000 Subject: [PATCH 11/12] broadcast NaNs across ensembles by reducing over members 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 --- doc/adr/ADR-003_ensemble-nan-masking.md | 55 ++++++ src/mxalign/alignment/nans.py | 110 ++++++++--- tests/test_nans.py | 253 ++++++++++++++++++++++++ tests/test_runner.py | 32 +++ 4 files changed, 423 insertions(+), 27 deletions(-) create mode 100644 doc/adr/ADR-003_ensemble-nan-masking.md create mode 100644 tests/test_nans.py diff --git a/doc/adr/ADR-003_ensemble-nan-masking.md b/doc/adr/ADR-003_ensemble-nan-masking.md new file mode 100644 index 0000000..2bfd4aa --- /dev/null +++ b/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. diff --git a/src/mxalign/alignment/nans.py b/src/mxalign/alignment/nans.py index 090d7c8..032f481 100644 --- a/src/mxalign/alignment/nans.py +++ b/src/mxalign/alignment/nans.py @@ -1,15 +1,23 @@ """Propagation of NaN masks across a collection of datasets.""" -import xarray as xr import itertools +import xarray as xr + def broadcast_nans( datasets: xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset], ) -> xr.Dataset | list[xr.Dataset] | dict[str, xr.Dataset]: """Propagate NaN masks so a NaN at a coordinate in any dataset is NaN in all. - Operates pairwise over deep copies; only shared coordinates and variables are considered. + Works pairwise on the variables two datasets share, at the coordinate + values they share; values outside that intersection are left as they are. + + The uncertainty dimension (``member`` for an ensemble, ``quantile`` for a + quantile dataset, taken from the uncertainty trait) is never matched + between datasets. Instead, a value counts as missing when any of its + members or quantiles is, and the mask is then applied to all of them. This + keeps the ensemble size constant over the verification sample. Parameters ---------- @@ -19,41 +27,89 @@ def broadcast_nans( Returns ------- xr.Dataset or list of xr.Dataset or dict of str to xr.Dataset - Deep copies of the input datasets with NaN masks propagated, in the - same container type as ``datasets``. + Copies of the input datasets with NaN masks propagated, in the same + container type as ``datasets``. + + Raises + ------ + KeyError + If a dataset has no uncertainty trait attribute. + ValueError + If a shared variable has different dimensions in two datasets, other + than their uncertainty dimensions. """ if isinstance(datasets, xr.Dataset): return datasets elif isinstance(datasets, dict): keys = list(datasets.keys()) - working = [ds.copy(deep=True) for ds in datasets.values()] + working = [ds.copy() for ds in datasets.values()] else: keys = None - working = [ds.copy(deep=True) for ds in datasets] + working = [ds.copy() for ds in datasets] - # Iterate over all pairs of datasets - for dsA, dsB in itertools.combinations(working, 2): - # Find the shared coordinates for all dimensions - common_coords = { - dim: sorted(set(dsA[dim].values) & set(dsB[dim].values)) for dim in dsA.dims - } + uncertainty_dims = [_uncertainty_dims(ds) for ds in working] - # Iterate over all variables + for i, j in itertools.combinations(range(len(working)), 2): + dsA, dsB = working[i], working[j] for var in dsA.data_vars: - if var in dsB: # Ensure both datasets have the variable - # Select the data at common coordinates - selA = dsA[var].sel(**common_coords) - selB = dsB[var].sel(**common_coords) - - # Compute NaN mask for shared coordinates - nan_mask = selA.isnull() | selB.isnull() - - # Apply NaN mask back to both datasets - dsA[var].loc[common_coords] = ( - dsA[var].sel(**common_coords).where(~nan_mask) - ) - dsB[var].loc[common_coords] = ( - dsB[var].sel(**common_coords).where(~nan_mask) + if var not in dsB: + continue + maskA = _missing(dsA[var], uncertainty_dims[i]) + maskB = _missing(dsB[var], uncertainty_dims[j]) + if set(maskA.dims) != set(maskB.dims): + raise ValueError( + f"Cannot broadcast NaNs for variable '{var}': dimensions " + f"{sorted(maskA.dims)} and {sorted(maskB.dims)} differ. Only " + "the uncertainty dimension (member, quantile) may differ " + "between datasets; align them in time and space first." ) + # Masks over the shared coordinate values, applied back to each + # dataset; outside the intersection nothing is masked + maskA, maskB = xr.align(maskA, maskB, join="inner") + mask = maskA | maskB + dsA[var] = dsA[var].where(~mask.reindex_like(dsA[var], fill_value=False)) + dsB[var] = dsB[var].where(~mask.reindex_like(dsB[var], fill_value=False)) return dict(zip(keys, working)) if keys else working + + +def _uncertainty_dims(ds: xr.Dataset) -> list[str]: + """Return the uncertainty dimensions of ``ds`` according to its trait. + + Parameters + ---------- + ds : xr.Dataset + Dataset with an uncertainty trait attribute. + + Returns + ------- + list of str + ``["member"]`` for an ensemble, ``["quantile"]`` for a quantile + dataset and ``[]`` for a deterministic one. + """ + if ds.mx.is_ensemble(): + return ["member"] + if ds.mx.is_quantile(): + return ["quantile"] + return [] + + +def _missing(da: xr.DataArray, uncertainty_dims: list[str]) -> xr.DataArray: + """Return where ``da`` is missing, reduced over its uncertainty dimensions. + + Parameters + ---------- + da : xr.DataArray + Variable to build the mask for. + uncertainty_dims : list of str + Dimensions over which a single missing value marks the whole + ensemble (or set of quantiles) as missing. + + Returns + ------- + xr.DataArray + Boolean mask without the uncertainty dimensions and without non-index + coordinates, so they cannot conflict when masks are combined. + """ + mask = da.isnull().any([d for d in uncertainty_dims if d in da.dims]) + return mask.reset_coords(drop=True) diff --git a/tests/test_nans.py b/tests/test_nans.py new file mode 100644 index 0000000..bba3049 --- /dev/null +++ b/tests/test_nans.py @@ -0,0 +1,253 @@ +"""Tests for ``broadcast_nans``: propagating NaN masks across datasets. + +The datasets are point observations over 4 valid times, with +``temp = 0, 1, 2, 3``. Ensembles add a ``member`` dimension, and member ``m`` +adds ``100 * m``. NaNs are placed by position. +""" + +from collections.abc import Sequence + +import dask.array as dda +import numpy as np +import pandas as pd +import pytest +import xarray as xr +from numpy.typing import ArrayLike + +from mxalign.alignment.nans import broadcast_nans + +TIMES = pd.date_range("2020-01-01", periods=4, freq="6h") +MEMBER_ATTRS = {"standard_name": "realization", "units": "1"} + + +def _traits(uncertainty: str) -> dict: + """Build the ``mlwp_*`` trait attributes for a point observation. + + Parameters + ---------- + uncertainty : str + Value of the ``mlwp_uncertainty_trait`` attribute. + + Returns + ------- + dict + Trait attributes to attach to a dataset's ``attrs``. + """ + return { + "mlwp_space_trait": "point", + "mlwp_time_trait": "observation", + "mlwp_uncertainty_trait": uncertainty, + } + + +def _det(nan_at: Sequence[int] = (), times: pd.DatetimeIndex = TIMES) -> xr.Dataset: + """Build a deterministic dataset with NaNs at the given time positions. + + Parameters + ---------- + nan_at : sequence of int, default () + Positions along ``valid_time`` to set to NaN. + times : pd.DatetimeIndex, default TIMES + Values of the ``valid_time`` coordinate. + + Returns + ------- + xr.Dataset + Dataset with a ``temp`` variable along ``valid_time``. + """ + values = np.arange(len(times), dtype=float) + values[list(nan_at)] = np.nan + return xr.Dataset( + {"temp": ("valid_time", values, {"units": "K"})}, + coords={"valid_time": times}, + attrs=_traits("deterministic"), + ) + + +def _ens( + nan_at: Sequence[tuple[int, int]] = (), members: ArrayLike = (0, 1, 2) +) -> xr.Dataset: + """Build an ensemble with NaNs at the given (member, time) positions. + + Parameters + ---------- + nan_at : sequence of tuple of int, default () + ``(member, valid_time)`` positions to set to NaN. + members : ArrayLike, default (0, 1, 2) + Values of the ``member`` coordinate. + + Returns + ------- + xr.Dataset + Ensemble with a ``temp`` variable over ``member`` and ``valid_time``. + """ + members = np.asarray(members) + values = 100.0 * members[:, np.newaxis] + np.arange(len(TIMES)) + for m, t in nan_at: + values[m, t] = np.nan + return xr.Dataset( + {"temp": (["member", "valid_time"], values, {"units": "K"})}, + coords={"member": ("member", members, MEMBER_ATTRS), "valid_time": TIMES}, + attrs=_traits("ensemble"), + ) + + +def _missing(ds: xr.Dataset) -> list[bool]: + """Return for each valid time whether ``temp`` is missing. + + Also asserts that an ensemble is missing either all its members or none. + + Parameters + ---------- + ds : xr.Dataset + Dataset with a ``temp`` variable along ``valid_time``. + + Returns + ------- + list of bool + Whether ``temp`` is NaN, per valid time. + """ + missing = ds["temp"].isnull() + if "member" in missing.dims: + xr.testing.assert_equal(missing.any("member"), missing.all("member")) + missing = missing.any("member") + return missing.values.tolist() + + +# --------------------------------------------------------------------------- +# Deterministic datasets +# --------------------------------------------------------------------------- + + +class TestDeterministic: + """Group tests of NaN propagation between deterministic datasets.""" + + def test_nan_propagates_both_ways(self) -> None: + """Check that a NaN in either dataset ends up in both.""" + a, b = broadcast_nans([_det(nan_at=[1]), _det(nan_at=[2])]) + + assert _missing(a) == _missing(b) == [False, True, True, False] + + def test_only_shared_times_are_masked(self) -> None: + """Check that values outside the shared coordinates are left alone.""" + a = _det(nan_at=[0]) + b = _det(nan_at=[0], times=TIMES[1:]) # NaN at TIMES[1] + + a, b = broadcast_nans([a, b]) + + assert _missing(a) == [True, True, False, False] + assert _missing(b) == [True, False, False] + + def test_inputs_are_not_modified(self) -> None: + """Check that the input datasets are left unchanged.""" + a, b = _det(nan_at=[1]), _det(nan_at=[2]) + a_before, b_before = a.copy(deep=True), b.copy(deep=True) + + broadcast_nans({"a": a, "b": b}) + + xr.testing.assert_identical(a, a_before) + xr.testing.assert_identical(b, b_before) + + def test_returns_same_container_type(self) -> None: + """Check that a dict, a list and a single dataset come back as given.""" + ds = _det() + + assert list(broadcast_nans({"x": ds, "y": _det()})) == ["x", "y"] + assert isinstance(broadcast_nans([ds, _det()]), list) + assert broadcast_nans(ds) is ds + + def test_variable_in_one_dataset_is_untouched(self) -> None: + """Check that a variable only one dataset has is not masked.""" + a = _det(nan_at=[1]).assign(rh=("valid_time", np.ones(len(TIMES)))) + + a, _ = broadcast_nans([a, _det(nan_at=[2])]) + + assert not a["rh"].isnull().any() + + def test_different_dims_raise(self) -> None: + """Check that a shared variable with other non-uncertainty dims raises.""" + a = _det().expand_dims(level=[500, 850]) + + with pytest.raises(ValueError, match="variable 'temp'"): + broadcast_nans([a, _det()]) + + def test_missing_uncertainty_trait_raises(self) -> None: + """Check that a dataset without an uncertainty trait raises.""" + a = _det() + del a.attrs["mlwp_uncertainty_trait"] + + with pytest.raises(KeyError, match="mlwp_uncertainty_trait"): + broadcast_nans([a, _det()]) + + +# --------------------------------------------------------------------------- +# Ensembles: members are reduced, never matched +# --------------------------------------------------------------------------- + + +class TestEnsemble: + """Group tests of NaN propagation with ensembles.""" + + @pytest.mark.parametrize("ensemble_first", [False, True]) + def test_observation_and_ensemble(self, ensemble_first: bool) -> None: + """Check that one missing member masks all members and the observation. + + Parameters + ---------- + ensemble_first : bool + Whether the ensemble comes before the observation in the input. + """ + datasets = {"obs": _det(nan_at=[1]), "ens": _ens(nan_at=[(0, 2)])} + if ensemble_first: + datasets = dict(reversed(datasets.items())) + + out = broadcast_nans(datasets) + + assert _missing(out["obs"]) == _missing(out["ens"]) + assert _missing(out["obs"]) == [False, True, True, False] + + def test_observation_deterministic_and_ensemble(self) -> None: + """Check that NaNs propagate across a mix of uncertainty types.""" + out = broadcast_nans( + [_det(nan_at=[0]), _det(nan_at=[1]), _ens(nan_at=[(2, 3)])] + ) + + assert all(_missing(ds) == [True, True, False, True] for ds in out) + + def test_members_are_not_matched_between_ensembles(self) -> None: + """Check that ensembles of different sizes keep all their members. + + Member labels of two ensembles are unrelated, so masks are compared + after reducing over ``member``, not per member. + """ + small = _ens(nan_at=[(0, 1)]) + large = _ens(nan_at=[(4, 2)], members=[10, 11, 12, 13, 14]) + + small, large = broadcast_nans([small, large]) + + np.testing.assert_array_equal(small["member"], [0, 1, 2]) + np.testing.assert_array_equal(large["member"], [10, 11, 12, 13, 14]) + assert _missing(small) == _missing(large) == [False, True, True, False] + + def test_dim_order_coords_and_attrs_are_kept(self) -> None: + """Check that a trailing member dim, coord attrs and var attrs survive.""" + ens = _ens(nan_at=[(1, 0)]).transpose("valid_time", "member") + + _, out = broadcast_nans([_det(), ens]) + + assert out["temp"].dims == ("valid_time", "member") + assert out["temp"].attrs == {"units": "K"} + assert out["member"].attrs == MEMBER_ATTRS + assert out.mx.is_ensemble() + assert _missing(out) == [True, False, False, False] + + def test_dask_inputs_stay_lazy(self) -> None: + """Check that dask-backed inputs give dask-backed, correct outputs.""" + obs = _det(nan_at=[1]).chunk({"valid_time": 2}) + ens = _ens(nan_at=[(0, 2)]).chunk({"member": 1, "valid_time": 2}) + + obs, ens = broadcast_nans([obs, ens]) + + assert isinstance(obs["temp"].data, dda.Array) + assert isinstance(ens["temp"].data, dda.Array) + assert _missing(obs) == _missing(ens) == [False, True, True, False] diff --git a/tests/test_runner.py b/tests/test_runner.py index 9bba149..d22bec4 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -50,6 +50,38 @@ def test_align_runs_time_alignment_without_time_section( assert runner.datasets["obs"].mx.is_forecast() +def test_align_mixes_deterministic_and_ensemble( + ds_fcst: xr.Dataset, ds_fcst_ens: xr.Dataset, ds_obs: xr.Dataset +) -> None: + """Check that ``align`` with NaN broadcasting handles an ensemble. + + The deterministic forecast and the observation must end up missing where + the ensemble is, and the ensemble must be missing all or none of its + members. + + Parameters + ---------- + ds_fcst : xr.Dataset + Point forecast fixture. + ds_fcst_ens : xr.Dataset + Point forecast ensemble fixture. + ds_obs : xr.Dataset + Point observation fixture, used as the reference. + """ + ens = ds_fcst_ens.copy(deep=True) + ens["temp"][{"member": 1, "lead_time": 1}] = float("nan") + runner = _runner({"obs": ds_obs, "det": ds_fcst, "ens": ens}, {"reference": "obs"}) + + runner.align() + + out = runner.datasets + missing = out["ens"]["temp"].isnull() + xr.testing.assert_equal(missing.any("member"), missing.all("member")) + for name in ("obs", "det"): + xr.testing.assert_equal(out[name]["temp"].isnull(), missing.any("member")) + assert missing.sel(lead_time="6h").all() + + def test_align_time_rejects_unknown_options() -> None: """Check that unknown ``alignment.time`` options are rejected.""" with pytest.raises(ValidationError, match="alignment.time.return_as"): From b793dccefb21f102b251748ecb323298f9d20f2d Mon Sep 17 00:00:00 2001 From: Michiel Van Ginderachter Date: Tue, 27 Oct 2026 13:27:12 +0000 Subject: [PATCH 12/12] update changelog Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 895c682..65d1be1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,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 @@ -27,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)