From 300e5c6db00f8c05245c68614381b4be3fc26564 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 11:48:21 +0100 Subject: [PATCH 1/5] 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 6d5aca535ca40847697d404d85bba217e3cb65dd Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 12:25:09 +0100 Subject: [PATCH 2/5] 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 80e6cd34a784eeb5b7066ca4ab03bb955bdd728b Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 12:26:26 +0100 Subject: [PATCH 3/5] 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 5986c748ebc4e22f45c1292980bfa2a16a2ce65d Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 13:43:35 +0100 Subject: [PATCH 4/5] 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 228f76c294b6fa50570509f1f387ab27f7c4288c Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Tue, 27 Oct 2026 13:38:57 +0100 Subject: [PATCH 5/5] 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)