From ceca8a11e31f1b4939f623ac9786fe3532268f14 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Mon, 1 Jun 2026 13:47:24 +0300 Subject: [PATCH 01/11] optimize loading --- src/mxalign/loaders/anemoi_inference.py | 142 +++++++++++++++++++++--- src/mxalign/loaders/base.py | 55 ++++++++- 2 files changed, 183 insertions(+), 14 deletions(-) diff --git a/src/mxalign/loaders/anemoi_inference.py b/src/mxalign/loaders/anemoi_inference.py index 68f8799..2c3eb3f 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,136 @@ 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 _open_mf_dataset(files, **kwargs): +def _load_nc_vars(path, var_names, engine): + """Load all named variables from one NC file. - times = xr.open_dataset(files[0], engine=kwargs["engine"], chunks=kwargs["chunks"])[ - "time" - ].values - lead_times = times - times[0] + 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. - ds = xr.open_mfdataset(files, preprocess=_preprocess, **kwargs) + 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 - ds_out = ( - ds.assign_coords({"lead_time": ("time", lead_times)}) - .rename_dims({"values": "grid_index"}) - .swap_dims({"time": "lead_time"}) - ) - return ds_out +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") + + # 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] + + 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"}) + ) + + # ------------------------------------------------------------------ + # 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 + + 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..d603eb2 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) From def9d0bec0110ada0055269931895debc75bad18 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Mon, 1 Jun 2026 13:48:26 +0300 Subject: [PATCH 02/11] vectorize wind speed transform --- src/mxalign/transformations/base.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/mxalign/transformations/base.py b/src/mxalign/transformations/base.py index 0638353..44a2840 100644 --- a/src/mxalign/transformations/base.py +++ b/src/mxalign/transformations/base.py @@ -33,6 +33,9 @@ def transform_kelvin_to_celcius(ds, variables, inverse=False): 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 From 41eed1c9d5f4bad0b6e016f4410624cf9848d4e2 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Mon, 1 Jun 2026 13:50:03 +0300 Subject: [PATCH 03/11] fix: add config change for loaders --- src/mxalign/utils/config.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) 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 From 0d196f970feebef44190d98629bc46324c7c42da Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Tue, 2 Jun 2026 11:04:21 +0300 Subject: [PATCH 04/11] optimize verification for use case --- src/mxalign/_progress.py | 170 ++++++ src/mxalign/loaders/anemoi_datasets.py | 20 + src/mxalign/loaders/anemoi_inference.py | 37 ++ src/mxalign/loaders/base.py | 8 + src/mxalign/runner.py | 121 ++++- src/mxalign/verification_fused.py | 673 ++++++++++++++++++++++++ 6 files changed, 1020 insertions(+), 9 deletions(-) create mode 100644 src/mxalign/_progress.py create mode 100644 src/mxalign/verification_fused.py 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..147e2d7 100644 --- a/src/mxalign/loaders/anemoi_datasets.py +++ b/src/mxalign/loaders/anemoi_datasets.py @@ -53,6 +53,26 @@ def _load(self): ) return ds_selected.to_dataset(dim="variable") + def fast_slice_recipe(self): + """Recipe for per-rt direct zarr region read (fused engine). + + Only the single-file zarr path is supported in v1. The leaf + computes valid_times = rt + lead_times and uses zarr's vectorised + indexing to fetch one slice; no xarray/dask lazy graph involved. + """ + if isinstance(self.files, list): + if len(self.files) != 1: + return None + path = self.files[0] + else: + path = self.files + return { + "kind": "anemoi-datasets-zarr", + "path": path, + "consolidated": False, + "drop_vars": list(DROP_VARS), + } + 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 2c3eb3f..bb10df3 100644 --- a/src/mxalign/loaders/anemoi_inference.py +++ b/src/mxalign/loaders/anemoi_inference.py @@ -69,6 +69,43 @@ def _load(self): return ds + def fast_slice_recipe(self): + """Per-reference_time loading recipe for the fused engine. + + 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) + if not files: + return None + if Path(files[0]).suffix.lower() == ".zarr": + return None + + engine = self.kwargs.get("engine", DEFAULTS_NETCDF["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 + + files_by_rt = {int(rt.astype("int64")): f for rt, f in zip(rt_values, files)} + return { + "kind": "anemoi-inference-nc", + "files_by_rt": files_by_rt, + "engine": engine, + } + + def _load_nc_vars(path, var_names, engine): """Load all named variables from one NC file. diff --git a/src/mxalign/loaders/base.py b/src/mxalign/loaders/base.py index d603eb2..2053a46 100644 --- a/src/mxalign/loaders/base.py +++ b/src/mxalign/loaders/base.py @@ -110,6 +110,14 @@ def _get_properties(self, ds): ) return properties + def fast_slice_recipe(self): + """Return a small picklable dict the fused verification engine can use + to load one reference_time slice directly from the underlying store, + bypassing the lazy xarray/dask graphs. Return None to opt out + (the fused engine will reject this loader for the dataset). + """ + return None + @register_loader class MxAlignLoader(BaseLoader): diff --git a/src/mxalign/runner.py b/src/mxalign/runner.py index aff4398..5ee3cc6 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"] @@ -63,6 +79,11 @@ def transform_datasets(self): self.datasets[name] = transform( name=transformation, datasets=ds, **config_trans ) + # Record (transform_name, kwargs) in application order for + # the fused engine to replay on per-rt slices. + self._transforms_by_ds.setdefault(name, []).append( + (transformation, dict(config_trans)) + ) def align(self): config = self.config["alignment"] @@ -116,7 +137,51 @@ 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: + 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 +191,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 +211,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 +247,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/verification_fused.py b/src/mxalign/verification_fused.py new file mode 100644 index 0000000..7a7175e --- /dev/null +++ b/src/mxalign/verification_fused.py @@ -0,0 +1,673 @@ +"""Fused verification engine (Phase 2 / lever B, recipe-based). + +For each reference_time, one client.submit task: + + 1. Loads the per-rt slice directly from the underlying store + (NetCDF per rt for forecasts; zarr region read for ERA5) + via the loader's `fast_slice_recipe`, bypassing xarray's lazy graphs. + 2. Replays the registered 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. + +Scope (v1): + - Sum-decomposable metrics with `reduce_dims` containing 'reference_time': + MSE, MAE, bias (mean error), mean(reference), mean(forecast). + - Loaders: anemoi-inference (per-rt NetCDF), anemoi-datasets (single zarr). + - Transformations: rename, kelvin_to_celcius, uv_to_speed (extend by + adding entries to `_TRANSFORM_IO`). + +Validation failures (unsupported metric/transform/loader, missing +reduce_dims, missing fast_slice_recipe) raise immediately. No silent +fallback. +""" +from __future__ import annotations + +import logging +import statistics +import time +import warnings +from collections import deque +from typing import Any, Callable + +import numpy as np +import xarray as xr + +LOG = logging.getLogger("mxalign") + + +# --------------------------------------------------------------------------- +# Metric kernels +# --------------------------------------------------------------------------- +# Each kernel takes (fcst, ref) numpy arrays of shape (n_var, n_lt, n_grid) +# and returns a per-sample partial of the same shape that is **summable** +# across reference_time. The finalize step (mean = sum / N, sum = sum) +# is applied after all leaves are reduced. + +def _kernel_squared_error(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: + diff = fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) + return diff * diff + + +def _kernel_abs_error(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: + return np.abs( + fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) + ) + + +def _kernel_error(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: + return fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) + + +def _kernel_identity_fcst(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: + return fcst.astype(np.float32, copy=False) + + +def _kernel_identity_ref(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: + return ref.astype(np.float32, copy=False) + + +# func_path -> (kernel, finalize_kind in {"mean", "sum"}) +_FUSED_KERNELS: dict[str, tuple[Callable, str]] = { + "scores.continuous.mse": (_kernel_squared_error, "mean"), + "scores.continuous.mae": (_kernel_abs_error, "mean"), + "scores.continuous.bias": (_kernel_error, "mean"), + "scores.continuous.mean_error": (_kernel_error, "mean"), +} + + +# --------------------------------------------------------------------------- +# Transformation source-variable bookkeeping +# --------------------------------------------------------------------------- +# Each entry returns (inputs, outputs) variable lists given the transform's +# kwargs (as recorded by Runner.transform_datasets). Used to walk the +# transformation chain backwards from `common_vars` to "what to load from +# the source". + +def _io_uv_to_speed(kwargs): + u = kwargs["u"]; v = kwargs["v"]; s = kwargs["speed"] + u = [u] if isinstance(u, str) else list(u) + v = [v] if isinstance(v, str) else list(v) + s = [s] if isinstance(s, str) else list(s) + return u + v, s + + +def _io_kelvin_to_celcius(kwargs): + v = kwargs["variables"] + v = [v] if isinstance(v, str) else list(v) + return v, v # in-place + + +def _io_rename(kwargs): + d = kwargs["rename_dict"] # new_name -> old_name(s) + outputs = list(d.keys()) + inputs: list[str] = [] + for v in d.values(): + inputs.extend(v if isinstance(v, list) else [v]) + return inputs, outputs + + +_TRANSFORM_IO: dict[str, Callable] = { + "uv_to_speed": _io_uv_to_speed, + "kelvin_to_celcius": _io_kelvin_to_celcius, + "rename": _io_rename, +} + + +def _derive_source_vars(common_vars, transforms_for_ds): + """Walk transformations backwards to derive the set of source variables + that need to be read from the store for one dataset.""" + needed = set(common_vars) + for tname, tkwargs in reversed(transforms_for_ds): + if tname not in _TRANSFORM_IO: + raise NotImplementedError( + f"fused engine: transformation {tname!r} has no input/output " + f"spec in _TRANSFORM_IO; add one or use engine=xarray" + ) + inputs, outputs = _TRANSFORM_IO[tname](tkwargs) + if any(o in needed for o in outputs): + needed -= set(outputs) + needed |= set(inputs) + return sorted(needed) + + +# --------------------------------------------------------------------------- +# Per-rt slice loaders (worker-side) +# --------------------------------------------------------------------------- + +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 _load_slice(recipe, rt_value, lead_times, var_names) -> xr.Dataset: + kind = recipe["kind"] + if kind == "anemoi-inference-nc": + return _load_anemoi_inference_slice(recipe, rt_value, lead_times, var_names) + if kind == "anemoi-datasets-zarr": + return _load_anemoi_datasets_slice(recipe, rt_value, lead_times, var_names) + raise NotImplementedError(f"fused engine: unknown recipe kind {kind!r}") + + +def _load_anemoi_inference_slice(recipe, rt_value, lead_times, var_names) -> xr.Dataset: + path = recipe["files_by_rt"][_rt_key(rt_value)] + engine = recipe["engine"] + with xr.open_dataset(path, engine=engine) as src: + # Subset variables + lead_times *lazily* and only then call .load(). + # Doing .load() up front (the previous behaviour) forces a read of the + # full time axis even when the file holds more steps than we need; it + # also turns the time-axis selection into an in-memory fancy index + # instead of a hyperslab read. Selecting first lets HDF5 issue a + # single contiguous read for the steady-state (cadence-1) case. + sub = src[list(var_names)] + 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(int(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"fused engine: missing lead_times in {path}: {bad[:5]}... " + f"(reference_time={rt_value})" + ) + # 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")) + ds = sub.load() + return ds + + +def _load_anemoi_datasets_slice(recipe, rt_value, lead_times, var_names) -> xr.Dataset: + path = recipe["path"] + src = xr.open_zarr(path, consolidated=recipe.get("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 var_names], dtype=np.int64) + except ValueError as e: + raise ValueError( + f"fused engine: variable not found in {path}: {e}" + ) from None + + rt = np.datetime64(rt_value, "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"fused engine: missing valid_times in {path}: {bad[:5]}... " + f"(reference_time={rt_value})" + ) + + 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"), + ) + loaded = arr_sel.load() + vals = np.asarray(loaded.values) # (n_lt, n_var, n_grid) + # `vals` lead_time axis matches the slice/index we asked for; for the + # contiguous-slice path it's already in the right order (and length). + ds = xr.Dataset( + { + v: (("lead_time", "grid_index"), vals[:, i, :]) + for i, v in enumerate(var_names) + } + ) + return ds + + +# --------------------------------------------------------------------------- +# Leaf task (runs on worker) +# --------------------------------------------------------------------------- + +def _leaf( + rt_value, + lead_times, + common_vars, + ref_name, + model_names, + recipes_by_ds, + source_vars_by_ds, + transforms_by_ds, + metric_kernels, +): + """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)}}, + } + """ + from mxalign.transformations.registry import get_transformation + + t0 = time.perf_counter() + timings: dict[str, float] = {} + + # 1. Load per-dataset slices. + slices: dict[str, xr.Dataset] = {} + for ds_name, recipe in recipes_by_ds.items(): + t = time.perf_counter() + slices[ds_name] = _load_slice( + recipe, rt_value, lead_times, 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. + ref = arrays[ref_name] + partials: dict[str, dict[str, np.ndarray]] = {} + t = time.perf_counter() + for m in model_names: + fcst = arrays[m] + partials[m] = { + mn: kernel(fcst, ref) for mn, (kernel, _) in metric_kernels.items() + } + 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 _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"], + recipes_by_ds=static["recipes_by_ds"], + source_vars_by_ds=static["source_vars_by_ds"], + transforms_by_ds=static["transforms_by_ds"], + metric_kernels=static["metric_kernels"], + ) + + +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + +def _validate(reference, datasets, loaders, transforms_by_ds, metrics_cfg, ref_name): + # 1. Every dataset must have a recipe. + recipes: dict[str, dict] = {} + for name, loader in loaders.items(): + if not hasattr(loader, "fast_slice_recipe"): + raise NotImplementedError( + f"fused engine: loader for dataset {name!r} has no " + f"fast_slice_recipe(); use engine=xarray or extend the loader." + ) + recipe = loader.fast_slice_recipe() + if recipe is None: + raise NotImplementedError( + f"fused engine: loader {type(loader).__name__!r} declined to " + f"produce a fast-slice recipe for dataset {name!r} (e.g. " + f"unsupported file layout); use engine=xarray." + ) + recipes[name] = recipe + + # 2. Every metric must be in the allow-list with reduce_dims=[reference_time]. + metric_kernels: dict[str, tuple[Callable, str]] = {} + for mn, mcfg in metrics_cfg.items(): + func_path = mcfg.get("function") + if func_path not in _FUSED_KERNELS: + raise NotImplementedError( + f"fused engine: metric {mn!r} uses function {func_path!r} which " + f"is not in the sum-decomposable allow-list " + f"({sorted(_FUSED_KERNELS)}); use engine=xarray." + ) + rd = mcfg.get("reduce_dims") or [] + rd = [rd] if isinstance(rd, str) else list(rd) + if "reference_time" not in rd: + raise ValueError( + f"fused engine: metric {mn!r} has reduce_dims={rd}; the fused " + f"engine requires 'reference_time' among reduce_dims." + ) + metric_kernels[mn] = _FUSED_KERNELS[func_path] + + # 3. Reference must be one of the datasets. + if ref_name not in datasets: + raise ValueError(f"fused engine: reference {ref_name!r} not in datasets") + + return recipes, metric_kernels + + +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). + """ + 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 + + # (model, metric, variable, lead_time, grid_index) + arr_by_metric: dict[str, np.ndarray] = {} + for mn in metric_order: + finalize = finalizers[mn] + stacked = np.stack( + [ + accums[m][mn] / (n_rt if finalize == "mean" else 1) + 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") + + +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: + # Collect per-stage timings across the 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 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()) + + recipes, metric_kernels = _validate( + reference, datasets, loaders, transforms_by_ds, metrics_cfg, reference_name + ) + + # Derive per-dataset source variables (walk transformations backwards). + source_vars_by_ds = { + name: _derive_source_vars(common_vars, transforms_by_ds.get(name, [])) + for name in datasets + } + + # Per-rt iteration: drive from the reference dataset's reference_time. + if "reference_time" not in reference.dims: + raise ValueError( + "fused engine: reference dataset has no 'reference_time' dim; " + "this engine requires forecast-shaped reference." + ) + 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) + finalizers = {mn: kind for mn, (_, kind) in metric_kernels.items()} + + # Pre-allocate driver-side accumulators (one per model+metric, ~1.5GB each). + n_var = len(common_vars) + n_lt = len(lead_times) + # We let the first arriving partial allocate via copy; saves a guess at n_grid. + accums: dict[str, dict[str, np.ndarray | None]] = { + m: {mn: None for mn in metric_order} for m in model_order + } + + # Try to get a Client; if none, run serial in-process. + 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 + + 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 recipes={%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}:{r['kind']}" for n, r in recipes.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 + partials = result["partials"] + for m, per_metric in 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, + recipes_by_ds=recipes, + source_vars_by_ds=source_vars_by_ds, + transforms_by_ds=transforms_by_ds, + metric_kernels=metric_kernels, + ) + + if client is None: + # Serial fallback (mainly for --cluster threads). + for i, rt in enumerate(rt_values): + 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 (large, identical-per-submit) static payload once and + # broadcast it to all workers. Each subsequent client.submit then ships + # only the per-leaf rt + lead_times + a pointer to the scattered + # bundle, keeping the per-submit graph size in the KB range. + # `lead_times_ns` is small (<=145 ints) but we scatter it too for + # symmetry. Broadcast=True ensures it's already on every worker before + # the first submit, so workers never pull from the scheduler at task + # start. + 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 + # Prime the window. + for _ in range(min(max_in_flight, n_rt)): + 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: + 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, finalizers, n_rt, common_vars, reference, model_order, metric_order + ) From afc39f53b9e37eeb1e2a40f4908edbbaeeeedac9 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Thu, 4 Jun 2026 15:58:55 +0300 Subject: [PATCH 05/11] set up prefetching data --- src/mxalign/verification_fused.py | 59 ++++++++++++++++++++++++++++++- 1 file changed, 58 insertions(+), 1 deletion(-) diff --git a/src/mxalign/verification_fused.py b/src/mxalign/verification_fused.py index 7a7175e..4e51e4e 100644 --- a/src/mxalign/verification_fused.py +++ b/src/mxalign/verification_fused.py @@ -29,6 +29,7 @@ import logging import statistics +import threading import time import warnings from collections import deque @@ -482,6 +483,47 @@ def _log_progress(done, total, t_start, timings_window, in_flight): LOG.info(" ".join(parts)) +def _prefetch_nc_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, + recipes: dict, + prefetch_ahead: int, +) -> None: + """Start a background prefetch daemon thread for the forecast NC file(s) + belonging to rt_values[idx + prefetch_ahead], if any. + Only fires for 'anemoi-inference-nc' recipes (not zarr). + """ + target_idx = idx + prefetch_ahead + if target_idx >= len(rt_values): + return + rt = rt_values[target_idx] + for name, recipe in recipes.items(): + if recipe.get("kind") != "anemoi-inference-nc": + continue + key = _rt_key(rt) + path = recipe.get("files_by_rt", {}).get(key) + if path: + threading.Thread( + target=_prefetch_nc_file, + args=(path,), + daemon=True, + name=f"mxalign-prefetch-{name}-{target_idx}", + ).start() + + def compute_metrics_fused( datasets, loaders, @@ -546,6 +588,15 @@ def compute_metrics_fused( else: max_in_flight = 1 + # Prefetch: background daemon threads warm the OS page cache for the next + # NC file(s) while the current leaf is being processed. Enabled via + # `prefetch: true` in the `verification:` yaml block. Only fires for + # anemoi-inference-nc recipes; zarr datasets are skipped. + prefetch_enabled = bool(engine_cfg.get("prefetch", False)) + # Look-ahead depth: start prefetching the file for leaf N+prefetch_ahead + # when leaf N is submitted/consumed. Default max_in_flight+1. + 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 recipes={%s}", @@ -586,6 +637,8 @@ def _consume(result): 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, recipes, prefetch_ahead) try: result = _leaf(rt, lead_times_ns, **leaf_kwargs) except Exception: @@ -622,8 +675,10 @@ def _consume(result): ac = as_completed() i_next = 0 in_flight = 0 - # Prime the window. + # Prime the window (and optionally prime the prefetch pipeline). for _ in range(min(max_in_flight, n_rt)): + if prefetch_enabled: + _schedule_prefetch(rt_values, i_next, recipes, prefetch_ahead) fut = client.submit(_leaf_bundled, rt_values[i_next], static_future, pure=False) fut._mxalign_rt_idx = i_next # informational @@ -647,6 +702,8 @@ def _consume(result): except Exception: pass if i_next < n_rt: + if prefetch_enabled: + _schedule_prefetch(rt_values, i_next, recipes, prefetch_ahead) fut2 = client.submit(_leaf_bundled, rt_values[i_next], static_future, pure=False) fut2._mxalign_rt_idx = i_next From 415be8aca747f9934b44bb8f3ef461a84e4447a8 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Thu, 4 Jun 2026 17:12:16 +0300 Subject: [PATCH 06/11] move slice to base --- src/mxalign/loaders/anemoi_datasets.py | 79 ++++++++++++++++++++ src/mxalign/loaders/anemoi_inference.py | 99 ++++++++++++++++++++++--- src/mxalign/loaders/base.py | 25 +++++++ 3 files changed, 192 insertions(+), 11 deletions(-) diff --git a/src/mxalign/loaders/anemoi_datasets.py b/src/mxalign/loaders/anemoi_datasets.py index 147e2d7..ff3bfcc 100644 --- a/src/mxalign/loaders/anemoi_datasets.py +++ b/src/mxalign/loaders/anemoi_datasets.py @@ -73,6 +73,85 @@ def fast_slice_recipe(self): "drop_vars": list(DROP_VARS), } + 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 bb10df3..8bd63f3 100644 --- a/src/mxalign/loaders/anemoi_inference.py +++ b/src/mxalign/loaders/anemoi_inference.py @@ -69,20 +69,20 @@ def _load(self): return ds - def fast_slice_recipe(self): - """Per-reference_time loading recipe for the fused engine. + def _files_by_rt(self): + """Return ({rt_ns: path}, engine) or (None, engine) if unsupported. - 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). + Shared by ``fast_slice_recipe`` and ``slice``: 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 + return None, engine if Path(files[0]).suffix.lower() == ".zarr": - return None - - engine = self.kwargs.get("engine", DEFAULTS_NETCDF["engine"]) + 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] @@ -96,15 +96,92 @@ def fast_slice_recipe(self): ) ) except ValueError: - return None + return None, engine + + return ( + {int(rt.astype("int64")): f for rt, f in zip(rt_values, files)}, + engine, + ) + + def fast_slice_recipe(self): + """Per-reference_time loading recipe for the fused engine. - files_by_rt = {int(rt.astype("int64")): f for rt, f in zip(rt_values, files)} + 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_by_rt, engine = self._files_by_rt() + if files_by_rt is None: + return None return { "kind": "anemoi-inference-nc", "files_by_rt": files_by_rt, "engine": engine, } + 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. diff --git a/src/mxalign/loaders/base.py b/src/mxalign/loaders/base.py index 2053a46..d56399e 100644 --- a/src/mxalign/loaders/base.py +++ b/src/mxalign/loaders/base.py @@ -118,6 +118,31 @@ def fast_slice_recipe(self): """ return None + 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): From f47b3876e6ccddeae0b074a25bc69e5a3396aaf2 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Thu, 4 Jun 2026 17:15:28 +0300 Subject: [PATCH 07/11] add variable awareness to transfos --- src/mxalign/transformations/base.py | 26 ++++++++++++++++++--- src/mxalign/transformations/external.py | 6 ++++- src/mxalign/transformations/registry.py | 31 ++++++++++++++++++++++++- 3 files changed, 58 insertions(+), 5 deletions(-) diff --git a/src/mxalign/transformations/base.py b/src/mxalign/transformations/base.py index 44a2840..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,7 +42,14 @@ 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 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..14e915f 100644 --- a/src/mxalign/transformations/registry.py +++ b/src/mxalign/transformations/registry.py @@ -1,9 +1,28 @@ _TRANSFORMATION_REGISTRY = {} +_SIGNATURE_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 +37,13 @@ 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) From 31fa402c631791a1525a20bedb2c632d94ea104e Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Tue, 9 Jun 2026 13:03:39 +0300 Subject: [PATCH 08/11] refactor fused --- src/mxalign/__init__.py | 5 + src/mxalign/loaders/anemoi_datasets.py | 20 - src/mxalign/loaders/anemoi_inference.py | 27 +- src/mxalign/loaders/base.py | 8 - src/mxalign/scores.py | 75 ++++ src/mxalign/verification.py | 50 +++ src/mxalign/verification_fused.py | 561 +++++++++--------------- 7 files changed, 351 insertions(+), 395 deletions(-) create mode 100644 src/mxalign/scores.py 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/loaders/anemoi_datasets.py b/src/mxalign/loaders/anemoi_datasets.py index ff3bfcc..d8da9bd 100644 --- a/src/mxalign/loaders/anemoi_datasets.py +++ b/src/mxalign/loaders/anemoi_datasets.py @@ -53,26 +53,6 @@ def _load(self): ) return ds_selected.to_dataset(dim="variable") - def fast_slice_recipe(self): - """Recipe for per-rt direct zarr region read (fused engine). - - Only the single-file zarr path is supported in v1. The leaf - computes valid_times = rt + lead_times and uses zarr's vectorised - indexing to fetch one slice; no xarray/dask lazy graph involved. - """ - if isinstance(self.files, list): - if len(self.files) != 1: - return None - path = self.files[0] - else: - path = self.files - return { - "kind": "anemoi-datasets-zarr", - "path": path, - "consolidated": False, - "drop_vars": list(DROP_VARS), - } - def slice(self, reference_time, lead_times, variables): """Eagerly read one (reference_time, lead_times, variables) slice. diff --git a/src/mxalign/loaders/anemoi_inference.py b/src/mxalign/loaders/anemoi_inference.py index 8bd63f3..873aee4 100644 --- a/src/mxalign/loaders/anemoi_inference.py +++ b/src/mxalign/loaders/anemoi_inference.py @@ -72,10 +72,10 @@ def _load(self): def _files_by_rt(self): """Return ({rt_ns: path}, engine) or (None, engine) if unsupported. - Shared by ``fast_slice_recipe`` and ``slice``: 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). + 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"]) @@ -103,21 +103,16 @@ def _files_by_rt(self): engine, ) - def fast_slice_recipe(self): - """Per-reference_time loading recipe for the fused engine. - - 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). + 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, engine = self._files_by_rt() + files_by_rt, _ = self._files_by_rt() if files_by_rt is None: return None - return { - "kind": "anemoi-inference-nc", - "files_by_rt": files_by_rt, - "engine": engine, - } + 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. diff --git a/src/mxalign/loaders/base.py b/src/mxalign/loaders/base.py index d56399e..83fa00e 100644 --- a/src/mxalign/loaders/base.py +++ b/src/mxalign/loaders/base.py @@ -110,14 +110,6 @@ def _get_properties(self, ds): ) return properties - def fast_slice_recipe(self): - """Return a small picklable dict the fused verification engine can use - to load one reference_time slice directly from the underlying store, - bypassing the lazy xarray/dask graphs. Return None to opt out - (the fused engine will reject this loader for the dataset). - """ - return None - def slice(self, reference_time, lead_times, variables): """Eagerly load a single per-reference_time slice. diff --git a/src/mxalign/scores.py b/src/mxalign/scores.py new file mode 100644 index 0000000..d14703a --- /dev/null +++ b/src/mxalign/scores.py @@ -0,0 +1,75 @@ +"""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 diff --git a/src/mxalign/verification.py b/src/mxalign/verification.py index 88ad172..00bbfd4 100644 --- a/src/mxalign/verification.py +++ b/src/mxalign/verification.py @@ -2,6 +2,56 @@ from functools import partial +# --------------------------------------------------------------------------- +# 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 diff --git a/src/mxalign/verification_fused.py b/src/mxalign/verification_fused.py index 4e51e4e..7c91d48 100644 --- a/src/mxalign/verification_fused.py +++ b/src/mxalign/verification_fused.py @@ -1,291 +1,99 @@ -"""Fused verification engine (Phase 2 / lever B, recipe-based). +"""Fused verification engine. For each reference_time, one client.submit task: - 1. Loads the per-rt slice directly from the underlying store - (NetCDF per rt for forecasts; zarr region read for ERA5) - via the loader's `fast_slice_recipe`, bypassing xarray's lazy graphs. - 2. Replays the registered transformations on that small per-rt Dataset. + 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 +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. - -Scope (v1): - - Sum-decomposable metrics with `reduce_dims` containing 'reference_time': - MSE, MAE, bias (mean error), mean(reference), mean(forecast). - - Loaders: anemoi-inference (per-rt NetCDF), anemoi-datasets (single zarr). - - Transformations: rename, kelvin_to_celcius, uv_to_speed (extend by - adding entries to `_TRANSFORM_IO`). - -Validation failures (unsupported metric/transform/loader, missing -reduce_dims, missing fast_slice_recipe) raise immediately. No silent -fallback. +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 -from typing import Any, Callable import numpy as np import xarray as xr -LOG = logging.getLogger("mxalign") - - -# --------------------------------------------------------------------------- -# Metric kernels -# --------------------------------------------------------------------------- -# Each kernel takes (fcst, ref) numpy arrays of shape (n_var, n_lt, n_grid) -# and returns a per-sample partial of the same shape that is **summable** -# across reference_time. The finalize step (mean = sum / N, sum = sum) -# is applied after all leaves are reduced. - -def _kernel_squared_error(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: - diff = fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) - return diff * diff - - -def _kernel_abs_error(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: - return np.abs( - fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) - ) - - -def _kernel_error(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: - return fcst.astype(np.float32, copy=False) - ref.astype(np.float32, copy=False) - - -def _kernel_identity_fcst(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: - return fcst.astype(np.float32, copy=False) +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 - -def _kernel_identity_ref(fcst: np.ndarray, ref: np.ndarray) -> np.ndarray: - return ref.astype(np.float32, copy=False) - - -# func_path -> (kernel, finalize_kind in {"mean", "sum"}) -_FUSED_KERNELS: dict[str, tuple[Callable, str]] = { - "scores.continuous.mse": (_kernel_squared_error, "mean"), - "scores.continuous.mae": (_kernel_abs_error, "mean"), - "scores.continuous.bias": (_kernel_error, "mean"), - "scores.continuous.mean_error": (_kernel_error, "mean"), -} +LOG = logging.getLogger("mxalign") # --------------------------------------------------------------------------- -# Transformation source-variable bookkeeping +# Helpers # --------------------------------------------------------------------------- -# Each entry returns (inputs, outputs) variable lists given the transform's -# kwargs (as recorded by Runner.transform_datasets). Used to walk the -# transformation chain backwards from `common_vars` to "what to load from -# the source". - -def _io_uv_to_speed(kwargs): - u = kwargs["u"]; v = kwargs["v"]; s = kwargs["speed"] - u = [u] if isinstance(u, str) else list(u) - v = [v] if isinstance(v, str) else list(v) - s = [s] if isinstance(s, str) else list(s) - return u + v, s - - -def _io_kelvin_to_celcius(kwargs): - v = kwargs["variables"] - v = [v] if isinstance(v, str) else list(v) - return v, v # in-place - -def _io_rename(kwargs): - d = kwargs["rename_dict"] # new_name -> old_name(s) - outputs = list(d.keys()) - inputs: list[str] = [] - for v in d.values(): - inputs.extend(v if isinstance(v, list) else [v]) - return inputs, outputs - - -_TRANSFORM_IO: dict[str, Callable] = { - "uv_to_speed": _io_uv_to_speed, - "kelvin_to_celcius": _io_kelvin_to_celcius, - "rename": _io_rename, -} +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 to derive the set of source variables - that need to be read from the store for one dataset.""" + """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): - if tname not in _TRANSFORM_IO: + sig = get_signature(tname) + if sig is None: raise NotImplementedError( - f"fused engine: transformation {tname!r} has no input/output " - f"spec in _TRANSFORM_IO; add one or use engine=xarray" + 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 = _TRANSFORM_IO[tname](tkwargs) + inputs, outputs = sig(**tkwargs) if any(o in needed for o in outputs): needed -= set(outputs) needed |= set(inputs) return sorted(needed) -# --------------------------------------------------------------------------- -# Per-rt slice loaders (worker-side) -# --------------------------------------------------------------------------- - -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 _load_slice(recipe, rt_value, lead_times, var_names) -> xr.Dataset: - kind = recipe["kind"] - if kind == "anemoi-inference-nc": - return _load_anemoi_inference_slice(recipe, rt_value, lead_times, var_names) - if kind == "anemoi-datasets-zarr": - return _load_anemoi_datasets_slice(recipe, rt_value, lead_times, var_names) - raise NotImplementedError(f"fused engine: unknown recipe kind {kind!r}") - - -def _load_anemoi_inference_slice(recipe, rt_value, lead_times, var_names) -> xr.Dataset: - path = recipe["files_by_rt"][_rt_key(rt_value)] - engine = recipe["engine"] - with xr.open_dataset(path, engine=engine) as src: - # Subset variables + lead_times *lazily* and only then call .load(). - # Doing .load() up front (the previous behaviour) forces a read of the - # full time axis even when the file holds more steps than we need; it - # also turns the time-axis selection into an in-memory fancy index - # instead of a hyperslab read. Selecting first lets HDF5 issue a - # single contiguous read for the steady-state (cadence-1) case. - sub = src[list(var_names)] - 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(int(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"fused engine: missing lead_times in {path}: {bad[:5]}... " - f"(reference_time={rt_value})" - ) - # 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")) - ds = sub.load() - return ds - - -def _load_anemoi_datasets_slice(recipe, rt_value, lead_times, var_names) -> xr.Dataset: - path = recipe["path"] - src = xr.open_zarr(path, consolidated=recipe.get("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 var_names], dtype=np.int64) - except ValueError as e: - raise ValueError( - f"fused engine: variable not found in {path}: {e}" - ) from None - - rt = np.datetime64(rt_value, "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"fused engine: missing valid_times in {path}: {bad[:5]}... " - f"(reference_time={rt_value})" - ) - - 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"), - ) - loaded = arr_sel.load() - vals = np.asarray(loaded.values) # (n_lt, n_var, n_grid) - # `vals` lead_time axis matches the slice/index we asked for; for the - # contiguous-slice path it's already in the right order (and length). - ds = xr.Dataset( - { - v: (("lead_time", "grid_index"), vals[:, i, :]) - for i, v in enumerate(var_names) - } - ) - return ds - - # --------------------------------------------------------------------------- # Leaf task (runs on worker) # --------------------------------------------------------------------------- def _leaf( rt_value, - lead_times, + lead_times_ns, common_vars, ref_name, model_names, - recipes_by_ds, + loaders, source_vars_by_ds, transforms_by_ds, - metric_kernels, + metric_specs, # {metric_name: (kernel_callable, inputs_or_None)} ): """One per-reference_time task. @@ -296,17 +104,15 @@ def _leaf( "partials": {model_name: {metric_name: np.ndarray(n_var, n_lt, n_grid)}}, } """ - from mxalign.transformations.registry import get_transformation - t0 = time.perf_counter() timings: dict[str, float] = {} - # 1. Load per-dataset slices. + # 1. Load per-dataset slices via the loader's eager slice() method. slices: dict[str, xr.Dataset] = {} - for ds_name, recipe in recipes_by_ds.items(): + for ds_name, loader in loaders.items(): t = time.perf_counter() - slices[ds_name] = _load_slice( - recipe, rt_value, lead_times, source_vars_by_ds[ds_name] + 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 @@ -330,15 +136,26 @@ def _leaf( axis=0, ) - # 4. Apply kernels. + # 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: - fcst = arrays[m] - partials[m] = { - mn: kernel(fcst, ref) for mn, (kernel, _) in metric_kernels.items() - } + 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 @@ -346,11 +163,12 @@ def _leaf( def _leaf_bundled(rt_value, static): - """Worker-side trampoline: unpack the scattered static bundle and call _leaf. + """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. + ``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, @@ -358,80 +176,134 @@ def _leaf_bundled(rt_value, static): common_vars=static["common_vars"], ref_name=static["ref_name"], model_names=static["model_names"], - recipes_by_ds=static["recipes_by_ds"], + loaders=static["loaders"], source_vars_by_ds=static["source_vars_by_ds"], transforms_by_ds=static["transforms_by_ds"], - metric_kernels=static["metric_kernels"], + metric_specs=static["metric_specs"], ) # --------------------------------------------------------------------------- -# Driver +# Driver: validation # --------------------------------------------------------------------------- -def _validate(reference, datasets, loaders, transforms_by_ds, metrics_cfg, ref_name): - # 1. Every dataset must have a recipe. - recipes: dict[str, dict] = {} - for name, loader in loaders.items(): - if not hasattr(loader, "fast_slice_recipe"): - raise NotImplementedError( - f"fused engine: loader for dataset {name!r} has no " - f"fast_slice_recipe(); use engine=xarray or extend the loader." +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)}." ) - recipe = loader.fast_slice_recipe() - if recipe is None: + 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"fused engine: loader {type(loader).__name__!r} declined to " - f"produce a fast-slice recipe for dataset {name!r} (e.g. " - f"unsupported file layout); use engine=xarray." + 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." ) - recipes[name] = recipe - # 2. Every metric must be in the allow-list with reduce_dims=[reference_time]. - metric_kernels: dict[str, tuple[Callable, str]] = {} + # 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 func_path not in _FUSED_KERNELS: + 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"fused engine: metric {mn!r} uses function {func_path!r} which " - f"is not in the sum-decomposable allow-list " - f"({sorted(_FUSED_KERNELS)}); use engine=xarray." + 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"fused engine: metric {mn!r} has reduce_dims={rd}; the fused " + f"engine=fused: metric {mn!r} has reduce_dims={rd}; the fused " f"engine requires 'reference_time' among reduce_dims." ) - metric_kernels[mn] = _FUSED_KERNELS[func_path] + inputs = _resolve_kernel_inputs(mn, func_path, kernel, mcfg.get("inputs")) + metric_specs[mn] = (kernel, inputs) + metric_finalizers[mn] = finalize - # 3. Reference must be one of the datasets. - if ref_name not in datasets: - raise ValueError(f"fused engine: reference {ref_name!r} not in datasets") + return metric_specs, metric_finalizers - return recipes, metric_kernels +# --------------------------------------------------------------------------- +# 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)``. -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). + 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 - # (model, metric, variable, lead_time, grid_index) arr_by_metric: dict[str, np.ndarray] = {} for mn in metric_order: finalize = finalizers[mn] stacked = np.stack( - [ - accums[m][mn] / (n_rt if finalize == "mean" else 1) - for m in model_order - ], + [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 @@ -457,6 +329,10 @@ def _make_xr_result(accums, finalizers, n_rt, common_vars, reference, model_orde ).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 @@ -469,7 +345,6 @@ def _log_progress(done, total, t_start, timings_window, in_flight): f"eta={eta:.0f}s", ] if timings_window: - # Collect per-stage timings across the window. keys = set().union(*(t.keys() for t in timings_window)) bits = [] for k in sorted(keys): @@ -483,10 +358,10 @@ def _log_progress(done, total, t_start, timings_window, in_flight): LOG.info(" ".join(parts)) -def _prefetch_nc_file(path: str) -> None: +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.""" + 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 @@ -496,34 +371,34 @@ def _prefetch_nc_file(path: str) -> None: pass -def _schedule_prefetch( - rt_values, - idx: int, - recipes: dict, - prefetch_ahead: int, -) -> None: - """Start a background prefetch daemon thread for the forecast NC file(s) - belonging to rt_values[idx + prefetch_ahead], if any. - Only fires for 'anemoi-inference-nc' recipes (not zarr). - """ +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, recipe in recipes.items(): - if recipe.get("kind") != "anemoi-inference-nc": + for name, loader in loaders.items(): + get_path = getattr(loader, "prefetch_path", None) + if get_path is None: continue - key = _rt_key(rt) - path = recipe.get("files_by_rt", {}).get(key) + path = get_path(rt) if path: threading.Thread( - target=_prefetch_nc_file, + target=_prefetch_file, args=(path,), daemon=True, name=f"mxalign-prefetch-{name}-{target_idx}", ).start() +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + def compute_metrics_fused( datasets, loaders, @@ -533,46 +408,38 @@ def compute_metrics_fused( metrics_cfg, engine_cfg, ): - """Driver entry point. Returns an xr.Dataset shaped - (model, metric, variable, lead_time, grid_index).""" + """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()) - recipes, metric_kernels = _validate( - reference, datasets, loaders, transforms_by_ds, metrics_cfg, reference_name + metric_specs, metric_finalizers = _validate( + reference, datasets, loaders, metrics_cfg, reference_name ) - # Derive per-dataset source variables (walk transformations backwards). + # 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 } - # Per-rt iteration: drive from the reference dataset's reference_time. - if "reference_time" not in reference.dims: - raise ValueError( - "fused engine: reference dataset has no 'reference_time' dim; " - "this engine requires forecast-shaped reference." - ) 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) - finalizers = {mn: kind for mn, (_, kind) in metric_kernels.items()} - - # Pre-allocate driver-side accumulators (one per model+metric, ~1.5GB each). n_var = len(common_vars) n_lt = len(lead_times) - # We let the first arriving partial allocate via copy; saves a guess at n_grid. + # 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 } - # Try to get a Client; if none, run serial in-process. + # Optional dask client. client = None try: from dask.distributed import default_client, as_completed @@ -588,21 +455,18 @@ def compute_metrics_fused( else: max_in_flight = 1 - # Prefetch: background daemon threads warm the OS page cache for the next - # NC file(s) while the current leaf is being processed. Enabled via - # `prefetch: true` in the `verification:` yaml block. Only fires for - # anemoi-inference-nc recipes; zarr datasets are skipped. + # 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)) - # Look-ahead depth: start prefetching the file for leaf N+prefetch_ahead - # when leaf N is submitted/consumed. Default max_in_flight+1. 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 recipes={%s}", + "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}:{r['kind']}" for n, r in recipes.items()), + ", ".join(f"{n}:{type(l).__name__}" for n, l in loaders.items()), ) timings_window: deque = deque(maxlen=64) @@ -613,8 +477,7 @@ def compute_metrics_fused( def _consume(result): nonlocal done, last_completion - partials = result["partials"] - for m, per_metric in partials.items(): + 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 @@ -628,17 +491,17 @@ def _consume(result): common_vars=common_vars, ref_name=reference_name, model_names=model_order, - recipes_by_ds=recipes, + loaders=loaders, source_vars_by_ds=source_vars_by_ds, transforms_by_ds=transforms_by_ds, - metric_kernels=metric_kernels, + 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, recipes, prefetch_ahead) + _schedule_prefetch(rt_values, i, loaders, prefetch_ahead) try: result = _leaf(rt, lead_times_ns, **leaf_kwargs) except Exception: @@ -650,14 +513,10 @@ def _consume(result): _log_progress(done, n_rt, t_start, list(timings_window), 0) last_progress_log = now else: - # Scatter the (large, identical-per-submit) static payload once and - # broadcast it to all workers. Each subsequent client.submit then ships - # only the per-leaf rt + lead_times + a pointer to the scattered - # bundle, keeping the per-submit graph size in the KB range. - # `lead_times_ns` is small (<=145 ints) but we scatter it too for - # symmetry. Broadcast=True ensures it's already on every worker before - # the first submit, so workers never pull from the scheduler at task - # start. + # 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) @@ -675,10 +534,9 @@ def _consume(result): ac = as_completed() i_next = 0 in_flight = 0 - # Prime the window (and optionally prime the prefetch pipeline). for _ in range(min(max_in_flight, n_rt)): if prefetch_enabled: - _schedule_prefetch(rt_values, i_next, recipes, prefetch_ahead) + _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 @@ -703,7 +561,7 @@ def _consume(result): pass if i_next < n_rt: if prefetch_enabled: - _schedule_prefetch(rt_values, i_next, recipes, prefetch_ahead) + _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 @@ -726,5 +584,6 @@ def _consume(result): _log_progress(done, n_rt, t_start, list(timings_window), 0) return _make_xr_result( - accums, finalizers, n_rt, common_vars, reference, model_order, metric_order + accums, metric_finalizers, n_rt, common_vars, reference, model_order, + metric_order, ) From cbc41db3176767867d74b37cf4b09d7522830394 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Tue, 9 Jun 2026 14:06:19 +0300 Subject: [PATCH 09/11] add gradient transfo --- src/mxalign/transformations/__init__.py | 2 + src/mxalign/transformations/cerra.py | 177 ++++++++++++++++++++++++ src/mxalign/verification.py | 131 ++++++++++++++++++ tests/test_cerra_gradient.py | 129 +++++++++++++++++ 4 files changed, 439 insertions(+) create mode 100644 src/mxalign/transformations/cerra.py create mode 100644 tests/test_cerra_gradient.py 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/cerra.py b/src/mxalign/transformations/cerra.py new file mode 100644 index 0000000..28770e6 --- /dev/null +++ b/src/mxalign/transformations/cerra.py @@ -0,0 +1,177 @@ +"""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 xarray as xr + +from .registry import register_transformation +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) + result = _compute_gradient(da_t.values, axis_xy, ny, nx, dx, dy) + out = xr.DataArray( + result, + dims=da_t.dims, + coords={k: v for k, v in da_t.coords.items() if set(v.dims).issubset(da_t.dims)}, + name=da.name, + attrs=dict(da.attrs), + ) + return out.transpose(*da.dims) + + +def _apply(ds: xr.Dataset, variables, outputs, axis_xy: str, *, + grid_dim: str = "grid_index", + ny: int = _NY_DEFAULT, nx: int = _NX_DEFAULT, + dx: float = _DX_DEFAULT, dy: float = _DY_DEFAULT) -> xr.Dataset: + vs = [variables] if isinstance(variables, str) else list(variables) + 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, **_): + v = [variables] if isinstance(variables, str) else list(variables) + o = [outputs] if isinstance(outputs, str) else list(outputs) + return v, o + + +@register_transformation("cerra_gradient_x", signature=_sig_cerra_grad) +def transform_cerra_gradient_x(ds, variables, outputs, **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, **grid_kwargs): + return _apply(ds, variables, outputs, axis_xy="y", **grid_kwargs) diff --git a/src/mxalign/verification.py b/src/mxalign/verification.py index 00bbfd4..3ce37e4 100644 --- a/src/mxalign/verification.py +++ b/src/mxalign/verification.py @@ -1,6 +1,83 @@ 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 @@ -55,6 +132,15 @@ def get_fused_kernel(fn): 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) @@ -82,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/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 + ) From 81c4a5be5915386e1c3c1e9390825bd3de5d2d11 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Wed, 10 Jun 2026 14:35:31 +0300 Subject: [PATCH 10/11] add nvv --- src/mxalign/scores.py | 110 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/src/mxalign/scores.py b/src/mxalign/scores.py index d14703a..60f5a12 100644 --- a/src/mxalign/scores.py +++ b/src/mxalign/scores.py @@ -73,3 +73,113 @@ def bias(fcst, obs, reduce_dims=None, **_): # so existing YAMLs that say ``function: scores.continuous.mean_error`` can # migrate to ``function: mxalign.scores.mean_error`` without semantic drift. mean_error = bias + + +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: + groups = { + grp_label: { + "components": grp_cfg["components"], + "eps": grp_cfg.get("eps", eps), + } + for grp_label, grp_cfg in variables.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) From 195efbd13eb7f3a07e8271caca387dccd6e69961 Mon Sep 17 00:00:00 2001 From: dietervdb-meteo Date: Mon, 15 Jun 2026 15:56:10 +0300 Subject: [PATCH 11/11] add nvv --- src/mxalign/runner.py | 35 ++- src/mxalign/scores.py | 97 ++++++- src/mxalign/transformations/cerra.py | 95 ++++++- src/mxalign/transformations/registry.py | 23 ++ src/mxalign/verification_fused.py | 361 ++++++++++++++++++++++++ 5 files changed, 591 insertions(+), 20 deletions(-) diff --git a/src/mxalign/runner.py b/src/mxalign/runner.py index 5ee3cc6..d6d2833 100644 --- a/src/mxalign/runner.py +++ b/src/mxalign/runner.py @@ -76,13 +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 ) - # Record (transform_name, kwargs) in application order for - # the fused engine to replay on per-rt slices. self._transforms_by_ds.setdefault(name, []).append( - (transformation, dict(config_trans)) + (transformation, dict(resolved)) ) def align(self): @@ -166,6 +169,30 @@ def verify(self): 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", diff --git a/src/mxalign/scores.py b/src/mxalign/scores.py index 60f5a12..7a3c8d1 100644 --- a/src/mxalign/scores.py +++ b/src/mxalign/scores.py @@ -75,6 +75,94 @@ def bias(fcst, obs, reduce_dims=None, **_): 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. @@ -147,12 +235,19 @@ def nvv(fcst, obs, reduce_dims=None, components=None, variables=None, # ---- 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 variables.items() + for grp_label, grp_cfg in resolved.items() } elif components: coord = label if label is not None else "+".join(components) diff --git a/src/mxalign/transformations/cerra.py b/src/mxalign/transformations/cerra.py index 28770e6..95eee43 100644 --- a/src/mxalign/transformations/cerra.py +++ b/src/mxalign/transformations/cerra.py @@ -21,9 +21,10 @@ from __future__ import annotations +import fnmatch import xarray as xr -from .registry import register_transformation +from .registry import register_transformation, register_expander from ..utils.projections import BUILTIN # Defaults sourced from the canonical CERRA grid description. @@ -126,23 +127,57 @@ def _cerra_gradient_axis(da: xr.DataArray, axis_xy: str, *, grid_dim: str, ) da_t = da.transpose(..., grid_dim) - result = _compute_gradient(da_t.values, axis_xy, ny, nx, dx, dy) - out = xr.DataArray( - result, - dims=da_t.dims, - coords={k: v for k, v in da_t.coords.items() if set(v.dims).issubset(da_t.dims)}, - name=da.name, - attrs=dict(da.attrs), + + # 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}, ) - return out.transpose(*da.dims) + 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, axis_xy: str, *, +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 = [variables] if isinstance(variables, str) else list(variables) - os_ = [outputs] if isinstance(outputs, str) else list(outputs) + 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, " @@ -161,17 +196,47 @@ def _apply(ds: xr.Dataset, variables, outputs, axis_xy: str, *, # --------------------------------------------------------------------------- -def _sig_cerra_grad(variables, outputs, **_): +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, **grid_kwargs): +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, **grid_kwargs): +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/registry.py b/src/mxalign/transformations/registry.py index 14e915f..1020fda 100644 --- a/src/mxalign/transformations/registry.py +++ b/src/mxalign/transformations/registry.py @@ -1,5 +1,6 @@ _TRANSFORMATION_REGISTRY = {} _SIGNATURE_REGISTRY = {} +_EXPANDER_REGISTRY = {} def register_transformation(name, signature=None): @@ -47,3 +48,25 @@ def get_signature(name): 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/verification_fused.py b/src/mxalign/verification_fused.py index 7c91d48..9c60554 100644 --- a/src/mxalign/verification_fused.py +++ b/src/mxalign/verification_fused.py @@ -587,3 +587,364 @@ def _consume(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)