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)