diff --git a/src/mxalign/__init__.py b/src/mxalign/__init__.py index 5a76c8d..652ac2e 100644 --- a/src/mxalign/__init__.py +++ b/src/mxalign/__init__.py @@ -7,11 +7,13 @@ from .interpolations.registry import available_interpolations, register_interpolator from .align.time import align_time from .align.space import align_space +from .verification import fused_metric, get_fused_kernel from . import accessors from . import loaders from . import transformations from . import interpolations +from . import scores __all__ = [ "Properties", @@ -29,8 +31,11 @@ "register_interpolator", "align_time", "align_space", + "fused_metric", + "get_fused_kernel", "accessors", "loaders", "transformations", "interpolations", + "scores", ] diff --git a/src/mxalign/_progress.py b/src/mxalign/_progress.py new file mode 100644 index 0000000..6feb6ba --- /dev/null +++ b/src/mxalign/_progress.py @@ -0,0 +1,170 @@ +"""Lightweight progress + diagnostics helpers for mxalign (Phase 0). + +All output goes to the 'mxalign' logger at INFO, single-line key=value +format so it is grep-friendly in SLURM logs. + +Helpers degrade silently if dask.distributed / psutil are unavailable. +""" +from __future__ import annotations + +import logging +import threading +import time + +LOG = logging.getLogger("mxalign") + + +def count_tasks(obj) -> int: + """Total task count across a dask-backed xarray/dask collection. + + Sums per-layer counts from the HighLevelGraph; avoids materializing the + full task dict (which can itself be slow for huge graphs). + """ + try: + graph = obj.__dask_graph__() + except AttributeError: + return 0 + try: + return sum(len(layer) for layer in graph.layers.values()) + except AttributeError: + try: + return len(dict(graph)) + except Exception: + return -1 + + +def _get_client(): + try: + from dask.distributed import default_client + return default_client() + except Exception: + return None + + +def _worker_rss_summary(client): + """Return (max_gb, mean_gb, n_workers) for current worker RSS, or None.""" + try: + import psutil # noqa: F401 + except ImportError: + return None + try: + rss = client.run( + lambda: __import__("psutil").Process().memory_info().rss + ) + except Exception: + return None + values = [v for v in rss.values() if isinstance(v, (int, float))] + if not values: + return None + n = len(values) + return max(values) / 1e9, (sum(values) / n) / 1e9, n + + +def _n_workers(client): + """Number of workers currently visible to the client (or -1 if unknown).""" + try: + info = client.scheduler_info() + except Exception: + return -1 + workers = info.get("workers") if isinstance(info, dict) else None + return len(workers) if workers is not None else -1 + + +def _pending_tasks(client): + try: + processing = client.processing() + except Exception: + return None + try: + return sum(len(v) for v in processing.values()) + except Exception: + return None + + +class ProgressTicker: + """Context manager: spawn a daemon thread emitting periodic status lines.""" + + def __init__(self, tag: str, interval: float = 15.0): + self.tag = tag + self.interval = interval + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._t0 = 0.0 + + def __enter__(self): + self._t0 = time.perf_counter() + self._thread = threading.Thread( + target=self._run, daemon=True, name=f"mxalign-tick-{self.tag}" + ) + self._thread.start() + return self + + def __exit__(self, exc_type, exc, tb): + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=2 * self.interval) + + def _run(self): + client = _get_client() + last_tick = time.perf_counter() + last_n_workers: int | None = None + warned_no_workers = False + while not self._stop.wait(self.interval): + now = time.perf_counter() + elapsed = now - self._t0 + delta = now - last_tick + last_tick = now + parts = [ + f"[mxalign] tick phase={self.tag}", + f"elapsed={elapsed:.1f}s", + f"since_last_tick={delta:.1f}s", + ] + n_workers = -1 + if client is not None: + pending = _pending_tasks(client) + if pending is not None: + parts.append(f"pending={pending}") + n_workers = _n_workers(client) + parts.append(f"workers={n_workers}") + rss = _worker_rss_summary(client) + if rss is not None: + max_gb, mean_gb, _ = rss + parts.append(f"rss_max_gb={max_gb:.2f}") + parts.append(f"rss_mean_gb={mean_gb:.2f}") + LOG.info(" ".join(parts)) + if ( + client is not None + and not warned_no_workers + and last_n_workers is not None + and last_n_workers > 0 + and n_workers == 0 + ): + LOG.warning( + "[mxalign] no workers visible to client (was %d); scheduler " + "likely lost the worker (e.g. heartbeat timeout). Subsequent " + "ticks will report stale state.", + last_n_workers, + ) + warned_no_workers = True + last_n_workers = n_workers + + +def log_phase_start(tag: str, **kv) -> None: + extras = " ".join(f"{k}={v}" for k, v in kv.items()) + LOG.info(f"[mxalign] phase={tag} status=start {extras}".rstrip()) + + +def log_phase_done(tag: str, elapsed: float, **kv) -> None: + extras = " ".join(f"{k}={v}" for k, v in kv.items()) + LOG.info( + f"[mxalign] phase={tag} status=done elapsed={elapsed:.2f}s {extras}".rstrip() + ) + + +def log_dashboard() -> None: + client = _get_client() + if client is None: + return + link = getattr(client, "dashboard_link", None) + if link: + LOG.info(f"[mxalign] dask dashboard={link}") diff --git a/src/mxalign/loaders/anemoi_datasets.py b/src/mxalign/loaders/anemoi_datasets.py index 3ccb646..d8da9bd 100644 --- a/src/mxalign/loaders/anemoi_datasets.py +++ b/src/mxalign/loaders/anemoi_datasets.py @@ -53,6 +53,85 @@ def _load(self): ) return ds_selected.to_dataset(dim="variable") + def slice(self, reference_time, lead_times, variables): + """Eagerly read one (reference_time, lead_times, variables) slice. + + See ``BaseLoader.slice`` for the contract. Only the single-file + zarr path is supported; multi-file returns ``None``. + """ + if isinstance(self.files, list): + if len(self.files) != 1: + return None + path = self.files[0] + else: + path = self.files + + src = xr.open_zarr(path, consolidated=False) + + # 'dates' coord on the 'time' dim is the canonical valid_time array. + valid_times = src["dates"].astype("datetime64[ns]").load().values + var_attr = list(src.attrs["variables"]) + try: + var_idx = np.array( + [var_attr.index(v) for v in variables], dtype=np.int64 + ) + except ValueError as e: + raise ValueError( + f"{type(self).__name__}.slice: variable not found in {path}: {e}" + ) from None + + rt = np.datetime64(reference_time, "ns") + requested_vts = np.array( + [rt + np.timedelta64(lt, "ns") for lt in lead_times], + dtype="datetime64[ns]", + ) + pos = np.searchsorted(valid_times, requested_vts) + if pos.max() >= valid_times.size or not np.all( + valid_times[pos] == requested_vts + ): + bad = requested_vts[ + (pos >= valid_times.size) + | (valid_times[pos.clip(max=valid_times.size - 1)] != requested_vts) + ] + raise ValueError( + f"{type(self).__name__}.slice: missing valid_times in {path}: " + f"{bad[:5]}... (reference_time={reference_time})" + ) + + arr = src["data"].isel(ensemble=0) + # If `pos` is strictly contiguous (the common case: 1 h cadence + # lead_times), issue a single slice read instead of a fancy index. + # Fancy indexing along `time` triggers one chunk read per requested + # step per variable, which on finely-time-chunked zarrs blows up + # into thousands of small reads. A slice is a single contiguous + # request and avoids that amplification entirely. + pos_arr = np.asarray(pos) + if pos_arr.size == 0: + contiguous = False + elif pos_arr.size == 1: + contiguous = True + else: + contiguous = bool(np.all(np.diff(pos_arr) == 1)) + if contiguous: + start = int(pos_arr[0]) + stop = int(pos_arr[-1]) + 1 + arr_sel = arr.isel( + time=slice(start, stop), + variable=xr.DataArray(var_idx, dims="variable_out"), + ) + else: + arr_sel = arr.isel( + time=xr.DataArray(pos_arr, dims="lead_time"), + variable=xr.DataArray(var_idx, dims="variable_out"), + ) + vals = np.asarray(arr_sel.load().values) # (n_lt, n_var, n_grid) + return xr.Dataset( + { + v: (("lead_time", "grid_index"), vals[:, i, :]) + for i, v in enumerate(variables) + } + ) + def _postprocess(dataset: xr.Dataset) -> xr.Dataset: """Post-process the dataset to add coordinates and drop unused variables. diff --git a/src/mxalign/loaders/anemoi_inference.py b/src/mxalign/loaders/anemoi_inference.py index 68f8799..873aee4 100644 --- a/src/mxalign/loaders/anemoi_inference.py +++ b/src/mxalign/loaders/anemoi_inference.py @@ -1,11 +1,17 @@ +from datetime import datetime from pathlib import Path + +import numpy as np import xarray as xr from .registry import register_loader from ..properties.properties import Space, Time, Uncertainty from .base import BaseLoader -DEFAULTS_NETCDF = {"chunks": "auto", "engine": "h5netcdf", "parallel": True} +DEFAULTS_NETCDF = {"chunks": "auto", "engine": "h5netcdf", "parallel": True, "identical_layout": True} + +# Variables that are static spatial fields, not per-timestep forecasts. +_SPATIAL_VARS = frozenset({"latitude", "longitude"}) DEFAULTS_ZARR = { "chunks": "auto", @@ -53,26 +59,245 @@ def _load(self): loader = _open_mf_dataset + # Pass reference_times hint to fast path so it doesn't need to parse + # filenames. Consumed (popped) inside _open_mf_dataset; ignored by + # _open_zarr. + if loader is _open_mf_dataset and self.reference_times is not None: + kwargs["_reference_times"] = np.asarray(self.reference_times) + ds = loader(files, **kwargs) return ds + def _files_by_rt(self): + """Return ({rt_ns: path}, engine) or (None, engine) if unsupported. + + Shared by ``slice`` and ``prefetch_path``: maps each reference_time + (np.datetime64[ns]) to the .nc file path containing that forecast. + Only the per-file netCDF path is supported (single-zarr forecasts + return None). + """ + files = [self.files] if isinstance(self.files, str) else list(self.files) + engine = self.kwargs.get("engine", DEFAULTS_NETCDF["engine"]) + if not files: + return None, engine + if Path(files[0]).suffix.lower() == ".zarr": + return None, engine + + if self.reference_times is not None and len(self.reference_times) == len(files): + rt_values = [np.datetime64(rt, "ns") for rt in self.reference_times] + else: + rt_values = [] + for f in files: + try: + rt_values.append( + np.datetime64( + datetime.strptime(Path(f).stem, "%Y-%m-%dT%H"), "ns" + ) + ) + except ValueError: + return None, engine + + return ( + {int(rt.astype("int64")): f for rt, f in zip(rt_values, files)}, + engine, + ) + + def prefetch_path(self, reference_time): + """Return the on-disk path the fused engine should warm into the OS + page cache for ``reference_time``, or ``None`` if not available + (single-zarr forecasts, unparseable filenames). + """ + files_by_rt, _ = self._files_by_rt() + if files_by_rt is None: + return None + rt_ns = int(np.datetime64(reference_time, "ns").astype("int64")) + return files_by_rt.get(rt_ns) + + def slice(self, reference_time, lead_times, variables): + """Eagerly read one (reference_time, lead_times, variables) slice. + + See ``BaseLoader.slice`` for the contract. Returns ``None`` for + single-zarr forecasts (use the regular loader path) or when + filename-based rt resolution fails. + """ + files_by_rt, engine = self._files_by_rt() + if files_by_rt is None: + return None + rt_ns = int(np.datetime64(reference_time, "ns").astype("int64")) + if rt_ns not in files_by_rt: + raise ValueError( + f"{type(self).__name__}.slice: reference_time {reference_time} " + f"not found in files mapping (have {len(files_by_rt)} files)." + ) + path = files_by_rt[rt_ns] + + with xr.open_dataset(path, engine=engine) as src: + # Subset variables + lead_times *lazily* and only then call + # .load(). Doing .load() up front forces a read of the full + # time axis and turns the time-axis selection into an + # in-memory fancy index instead of a hyperslab read. + sub = src[list(variables)] + if "time" in sub.dims: + times = sub["time"].values + lts = (times - times[0]).astype("timedelta64[ns]") + sub = sub.assign_coords({"lead_time": ("time", lts)}).swap_dims( + {"time": "lead_time"} + ) + if "values" in sub.dims: + sub = sub.rename_dims({"values": "grid_index"}) + requested = np.asarray( + [np.timedelta64(lt, "ns") for lt in lead_times], + dtype="timedelta64[ns]", + ) + file_lts = sub["lead_time"].values.astype("timedelta64[ns]") + pos = np.searchsorted(file_lts, requested) + if pos.max() >= file_lts.size or not np.all(file_lts[pos] == requested): + bad = requested[ + (pos >= file_lts.size) + | (file_lts[pos.clip(max=file_lts.size - 1)] != requested) + ] + raise ValueError( + f"{type(self).__name__}.slice: missing lead_times in " + f"{path}: {bad[:5]}... (reference_time={reference_time})" + ) + # Contiguous fast path → hyperslab; otherwise fancy index. + pos_arr = np.asarray(pos) + if pos_arr.size == 0: + contiguous = False + elif pos_arr.size == 1: + contiguous = True + else: + contiguous = bool(np.all(np.diff(pos_arr) == 1)) + if contiguous: + sub = sub.isel( + lead_time=slice(int(pos_arr[0]), int(pos_arr[-1]) + 1) + ) + else: + sub = sub.isel(lead_time=xr.DataArray(pos_arr, dims="lead_time")) + return sub.load() + + +def _load_nc_vars(path, var_names, engine): + """Load all named variables from one NC file. + + Executed by dask workers at compute-time (not graph-build time), so 358 + files are opened in parallel across dask threads rather than serially + during graph construction. + + Returns a dict {var_name: np.ndarray shape (n_lt, n_grid)}. + """ + ds = xr.open_dataset(path, engine=engine) + result = {v: ds[v].values for v in var_names} + ds.close() + return result + + +def _load_nc_var(path, var_name, engine): + """Load a single variable from one NC file. + + One delayed task per (file, variable) pair: each result is ~23 MB instead + of ~1.5 GB per file. Without a shared intermediate dict there is no + dependency forcing all 65 variable results to stay in memory at once, + so peak worker memory scales with concurrency (O(n_threads × chunk_size)) + rather than with n_files × file_size. + """ + ds = xr.open_dataset(path, engine=engine) + result = ds[var_name].values + ds.close() + return result + + def _open_mf_dataset(files, **kwargs): + identical_layout = kwargs.pop("identical_layout", True) + # Reference times from the blueprint config (sorted datetime64 array, + # index-aligned with the sorted files list). Preferred over filename + # parsing; absent when the loader is called outside the blueprint system. + reference_times_hint = kwargs.pop("_reference_times", None) + engine = kwargs.get("engine", "h5netcdf") - times = xr.open_dataset(files[0], engine=kwargs["engine"], chunks=kwargs["chunks"])[ - "time" - ].values - lead_times = times - times[0] + # Always open file 0: needed for lead_times (and schema in fast path). + ds0 = xr.open_dataset(files[0], engine=engine) + times0 = ds0["time"].values + lead_times = times0 - times0[0] - ds = xr.open_mfdataset(files, preprocess=_preprocess, **kwargs) + if not identical_layout or len(files) == 1: + ds0.close() + ds = xr.open_mfdataset(files, preprocess=_preprocess, **kwargs) + return ( + ds.assign_coords({"lead_time": ("time", lead_times)}) + .rename_dims({"values": "grid_index"}) + .swap_dims({"time": "lead_time"}) + ) - ds_out = ( - ds.assign_coords({"lead_time": ("time", lead_times)}) - .rename_dims({"values": "grid_index"}) - .swap_dims({"time": "lead_time"}) - ) + # ------------------------------------------------------------------ + # Fast path: identical_layout=True + # Build a lazy dataset without opening files[1:]. Each file's data + # becomes a dask.delayed task executed at compute-time. Only the + # schema (shape, dtype, coords) is read here, from file 0 only. + # ------------------------------------------------------------------ + import dask + import dask.array as dsa - return ds_out + data_vars = tuple(v for v in ds0.data_vars if v not in _SPATIAL_VARS) + lat = ds0["latitude"].values + lon = ds0["longitude"].values + n_lt, n_grid = ds0[data_vars[0]].shape # (time, values) + dtype = ds0[data_vars[0]].dtype + ds0.close() + + # Resolve a reference_time for each file. + # Primary: use the blueprint-provided array (format-agnostic, no I/O). + # Fallback: parse from filename stem (ISO-8601: 2023-01-01T00.nc). + # If neither works, abort the fast path. + if reference_times_hint is not None and len(reference_times_hint) == len(files): + ref_time_list = [np.datetime64(rt, "ns") for rt in reference_times_hint] + else: + ref_time_list = [] + for f in files: + try: + ref_time_list.append( + np.datetime64(datetime.strptime(Path(f).stem, "%Y-%m-%dT%H"), "ns") + ) + except ValueError: + import warnings + warnings.warn( + f"identical_layout=True: cannot parse reference_time from " + f"{Path(f).name!r}; falling back to open_mfdataset", + stacklevel=2, + ) + return _open_mf_dataset(files, identical_layout=False, **kwargs) + + individual_dss = [] + for f, ref_time in zip(files, ref_time_list): + # One delayed task per (file, variable): no shared intermediate dict, + # so dask can free each ~23 MB result immediately after its consumer + # finishes instead of holding a ~1.5 GB per-file dict until all 65 + # getitem tasks complete. + ds_vars = { + v: xr.DataArray( + dsa.from_delayed( + dask.delayed(_load_nc_var)(f, v, engine), + shape=(n_lt, n_grid), + dtype=dtype, + ), + dims=["lead_time", "grid_index"], + ) + for v in data_vars + } + + ds_individual = ( + xr.Dataset(ds_vars) + .assign_coords({ + "lead_time": lead_times, + "latitude": ("grid_index", lat), + "longitude": ("grid_index", lon), + }) + .expand_dims({"reference_time": [ref_time]}) + ) + individual_dss.append(ds_individual) + + return xr.concat(individual_dss, dim="reference_time", coords="minimal", join="override") def _open_zarr(files, **kwargs): diff --git a/src/mxalign/loaders/base.py b/src/mxalign/loaders/base.py index c7020b5..83fa00e 100644 --- a/src/mxalign/loaders/base.py +++ b/src/mxalign/loaders/base.py @@ -1,5 +1,7 @@ from abc import ABC, abstractmethod +import numpy as np + from .registry import register_loader from ..properties.properties import Properties, Space, Time, Uncertainty from ..properties.validation import validate_dataset @@ -15,10 +17,21 @@ class BaseLoader(ABC): time: Time | None = None uncertainty: Uncertainty | None = None - def __init__(self, files, variables=None, grid_mapping=None, **kwargs): + def __init__(self, files, variables=None, grid_mapping=None, + valid_times=None, reference_times=None, lead_times=None, + **kwargs): self.files = files self.variables = [variables] if isinstance(variables, str) else variables self.grid_mapping = grid_mapping + # Optional pre-pruning hints; consumed in load(), not forwarded + # to backend kwargs (which would explode for unknown args). + # - valid_times: 1D datetime64 set, used to prune observation + # datasets carrying a `valid_time` dim. + # - reference_times / lead_times: 1D arrays defining the + # allowed rectangular (rt, lt) window for forecast datasets. + self.valid_times = valid_times + self.reference_times = reference_times + self.lead_times = lead_times self.kwargs = kwargs def load(self): @@ -26,6 +39,46 @@ def load(self): if self.variables: ds = self._select_variables(ds) + # Generic time pre-pruning. Applied here (after _load / variable + # selection, before properties/validation) so every loader benefits + # without needing to know about the time hints. Dask's culling will + # drop the unused upstream chunks at execution time. + # + # Two cases: + # - Observation datasets carry `valid_time` as a 1D dim, pruned + # against `self.valid_times`. + # - Forecast datasets carry `reference_time` and `lead_time` as + # dims; pruned rectangularly against `self.reference_times` and + # `self.lead_times`. This enforces the blueprint's `dates.range` + # (max lead time) and `dates.period` (rt spacing) without the + # spurious over-keep that an axis-independent mask derived from + # `valid_times` would produce for commensurate spacings. + if "valid_time" in ds.dims and self.valid_times is not None: + wanted = np.asarray(self.valid_times) + keep = np.intersect1d(wanted, ds["valid_time"].values) + if keep.size and keep.size < ds["valid_time"].size: + all_vt = ds["valid_time"].values + positions = np.searchsorted(all_vt, keep) + if positions[-1] - positions[0] == len(positions) - 1: + # contiguous block — isel with a slice keeps the dask graph + # small (only the needed chunks, not all 403K zarr tasks) + ds = ds.isel(valid_time=slice(int(positions[0]), int(positions[-1]) + 1)) + else: + ds = ds.sel(valid_time=keep) + elif {"reference_time", "lead_time"} <= set(ds.dims): + if self.lead_times is not None: + lt = ds["lead_time"].values + wanted_lt = np.asarray(self.lead_times).astype(lt.dtype) + keep_lt = np.isin(lt, wanted_lt) + if keep_lt.any() and keep_lt.sum() < lt.size: + ds = ds.isel(lead_time=keep_lt) + if self.reference_times is not None: + rt = ds["reference_time"].values + wanted_rt = np.asarray(self.reference_times).astype(rt.dtype) + keep_rt = np.isin(rt, wanted_rt) + if keep_rt.any() and keep_rt.sum() < rt.size: + ds = ds.isel(reference_time=keep_rt) + properties = self._get_properties(ds) validate_dataset(ds, properties) @@ -57,6 +110,31 @@ def _get_properties(self, ds): ) return properties + def slice(self, reference_time, lead_times, variables): + """Eagerly load a single per-reference_time slice. + + Returns an in-memory ``xr.Dataset`` with dims ``(lead_time, grid_index)`` + and one data variable per name in ``variables``. No ``reference_time`` + dim (callers iterate reference_times). Bypasses xarray/dask lazy + graphs by reading directly from the underlying store. + + Parameters + ---------- + reference_time + datetime64-coercible. Forecast initial time to read. + lead_times + Sequence of timedelta64-coercible offsets from reference_time. + variables + List of variable names to read. Required. + + Returns + ------- + xr.Dataset | None + None if this loader cannot serve a per-rt slice (e.g. unsupported + file layout). The fused engine treats None as a hard error. + """ + return None + @register_loader class MxAlignLoader(BaseLoader): diff --git a/src/mxalign/runner.py b/src/mxalign/runner.py index aff4398..d6d2833 100644 --- a/src/mxalign/runner.py +++ b/src/mxalign/runner.py @@ -1,20 +1,34 @@ import os +import time import xarray as xr from .utils.config import Config -from .loaders.loader import load +from .loaders.loader import load # noqa: F401 (kept for external API back-compat) +from .loaders.registry import get_loader from .transformations.transform import transform from .align.time import align_time from .align.space import align_space from .align.nans import broadcast_nans from .utils.save import save_dataset, save_metrics from .verification import Metric +from ._progress import ( + ProgressTicker, + count_tasks, + log_dashboard, + log_phase_done, + log_phase_start, +) class Runner: def __init__(self, config: str | dict): self.config = Config(config) self.datasets = {} + # Bookkeeping required by the fused verification engine. + # Populated by load_datasets / transform_datasets and ignored by + # the legacy xarray engine. + self.loaders: dict[str, object] = {} + self._transforms_by_ds: dict[str, list[tuple[str, dict]]] = {} def run(self): # 1. Load the datasets @@ -32,7 +46,7 @@ def load_datasets(self): for name, config_ds in config.items(): config_ds = config_ds.copy() # Check if all the files exist - loader = config_ds.pop("loader") + loader_name = config_ds.pop("loader") variables = config_ds.pop("variables", None) grid_mapping = config_ds.pop("grid_mapping", None) files = [] @@ -42,13 +56,15 @@ def load_datasets(self): files.append(file) else: print(f"File: {file} is missing, skipping.") - self.datasets[name] = load( - name=loader, - files=files, + loader_cls = get_loader(loader_name) + loader_inst = loader_cls( + files, variables=variables, grid_mapping=grid_mapping, **config_ds, ) + self.datasets[name] = loader_inst.load() + self.loaders[name] = loader_inst def transform_datasets(self): config = self.config["transformations"] @@ -60,8 +76,16 @@ def transform_datasets(self): names_ds = config_trans.pop("datasets", self.datasets.keys()) for name in names_ds: ds = self.datasets[name] + # Expand glob patterns / fill defaults before executing and + # recording, so downstream engines always see concrete names. + from .transformations.registry import get_expander + expander = get_expander(transformation) + resolved = expander(ds, config_trans) if expander else config_trans self.datasets[name] = transform( - name=transformation, datasets=ds, **config_trans + name=transformation, datasets=ds, **resolved + ) + self._transforms_by_ds.setdefault(name, []).append( + (transformation, dict(resolved)) ) def align(self): @@ -116,7 +140,75 @@ def verify(self): common_vars.intersection_update(set(ds.data_vars)) common_vars = list(common_vars) - if config_metrics: + rechunk_lead_time = config.get("rechunk_lead_time", True) + engine = config.get("engine", "xarray") + + if config_metrics and engine == "fused": + from .verification_fused import compute_metrics_fused + log_phase_start( + "verify-build", + engine="fused", + n_models=len(self.datasets) - 1, + n_metrics=len(config_metrics), + n_vars=len(common_vars), + n_rt=int(reference.sizes.get("reference_time", -1)), + n_lt=int(reference.sizes.get("lead_time", -1)), + ) + t_build = time.perf_counter() + self.metrics = compute_metrics_fused( + datasets=self.datasets, + loaders=self.loaders, + transforms_by_ds=self._transforms_by_ds, + reference_name=config["reference"], + common_vars=common_vars, + metrics_cfg=config["metrics"], + engine_cfg=config, + ) + log_phase_done( + "verify-build+exec", + time.perf_counter() - t_build, + engine="fused", + ) + elif config_metrics and engine == "fused_collect": + from .verification_fused import compute_metrics_collect + log_phase_start( + "verify-build", + engine="fused_collect", + n_models=len(self.datasets) - 1, + n_metrics=len(config_metrics), + n_rt=int(reference.sizes.get("reference_time", -1)), + n_lt=int(reference.sizes.get("lead_time", -1)), + ) + t_build = time.perf_counter() + self.metrics = compute_metrics_collect( + datasets=self.datasets, + loaders=self.loaders, + transforms_by_ds=self._transforms_by_ds, + reference_name=config["reference"], + metrics_cfg=config["metrics"], + engine_cfg=config, + ) + log_phase_done( + "verify-build+exec", + time.perf_counter() - t_build, + engine="fused_collect", + ) + elif config_metrics: + log_phase_start( + "verify-build", + n_models=len(self.datasets) - 1, + n_metrics=len(config_metrics), + n_vars=len(common_vars), + n_rt=int(reference.sizes.get("reference_time", -1)), + n_lt=int(reference.sizes.get("lead_time", -1)), + rechunk_lead_time=rechunk_lead_time, + ) + t_build = time.perf_counter() + ds_ref_for_metric = ( + _rechunk_for_metric(reference[common_vars]) + if rechunk_lead_time + else reference[common_vars] + ) metrics = {} for metric_name, config_metric in config["metrics"].items(): config_metric = config_metric.copy() @@ -126,14 +218,19 @@ def verify(self): metric = Metric( name=metric_name, func_path=func_path, - ds_ref=reference[common_vars], + ds_ref=ds_ref_for_metric, inputs=inputs, **config_metric, ) models = {} for ds_name, ds in self.datasets.items(): if ds_name != config["reference"]: - models[ds_name] = metric.compute(ds[common_vars]) + ds_slice = ( + _rechunk_for_metric(ds[common_vars]) + if rechunk_lead_time + else ds[common_vars] + ) + models[ds_name] = metric.compute(ds_slice) models = xr.concat( models.values(), dim=xr.Variable("model", list(models.keys())) ) @@ -141,7 +238,19 @@ def verify(self): metrics = xr.concat( metrics.values(), dim=xr.Variable("metric", list(metrics.keys())) ) - self.metrics = metrics.transpose("model", "metric", ...).compute() + metrics_lazy = metrics.transpose("model", "metric", ...) + n_tasks = count_tasks(metrics_lazy) + log_phase_done( + "verify-build", + time.perf_counter() - t_build, + n_tasks=n_tasks, + ) + log_dashboard() + log_phase_start("verify-exec") + t_exec = time.perf_counter() + with ProgressTicker("verify-exec"): + self.metrics = metrics_lazy.compute() + log_phase_done("verify-exec", time.perf_counter() - t_exec) if config_save_metrics: config = config_save_metrics.copy() @@ -165,3 +274,24 @@ def get_spatial_alignment(ds, reference): if reference.space.is_grid() and ds.space.is_grid(): return "regrid" return "null" + + +def _rechunk_for_metric(ds: xr.Dataset) -> xr.Dataset: + """Rechunk to (reference_time=1, lead_time=-1, ...) before metric graph build. + + This aligns the ERA5 observation chunks (typically (1,1,40320) after time + alignment) with the forecast chunks (1, n_lt, n_grid) produced by the + anemoi-inference loader. Without this, xarray/dask fans out 144 tasks per + (reference_time, variable) cell when it tries to broadcast mismatched + lead_time chunks, turning an O(N_rt) graph into an O(N_rt * N_lt) one. + + Only rechunks dims that are present; leaves grid_index at its natural size. + """ + chunks: dict[str, int] = {} + if "reference_time" in ds.dims: + chunks["reference_time"] = 1 + if "lead_time" in ds.dims: + chunks["lead_time"] = -1 + if not chunks: + return ds + return ds.chunk(chunks) diff --git a/src/mxalign/scores.py b/src/mxalign/scores.py new file mode 100644 index 0000000..7a3c8d1 --- /dev/null +++ b/src/mxalign/scores.py @@ -0,0 +1,280 @@ +"""Built-in metric implementations with a fused-engine fast path. + +Each metric here is a regular Python function that operates on xarray +objects — usable under ``engine: xarray`` exactly like ``scores.continuous.*`` +or ``xskillscore.*``. They are additionally decorated with +:func:`mxalign.fused_metric`, attaching a numpy kernel + finalizer that the +fused verification engine discovers and uses on the per-reference_time +fast path. + +YAML usage:: + + metrics: + mse: + function: mxalign.scores.mse + inputs: {fcst: forecast, obs: reference} + reduce_dims: [reference_time] + +Switching to a different backend (``scores.continuous.mse``, +``xskillscore.mse``, ...) stays a one-line config change; the fused fast +path is only available when ``function:`` resolves to a function decorated +with ``@fused_metric``. +""" +from __future__ import annotations + +import numpy as np + +from .verification import fused_metric + + +def _np32(a): + return np.asarray(a, dtype=np.float32) if not isinstance(a, np.ndarray) else ( + a if a.dtype == np.float32 else a.astype(np.float32, copy=False) + ) + + +def _kernel_squared_error(fcst, obs): + diff = _np32(fcst) - _np32(obs) + return diff * diff + + +def _kernel_absolute_error(fcst, obs): + return np.abs(_np32(fcst) - _np32(obs)) + + +def _kernel_error(fcst, obs): + return _np32(fcst) - _np32(obs) + + +def _finalize_mean(partial_sum, n): + return partial_sum / n + + +@fused_metric(kernel=_kernel_squared_error, finalize=_finalize_mean) +def mse(fcst, obs, reduce_dims=None, **_): + """Mean squared error along ``reduce_dims``.""" + diff = fcst - obs + return (diff * diff).mean(dim=reduce_dims) + + +@fused_metric(kernel=_kernel_absolute_error, finalize=_finalize_mean) +def mae(fcst, obs, reduce_dims=None, **_): + """Mean absolute error along ``reduce_dims``.""" + return abs(fcst - obs).mean(dim=reduce_dims) + + +@fused_metric(kernel=_kernel_error, finalize=_finalize_mean) +def bias(fcst, obs, reduce_dims=None, **_): + """Mean error (forecast minus observation) along ``reduce_dims``.""" + return (fcst - obs).mean(dim=reduce_dims) + + +# Alias: ``mean_error`` is the same thing as ``bias`` in this context, kept +# so existing YAMLs that say ``function: scores.continuous.mean_error`` can +# migrate to ``function: mxalign.scores.mean_error`` without semantic drift. +mean_error = bias + + +# --------------------------------------------------------------------------- +# NVV group-pattern helpers (also imported by the fused_collect engine) +# --------------------------------------------------------------------------- + +def _extract_captures(pattern, ds_vars): + """Return the set of ``*`` captures for a single-wildcard pattern. + + Given a pattern such as ``"*_x"`` and a collection of variable names, + returns every string ``c`` such that ``pattern.replace("*", c)`` is + present in ``ds_vars``. Literal patterns (no ``*``) return + ``{pattern}`` if the name exists, else an empty set. + """ + if "*" not in pattern: + return {pattern} if pattern in ds_vars else set() + idx = pattern.index("*") + prefix, suffix = pattern[:idx], pattern[idx + 1:] + if "*" in suffix: + raise ValueError( + f"Multiple wildcards are not supported in component patterns: {pattern!r}" + ) + captures: set[str] = set() + for v in ds_vars: + if not v.startswith(prefix): + continue + rest = v[len(prefix):] + if suffix: + if not rest.endswith(suffix): + continue + captures.add(rest[: -len(suffix)]) + else: + captures.add(rest) + return captures + + +def _expand_group_patterns(variables_cfg, ds_vars): + """Expand glob-keyed entries in an NVV ``variables:`` config dict. + + A group is treated as a glob entry when its key or any component + pattern contains ``*``. For each such entry, all values of the + wildcard are found such that **every** component pattern evaluates to + an existing variable name (present in ``ds_vars``); a concrete group + is created for each valid capture. + + Literal-keyed entries (no ``*`` in key or components) are kept as-is + and override any glob-generated group that has the same resolved name. + + Parameters + ---------- + variables_cfg : dict + Raw ``variables:`` dict from the YAML metric config. + ds_vars : iterable of str + Variable names available in the dataset at evaluation time. + + Returns + ------- + dict + Fully-resolved ``{group_name: {"components": [...], ...}}`` dict + suitable for direct use in :func:`nvv`. + """ + ds_set = set(ds_vars) + literal: dict = {} + globs: list = [] + for key, grp_cfg in variables_cfg.items(): + comps = grp_cfg.get("components", []) + if "*" in key or any("*" in c for c in comps): + globs.append((key, grp_cfg)) + else: + literal[key] = grp_cfg + + expanded: dict = {} + for key_tmpl, grp_cfg in globs: + comps = grp_cfg.get("components", []) + # Intersection of captures that satisfy ALL component patterns. + valid: set[str] | None = None + for comp in comps: + caps = _extract_captures(comp, ds_set) + valid = caps if valid is None else valid & caps + if not valid: + continue + extra = {k: v for k, v in grp_cfg.items() if k != "components"} + for cap in sorted(valid): + name = key_tmpl.replace("*", cap) + expanded[name] = {"components": [c.replace("*", cap) for c in comps], **extra} + + expanded.update(literal) # literal entries override glob-generated ones + return expanded + + +def nvv(fcst, obs, reduce_dims=None, components=None, variables=None, + eps=0.0, label=None, **_): + """Normalized Vector Variance. + + Compares the temporal spread of one or more vector fields in forecast vs + observation. For each group of components ``i``, the vector variance is:: + + VV = sqrt( sum_i Var(x_i) ) + + where ``Var(x_i)`` is the variance of component ``i`` along + ``reduce_dims``. The NVV is:: + + NVV = VV_fcst / max(VV_obs, eps) + + Two calling forms are supported: + + **Single group** (``components`` key):: + + metrics: + nvv_wind10m: + function: mxalign.scores.nvv + inputs: {fcst: forecast, obs: reference} + reduce_dims: [reference_time] + components: [10u, 10v] + label: wind10m # optional; default "10u+10v" + + **Multiple groups** (``variables`` key):: + + metrics: + nvv: + function: mxalign.scores.nvv + inputs: {fcst: forecast, obs: reference} + reduce_dims: [reference_time, grid_index] + variables: + z_500_grad: + components: [z_500_grad_x, z_500_grad_y] + q_500_grad: + components: [q_500_grad_x, q_500_grad_y] + eps: 1.0e-5 # per-group override (optional) + + Parameters + ---------- + fcst, obs : xr.Dataset + Forecast and observation datasets. + reduce_dims : str or list of str + Dimension(s) to reduce over (e.g. ``["reference_time"]``). + components : list of str + Single-group form: variable names forming the vector. + variables : dict + Multi-group form: ``{label: {components: [...], eps: ...}, ...}``. + Top-level ``eps`` is the default for any group that omits it. + eps : float + Zero-guard threshold, in the units of VV (not VV squared). NVV is set + to NaN wherever ``VV_obs <= eps``. Default ``0.0`` masks only an + exactly-zero observation spread. Note: this is *not* a floor added to + the denominator (which would corrupt small-magnitude fields such as + specific-humidity gradients); it only masks undefined ratios. + label : str, optional + Coordinate value for the synthetic ``variable`` dim in single-group + form. Defaults to ``"+".join(components)``. + + Returns + ------- + xr.DataArray + Shape ``(variable, ...)`` where ``variable`` has one entry per group + and ``...`` are whatever dims remain after reducing over + ``reduce_dims``. + """ + import xarray as xr + + # ---- normalise to the grouped form ------------------------------------ + if variables is not None: + ds_vars = list(fcst.data_vars) if isinstance(fcst, xr.Dataset) else [] + # Expand glob patterns (e.g. "*": {components: ["*_x", "*_y"]}). + needs_expansion = any( + "*" in k or any("*" in c for c in v.get("components", [])) + for k, v in variables.items() + ) + resolved = _expand_group_patterns(variables, ds_vars) if needs_expansion else variables + groups = { + grp_label: { + "components": grp_cfg["components"], + "eps": grp_cfg.get("eps", eps), + } + for grp_label, grp_cfg in resolved.items() + } + elif components: + coord = label if label is not None else "+".join(components) + groups = {coord: {"components": components, "eps": eps}} + else: + raise ValueError( + "mxalign.scores.nvv requires either `components:` (single group) " + "or `variables:` (multiple groups) to be set." + ) + + rd = [reduce_dims] if isinstance(reduce_dims, str) else list(reduce_dims or []) + + results = {} + for grp_label, grp_cfg in groups.items(): + comps = grp_cfg["components"] + grp_eps = grp_cfg["eps"] + + vv_sq_fcst = None + vv_sq_obs = None + for var in comps: + s2_fcst = (fcst[var] if isinstance(fcst, xr.Dataset) else fcst).var(dim=rd) + s2_obs = (obs[var] if isinstance(obs, xr.Dataset) else obs).var(dim=rd) + vv_sq_fcst = s2_fcst if vv_sq_fcst is None else vv_sq_fcst + s2_fcst + vv_sq_obs = s2_obs if vv_sq_obs is None else vv_sq_obs + s2_obs + + vv_fcst = vv_sq_fcst ** 0.5 + vv_obs = vv_sq_obs ** 0.5 + results[grp_label] = (vv_fcst / vv_obs).where(vv_obs > grp_eps) + + return xr.Dataset(results) diff --git a/src/mxalign/transformations/__init__.py b/src/mxalign/transformations/__init__.py index c337718..404fe38 100644 --- a/src/mxalign/transformations/__init__.py +++ b/src/mxalign/transformations/__init__.py @@ -1,7 +1,9 @@ from . import base +from . import cerra from . import external __all__ = [ "base", + "cerra", "external", ] diff --git a/src/mxalign/transformations/base.py b/src/mxalign/transformations/base.py index 0638353..54f997c 100644 --- a/src/mxalign/transformations/base.py +++ b/src/mxalign/transformations/base.py @@ -1,7 +1,15 @@ from .registry import register_transformation -@register_transformation("rename") +def _sig_rename(rename_dict): + outputs = list(rename_dict.keys()) + inputs: list[str] = [] + for v in rename_dict.values(): + inputs.extend(v if isinstance(v, list) else [v]) + return inputs, outputs + + +@register_transformation("rename", signature=_sig_rename) def transform_rename(ds, rename_dict): new_dict = {} for new_name, old_names in rename_dict.items(): @@ -13,7 +21,12 @@ def transform_rename(ds, rename_dict): return ds.rename(new_dict) -@register_transformation("kelvin_to_celcius") +def _sig_kelvin_to_celcius(variables, inverse=False): + v = [variables] if isinstance(variables, str) else list(variables) + return v, v # in-place + + +@register_transformation("kelvin_to_celcius", signature=_sig_kelvin_to_celcius) def transform_kelvin_to_celcius(ds, variables, inverse=False): T_C2K = 273.15 if isinstance(variables, str): @@ -29,10 +42,20 @@ def transform_kelvin_to_celcius(ds, variables, inverse=False): return ds -@register_transformation("uv_to_speed") +def _sig_uv_to_speed(u, v, speed): + us = [u] if isinstance(u, str) else list(u) + vs = [v] if isinstance(v, str) else list(v) + ss = [speed] if isinstance(speed, str) else list(speed) + return us + vs, ss + + +@register_transformation("uv_to_speed", signature=_sig_uv_to_speed) def transform(ds, u, v, speed): import numpy as np - result = np.sqrt(ds[u] ** 2 + ds[v] ** 2) - ds[speed] = result + us = [u] if isinstance(u, str) else u + vs = [v] if isinstance(v, str) else v + speeds = [speed] if isinstance(speed, str) else speed + for u_var, v_var, s_var in zip(us, vs, speeds): + ds[s_var] = np.sqrt(ds[u_var] ** 2 + ds[v_var] ** 2) return ds diff --git a/src/mxalign/transformations/cerra.py b/src/mxalign/transformations/cerra.py new file mode 100644 index 0000000..95eee43 --- /dev/null +++ b/src/mxalign/transformations/cerra.py @@ -0,0 +1,242 @@ +"""CERRA-grid spatial-gradient transformations. + +Registers ``cerra_gradient_x`` and ``cerra_gradient_y``: discrete first-order +gradients along the projection x and y axes of the CERRA LCC grid, using +central differences in the interior and one-sided differences at the borders. +Units are [variable] / m (projection-plane meters; no map-scale-factor +correction). + +The transformation operates on a flat spatial dim (default ``grid_index``) +of length ``ny * nx``. It reshapes to 2D for the stencil and flattens back, +so the resulting variable has the same dims/coords as the input. + +Backend selection (per call): + +* If PyTorch + CUDA is available, the gradient runs on GPU. +* Otherwise it runs as a vectorised NumPy stencil. + +Both backends use the exact same slicing expression and produce identical +numerical output (to floating-point round-off). +""" + +from __future__ import annotations + +import fnmatch +import xarray as xr + +from .registry import register_transformation, register_expander +from ..utils.projections import BUILTIN + +# Defaults sourced from the canonical CERRA grid description. +_GRID = BUILTIN["cerra"]["kws_grid"] +_NY_DEFAULT: int = int(_GRID["ny"]) +_NX_DEFAULT: int = int(_GRID["nx"]) +_DX_DEFAULT: float = float(_GRID["dx"]) +_DY_DEFAULT: float = float(_GRID["dy"]) + + +# --------------------------------------------------------------------------- +# Backend +# --------------------------------------------------------------------------- + + +def _torch_cuda(): + """Return the ``torch`` module if CUDA is available, else ``None``. + + Import is lazy so the module remains importable without torch installed. + """ + try: + import torch + except ImportError: + return None + if not torch.cuda.is_available(): + return None + return torch + + +def _grad_axis(arr2d, axis_xy: str, dx: float, dy: float, empty_like): + """Discrete gradient on a 2D array, central interior + one-sided edges. + + ``arr2d`` has shape ``(..., ny, nx)`` in image orientation (row 0 = North, + col 0 = West). ``axis_xy`` is ``"x"`` (along columns / west-east) or + ``"y"`` (along rows / south-north). ``empty_like`` is the array + library's ``empty_like`` constructor; the same slicing expression works + for NumPy and PyTorch tensors. + """ + out = empty_like(arr2d) + if axis_xy == "x": + out[..., :, 1:-1] = (arr2d[..., :, 2:] - arr2d[..., :, :-2]) / (2.0 * dx) + out[..., :, 0] = (arr2d[..., :, 1] - arr2d[..., :, 0]) / dx + out[..., :, -1] = (arr2d[..., :, -1] - arr2d[..., :, -2]) / dx + elif axis_xy == "y": + # row 0 = North, so +y (north) corresponds to decreasing row index. + out[..., 1:-1, :] = (arr2d[..., :-2, :] - arr2d[..., 2:, :]) / (2.0 * dy) + out[..., 0, :] = (arr2d[..., 0, :] - arr2d[..., 1, :]) / dy + out[..., -1, :] = (arr2d[..., -2, :] - arr2d[..., -1, :]) / dy + else: + raise ValueError(f"axis_xy must be 'x' or 'y', got {axis_xy!r}") + return out + + +def _compute_gradient(arr_flat, axis_xy: str, ny: int, nx: int, + dx: float, dy: float): + """Compute the gradient of a NumPy array shaped ``(..., ny*nx)``. + + Reshapes to 2D image orientation, dispatches to the GPU backend if + available, and returns a NumPy array of the original shape. + """ + import numpy as np + + arr = np.ascontiguousarray(arr_flat) + lead_shape = arr.shape[:-1] + arr2d_image = arr.reshape(*lead_shape, ny, nx)[..., ::-1, :] + + torch = _torch_cuda() + if torch is not None: + device = torch.device("cuda") + t = torch.from_numpy(np.ascontiguousarray(arr2d_image)).to( + device, non_blocking=True + ) + out_t = _grad_axis(t, axis_xy, dx, dy, torch.empty_like) + out2d_image = out_t.detach().cpu().numpy() + else: + out2d_image = _grad_axis(arr2d_image, axis_xy, dx, dy, np.empty_like) + + # Reverse the image flip and flatten back. + out2d = out2d_image[..., ::-1, :] + return np.ascontiguousarray(out2d).reshape(*lead_shape, ny * nx) + + +# --------------------------------------------------------------------------- +# DataArray wrapper +# --------------------------------------------------------------------------- + + +def _cerra_gradient_axis(da: xr.DataArray, axis_xy: str, *, grid_dim: str, + ny: int, nx: int, dx: float, dy: float) -> xr.DataArray: + """Apply the gradient to a single ``DataArray`` and return a new one.""" + if grid_dim not in da.dims: + raise ValueError( + f"DataArray has no dim '{grid_dim}'. dims={da.dims}" + ) + n_expected = ny * nx + if da.sizes[grid_dim] != n_expected: + raise ValueError( + f"DataArray dim '{grid_dim}' has size {da.sizes[grid_dim]}, " + f"expected ny*nx = {n_expected} (ny={ny}, nx={nx})." + ) + + da_t = da.transpose(..., grid_dim) + + # Use apply_ufunc so the computation is lazy when the input DataArray is + # backed by dask (e.g. when transform_datasets runs on the full loaded + # dataset before alignment). The grid_dim is a core dimension passed as + # the last axis to _compute_gradient; allow_rechunk ensures the spatial + # axis is never split across chunks (required for the 2-D reshape). + result = xr.apply_ufunc( + _compute_gradient, + da_t, + kwargs=dict(axis_xy=axis_xy, ny=ny, nx=nx, dx=dx, dy=dy), + input_core_dims=[[grid_dim]], + output_core_dims=[[grid_dim]], + dask="parallelized", + output_dtypes=[da_t.dtype], + dask_gufunc_kwargs={"allow_rechunk": True}, + ) + result.name = da.name + result.attrs.update(da.attrs) + return result.transpose(*da.dims) + + +def _expand_vars(patterns, ds_vars): + """Expand a list of variable name patterns (may contain ``*`` / ``?`` globs) + against the concrete variable names in ``ds_vars``. + + Literal names that contain no wildcards are kept as-is (and will cause an + error later if they are absent from the dataset, which is the intended + behaviour). Glob patterns that match nothing are silently dropped. + """ + result = [] + for p in (patterns if not isinstance(patterns, str) else [patterns]): + if any(c in p for c in ("*", "?", "[")): + result.extend(sorted(v for v in ds_vars if fnmatch.fnmatch(v, p))) + else: + result.append(p) + return result + + +def _apply(ds: xr.Dataset, variables, outputs=None, axis_xy: str = "x", *, + grid_dim: str = "grid_index", + ny: int = _NY_DEFAULT, nx: int = _NX_DEFAULT, + dx: float = _DX_DEFAULT, dy: float = _DY_DEFAULT) -> xr.Dataset: + vs = _expand_vars( + [variables] if isinstance(variables, str) else list(variables), + list(ds.data_vars), + ) + if outputs is None: + suffix = f"_grad_{axis_xy}" + os_ = [f"{v}{suffix}" for v in vs] + else: + os_ = [outputs] if isinstance(outputs, str) else list(outputs) + if len(vs) != len(os_): + raise ValueError( + f"variables and outputs must have the same length, " + f"got {len(vs)} vs {len(os_)}." + ) + for in_name, out_name in zip(vs, os_): + ds[out_name] = _cerra_gradient_axis( + ds[in_name], axis_xy, + grid_dim=grid_dim, ny=ny, nx=nx, dx=dx, dy=dy, + ) + return ds + + +# --------------------------------------------------------------------------- +# Registry entry points +# --------------------------------------------------------------------------- + + +def _sig_cerra_grad(variables, outputs=None, **_): + v = [variables] if isinstance(variables, str) else list(variables) + if outputs is None: + # Outputs can't be derived without the dataset when globs are present; + # the expander will have resolved them to concrete names before this + # is called by _derive_source_vars, so outputs will not be None there. + raise ValueError( + "cerra_gradient signature called with outputs=None; ensure the " + "transformation expander ran before recording kwargs." + ) + o = [outputs] if isinstance(outputs, str) else list(outputs) + return v, o + + +def _make_expander(axis_xy: str): + suffix = f"_grad_{axis_xy}" + + def expander(ds, kwargs: dict) -> dict: + kw = dict(kwargs) + vars_raw = kw.get("variables", []) + expanded = _expand_vars( + [vars_raw] if isinstance(vars_raw, str) else list(vars_raw), + list(ds.data_vars), + ) + kw["variables"] = expanded + if kw.get("outputs") is None: + kw["outputs"] = [f"{v}{suffix}" for v in expanded] + return kw + + return expander + + +register_expander("cerra_gradient_x")(_make_expander("x")) +register_expander("cerra_gradient_y")(_make_expander("y")) + + +@register_transformation("cerra_gradient_x", signature=_sig_cerra_grad) +def transform_cerra_gradient_x(ds, variables, outputs=None, **grid_kwargs): + return _apply(ds, variables, outputs, axis_xy="x", **grid_kwargs) + + +@register_transformation("cerra_gradient_y", signature=_sig_cerra_grad) +def transform_cerra_gradient_y(ds, variables, outputs=None, **grid_kwargs): + return _apply(ds, variables, outputs, axis_xy="y", **grid_kwargs) diff --git a/src/mxalign/transformations/external.py b/src/mxalign/transformations/external.py index aa97e69..a32a849 100644 --- a/src/mxalign/transformations/external.py +++ b/src/mxalign/transformations/external.py @@ -1,7 +1,11 @@ from .registry import register_transformation -@register_transformation("external") +def _sig_external(func_path, inputs, output, **_): + return list(inputs.values()), [output] + + +@register_transformation("external", signature=_sig_external) def transform(ds, func_path, inputs, output, **kwargs): func = _resolve_function(func_path) diff --git a/src/mxalign/transformations/registry.py b/src/mxalign/transformations/registry.py index be6378d..1020fda 100644 --- a/src/mxalign/transformations/registry.py +++ b/src/mxalign/transformations/registry.py @@ -1,9 +1,29 @@ _TRANSFORMATION_REGISTRY = {} +_SIGNATURE_REGISTRY = {} +_EXPANDER_REGISTRY = {} -def register_transformation(name): +def register_transformation(name, signature=None): + """Register a transformation function under ``name``. + + Parameters + ---------- + name + Registry key, also the value used in YAML ``transformations:`` blocks. + signature + Optional callable ``(**kwargs) -> (inputs, outputs)`` returning the + lists of source and sink variable names for the transformation given + its YAML kwargs. Used by consumers (e.g. the fused verification + engine) to determine, without executing, which variables a + transformation reads from and writes to a dataset. Transformations + whose I/O cannot be derived from kwargs alone may omit it; callers + that need the information must then either fall back or fail. + """ + def decorator(func): _TRANSFORMATION_REGISTRY[name] = func + if signature is not None: + _SIGNATURE_REGISTRY[name] = signature return func return decorator @@ -18,3 +38,35 @@ def get_transformation(name): return _TRANSFORMATION_REGISTRY[name] except KeyError: raise ValueError(f"Unknown transformation: {name}") + + +def get_signature(name): + """Return the variable I/O signature callable for ``name``, or ``None``. + + The callable, when invoked with the transformation's YAML kwargs, + returns ``(inputs, outputs)`` — two lists of variable names. ``None`` + means the transformation did not declare a signature. + """ + return _SIGNATURE_REGISTRY.get(name) + + +def register_expander(name): + """Register a pre-execution expander for transformation ``name``. + + An expander is a callable ``(ds, kwargs: dict) -> dict`` that receives + the current dataset and the raw YAML kwargs, and returns a new kwargs + dict with glob patterns expanded and optional defaults filled in. The + runner calls the expander (if present) before both executing the + transformation and recording its kwargs in ``_transforms_by_ds``, so + that downstream engines always see concrete, fully-resolved variable + names. + """ + def decorator(func): + _EXPANDER_REGISTRY[name] = func + return func + return decorator + + +def get_expander(name): + """Return the expander callable for ``name``, or ``None``.""" + return _EXPANDER_REGISTRY.get(name) diff --git a/src/mxalign/utils/config.py b/src/mxalign/utils/config.py index 5deb451..1c5f376 100644 --- a/src/mxalign/utils/config.py +++ b/src/mxalign/utils/config.py @@ -1,3 +1,4 @@ +import numpy as np import yaml from .dates import Dates @@ -52,4 +53,25 @@ def _init_datasets(self): if dates: dates = Dates(**dates) loader["files"] = dates.substitute(loader["files"]) + # Propagate declarative time hints to every loader. + # BaseLoader.load() uses these to pre-prune datasets: + # - `valid_times` prunes observation datasets (1D dim). + # - `reference_times` + `lead_times` prune forecast + # datasets rectangularly, enforcing `dates.range` + # (max lead) and `dates.period` (rt spacing). + loader.setdefault( + "valid_times", + np.sort(np.array(dates.valid_times)), + ) + loader.setdefault( + "reference_times", + np.sort(np.array(dates.reference_times)), + ) + loader.setdefault( + "lead_times", + # dates.lead_times strips unit info (stored as plain ints). + # Reconstruct from _step/_range to keep the timedelta64 unit + # so that BaseLoader can cast correctly to the dataset dtype. + np.arange(int(dates._range / dates._step) + 1) * dates._step, + ) self.config["datasets"][key] = loader diff --git a/src/mxalign/verification.py b/src/mxalign/verification.py index 88ad172..3ce37e4 100644 --- a/src/mxalign/verification.py +++ b/src/mxalign/verification.py @@ -1,10 +1,146 @@ from .transformations.external import _resolve_function from functools import partial +import numpy as np + + +# --------------------------------------------------------------------------- +# Sum-decomposable kernels & finalizers (opt-in, config-driven) +# --------------------------------------------------------------------------- +# A metric is "sum-decomposable along reference_time" if its full result can +# be obtained by: +# 1. computing a per-sample partial array via ``kernel(fcst, ref)``, +# 2. summing partials across reference_time, +# 3. applying ``finalize(partial_sum, n_samples)`` once at the end. +# +# Kernels and finalizers are independent registries keyed by short +# mathematical names. Whether (and how) a metric uses them is decided in the +# YAML per metric, e.g.: +# +# metrics: +# mse: +# function: scores.continuous.mse # backend choice — config-level +# kernel: squared_error # opt-in to the fused fast path +# finalize: mean # default if omitted +# +# Switching backends (e.g. ``xskillscore.mse`` ↔ ``scores.continuous.mse``) +# does not require touching the registry: the kernel is determined by the +# math, not the implementation. +# +# Both registries are public; downstream code may register custom kernels or +# finalizers (e.g. ``register_finalize("rms", lambda s, n: np.sqrt(s/n))``). + +_KERNEL_REGISTRY: dict[str, "callable"] = {} +_FINALIZE_REGISTRY: dict[str, "callable"] = {} + + +def register_kernel(name, fn): + """Register a per-sample kernel ``fn(fcst, ref) -> np.ndarray`` under + ``name``. The output must be summable along the reduction dimension.""" + _KERNEL_REGISTRY[name] = fn + + +def register_finalize(name, fn): + """Register a finalizer ``fn(partial_sum, n_samples) -> result`` under + ``name``.""" + _FINALIZE_REGISTRY[name] = fn + + +def get_kernel(name): + """Return the kernel callable for ``name``, or ``None`` if unknown.""" + return _KERNEL_REGISTRY.get(name) + + +def get_finalize(name): + """Return the finalizer callable for ``name``, or ``None`` if unknown.""" + return _FINALIZE_REGISTRY.get(name) + + +def _kernel_squared_error(fcst, ref): + diff = fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) + return diff * diff + + +def _kernel_absolute_error(fcst, ref): + return np.abs( + fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) + ) + + +def _kernel_error(fcst, ref): + return fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) + + +register_kernel("squared_error", _kernel_squared_error) +register_kernel("absolute_error", _kernel_absolute_error) +register_kernel("error", _kernel_error) + +register_finalize("mean", lambda partial_sum, n: partial_sum / n) +register_finalize("sum", lambda partial_sum, n: partial_sum) + + +# --------------------------------------------------------------------------- +# Fused fast-path marker +# --------------------------------------------------------------------------- +# A metric function may declare a sum-decomposable fast path by being +# decorated with ``@fused_metric(kernel=..., finalize=...)``. The decorator +# attaches two attributes to the function object: +# +# * ``_fused_kernel(fcst, ref) -> np.ndarray`` — per-sample partial, +# summable along the reduction dimension. +# * ``_fused_finalize(partial_sum, n_samples) -> result`` — final +# reduction (e.g. divide by ``n`` for means). +# +# The fused verification engine discovers the fast path by inspecting these +# attributes on the function resolved from the YAML's ``function:`` field; +# no registry is involved. The decorated function itself must remain a +# valid xarray-side metric so it works under ``engine: xarray`` too. + +def fused_metric(*, kernel, finalize): + """Mark a metric function as having a fused-engine fast path. + + Parameters + ---------- + kernel : callable + ``kernel(fcst, ref) -> np.ndarray`` returning a per-sample partial + that is summable across the reduction dimension. Operates on plain + numpy arrays of identical shape. + finalize : callable + ``finalize(partial_sum, n_samples) -> result`` reducing the + accumulated sum (e.g. ``lambda s, n: s / n`` for means). + """ + + def decorator(fn): + fn._fused_kernel = kernel + fn._fused_finalize = finalize + return fn + + return decorator + + +def get_fused_kernel(fn): + """Return ``(kernel, finalize)`` for a function decorated with + :func:`fused_metric`, or ``None`` if the function has no fast path. + """ + kernel = getattr(fn, "_fused_kernel", None) + finalize = getattr(fn, "_fused_finalize", None) + if kernel is None or finalize is None: + return None + return kernel, finalize + class Metric: def __init__(self, name, func_path, ds_ref, inputs, **kwargs): self.name = name + self.func_path = func_path + # Opt-in sum-decomposable fields. Popped from kwargs so they are + # never forwarded to the metric function itself. Validation of the + # registered names is deferred until ``.kernel`` / ``.finalize`` is + # actually read, so the legacy path (which doesn't care) keeps + # working regardless of typos. + self._kernel_name = kwargs.pop("kernel", None) + self._finalize_name = kwargs.pop("finalize", "mean") + func = _resolve_function(func_path) self._is_xskillscore = func.__module__.startswith("xskillscore") self._dim = kwargs.get("dim", None) @@ -32,6 +168,51 @@ def compute(self, ds): kwarg_ds = {self._kwarg_ds: ds} return self._func(**kwarg_ds) + @property + def is_decomposable(self): + """True if the metric config opted into a fused-engine kernel.""" + return self._kernel_name is not None + + @property + def kernel_name(self): + return self._kernel_name + + @property + def finalize_name(self): + return self._finalize_name + + @property + def kernel(self): + """The kernel callable, or ``None`` if not opted in. + + Raises ``KeyError`` if the configured name is unknown. + """ + if self._kernel_name is None: + return None + fn = get_kernel(self._kernel_name) + if fn is None: + raise KeyError( + f"metric {self.name!r}: unknown kernel " + f"{self._kernel_name!r} (known: {sorted(_KERNEL_REGISTRY)})" + ) + return fn + + @property + def finalize(self): + """The finalizer callable, or ``None`` if not opted in. + + Raises ``KeyError`` if the configured name is unknown. + """ + if self._kernel_name is None: + return None + fn = get_finalize(self._finalize_name) + if fn is None: + raise KeyError( + f"metric {self.name!r}: unknown finalize " + f"{self._finalize_name!r} (known: {sorted(_FINALIZE_REGISTRY)})" + ) + return fn + def _rechunk(self, ds): if self._dim is None: return ds diff --git a/src/mxalign/verification_fused.py b/src/mxalign/verification_fused.py new file mode 100644 index 0000000..9c60554 --- /dev/null +++ b/src/mxalign/verification_fused.py @@ -0,0 +1,950 @@ +"""Fused verification engine. + +For each reference_time, one client.submit task: + + 1. Loads the per-rt slice directly from the underlying store via + ``loader.slice(rt, lead_times, source_vars)`` — bypassing xarray's + lazy graphs. + 2. Replays the recorded transformations on that small per-rt Dataset. + 3. Applies a sum-decomposable kernel (e.g. squared error for MSE). + 4. Returns numpy partials + per-stage timings. + +Driver runs an ``as_completed`` loop with a bounded submission window +("backpressure"), accumulating partials in driver memory (~few GB total). +After all leaves complete it finalises (e.g. divides by N_rt for means) and +wraps the result into an ``xr.Dataset`` matching the legacy engine shape. + +Required abstractions (extend these to add capability): + + * Loaders override :meth:`mxalign.loaders.base.BaseLoader.slice` for the + per-rt fast read. May optionally provide ``prefetch_path(rt)`` to + enable OS-page-cache prefetch. + * Metric functions are decorated with + :func:`mxalign.verification.fused_metric` to expose a numpy kernel + + finalizer. The bundled :mod:`mxalign.scores` ships + ``mse`` / ``mae`` / ``bias`` / ``mean_error``. + * Transformations declare an I/O signature via + ``register_transformation(..., signature=...)`` so the engine can + derive which source variables to load. + +Validation failures raise immediately with a message pointing at the +abstraction to extend; no silent fallback. +""" +from __future__ import annotations + +import inspect +import logging +import statistics +import threading +import time +import warnings +from collections import deque + +import numpy as np +import xarray as xr + +from .loaders.base import BaseLoader +from .transformations.external import _resolve_function +from .transformations.registry import get_signature, get_transformation +from .verification import get_fused_kernel + +LOG = logging.getLogger("mxalign") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _rt_key(rt) -> int: + """Canonical hashable key for a reference_time: ns-since-epoch int.""" + return int(np.datetime64(rt, "ns").astype("int64")) + + +def _derive_source_vars(common_vars, transforms_for_ds): + """Walk transformations backwards (using their declared signatures) to + determine which source variables must be read from the store for one + dataset.""" + needed = set(common_vars) + for tname, tkwargs in reversed(transforms_for_ds): + sig = get_signature(tname) + if sig is None: + raise NotImplementedError( + f"engine=fused: transformation {tname!r} has no declared I/O " + f"signature; add `signature=...` to its " + f"`register_transformation(...)` call or use engine=xarray." + ) + inputs, outputs = sig(**tkwargs) + if any(o in needed for o in outputs): + needed -= set(outputs) + needed |= set(inputs) + return sorted(needed) + + +# --------------------------------------------------------------------------- +# Leaf task (runs on worker) +# --------------------------------------------------------------------------- + +def _leaf( + rt_value, + lead_times_ns, + common_vars, + ref_name, + model_names, + loaders, + source_vars_by_ds, + transforms_by_ds, + metric_specs, # {metric_name: (kernel_callable, inputs_or_None)} +): + """One per-reference_time task. + + Returns: + { + "rt_value": rt_value, + "timings": {load_: float, transform_: float, kernel: float, total: float}, + "partials": {model_name: {metric_name: np.ndarray(n_var, n_lt, n_grid)}}, + } + """ + t0 = time.perf_counter() + timings: dict[str, float] = {} + + # 1. Load per-dataset slices via the loader's eager slice() method. + slices: dict[str, xr.Dataset] = {} + for ds_name, loader in loaders.items(): + t = time.perf_counter() + slices[ds_name] = loader.slice( + rt_value, lead_times_ns, source_vars_by_ds[ds_name] + ) + timings[f"load_{ds_name}"] = time.perf_counter() - t + + # 2. Replay transformations in recorded order. + for ds_name, ds in list(slices.items()): + t = time.perf_counter() + for tname, tkwargs in transforms_by_ds.get(ds_name, []): + func = get_transformation(tname) + ds = func(ds.copy(), **tkwargs) + slices[ds_name] = ds + timings[f"transform_{ds_name}"] = time.perf_counter() - t + + # 3. Stack to canonical (n_var, n_lt, n_grid) float32 numpy. + arrays: dict[str, np.ndarray] = {} + for ds_name, ds in slices.items(): + arrays[ds_name] = np.stack( + [ + np.ascontiguousarray(ds[v].values, dtype=np.float32) + for v in common_vars + ], + axis=0, + ) + + # 4. Apply kernels. Each metric binds its declared `inputs:` roles to the + # kernel's parameters by name: role 'reference' -> reference array, + # any other role -> the model array being scored. Metrics without an + # `inputs:` block fall back to positional (forecast, reference). + ref = arrays[ref_name] + partials: dict[str, dict[str, np.ndarray]] = {} + t = time.perf_counter() + for m in model_names: + model_arr = arrays[m] + out: dict[str, np.ndarray] = {} + for mn, (kern, inputs) in metric_specs.items(): + if inputs: + kwargs = { + arg: (ref if role == "reference" else model_arr) + for arg, role in inputs.items() + } + out[mn] = kern(**kwargs) + else: + out[mn] = kern(model_arr, ref) + partials[m] = out + timings["kernel"] = time.perf_counter() - t + + timings["total"] = time.perf_counter() - t0 + return {"rt_value": rt_value, "timings": timings, "partials": partials} + + +def _leaf_bundled(rt_value, static): + """Worker-side trampoline: unpack the scattered static bundle and call + :func:`_leaf`. + + ``static`` is a plain dict that was shipped to every worker once via + ``client.scatter(..., broadcast=True)``. Dask resolves the Future to + its materialized value before invoking this function. + """ + return _leaf( + rt_value, + static["lead_times_ns"], + common_vars=static["common_vars"], + ref_name=static["ref_name"], + model_names=static["model_names"], + loaders=static["loaders"], + source_vars_by_ds=static["source_vars_by_ds"], + transforms_by_ds=static["transforms_by_ds"], + metric_specs=static["metric_specs"], + ) + + +# --------------------------------------------------------------------------- +# Driver: validation +# --------------------------------------------------------------------------- + +def _resolve_kernel_inputs(metric_name, func_path, kernel, inputs): + """Validate a metric's YAML ``inputs:`` map against its fused kernel. + + ``inputs`` maps the metric function's argument names to roles + (``reference`` for the reference dataset, any other role for the model + being scored). The fused leaf binds these names directly onto the + kernel, so the kernel must accept them. Returns the inputs dict (copy), + or ``None`` when no ``inputs:`` block was given (positional fallback). + """ + if not inputs: + return None + roles = list(inputs.values()) + n_ref = sum(1 for r in roles if r == "reference") + if n_ref != 1: + raise ValueError( + f"engine=fused: metric {metric_name!r} inputs={inputs} must map " + f"exactly one argument to role 'reference' (got {n_ref})." + ) + sig = inspect.signature(kernel) + has_var_kw = any( + p.kind is inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) + if not has_var_kw: + valid = { + p.name + for p in sig.parameters.values() + if p.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + bad = [a for a in inputs if a not in valid] + if bad: + raise ValueError( + f"engine=fused: metric {metric_name!r} function {func_path!r} " + f"fused kernel does not accept input argument(s) {bad}; kernel " + f"parameters are {sorted(valid)}." + ) + return dict(inputs) + + +def _validate(reference, datasets, loaders, metrics_cfg, ref_name): + # 1. Reference must be one of the datasets and forecast-shaped. + if ref_name not in datasets: + raise ValueError(f"engine=fused: reference {ref_name!r} not in datasets") + if "reference_time" not in reference.dims: + raise ValueError( + "engine=fused: reference dataset has no 'reference_time' dim; " + "this engine requires forecast-shaped reference." + ) + + # 2. Every loader must override BaseLoader.slice. + for name, loader in loaders.items(): + if type(loader).slice is BaseLoader.slice: + raise NotImplementedError( + f"engine=fused: loader {type(loader).__name__!r} for dataset " + f"{name!r} does not override BaseLoader.slice(); add a " + f"slice() method or use engine=xarray." + ) + + # 3. Every metric must resolve to a function with a fused fast path, + # and must include 'reference_time' among its reduce_dims. + metric_specs: dict[str, tuple] = {} + metric_finalizers: dict[str, "callable"] = {} + for mn, mcfg in metrics_cfg.items(): + func_path = mcfg.get("function") + if not func_path: + raise ValueError( + f"engine=fused: metric {mn!r} has no 'function:' entry." + ) + fn = _resolve_function(func_path) + fast = get_fused_kernel(fn) + if fast is None: + raise NotImplementedError( + f"engine=fused: metric {mn!r} uses function {func_path!r} " + f"which has no fused fast path. Decorate the function with " + f"@fused_metric or use one of mxalign.scores.* " + f"(e.g. mxalign.scores.mse)." + ) + kernel, finalize = fast + rd = mcfg.get("reduce_dims") or [] + rd = [rd] if isinstance(rd, str) else list(rd) + if "reference_time" not in rd: + raise ValueError( + f"engine=fused: metric {mn!r} has reduce_dims={rd}; the fused " + f"engine requires 'reference_time' among reduce_dims." + ) + inputs = _resolve_kernel_inputs(mn, func_path, kernel, mcfg.get("inputs")) + metric_specs[mn] = (kernel, inputs) + metric_finalizers[mn] = finalize + + return metric_specs, metric_finalizers + + +# --------------------------------------------------------------------------- +# Driver: result wrap +# --------------------------------------------------------------------------- + +def _make_xr_result(accums, finalizers, n_rt, common_vars, reference, + model_order, metric_order): + """Wrap accumulated partials into an ``xr.Dataset`` matching the legacy + shape: dims = ``(model, metric, variable, lead_time, grid_index)``. + + Coords: ``model``, ``metric``, ``variable``, ``lead_time`` + (+ ``latitude`` / ``longitude`` on ``grid_index`` when present on the + reference). + """ + lead_time = reference["lead_time"].values + lat = reference["latitude"].values if "latitude" in reference.coords else None + lon = reference["longitude"].values if "longitude" in reference.coords else None + + arr_by_metric: dict[str, np.ndarray] = {} + for mn in metric_order: + finalize = finalizers[mn] + stacked = np.stack( + [finalize(accums[m][mn], n_rt) for m in model_order], + axis=0, + ) # (n_model, n_var, n_lt, n_grid) + arr_by_metric[mn] = stacked + + full = np.stack([arr_by_metric[mn] for mn in metric_order], axis=1) + # full: (n_model, n_metric, n_var, n_lt, n_grid) + + coords = { + "model": list(model_order), + "metric": list(metric_order), + "variable": list(common_vars), + "lead_time": lead_time, + } + if lat is not None: + coords["latitude"] = ("grid_index", lat) + if lon is not None: + coords["longitude"] = ("grid_index", lon) + + return xr.DataArray( + full, + dims=("model", "metric", "variable", "lead_time", "grid_index"), + coords=coords, + ).to_dataset(name="metrics") + + +# --------------------------------------------------------------------------- +# Driver: progress + prefetch +# --------------------------------------------------------------------------- + +def _log_progress(done, total, t_start, timings_window, in_flight): + elapsed = time.perf_counter() - t_start + throughput = done / elapsed if elapsed > 0 else 0 + eta = (total - done) / throughput if throughput > 0 else float("nan") + parts = [ + f"[mxalign] fused progress done={done}/{total}", + f"inflight={in_flight}", + f"elapsed={elapsed:.1f}s", + f"throughput={throughput:.2f}leaf/s", + f"eta={eta:.0f}s", + ] + if timings_window: + keys = set().union(*(t.keys() for t in timings_window)) + bits = [] + for k in sorted(keys): + vals = [t[k] for t in timings_window if k in t] + if not vals: + continue + p50 = statistics.median(vals) + p95 = sorted(vals)[max(0, int(0.95 * len(vals)) - 1)] + bits.append(f"{k}(p50={p50:.2f}s,p95={p95:.2f}s)") + parts.append("timings=[" + " ".join(bits) + "]") + LOG.info(" ".join(parts)) + + +def _prefetch_file(path: str) -> None: + """Read *path* sequentially in a daemon thread to populate the OS page + cache. Errors are silently swallowed — a failed prefetch just means + the next leaf reads cold, which is no worse than before.""" + try: + with open(path, "rb") as fh: + buf = bytearray(8 << 20) # 8 MB read buffer + while fh.readinto(buf): + pass + except OSError: + pass + + +def _schedule_prefetch(rt_values, idx: int, loaders: dict, + prefetch_ahead: int) -> None: + """Start a background prefetch daemon thread for each loader's file(s) + at ``rt_values[idx + prefetch_ahead]``. Loaders that do not expose a + ``prefetch_path(rt)`` method are skipped (typically zarr-backed + loaders, where OS-level prefetch is not useful).""" + target_idx = idx + prefetch_ahead + if target_idx >= len(rt_values): + return + rt = rt_values[target_idx] + for name, loader in loaders.items(): + get_path = getattr(loader, "prefetch_path", None) + if get_path is None: + continue + path = get_path(rt) + if path: + threading.Thread( + target=_prefetch_file, + args=(path,), + daemon=True, + name=f"mxalign-prefetch-{name}-{target_idx}", + ).start() + + +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + +def compute_metrics_fused( + datasets, + loaders, + transforms_by_ds, + reference_name, + common_vars, + metrics_cfg, + engine_cfg, +): + """Driver entry point. Returns an ``xr.Dataset`` shaped + ``(model, metric, variable, lead_time, grid_index)``.""" + common_vars = sorted(common_vars) + reference = datasets[reference_name] + model_order = sorted(n for n in datasets if n != reference_name) + metric_order = list(metrics_cfg.keys()) + + metric_specs, metric_finalizers = _validate( + reference, datasets, loaders, metrics_cfg, reference_name + ) + + # Per-dataset source variables (walk transformation signatures backwards). + source_vars_by_ds = { + name: _derive_source_vars(common_vars, transforms_by_ds.get(name, [])) + for name in datasets + } + + rt_values = reference["reference_time"].values + lead_times = reference["lead_time"].values # timedelta64[ns] + # Convert lead_times to integer ns for stable pickling. + lead_times_ns = [int(np.timedelta64(lt, "ns").astype("int64")) for lt in lead_times] + + n_rt = len(rt_values) + n_var = len(common_vars) + n_lt = len(lead_times) + # Driver-side accumulators (one per model+metric). The first arriving + # partial allocates via copy — avoids guessing n_grid up front. + accums: dict[str, dict[str, np.ndarray | None]] = { + m: {mn: None for mn in metric_order} for m in model_order + } + + # Optional dask client. + client = None + try: + from dask.distributed import default_client, as_completed + client = default_client() + except Exception: + client = None + + max_in_flight_cfg = engine_cfg.get("max_in_flight") + if client is not None: + n_workers = max(1, len(client.scheduler_info().get("workers", {}))) + default_window = 2 * n_workers + max_in_flight = int(max_in_flight_cfg) if max_in_flight_cfg else default_window + else: + max_in_flight = 1 + + # Prefetch: background daemon threads warm the OS page cache for the + # next forecast file(s) while the current leaf is being processed. + # Enabled via `prefetch: true` in the `verification:` yaml block. + prefetch_enabled = bool(engine_cfg.get("prefetch", False)) + prefetch_ahead = max(1, int(engine_cfg.get("prefetch_ahead", max_in_flight + 1))) + + LOG.info( + "[mxalign] fused start n_rt=%d n_models=%d n_metrics=%d n_vars=%d " + "n_lt=%d max_in_flight=%d client=%s loaders={%s}", + n_rt, len(model_order), len(metric_order), n_var, n_lt, max_in_flight, + "yes" if client is not None else "no (serial)", + ", ".join(f"{n}:{type(l).__name__}" for n, l in loaders.items()), + ) + + timings_window: deque = deque(maxlen=64) + last_progress_log = time.perf_counter() + last_completion = time.perf_counter() + t_start = time.perf_counter() + done = 0 + + def _consume(result): + nonlocal done, last_completion + for m, per_metric in result["partials"].items(): + for mn, arr in per_metric.items(): + if accums[m][mn] is None: + accums[m][mn] = arr # take ownership + else: + accums[m][mn] += arr + timings_window.append(result["timings"]) + done += 1 + last_completion = time.perf_counter() + + leaf_kwargs = dict( + common_vars=common_vars, + ref_name=reference_name, + model_names=model_order, + loaders=loaders, + source_vars_by_ds=source_vars_by_ds, + transforms_by_ds=transforms_by_ds, + metric_specs=metric_specs, + ) + + if client is None: + # Serial fallback (mainly for --cluster threads). + for i, rt in enumerate(rt_values): + if prefetch_enabled: + _schedule_prefetch(rt_values, i, loaders, prefetch_ahead) + try: + result = _leaf(rt, lead_times_ns, **leaf_kwargs) + except Exception: + LOG.exception("[mxalign] fused leaf-failed rt_idx=%d rt=%s", i, rt) + raise + _consume(result) + now = time.perf_counter() + if now - last_progress_log >= 15.0: + _log_progress(done, n_rt, t_start, list(timings_window), 0) + last_progress_log = now + else: + # Scatter the (identical-per-submit) static payload once and broadcast + # to all workers. Each subsequent client.submit then ships only the + # per-leaf rt + a pointer to the scattered bundle, keeping the + # per-submit graph size in the KB range. + static_bundle = dict(leaf_kwargs) + static_bundle["lead_times_ns"] = lead_times_ns + static_future = client.scatter(static_bundle, broadcast=True, hash=False) + + # Suppress the (now-spurious) per-submit "Sending large graph" warning; + # with the scattered bundle each submit ships only ~hundreds of bytes. + warnings.filterwarnings( + "ignore", + message="Sending large graph of size", + category=UserWarning, + module=r"distributed\.client", + ) + + # Streaming as_completed with a sliding submission window. + ac = as_completed() + i_next = 0 + in_flight = 0 + for _ in range(min(max_in_flight, n_rt)): + if prefetch_enabled: + _schedule_prefetch(rt_values, i_next, loaders, prefetch_ahead) + fut = client.submit(_leaf_bundled, rt_values[i_next], static_future, + pure=False) + fut._mxalign_rt_idx = i_next # informational + ac.add(fut) + i_next += 1 + in_flight += 1 + for fut in ac: + try: + result = fut.result() + except Exception: + LOG.exception( + "[mxalign] fused leaf-failed rt_idx=%d", + getattr(fut, "_mxalign_rt_idx", -1), + ) + raise + _consume(result) + in_flight -= 1 + # Release the future (and its scheduler-held result) ASAP. + try: + fut.release() + except Exception: + pass + if i_next < n_rt: + if prefetch_enabled: + _schedule_prefetch(rt_values, i_next, loaders, prefetch_ahead) + fut2 = client.submit(_leaf_bundled, rt_values[i_next], + static_future, pure=False) + fut2._mxalign_rt_idx = i_next + ac.add(fut2) + i_next += 1 + in_flight += 1 + now = time.perf_counter() + if now - last_progress_log >= 15.0: + _log_progress(done, n_rt, t_start, list(timings_window), in_flight) + last_progress_log = now + if now - last_completion >= 60.0 and in_flight > 0: + LOG.warning( + "[mxalign] fused stall: no leaf completion for %.0fs " + "(done=%d/%d, inflight=%d)", + now - last_completion, done, n_rt, in_flight, + ) + last_completion = now # de-spam + + # Final progress line. + _log_progress(done, n_rt, t_start, list(timings_window), 0) + + return _make_xr_result( + accums, metric_finalizers, n_rt, common_vars, reference, model_order, + metric_order, + ) + + +# =========================================================================== +# Collect engine (fused_collect) +# =========================================================================== +# +# For metrics that reduce over spatial dims only (e.g. NVV with +# reduce_dims=[grid_index]), each reference-time task is independent: no +# partial accumulation across rts is needed. The pattern is: +# +# for each rt (in parallel): +# load → transform → call metric(fcst_ds, obs_ds, **kwargs) → xr.Dataset +# collect all per-rt results → concat along reference_time → return +# +# The leaf passes xr.Dataset objects directly to the metric function, so the +# function can access variables by name (required for grouped vector metrics). +# =========================================================================== + + +def _collect_component_vars(metrics_cfg, ds_vars=None): + """All variable names used as vector components across all collect metrics. + + When ``ds_vars`` is provided, glob patterns in group keys or component + lists are expanded against it before collecting variable names. + """ + needed: set[str] = set() + for mcfg in metrics_cfg.values(): + vars_cfg = mcfg.get("variables") or {} + # Expand glob patterns if dataset variable names are available. + if ds_vars is not None and vars_cfg and ( + any("*" in k for k in vars_cfg) + or any("*" in c for v in vars_cfg.values() for c in v.get("components", [])) + ): + from .scores import _expand_group_patterns + vars_cfg = _expand_group_patterns(vars_cfg, ds_vars) + for grp_cfg in vars_cfg.values(): + needed.update(grp_cfg.get("components", [])) + comps = mcfg.get("components") + if comps: + needed.update(comps) + return needed + + +def _leaf_collect( + rt_value, + lead_times_ns, + ref_name, + model_names, + loaders, + source_vars_by_ds, + transforms_by_ds, + metric_specs, # {metric_name: (func_path, inputs_map_or_None, extra_kwargs)} +): + """Per-rt leaf for the collect engine. + + Returns: + { + "rt_value": rt_value, + "timings": {load_: float, transform_: float, kernel: float, total: float}, + "results": {model_name: {metric_name: xr.Dataset}}, + } + """ + t0 = time.perf_counter() + timings: dict[str, float] = {} + + # 1. Load + slices: dict[str, xr.Dataset] = {} + for ds_name, loader in loaders.items(): + t = time.perf_counter() + slices[ds_name] = loader.slice(rt_value, lead_times_ns, source_vars_by_ds[ds_name]) + timings[f"load_{ds_name}"] = time.perf_counter() - t + + # 2. Transform + for ds_name, ds in list(slices.items()): + t = time.perf_counter() + for tname, tkwargs in transforms_by_ds.get(ds_name, []): + func = get_transformation(tname) + ds = func(ds.copy(), **tkwargs) + slices[ds_name] = ds + timings[f"transform_{ds_name}"] = time.perf_counter() - t + + # 3. Apply metric functions directly on xr.Dataset slices. + ref_ds = slices[ref_name] + results: dict[str, dict[str, xr.Dataset]] = {} + t = time.perf_counter() + for m in model_names: + model_ds = slices[m] + per_metric: dict[str, xr.Dataset] = {} + for mn, (func_path, inputs, extra_kwargs) in metric_specs.items(): + fn = _resolve_function(func_path) + if inputs: + call_kwargs: dict = { + arg: (ref_ds if role == "reference" else model_ds) + for arg, role in inputs.items() + } + else: + call_kwargs = {"fcst": model_ds, "obs": ref_ds} + call_kwargs.update(extra_kwargs) + per_metric[mn] = fn(**call_kwargs) + results[m] = per_metric + timings["kernel"] = time.perf_counter() - t + timings["total"] = time.perf_counter() - t0 + + return {"rt_value": rt_value, "timings": timings, "results": results} + + +def _leaf_collect_bundled(rt_value, static): + return _leaf_collect( + rt_value, + static["lead_times_ns"], + ref_name=static["ref_name"], + model_names=static["model_names"], + loaders=static["loaders"], + source_vars_by_ds=static["source_vars_by_ds"], + transforms_by_ds=static["transforms_by_ds"], + metric_specs=static["metric_specs"], + ) + + +def _make_xr_collect_result(collected, reference, model_order, metric_order): + """Build final xr.DataArray from a list of per-rt result dicts (sorted by rt). + + Each result["results"][model][metric] is an xr.Dataset with one data + variable per vector group and dim ``(lead_time,)``. + + Output: xr.DataArray with dims ``(model, metric, variable, reference_time, + lead_time)``, matching the format produced by the xarray engine so that + downstream code and saved netCDF files have a consistent ``metric`` + coordinate dimension. + """ + rt_values = np.array([r["rt_value"] for r in collected]) + lead_times = reference["lead_time"].values + + out: dict[str, xr.DataArray] = {} + for mn in metric_order: + model_arrays = [] + group_names = None + for m in model_order: + rt_datasets = [r["results"][m][mn] for r in collected] + if group_names is None: + group_names = list(rt_datasets[0].data_vars) + + # Concat per-rt xr.Datasets along a new reference_time dimension. + combined = xr.concat( + rt_datasets, + dim=xr.DataArray(rt_values, dims="reference_time", name="reference_time"), + ) + # combined: xr.Dataset({group: DataArray(reference_time, lead_time)}) + + # Stack vector groups into a "variable" dimension. + group_stack = xr.concat( + [combined[g] for g in group_names], + dim=xr.DataArray(group_names, dims="variable", name="variable"), + ) + # Assign the canonical lead_time coordinate values. + group_stack = group_stack.assign_coords(lead_time=lead_times) + model_arrays.append(group_stack) + + metric_da = xr.concat( + model_arrays, + dim=xr.DataArray(model_order, dims="model", name="model"), + ) + out[mn] = metric_da # (model, variable, reference_time, lead_time) + + # Stack metrics along a "metric" dimension and wrap as a Dataset named + # "metrics", matching the fused engine's output format (_make_xr_result). + return xr.concat( + list(out.values()), + dim=xr.Variable("metric", metric_order), + ).transpose("model", "metric", ...).to_dataset(name="metrics") + + +def compute_metrics_collect( + datasets, + loaders, + transforms_by_ds, + reference_name, + metrics_cfg, + engine_cfg, +): + """Driver for the *fused_collect* engine. + + For each reference_time, loads + transforms + calls the metric function + directly on ``xr.Dataset`` slices, then returns the per-rt result without + accumulation. This is correct for metrics that reduce over spatial + dimensions only (e.g. NVV with ``reduce_dims: [grid_index]``). + + Returns an ``xr.DataArray`` with dims + ``(model, metric, variable, reference_time, lead_time)``, matching the + format produced by the xarray engine. + """ + reference = datasets[reference_name] + model_order = sorted(n for n in datasets if n != reference_name) + metric_order = list(metrics_cfg.keys()) + + # -- validate loaders --------------------------------------------------- + for name, loader in loaders.items(): + if type(loader).slice is BaseLoader.slice: + raise NotImplementedError( + f"engine=fused_collect: loader {type(loader).__name__!r} for " + f"dataset {name!r} does not override BaseLoader.slice(); add a " + f"slice() method or use engine=xarray." + ) + + # -- build metric specs ------------------------------------------------- + metric_specs: dict[str, tuple] = {} + for mn, mcfg in metrics_cfg.items(): + func_path = mcfg.get("function") + if not func_path: + raise ValueError( + f"engine=fused_collect: metric {mn!r} has no 'function:' entry." + ) + _resolve_function(func_path) # fail-fast on bad path + inputs = mcfg.get("inputs") + extra_kwargs = {k: v for k, v in mcfg.items() if k not in ("function", "inputs")} + metric_specs[mn] = (func_path, inputs, extra_kwargs) + + # -- derive source vars per dataset ------------------------------------ + needed_vars = _collect_component_vars( + metrics_cfg, + ds_vars=datasets[reference_name].data_vars, + ) + source_vars_by_ds = { + name: _derive_source_vars(sorted(needed_vars), transforms_by_ds.get(name, [])) + for name in datasets + } + + rt_values = reference["reference_time"].values + lead_times = reference["lead_time"].values + lead_times_ns = [int(np.timedelta64(lt, "ns").astype("int64")) for lt in lead_times] + n_rt = len(rt_values) + + # -- Dask / serial setup (identical pattern to compute_metrics_fused) --- + client = None + try: + from dask.distributed import default_client, as_completed + client = default_client() + except Exception: + client = None + + max_in_flight_cfg = engine_cfg.get("max_in_flight") + if client is not None: + n_workers = max(1, len(client.scheduler_info().get("workers", {}))) + default_window = 2 * n_workers + max_in_flight = int(max_in_flight_cfg) if max_in_flight_cfg else default_window + else: + max_in_flight = 1 + + prefetch_enabled = bool(engine_cfg.get("prefetch", False)) + prefetch_ahead = max(1, int(engine_cfg.get("prefetch_ahead", max_in_flight + 1))) + + LOG.info( + "[mxalign] fused_collect start n_rt=%d n_models=%d n_metrics=%d " + "n_lt=%d max_in_flight=%d client=%s", + n_rt, len(model_order), len(metric_order), len(lead_times), max_in_flight, + "yes" if client is not None else "no (serial)", + ) + + timings_window: deque = deque(maxlen=64) + last_progress_log = time.perf_counter() + last_completion = time.perf_counter() + t_start = time.perf_counter() + done = 0 + collected: list[dict] = [] + + def _consume(result): + nonlocal done, last_completion + collected.append(result) + timings_window.append(result["timings"]) + done += 1 + last_completion = time.perf_counter() + + leaf_kwargs = dict( + ref_name=reference_name, + model_names=model_order, + loaders=loaders, + source_vars_by_ds=source_vars_by_ds, + transforms_by_ds=transforms_by_ds, + metric_specs=metric_specs, + ) + + if client is None: + for i, rt in enumerate(rt_values): + if prefetch_enabled: + _schedule_prefetch(rt_values, i, loaders, prefetch_ahead) + try: + result = _leaf_collect(rt, lead_times_ns, **leaf_kwargs) + except Exception: + LOG.exception("[mxalign] fused_collect leaf-failed rt_idx=%d rt=%s", i, rt) + raise + _consume(result) + now = time.perf_counter() + if now - last_progress_log >= 15.0: + _log_progress(done, n_rt, t_start, list(timings_window), 0) + last_progress_log = now + else: + static_bundle = dict(leaf_kwargs) + static_bundle["lead_times_ns"] = lead_times_ns + static_future = client.scatter(static_bundle, broadcast=True, hash=False) + + warnings.filterwarnings( + "ignore", + message="Sending large graph of size", + category=UserWarning, + module=r"distributed\.client", + ) + + ac = as_completed() + i_next = 0 + in_flight = 0 + for _ in range(min(max_in_flight, n_rt)): + if prefetch_enabled: + _schedule_prefetch(rt_values, i_next, loaders, prefetch_ahead) + fut = client.submit(_leaf_collect_bundled, rt_values[i_next], static_future, + pure=False) + fut._mxalign_rt_idx = i_next + ac.add(fut) + i_next += 1 + in_flight += 1 + for fut in ac: + try: + result = fut.result() + except Exception: + LOG.exception( + "[mxalign] fused_collect leaf-failed rt_idx=%d", + getattr(fut, "_mxalign_rt_idx", -1), + ) + raise + _consume(result) + in_flight -= 1 + try: + fut.release() + except Exception: + pass + if i_next < n_rt: + if prefetch_enabled: + _schedule_prefetch(rt_values, i_next, loaders, prefetch_ahead) + fut2 = client.submit(_leaf_collect_bundled, rt_values[i_next], + static_future, pure=False) + fut2._mxalign_rt_idx = i_next + ac.add(fut2) + i_next += 1 + in_flight += 1 + now = time.perf_counter() + if now - last_progress_log >= 15.0: + _log_progress(done, n_rt, t_start, list(timings_window), in_flight) + last_progress_log = now + if now - last_completion >= 60.0 and in_flight > 0: + LOG.warning( + "[mxalign] fused_collect stall: no leaf completion for %.0fs " + "(done=%d/%d, inflight=%d)", + now - last_completion, done, n_rt, in_flight, + ) + last_completion = now + + _log_progress(done, n_rt, t_start, list(timings_window), 0) + + # Sort collected results by rt_value for deterministic output ordering. + collected.sort(key=lambda r: r["rt_value"]) + + return _make_xr_collect_result(collected, reference, model_order, metric_order) diff --git a/tests/test_cerra_gradient.py b/tests/test_cerra_gradient.py new file mode 100644 index 0000000..8e9e09d --- /dev/null +++ b/tests/test_cerra_gradient.py @@ -0,0 +1,129 @@ +"""Unit tests for the CERRA gradient transformations. + +These tests run without the real CERRA dataset by constructing a small +synthetic grid (still 1069x1069 to match the registered defaults) with +analytically-known gradients. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import xarray as xr + +from mxalign.transformations import cerra # noqa: F401 — registers transforms +from mxalign.transformations.registry import get_transformation +from mxalign.utils.projections import BUILTIN + + +GRID = BUILTIN["cerra"]["kws_grid"] +NY = int(GRID["ny"]) +NX = int(GRID["nx"]) +DX = float(GRID["dx"]) +DY = float(GRID["dy"]) +N = NY * NX + + +def _flat_field_from_image(image_2d: np.ndarray) -> np.ndarray: + """Encode a 2D image (row 0 = North) as the 1D anemoi flatten order + (row 0 of the raw reshape = South). + """ + return image_2d[::-1, :].reshape(-1) + + +def _make_dataset(image_2d: np.ndarray, name: str = "f") -> xr.Dataset: + flat = _flat_field_from_image(image_2d) + da = xr.DataArray(flat[None, :], dims=("valid_time", "grid_index")) + return xr.Dataset({name: da}) + + +def test_cerra_gradient_x_linear_field(): + """f(col) = col * DX => df/dx == 1 everywhere (also at the edges).""" + cols = np.arange(NX, dtype=np.float64) + image = np.broadcast_to(cols * DX, (NY, NX)).copy() + ds = _make_dataset(image, name="f") + + fn = get_transformation("cerra_gradient_x") + ds_out = fn(ds.copy(), variables=["f"], outputs=["fx"]) + + gx_flat = ds_out["fx"].values[0] + gx_image = gx_flat.reshape(NY, NX)[::-1, :] + np.testing.assert_allclose(gx_image, np.ones_like(gx_image), atol=1e-9) + + +def test_cerra_gradient_y_linear_field_north_positive(): + """f(row_image) = (NY-1-row_image) * DY -> values grow going North, + so df/dy == 1 everywhere (also at the edges). + """ + rows_image = np.arange(NY, dtype=np.float64) + f_per_row = (NY - 1 - rows_image) * DY # increases northward + image = np.broadcast_to(f_per_row[:, None], (NY, NX)).copy() + ds = _make_dataset(image, name="f") + + fn = get_transformation("cerra_gradient_y") + ds_out = fn(ds.copy(), variables=["f"], outputs=["fy"]) + + gy_flat = ds_out["fy"].values[0] + gy_image = gy_flat.reshape(NY, NX)[::-1, :] + np.testing.assert_allclose(gy_image, np.ones_like(gy_image), atol=1e-9) + + +def test_cerra_gradient_y_constant_in_x(): + """A field constant in the x-direction has zero x-gradient.""" + rows_image = np.arange(NY, dtype=np.float64) + image = np.broadcast_to(rows_image[:, None] * DY, (NY, NX)).copy() + ds = _make_dataset(image, name="f") + + fn = get_transformation("cerra_gradient_x") + ds_out = fn(ds.copy(), variables=["f"], outputs=["fx"]) + assert np.max(np.abs(ds_out["fx"].values)) < 1e-9 + + +def test_cerra_gradient_preserves_dims_and_coords(): + rng = np.random.default_rng(0) + image = rng.standard_normal((NY, NX)).astype(np.float32) + flat = _flat_field_from_image(image) + times = np.array(["2024-01-01", "2024-01-02"], dtype="datetime64[ns]") + da = xr.DataArray( + np.stack([flat, flat]), + dims=("valid_time", "grid_index"), + coords={"valid_time": times}, + name="f", + ) + ds = xr.Dataset({"f": da}) + + fn = get_transformation("cerra_gradient_x") + ds_out = fn(ds.copy(), variables=["f"], outputs=["fx"]) + + assert ds_out["fx"].dims == ("valid_time", "grid_index") + assert ds_out["fx"].shape == (2, N) + np.testing.assert_array_equal( + ds_out["fx"].coords["valid_time"].values, times + ) + + +def test_cerra_gradient_backend_parity(monkeypatch): + """Numpy and torch-CUDA paths must produce identical numerical output. + + Skipped if torch+CUDA is not available. + """ + torch_mod = pytest.importorskip("torch") + if not torch_mod.cuda.is_available(): + pytest.skip("CUDA not available") + + rng = np.random.default_rng(1) + image = rng.standard_normal((NY, NX)).astype(np.float32) + ds = _make_dataset(image, name="f") + + fn = get_transformation("cerra_gradient_x") + + # GPU path (default since torch+CUDA is present). + ds_gpu = fn(ds.copy(), variables=["f"], outputs=["fx"]) + + # Force numpy path. + monkeypatch.setattr(cerra, "_torch_cuda", lambda: None) + ds_cpu = fn(ds.copy(), variables=["f"], outputs=["fx"]) + + np.testing.assert_allclose( + ds_gpu["fx"].values, ds_cpu["fx"].values, rtol=0, atol=1e-6 + )