From d40c9050a8e73a2bf438d1a04fbd155b4295ba5b Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Fri, 5 Jun 2026 17:47:22 +0200 Subject: [PATCH 01/21] add power-spectra QC diagnostic --- pyproject.toml | 3 +- src/evalml/config.py | 59 ++++++- src/spectra/__init__.py | 7 + src/spectra/compute.py | 119 ++++++++++++++ src/spectra/core.py | 201 ++++++++++++++++++++++++ src/spectra/io.py | 69 +++++++++ src/spectra/regrid.py | 90 +++++++++++ tests/unit/test_config.py | 53 ++++++- tests/unit/test_spectra_core.py | 103 +++++++++++++ tests/unit/test_spectra_io.py | 87 +++++++++++ tests/unit/test_spectra_regrid.py | 82 ++++++++++ workflow/Snakefile | 6 + workflow/rules/spectra.smk | 213 ++++++++++++++++++++++++++ workflow/scripts/spectra_aggregate.py | 28 ++++ workflow/scripts/spectra_compute.py | 79 ++++++++++ workflow/scripts/spectra_plot.py | 32 ++++ workflow/tools/config.schema.json | 65 ++++++++ 17 files changed, 1292 insertions(+), 4 deletions(-) create mode 100644 src/spectra/__init__.py create mode 100644 src/spectra/compute.py create mode 100644 src/spectra/core.py create mode 100644 src/spectra/io.py create mode 100644 src/spectra/regrid.py create mode 100644 tests/unit/test_spectra_core.py create mode 100644 tests/unit/test_spectra_io.py create mode 100644 tests/unit/test_spectra_regrid.py create mode 100644 workflow/rules/spectra.smk create mode 100644 workflow/scripts/spectra_aggregate.py create mode 100644 workflow/scripts/spectra_compute.py create mode 100644 workflow/scripts/spectra_plot.py diff --git a/pyproject.toml b/pyproject.toml index d3279e91..e3033ad2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,8 @@ markers = [ packages = [ "src/evalml", "src/verification", - "src/data_input" + "src/data_input", + "src/spectra" ] [tool.uv.sources] diff --git a/src/evalml/config.py b/src/evalml/config.py index c36ea4df..c2822835 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -1,7 +1,7 @@ from pathlib import Path -from typing import Dict, List, Any, ClassVar, FrozenSet, Optional +from typing import Dict, List, Any, ClassVar, FrozenSet, Literal, Optional -from pydantic import BaseModel, Field, RootModel, field_validator +from pydantic import BaseModel, Field, RootModel, field_validator, model_validator PROJECT_ROOT = Path(__file__).parents[2] @@ -302,6 +302,57 @@ class ExperimentScorecardConfig(BaseModel): model_config = {"extra": "forbid"} +class SpectraConfig(BaseModel): + """Configuration for power-spectra QC plots in the experiment pipeline.""" + + enabled: bool = Field( + default=False, + description="Whether to compute and plot power spectra.", + ) + method: Literal["dct", "fft"] = Field( + default="dct", + description="Spectral method: 'dct' (default, recommended for LAM) or 'fft'.", + ) + lead_times: List[int] = Field( + default_factory=list, + description="Representative lead times (hours) at which to compute spectra.", + ) + variables: List[str] = Field( + default=["T_2M", "WIND_KE", "TOT_PREC"], + description="Spectra variables. Supported: T_2M, WIND_KE (from U/V_10M), TOT_PREC.", + ) + init_hours: Optional[List[int]] = Field( + default=None, + description="Optional subset of init hours to average over. None = all.", + ) + + @field_validator("variables") + @classmethod + def validate_variables(cls, v: List[str]) -> List[str]: + if not v: + raise ValueError( + "`variables` must list at least one variable when spectra is configured." + ) + allowed = {"T_2M", "WIND_KE", "TOT_PREC"} + invalid = set(v) - allowed + if invalid: + raise ValueError( + f"Unsupported spectra variable(s) {invalid!r}. Must be subset of {allowed}." + ) + return v + + @model_validator(mode="after") + def check_enabled_requirements(self): + if self.enabled and not self.lead_times: + raise ValueError( + "experiment.spectra.enabled is true but `lead_times` is empty; " + "specify at least one lead time (hours)." + ) + return self + + model_config = {"extra": "forbid"} + + class ShowcaseConfig(BaseModel): """Configuration for the showcase workflow.""" @@ -373,6 +424,10 @@ class ExperimentConfig(BaseModel): default=None, description="Scorecard generation configuration. Omit or set enabled: false to disable.", ) + spectra: SpectraConfig = Field( + default_factory=SpectraConfig, + description="Power-spectra QC configuration. Disabled by default.", + ) @field_validator("thresholds") @classmethod diff --git a/src/spectra/__init__.py b/src/spectra/__init__.py new file mode 100644 index 00000000..7599fd4a --- /dev/null +++ b/src/spectra/__init__.py @@ -0,0 +1,7 @@ +"""Variance power-spectra diagnostics for evalml. + +Public surface re-exported for convenience; submodules can also be imported +directly. The spectral `core` is numpy-only and reusable for any model. +""" + +from spectra import compute, core, io, regrid # noqa: F401 diff --git a/src/spectra/compute.py b/src/spectra/compute.py new file mode 100644 index 00000000..336788fa --- /dev/null +++ b/src/spectra/compute.py @@ -0,0 +1,119 @@ +"""Orchestration: per-source spectra computation, over-init aggregation, and the +experiment overlay plot. Defines the spectra.nc schema.""" + +from __future__ import annotations + +from itertools import cycle +from pathlib import Path + +import numpy as np +import xarray as xr + +from spectra import core, io, regrid + +_PALETTE = ["#1f3b73", "#c1272d", "#1b7837", "#6a3d9a", "#e08214", "#4d4d4d"] + + +def compute_source_spectra(ds: xr.Dataset, variables, lead_times, method, label): + """Compute spectra for one source (already-loaded native dataset). + + Returns an x.Dataset with dims (variable, leadtime, wavenumber) and a shared + `wavelength` coordinate. The native grid is detected from the field length. + """ + if not variables or not lead_times: + raise ValueError( + "compute_source_spectra requires non-empty `variables` and `lead_times`." + ) + spectrum_fn = core.SPECTRUM_FUNCS.get(method) + if spectrum_fn is None: + raise ValueError( + f"Unknown spectrum method {method!r}. Valid: {list(core.SPECTRUM_FUNCS)}." + ) + npoints = io.native_field( + ds, io.required_params(variables)[0], lead_times[0] + ).shape[0] + matrix, ny, nx, dx_km = regrid.load_regridder(npoints) + + wavelength = None + power = np.full((len(variables), len(lead_times), min(ny, nx) // 2), np.nan) + for vi, var in enumerate(variables): + for li, step in enumerate(lead_times): + comps, factor = io.native_components(ds, var, step) + grids = [regrid.regrid(c, matrix, ny, nx) for c in comps] + wl, p = core.combined_spectrum(spectrum_fn, grids, dx_km, factor=factor) + wavelength = wl + power[vi, li, :] = p + + return xr.Dataset( + {"power": (("variable", "leadtime", "wavenumber"), power)}, + coords={ + "variable": list(variables), + "leadtime": list(lead_times), + "wavenumber": np.arange(power.shape[-1]), + "wavelength": ("wavenumber", wavelength), + }, + attrs={ + "dx_km": float(dx_km), + "npoints": int(npoints), + "label": label, + "method": method, + }, + ) + + +def aggregate_spectra(spectra_files) -> xr.Dataset: + """Average power over init times (nanmean). All inputs share one grid.""" + datasets = [xr.open_dataset(f) for f in spectra_files] + try: + stacked = xr.concat(datasets, dim="init") + agg = stacked.mean(dim="init", skipna=True) + agg.attrs = datasets[0].attrs + agg["wavelength"] = datasets[0]["wavelength"].copy() + agg.load() + finally: + for ds in datasets: + ds.close() + return agg + + +def plot_experiment_spectra( + truth_file, participant_files, out_dir, variables, lead_times +): + """One overlay figure per (variable, lead time): truth + all participants.""" + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + participants = [xr.open_dataset(f) for f in participant_files] + try: + with xr.open_dataset(truth_file) as truth: + for var in variables: + for step in lead_times: + spectra = {} + t_wl = truth["wavelength"].values + t_p = truth["power"].sel(variable=var, leadtime=step).values + spectra[f"{truth.attrs['label']} (truth)"] = ( + t_wl, + t_p, + "k", + "-", + 1.0, + 2.0, + ) + color = cycle(_PALETTE) + for ds in participants: + wl = ds["wavelength"].values + p = ds["power"].sel(variable=var, leadtime=step).values + spectra[ds.attrs["label"]] = (wl, p, next(color), "-", 0.9, 1.6) + out = out_dir / f"spectrum_{var}_{step:03d}.png" + # Reference grid/eff-res/Nyquist lines use the TRUTH grid's dx (labelled + # "truth"); each participant curve keeps its own wavelength axis. + core.plot_power_spectra( + spectra, + out, + f"Power spectrum — {var} @ +{step}h", + grid_dx_km=float(truth.attrs["dx_km"]), + model_short="truth", + model_color="k", + ) + finally: + for ds in participants: + ds.close() diff --git a/src/spectra/core.py b/src/spectra/core.py new file mode 100644 index 00000000..62ce95e3 --- /dev/null +++ b/src/spectra/core.py @@ -0,0 +1,201 @@ +"""Numpy-only 2D variance (power) spectra of regular-grid fields. + +Ported from the `icon-power-spectra` skill. The only behavioural change vs. the +skill is that `radial_spectrum` returns FIXED-LENGTH arrays (length `nbins`, +`NaN` for empty bins) on a deterministic, edge-centred wavenumber grid, so +spectra from different init times share an identical wavelength coordinate and +can be stacked/averaged. Variance-conserving normalization is unchanged. +""" + +from __future__ import annotations + +from itertools import cycle + +import matplotlib.pyplot as plt +import numpy as np +import scipy.fft + +KM_PER_DEG = 111.2 + +REFERENCE_GRIDS = { + "O96 (~104 km)": 90.0 / 96.0 * KM_PER_DEG, + "N320 (~31 km)": 0.28125 * KM_PER_DEG, + "CERRA (~9 km)": 0.10 * KM_PER_DEG, +} + +EFF_RES_FACTOR = 6.0 +TRANSITION_KM = 400.0 + + +def _hann2d(ny: int, nx: int) -> np.ndarray: + win = np.hanning(ny)[:, None] * np.hanning(nx)[None, :] + return win / np.sqrt((win**2).mean()) + + +def radial_spectrum(coeff, kx, ky, nbins, norm, oro=None): + """Radially bin |coeff|^2 over isotropic |k| [cycles/km] into a density. + + Returns (wavelength_km, power) each of length `nbins`, sorted by ascending + wavelength, with `NaN` in bins that contain no samples. `norm` scales the + summed coefficients to a variance-conserving spectrum so spectra from grids + of different resolution overlap where both resolve. `oro` (transform of the + orography on the same grid), if given, removes the terrain-coherent + component per band (unused in v1). + """ + kk = np.sqrt(kx[None, :] ** 2 + ky[:, None] ** 2).ravel() + c = coeff.ravel() + h = oro.ravel() if oro is not None else None + + k_max = min(kx.max(), ky.max()) + edges = np.linspace(0.0, k_max, nbins + 1) + dk = edges[1] - edges[0] + which = np.digitize(kk, edges) + + power = np.full(nbins, np.nan) + for b in range(1, nbins + 1): + sel = which == b + if not np.any(sel): + continue + cb = c[sel] + s_ff = float(np.vdot(cb, cb).real) + if h is not None: + hb = h[sel] + s_hh = float(np.vdot(hb, hb).real) + if s_hh > 0: + s_fh = np.vdot(hb, cb) + s_ff -= abs(s_fh) ** 2 / s_hh + val = max(s_ff, 0.0) * norm / dk + power[b - 1] = val if val > 0.0 else np.nan + + centre_k = 0.5 * (edges[:-1] + edges[1:]) + wavelength = 1.0 / centre_k + order = np.argsort(wavelength) + return wavelength[order], power[order] + + +def dct_power_spectrum(field, dx_km, nbins=None, oro_coeff=None): + """2D DCT (Denis et al. 2002) variance spectrum of a regular-grid field.""" + field = field - field.mean() + ny, nx = field.shape + nbins = nbins or min(ny, nx) // 2 + coeff = scipy.fft.dctn(field, type=2, norm="ortho") + kx = np.arange(nx) / (2.0 * nx) / dx_km + ky = np.arange(ny) / (2.0 * ny) / dx_km + return radial_spectrum(coeff, kx, ky, nbins, 1.0 / (nx * ny), oro_coeff) + + +def fft_power_spectrum(field, dx_km, nbins=None, oro_coeff=None): + """2D windowed-FFT (Hann) spectrum, normalized to match the DCT.""" + # NOTE: ported verbatim from the icon-power-spectra skill. FFT is the + # optional (non-default) comparison method; the magnitude binning follows + # the reference implementation exactly. Prefer DCT for quantitative use. + field = field - field.mean() + ny, nx = field.shape + nbins = nbins or min(ny, nx) // 2 + coeff = scipy.fft.fft2(field * _hann2d(ny, nx)) + kx = scipy.fft.fftfreq(nx, d=dx_km) + ky = scipy.fft.fftfreq(ny, d=dx_km) + return radial_spectrum(coeff, kx, ky, nbins, 1.0 / (nx * ny) ** 2, oro_coeff) + + +SPECTRUM_FUNCS = {"dct": dct_power_spectrum, "fft": fft_power_spectrum} + + +def combined_spectrum( + spectrum_fn, fields_2d, dx_km, nbins=None, factor=1.0, oro_coeff=None +): + """Sum the spectra of one or more component fields, scaled by `factor`. + + For wind kinetic energy pass [u, v] and factor=0.5 -> 0.5*(|u_hat|^2+|v_hat|^2). + NaN bins are treated as zero in the sum so the common coordinate is preserved. + """ + if not fields_2d: + raise ValueError("combined_spectrum requires at least one field") + wl = total = None + for f in fields_2d: + w, p = spectrum_fn(f, dx_km, nbins, oro_coeff=oro_coeff) + wl = w + total = np.nan_to_num(p) if total is None else total + np.nan_to_num(p) + return wl, factor * total + + +def _add_slope_guide(ax, n, label, wl_range, anchor): + wl = np.array(wl_range, dtype=float) + wl0, p0 = anchor + power = p0 * (wl / wl0) ** n + ax.plot(wl, power, ls=":", color="0.35", lw=1.3) + ax.annotate( + label, + xy=(wl.min(), power[np.argmin(wl)]), + xytext=(3, 3), + textcoords="offset points", + color="0.35", + fontsize=11, + style="italic", + ) + + +def plot_power_spectra( + spectra, + out_path, + title, + grid_dx_km=None, + model_short="model", + model_color="#1f3b73", + extra_lines=None, +): + """Plot one or more spectra vs wavelength with the standard annotations. + + spectra : {label: (wavelength_km, power, color, linestyle, alpha, lw)} + """ + fig, ax = plt.subplots(figsize=(10, 6)) + for label, (wl, power, color, ls, alpha, lw) in spectra.items(): + ax.loglog(wl, power, lw=lw, label=label, color=color, ls=ls, alpha=alpha) + + lines = [(lbl, wl, "0.55") for lbl, wl in REFERENCE_GRIDS.items()] + if grid_dx_km is not None: + lines += [ + (f"{model_short} grid", grid_dx_km, model_color), + (f"{model_short} eff.", EFF_RES_FACTOR * grid_dx_km, model_color), + (f"{model_short} Nyq.", 2.0 * grid_dx_km, model_color), + ] + lines += list(extra_lines or []) + _, ymax = ax.get_ylim() + label_y = cycle([1.0, 0.6, 0.35]) + for label, wl, color in sorted(lines, key=lambda t: t[1]): + ax.axvline(wl, color=color, ls="--", lw=1.0, alpha=0.8) + ax.text( + wl, + ymax * next(label_y), + " " + label, + rotation=90, + va="top", + ha="right", + color=color, + fontsize=8.5, + ) + + ref_wl, ref_p = next(iter(spectra.values()))[:2] + ref_wl = np.asarray(ref_wl) + ref_p = np.asarray(ref_p) + finite = np.isfinite(ref_wl) & np.isfinite(ref_p) + ref_wl, ref_p = ref_wl[finite], ref_p[finite] + meso = (ref_wl >= 5) & (ref_wl <= TRANSITION_KM) + if meso.any(): + _add_slope_guide( + ax, + 5.0 / 3.0, + r"$k^{-5/3}$", + (ref_wl[meso].min(), ref_wl[meso].max()), + (np.median(ref_wl[meso]), 5 * np.median(ref_p[meso])), + ) + + ax.set_xlabel("Wavelength [km]") + ax.set_ylabel("Power (variance density)") + ax.set_title(title) + ax.invert_xaxis() + ax.grid(True, which="both", ls=":", alpha=0.3) + ax.legend(loc="lower left", frameon=True) + fig.tight_layout() + fig.savefig(out_path, dpi=200, bbox_inches="tight") + plt.close(fig) diff --git a/src/spectra/io.py b/src/spectra/io.py new file mode 100644 index 00000000..a289eb59 --- /dev/null +++ b/src/spectra/io.py @@ -0,0 +1,69 @@ +"""Thin evalml I/O for spectra: variable->component mapping and native-field +extraction. Field loading is delegated to data_input loaders.""" + +from __future__ import annotations + +import numpy as np +import xarray as xr + +# Logical spectra variable -> (component GRIB params, scale factor). +# WIND_KE combines U/V as kinetic energy E = 0.5*(|u_hat|^2 + |v_hat|^2). +VARIABLE_COMPONENTS: dict[str, tuple[list[str], float]] = { + "T_2M": (["T_2M"], 1.0), + "WIND_KE": (["U_10M", "V_10M"], 0.5), + "TOT_PREC": (["TOT_PREC"], 1.0), +} + + +def required_params(variables: list[str]) -> list[str]: + """All GRIB params needed to compute the requested spectra variables.""" + params: list[str] = [] + for var in variables: + if var not in VARIABLE_COMPONENTS: + raise KeyError( + f"Unknown spectra variable {var!r}. " + f"Known: {sorted(VARIABLE_COMPONENTS)}." + ) + for p in VARIABLE_COMPONENTS[var][0]: + if p not in params: + params.append(p) + return params + + +def native_field(ds: xr.Dataset, param: str, step: int) -> np.ndarray: + """1-D native field for `param` at lead-time `step` (hours). + + Handles both timedelta64 step coordinates (as produced by the data_input + GRIB/baseline loaders) and plain integer-hour step coordinates. + """ + da = ds[param] + if "step" in da.dims: + step_coord = da.coords["step"] + if np.issubdtype(step_coord.dtype, np.timedelta64): + da = da.sel(step=np.timedelta64(step, "h")) + else: + da = da.sel(step=step) + # Drop singleton non-spatial dims left over from the loaders (e.g. a size-1 + # forecast_reference_time, or a single ensemble `number`), keeping only the + # native grid axis. A genuine multi-member dim survives and trips the guard. + da = da.squeeze(drop=True) + if da.ndim != 1: + raise ValueError( + f"Expected a 1-D native field for {param!r} after step selection, " + f"got shape {da.shape} with dims {da.dims}. Power spectra require " + f"data on the ICON native (unstructured) grid." + ) + return np.asarray(da.values).reshape(-1) + + +def native_components( + ds: xr.Dataset, variable: str, step: int +) -> tuple[list[np.ndarray], float]: + """Native component arrays + scale factor for a spectra variable at `step`.""" + if variable not in VARIABLE_COMPONENTS: + raise KeyError( + f"Unknown spectra variable {variable!r}. " + f"Known: {sorted(VARIABLE_COMPONENTS)}." + ) + params, factor = VARIABLE_COMPONENTS[variable] + return [native_field(ds, p, step) for p in params], factor diff --git a/src/spectra/regrid.py b/src/spectra/regrid.py new file mode 100644 index 00000000..3aa9ce39 --- /dev/null +++ b/src/spectra/regrid.py @@ -0,0 +1,90 @@ +# src/spectra/regrid.py +"""ICON native (unstructured) -> regular rotlatlon regridding via iconremap RBF +weights. The weights file is selected at runtime from the native point count.""" + +from __future__ import annotations + +from functools import lru_cache +from pathlib import Path + +import numpy as np +import scipy.sparse +import xarray as xr + +from spectra import core + +# Hardcoded iconremap RBF weight files on Balfrin, keyed by native cell count. +# CH1 native count is 1,147,980 (per the icon-power-spectra skill). +# TODO confirm CH2 count (see plan Task 2 Step 1). +_WEIGHTS_ROOT = Path("/store_new/mch/msopr/icon_workflow_2/iconremap-weights") +CH1_NPOINTS = 1_147_980 # confirmed vs ICON-CH1-EPS native T_2M field +CH2_NPOINTS = 283_876 # confirmed vs ICON-CH2-EPS native T_2M field +ICON_WEIGHTS = { + CH1_NPOINTS: _WEIGHTS_ROOT / "icon-ch1-eps-rotlatlon.nc", + CH2_NPOINTS: _WEIGHTS_ROOT / "icon-ch2-eps-rotlatlon.nc", +} +# TODO: prefer uuidOfHGrid as the registry key if point counts ever collide. +# TODO: eventually consume precomputed .npz sparse matrices directly. + + +def weights_path_for_npoints(npoints: int) -> Path: + """Resolve the iconremap weights file for a native grid of `npoints` cells.""" + try: + return ICON_WEIGHTS[npoints] + except KeyError as exc: + raise KeyError( + f"No iconremap weights registered for native grid with {npoints} " + f"cells. Known grids: {sorted(ICON_WEIGHTS)}." + ) from exc + + +def build_truncation_matrix(iconremap_nc, n_source): + """Sparse native(n_source) -> regular(ny*nx) matrix from iconremap weights.""" + with xr.open_dataset(iconremap_nc) as ds: + idx = ds["rbf_B_glbidx"].values + wgt = ds["rbf_B_wgt"].values + nx, ny = int(ds.attrs["nx"]), int(ds.attrs["ny"]) + npoints = idx.shape[-1] + idx = idx.reshape(ny * nx, npoints) + wgt = wgt.reshape(ny * nx, npoints) + n_targets = ny * nx + + # iconremap indices are 0-based and used as-is (cf. meteodata-lab's + # `_icon2regular`, which indexes with `np.take(field, indices)`). + rows = np.broadcast_to(np.arange(n_targets)[:, None], idx.shape).ravel() + cols = idx.ravel() + vals = wgt.ravel().astype(np.float32) + + M = scipy.sparse.coo_matrix( + (vals, (rows, cols)), shape=(n_targets, n_source) + ).tocsr() + return M, (ny, nx) + + +def grid_geometry(iconremap_nc): + """(ny, nx, dx_km) of the regular target grid from the iconremap weights.""" + with xr.open_dataset(iconremap_nc) as ds: + return ( + int(ds.attrs["ny"]), + int(ds.attrs["nx"]), + float(ds.attrs["dx"]) * core.KM_PER_DEG, + ) + + +def regrid(native_1d, matrix, ny, nx): + """Apply a truncation matrix to a native field -> (ny, nx) regular array.""" + return (matrix @ np.asarray(native_1d)).reshape(ny, nx) + + +@lru_cache(maxsize=4) +def load_regridder(npoints: int): + """Cached regridder for a native grid: (matrix, ny, nx, dx_km).""" + weights = weights_path_for_npoints(npoints) + matrix, (ny, nx) = build_truncation_matrix(weights, n_source=npoints) + g_ny, g_nx, dx_km = grid_geometry(weights) + if (g_ny, g_nx) != (ny, nx): + raise ValueError( + f"Grid geometry mismatch for {npoints} cells: " + f"matrix target {(ny, nx)} vs grid_geometry {(g_ny, g_nx)}." + ) + return matrix, ny, nx, dx_km diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 4c715931..e34b87a4 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1,6 +1,7 @@ import pytest +from pydantic import ValidationError -from evalml.config import ConfigModel +from evalml.config import ConfigModel, SpectraConfig def test_example_forecasters_config(example_forecasters_config): @@ -39,3 +40,53 @@ def test_legacy_top_level_baselines_still_supported(example_forecasters_config): ] _ = ConfigModel.model_validate(cfg) + + +def test_spectra_config_defaults(): + cfg = SpectraConfig() + assert cfg.enabled is False + assert cfg.method == "dct" + assert cfg.variables == ["T_2M", "WIND_KE", "TOT_PREC"] + assert cfg.lead_times == [] + + +def test_spectra_config_accepts_valid(): + cfg = SpectraConfig( + enabled=True, + method="fft", + lead_times=[6, 48, 120], + variables=["T_2M"], + init_hours=[0, 12], + ) + assert cfg.method == "fft" + assert cfg.lead_times == [6, 48, 120] + + +def test_spectra_config_rejects_bad_method(): + with pytest.raises(ValidationError): + SpectraConfig(method="wavelet") + + +def test_spectra_config_rejects_unknown_variable(): + with pytest.raises(ValidationError): + SpectraConfig(variables=["GEOPOTENTIAL"]) + + +def test_spectra_config_rejects_empty_variables(): + with pytest.raises(ValidationError): + SpectraConfig(variables=[]) + + +def test_spectra_config_forbids_extra(): + with pytest.raises(ValidationError): + SpectraConfig(enabled=True, bogus=1) + + +def test_spectra_config_enabled_requires_lead_times(): + with pytest.raises(ValidationError): + SpectraConfig(enabled=True, lead_times=[]) + + +def test_spectra_config_disabled_allows_empty_lead_times(): + cfg = SpectraConfig(enabled=False) + assert cfg.lead_times == [] diff --git a/tests/unit/test_spectra_core.py b/tests/unit/test_spectra_core.py new file mode 100644 index 00000000..19f9deb9 --- /dev/null +++ b/tests/unit/test_spectra_core.py @@ -0,0 +1,103 @@ +import numpy as np +import pytest +import xarray as xr + +from spectra import compute, core + + +def _cosine_field(n=128, wavelength_px=16): + x = np.arange(n) + return np.cos(2 * np.pi * x[None, :] / wavelength_px) * np.ones((n, 1)) + + +def test_radial_spectrum_returns_fixed_length_with_nan_padding(): + # A coeff array with energy only in a couple of modes -> most bins empty (NaN). + coeff = np.zeros((8, 8)) + coeff[0, 3] = 5.0 + kx = np.arange(8) / (2.0 * 8) / 1.0 + ky = np.arange(8) / (2.0 * 8) / 1.0 + wl, power = core.radial_spectrum(coeff, kx, ky, nbins=4, norm=1.0 / 64) + assert wl.shape == (4,) and power.shape == (4,) + assert np.isnan(power).any() # empty bins are NaN, not dropped + assert np.nanmax(power) > 0 + assert np.all(np.diff(wl) > 0) # wavelengths strictly ascending + + +def test_dct_spectrum_peaks_at_input_wavelength(): + dx = 2.0 # km per pixel + wavelength_px = 16 + field = _cosine_field(n=128, wavelength_px=wavelength_px) + wl, power = core.dct_power_spectrum(field, dx) + peak_wl = wl[np.nanargmax(power)] + expected = wavelength_px * dx # 32 km + assert peak_wl == pytest.approx(expected, rel=0.25) + + +def test_constant_field_has_no_power(): + field = np.full((64, 64), 7.3) + wl, power = core.dct_power_spectrum(field, 1.0) + assert np.nansum(power) == pytest.approx(0.0, abs=1e-9) + + +def test_combined_spectrum_scales_and_sums(): + field = _cosine_field(64, 8) + wl1, p1 = core.dct_power_spectrum(field, 1.0) + # two identical fields with factor 0.5 -> 0.5*(p+p) == p + wl2, p2 = core.combined_spectrum( + core.dct_power_spectrum, [field, field], 1.0, factor=0.5 + ) + np.testing.assert_allclose(np.nan_to_num(p2), np.nan_to_num(p1), rtol=1e-6) + + +def test_aggregate_spectra_averages_over_init(tmp_path): + wl = np.array([100.0, 50.0, 25.0]) + + def mk(power, init): + ds = xr.Dataset( + {"power": (("variable", "leadtime", "wavenumber"), np.array([[power]]))}, + coords={ + "variable": ["T_2M"], + "leadtime": [6], + "wavenumber": np.arange(3), + "wavelength": ("wavenumber", wl), + }, + attrs={"dx_km": 1.1, "npoints": 10, "label": "m"}, + ) + p = tmp_path / f"s_{init}.nc" + ds.to_netcdf(p) + return p + + a = mk([1.0, 2.0, 3.0], 0) + b = mk([3.0, 4.0, 5.0], 1) + agg = compute.aggregate_spectra([a, b]) + np.testing.assert_allclose( + agg["power"].sel(variable="T_2M", leadtime=6).values, [2.0, 3.0, 4.0] + ) + assert agg.attrs["label"] == "m" + np.testing.assert_array_equal(agg["wavelength"].values, wl) + + +def test_aggregate_spectra_nanmean_ignores_missing(tmp_path): + wl = np.array([100.0, 50.0]) + + def mk(power, init): + ds = xr.Dataset( + {"power": (("variable", "leadtime", "wavenumber"), np.array([[power]]))}, + coords={ + "variable": ["T_2M"], + "leadtime": [6], + "wavenumber": np.arange(2), + "wavelength": ("wavenumber", wl), + }, + attrs={"dx_km": 1.1, "npoints": 10, "label": "m"}, + ) + p = tmp_path / f"s_{init}.nc" + ds.to_netcdf(p) + return p + + a = mk([np.nan, 2.0], 0) + b = mk([4.0, 4.0], 1) + agg = compute.aggregate_spectra([a, b]) + np.testing.assert_allclose( + agg["power"].sel(variable="T_2M", leadtime=6).values, [4.0, 3.0] + ) diff --git a/tests/unit/test_spectra_io.py b/tests/unit/test_spectra_io.py new file mode 100644 index 00000000..a9359b68 --- /dev/null +++ b/tests/unit/test_spectra_io.py @@ -0,0 +1,87 @@ +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from spectra import io + + +def _toy_dataset(): + # native grid of 6 cells, two steps + return xr.Dataset( + { + "T_2M": (("step", "values"), np.arange(12.0).reshape(2, 6)), + "U_10M": (("step", "values"), np.ones((2, 6))), + "V_10M": (("step", "values"), 2 * np.ones((2, 6))), + }, + coords={"step": [0, 6], "values": np.arange(6)}, + ) + + +def test_variable_components_known(): + assert io.VARIABLE_COMPONENTS["T_2M"] == (["T_2M"], 1.0) + assert io.VARIABLE_COMPONENTS["WIND_KE"] == (["U_10M", "V_10M"], 0.5) + assert io.VARIABLE_COMPONENTS["TOT_PREC"] == (["TOT_PREC"], 1.0) + + +def test_native_field_squeezes_singleton_reftime_dim(): + # Real GRIB loads carry a size-1 forecast_reference_time dim; native_field + # must squeeze it so the field reduces to the 1-D native grid axis. + ds = xr.Dataset( + { + "T_2M": ( + ("forecast_reference_time", "step", "values"), + np.arange(12.0).reshape(1, 2, 6), + ) + }, + coords={"step": [0, 6], "values": np.arange(6)}, + ) + arr = io.native_field(ds, "T_2M", step=6) + assert arr.shape == (6,) + np.testing.assert_array_equal(arr, np.arange(6, 12)) + + +def test_native_field_extracts_1d_for_step(): + ds = _toy_dataset() + arr = io.native_field(ds, "T_2M", step=6) + assert arr.shape == (6,) + np.testing.assert_array_equal(arr, np.arange(6, 12)) + + +def test_native_components_returns_list_and_factor(): + ds = _toy_dataset() + comps, factor = io.native_components(ds, "WIND_KE", step=0) + assert factor == 0.5 + assert len(comps) == 2 + assert comps[0].shape == (6,) + + +def test_unknown_variable_raises(): + with pytest.raises(KeyError, match="Unknown spectra variable"): + io.native_components(_toy_dataset(), "NOPE", step=0) + + +def test_native_field_handles_timedelta_step(): + ds = xr.Dataset( + {"T_2M": (("step", "values"), np.arange(12.0).reshape(2, 6))}, + coords={"step": pd.to_timedelta([0, 6], unit="h"), "values": np.arange(6)}, + ) + arr = io.native_field(ds, "T_2M", step=6) + np.testing.assert_array_equal(arr, np.arange(6, 12)) + + +def test_native_field_rejects_non_1d(): + ds = xr.Dataset( + {"T_2M": (("step", "y", "x"), np.zeros((1, 3, 4)))}, + coords={"step": [0], "y": np.arange(3), "x": np.arange(4)}, + ) + with pytest.raises(ValueError, match="1-D native field"): + io.native_field(ds, "T_2M", step=0) + + +def test_required_params_deduplicates(): + assert io.required_params(["WIND_KE", "T_2M", "WIND_KE"]) == [ + "U_10M", + "V_10M", + "T_2M", + ] diff --git a/tests/unit/test_spectra_regrid.py b/tests/unit/test_spectra_regrid.py new file mode 100644 index 00000000..bf439796 --- /dev/null +++ b/tests/unit/test_spectra_regrid.py @@ -0,0 +1,82 @@ +# tests/unit/test_spectra_regrid.py +import numpy as np +import pytest +import xarray as xr + +from spectra import regrid + + +def _fake_weights(tmp_path, ny=4, nx=5, n_source=20, npoints=3): + # Realistic iconremap file: 0-based source indices, used as-is. Index 0 is a + # valid source cell (here paired with a nonzero weight), not padding. + rng = np.arange(ny * nx * npoints) % n_source + idx = rng.reshape(ny * nx, npoints).astype("int64") # 0-based, 0..n_source-1 + wgt = np.full((ny * nx, npoints), 1.0 / npoints, dtype="float64") + ds = xr.Dataset( + { + "rbf_B_glbidx": (("cell", "stencil"), idx), + "rbf_B_wgt": (("cell", "stencil"), wgt), + }, + attrs={"nx": nx, "ny": ny, "dx": 0.01}, + ) + path = tmp_path / "weights.nc" + ds.to_netcdf(path) + return path, ny, nx, n_source + + +def test_build_truncation_matrix_applies_weights(tmp_path): + path, ny, nx, n_source = _fake_weights(tmp_path) + M, (rny, rnx) = regrid.build_truncation_matrix(path, n_source=n_source) + assert (rny, rnx) == (ny, nx) + native = np.ones(n_source) + out = regrid.regrid(native, M, ny, nx) + assert out.shape == (ny, nx) + np.testing.assert_allclose(out, 1.0, rtol=1e-6) # average of ones is one + + +def test_grid_geometry_reads_dx_in_km(tmp_path): + path, ny, nx, _ = _fake_weights(tmp_path) + g_ny, g_nx, dx_km = regrid.grid_geometry(path) + assert (g_ny, g_nx) == (ny, nx) + assert dx_km == pytest.approx(0.01 * regrid.core.KM_PER_DEG) + + +def test_weights_path_for_npoints_unknown_raises(): + with pytest.raises(KeyError, match="No iconremap weights"): + regrid.weights_path_for_npoints(999_999) + + +def test_weights_path_for_npoints_known(): + assert "ch1" in str(regrid.weights_path_for_npoints(1_147_980)).lower() + + +def _fake_weights_zero_based_no_zeros(tmp_path, ny=4, nx=5, n_source=20): + # 0-based file (as the real ICON-CH1 weights): no zero entries at all, and + # the smallest *referenced* cell is >= 1, yet max index == n_source - 1. A + # min-index heuristic wrongly calls this 1-based and shifts every reference + # by one -- the off-by-one that injected grid-scale noise into the spectra. + # Single-point stencils (weight 1) so the regridded value IS the referenced + # native cell, making any shift directly observable. + idx = ((np.arange(ny * nx) % (n_source - 5)) + 5).reshape(ny * nx, 1) + assert idx.min() >= 1 and idx.max() == n_source - 1 and (idx != 0).all() + wgt = np.ones((ny * nx, 1), dtype="float64") + ds = xr.Dataset( + { + "rbf_B_glbidx": (("cell", "stencil"), idx.astype("int64")), + "rbf_B_wgt": (("cell", "stencil"), wgt), + }, + attrs={"nx": nx, "ny": ny, "dx": 0.01}, + ) + path = tmp_path / "weights_zero_no_pad.nc" + ds.to_netcdf(path) + return path, ny, nx, n_source, idx.ravel() + + +def test_build_truncation_matrix_zero_based_no_zeros_not_shifted(tmp_path): + # Regression: a 0-based, zero-padding-free file must NOT be treated as 1-based. + path, ny, nx, n_source, ref_cells = _fake_weights_zero_based_no_zeros(tmp_path) + M, _ = regrid.build_truncation_matrix(path, n_source=n_source) + native = np.arange(n_source, dtype="float64") # value == cell index + out = regrid.regrid(native, M, ny, nx).ravel() + # No off-by-one: each target picks up its referenced native cell as-is. + np.testing.assert_array_equal(out, ref_cells) diff --git a/workflow/Snakefile b/workflow/Snakefile index 4a3ec692..9d446892 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -19,6 +19,7 @@ include: "rules/inference.smk" include: "rules/verification.smk" include: "rules/report.smk" include: "rules/plot.smk" +include: "rules/spectra.smk" # about workflow @@ -148,6 +149,11 @@ rule experiment_all: ) else [] ), + ( + expand(rules.spectra_plot.output, experiment=EXPERIMENT_NAME) + if config.get("experiment", {}).get("spectra", {}).get("enabled", False) + else [] + ), rule showcase_all: diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk new file mode 100644 index 00000000..b7f152ca --- /dev/null +++ b/workflow/rules/spectra.smk @@ -0,0 +1,213 @@ +# ----------------------------------------------------- # +# POWER SPECTRA WORKFLOW # +# ----------------------------------------------------- # + + +include: "common.smk" + + +SPECTRA = config["experiment"].get("spectra", {}) +SPECTRA_LEAD_TIMES = ",".join(str(s) for s in SPECTRA.get("lead_times", [])) +SPECTRA_VARIABLES = ",".join(SPECTRA.get("variables", ["T_2M", "WIND_KE", "TOT_PREC"])) +SPECTRA_METHOD = SPECTRA.get("method", "dct") + + +def _spectra_reftimes(): + hours = SPECTRA.get("init_hours") + if hours is None: + return [t.strftime("%Y%m%d%H%M") for t in REFTIMES] + return [t.strftime("%Y%m%d%H%M") for t in REFTIMES if t.hour in hours] + + +def _spectra_truth_steps(): + candidates = list(collect_all_candidates()) + if candidates: + return RUN_CONFIGS[candidates[0]]["steps"] + # no candidate runs (baseline-only experiment): fall back to any baseline's steps + if BASELINE_CONFIGS: + return next(iter(BASELINE_CONFIGS.values()))["steps"] + return "0/120/6" + + +def spectra_participants(): + """participant key -> aggregated spectra path (runs + baselines).""" + out = {} + for base in BASELINE_CONFIGS: + out[base] = OUT_ROOT / f"data/baselines/{base}/spectra_aggregated.nc" + for run_id, cfg in RUN_CONFIGS.items(): + if cfg.get("_is_candidate", False): + out[run_id] = OUT_ROOT / f"data/runs/{run_id}/spectra_aggregated.nc" + return out + + +SPECTRA_PARTICIPANTS = spectra_participants() + + +rule spectra_compute: + input: + "src/spectra/__init__.py", + "src/data_input/__init__.py", + script="workflow/scripts/spectra_compute.py", + inference_okfile=rules.inference_execute.output.okfile, + eckit_grids=rules.data_download_eckit_geo_grids.output, + output: + OUT_ROOT / "data/runs/{run_id}/{init_time}/spectra.nc", + log: + OUT_ROOT / "logs/spectra_compute/{run_id}-{init_time}.log", + wildcard_constraints: + run_id=r"[^/]+/[^/]+", + init_time=r"\d+", + resources: + cpus_per_task=8, + mem_mb=50_000, + runtime="30m", + params: + grib=lambda wc: ( + Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" + ).resolve(), + steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], + label=lambda wc: RUN_CONFIGS[wc.run_id].get("label", wc.run_id), + shell: + """ + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + uv run {input.script} \ + --forecast {params.grib} --reftime {wildcards.init_time} \ + --steps "{params.steps}" --lead_times "{SPECTRA_LEAD_TIMES}" \ + --variables "{SPECTRA_VARIABLES}" --method "{SPECTRA_METHOD}" \ + --label "{params.label}" --output {output} >{log} 2>&1 + """ + + +rule spectra_compute_baseline: + input: + "src/spectra/__init__.py", + "src/data_input/__init__.py", + script="workflow/scripts/spectra_compute.py", + forecast=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["root"], + eckit_grids=rules.data_download_eckit_geo_grids.output, + output: + OUT_ROOT / "data/baselines/{baseline_id}/{init_time}/spectra.nc", + log: + OUT_ROOT / "logs/spectra_compute_baseline/{baseline_id}-{init_time}.log", + resources: + cpus_per_task=8, + mem_mb=50_000, + runtime="30m", + params: + steps=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["steps"], + label=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("label", wc.baseline_id), + shell: + """ + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + uv run {input.script} \ + --forecast {input.forecast} --reftime {wildcards.init_time} \ + --steps "{params.steps}" --lead_times "{SPECTRA_LEAD_TIMES}" \ + --variables "{SPECTRA_VARIABLES}" --method "{SPECTRA_METHOD}" \ + --label "{params.label}" --output {output} >{log} 2>&1 + """ + + +rule spectra_compute_truth: + input: + "src/spectra/__init__.py", + "src/data_input/__init__.py", + script="workflow/scripts/spectra_compute.py", + truth=config["truth"]["root"], + eckit_grids=rules.data_download_eckit_geo_grids.output, + output: + OUT_ROOT / "data/truth/{init_time}/spectra.nc", + log: + OUT_ROOT / "logs/spectra_truth/{init_time}.log", + resources: + cpus_per_task=8, + mem_mb=50_000, + runtime="30m", + params: + steps=_spectra_truth_steps(), + label=config["truth"]["label"], + shell: + """ + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + uv run {input.script} \ + --truth {input.truth} --reftime {wildcards.init_time} \ + --steps "{params.steps}" --lead_times "{SPECTRA_LEAD_TIMES}" \ + --variables "{SPECTRA_VARIABLES}" --method "{SPECTRA_METHOD}" \ + --label "{params.label}" --output {output} >{log} 2>&1 + """ + + +rule spectra_aggregate: + input: + script="workflow/scripts/spectra_aggregate.py", + spectra_nc=lambda wc: expand( + rules.spectra_compute.output, + init_time=_spectra_reftimes(), + allow_missing=True, + ), + output: + OUT_ROOT / "data/runs/{run_id}/spectra_aggregated.nc", + log: + OUT_ROOT / "logs/spectra_aggregate/{run_id}.log", + wildcard_constraints: + run_id=r"[^/]+/[^/]+", + localrule: True + resources: + cpus_per_task=2, + mem_mb=20_000, + runtime="20m", + shell: + "uv run {input.script} {input.spectra_nc} --output {output} >{log} 2>&1" + + +use rule spectra_aggregate as spectra_aggregate_baseline with: + input: + script="workflow/scripts/spectra_aggregate.py", + spectra_nc=lambda wc: expand( + rules.spectra_compute_baseline.output, + init_time=_spectra_reftimes(), + allow_missing=True, + ), + output: + OUT_ROOT / "data/baselines/{baseline_id}/spectra_aggregated.nc", + log: + OUT_ROOT / "logs/spectra_aggregate_baseline/{baseline_id}.log", + localrule: True + + +use rule spectra_aggregate as spectra_aggregate_truth with: + input: + script="workflow/scripts/spectra_aggregate.py", + spectra_nc=expand( + rules.spectra_compute_truth.output, + init_time=_spectra_reftimes(), + ), + output: + OUT_ROOT / "data/truth/spectra_aggregated.nc", + log: + OUT_ROOT / "logs/spectra_aggregate_truth/truth.log", + localrule: True + + +rule spectra_plot: + input: + script="workflow/scripts/spectra_plot.py", + truth=OUT_ROOT / "data/truth/spectra_aggregated.nc", + participants=list(SPECTRA_PARTICIPANTS.values()), + output: + report( + directory(OUT_ROOT / "results/{experiment}/spectra"), + patterns=["{name}.png"], + ), + log: + OUT_ROOT / "logs/spectra_plot/{experiment}.log", + resources: + cpus_per_task=2, + mem_mb=20_000, + runtime="20m", + shell: + """ + uv run {input.script} \ + --truth {input.truth} --participants {input.participants} \ + --variables "{SPECTRA_VARIABLES}" --lead_times "{SPECTRA_LEAD_TIMES}" \ + --output_dir {output} >{log} 2>&1 + """ diff --git a/workflow/scripts/spectra_aggregate.py b/workflow/scripts/spectra_aggregate.py new file mode 100644 index 00000000..5cf1a39b --- /dev/null +++ b/workflow/scripts/spectra_aggregate.py @@ -0,0 +1,28 @@ +"""Average per-init spectra over init times for one participant.""" + +import logging +from argparse import ArgumentParser +from pathlib import Path + +from spectra.compute import aggregate_spectra + +LOG = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" +) + + +def main(): + ap = ArgumentParser() + ap.add_argument("inputs", nargs="+", help="per-init spectra.nc files.") + ap.add_argument("--output", required=True) + args = ap.parse_args() + agg = aggregate_spectra(args.inputs) + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + agg.to_netcdf(out) + LOG.info("Saved aggregated spectra to %s", out) + + +if __name__ == "__main__": + main() diff --git a/workflow/scripts/spectra_compute.py b/workflow/scripts/spectra_compute.py new file mode 100644 index 00000000..327ab59f --- /dev/null +++ b/workflow/scripts/spectra_compute.py @@ -0,0 +1,79 @@ +"""Compute power spectra for one source (run, baseline, or truth) at one init.""" + +import logging +from argparse import ArgumentParser +from datetime import datetime +from pathlib import Path + +import pandas as pd + +from data_input import load_forecast_data, load_truth_data, parse_steps +from spectra import io +from spectra.compute import compute_source_spectra + +LOG = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) + + +def main(): + ap = ArgumentParser() + ap.add_argument("--forecast", help="Forecast root (grib dir or baseline archive).") + ap.add_argument("--truth", help="Truth zarr root (use instead of --forecast).") + ap.add_argument("--reftime", required=True) + ap.add_argument("--steps", required=True, help="start/stop/step in hours.") + ap.add_argument("--lead_times", required=True, help="comma-separated hours.") + ap.add_argument("--variables", required=True, help="comma-separated names.") + ap.add_argument("--method", default="dct") + ap.add_argument("--label", required=True) + ap.add_argument("--output", required=True) + args = ap.parse_args() + + if bool(args.forecast) == bool(args.truth): + ap.error("Exactly one of --forecast / --truth must be supplied.") + + reftime = datetime.strptime(args.reftime, "%Y%m%d%H%M") + all_steps = parse_steps(args.steps) + lead_times = [int(s) for s in args.lead_times.split(",")] + missing = [s for s in lead_times if s not in all_steps] + if missing: + LOG.warning("Requested lead times %s not in source steps; dropping.", missing) + lead_times = [s for s in lead_times if s in all_steps] + if not lead_times: + LOG.error( + "None of the requested lead times are available in source steps %s; aborting.", + all_steps, + ) + raise SystemExit(1) + variables = args.variables.split(",") + params = io.required_params(variables) + + if args.truth: + ds = load_truth_data(Path(args.truth), reftime, all_steps, params) + if "time" in ds.dims and "step" not in ds.dims: + valid = pd.to_datetime(ds["time"].values) + steps_h = ( + ((valid - pd.Timestamp(reftime)) / pd.Timedelta(hours=1)) + .round() + .astype(int) + ) + if pd.Series(steps_h).duplicated().any(): + raise ValueError( + f"Truth timestamps rounded to integer hours produced duplicate " + f"steps: {steps_h.tolist()}. Sub-hourly truth is unsupported for spectra." + ) + ds = ds.assign_coords(step=("time", steps_h)).swap_dims({"time": "step"}) + else: + ds = load_forecast_data(Path(args.forecast), reftime, all_steps, params) + + result = compute_source_spectra(ds, variables, lead_times, args.method, args.label) + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + result.to_netcdf(out) + LOG.info("Saved spectra to %s", out) + + +if __name__ == "__main__": + main() diff --git a/workflow/scripts/spectra_plot.py b/workflow/scripts/spectra_plot.py new file mode 100644 index 00000000..b81cc1a8 --- /dev/null +++ b/workflow/scripts/spectra_plot.py @@ -0,0 +1,32 @@ +"""Overlay truth + all participants into one figure per (variable, lead time).""" + +import logging +from argparse import ArgumentParser + +from spectra.compute import plot_experiment_spectra + +LOG = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" +) + + +def main(): + ap = ArgumentParser() + ap.add_argument("--truth", required=True) + ap.add_argument("--participants", nargs="+", required=True) + ap.add_argument("--variables", required=True) + ap.add_argument("--lead_times", required=True) + ap.add_argument("--output_dir", required=True) + args = ap.parse_args() + plot_experiment_spectra( + args.truth, + args.participants, + args.output_dir, + variables=args.variables.split(","), + lead_times=[int(s) for s in args.lead_times.split(",")], + ) + + +if __name__ == "__main__": + main() diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 5bc95b48..de1de588 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -269,6 +269,10 @@ ], "default": null, "description": "Scorecard generation configuration. Omit or set enabled: false to disable." + }, + "spectra": { + "$ref": "#/$defs/SpectraConfig", + "description": "Power-spectra QC configuration. Disabled by default." } }, "required": [ @@ -734,6 +738,67 @@ "title": "ShowcaseConfig", "type": "object" }, + "SpectraConfig": { + "additionalProperties": false, + "description": "Configuration for power-spectra QC plots in the experiment pipeline.", + "properties": { + "enabled": { + "default": false, + "description": "Whether to compute and plot power spectra.", + "title": "Enabled", + "type": "boolean" + }, + "method": { + "default": "dct", + "description": "Spectral method: 'dct' (default, recommended for LAM) or 'fft'.", + "enum": [ + "dct", + "fft" + ], + "title": "Method", + "type": "string" + }, + "lead_times": { + "description": "Representative lead times (hours) at which to compute spectra.", + "items": { + "type": "integer" + }, + "title": "Lead Times", + "type": "array" + }, + "variables": { + "default": [ + "T_2M", + "WIND_KE", + "TOT_PREC" + ], + "description": "Spectra variables. Supported: T_2M, WIND_KE (from U/V_10M), TOT_PREC.", + "items": { + "type": "string" + }, + "title": "Variables", + "type": "array" + }, + "init_hours": { + "anyOf": [ + { + "items": { + "type": "integer" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional subset of init hours to average over. None = all.", + "title": "Init Hours" + } + }, + "title": "SpectraConfig", + "type": "object" + }, "Stratification": { "description": "Stratification settings for the analysis.", "properties": { From 2c5e0d1c897e365cbd74238c7d324cc430fd4e34 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:18:11 +0100 Subject: [PATCH 02/21] read spectra config from optional experiment block --- workflow/Snakefile | 2 +- workflow/rules/spectra.smk | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/workflow/Snakefile b/workflow/Snakefile index a3251105..22fb477c 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -216,7 +216,7 @@ rule experiment_all: ), ( expand(rules.spectra_plot.output, experiment=EXPERIMENT_NAME) - if config.get("experiment", {}).get("spectra", {}).get("enabled", False) + if SPECTRA.get("enabled", False) else [] ), # SAL precipitation scores (compute-only: the CSVs themselves are the diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index 66fc6fcd..08d40e5e 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -6,7 +6,7 @@ include: "common.smk" -SPECTRA = config["experiment"].get("spectra", {}) +SPECTRA = EXPERIMENT_CONFIG.get("spectra") or {} SPECTRA_LEAD_TIMES = ",".join(str(s) for s in SPECTRA.get("lead_times", [])) SPECTRA_VARIABLES = ",".join(SPECTRA.get("variables", ["T_2M", "WIND_KE", "TOT_PREC"])) SPECTRA_METHOD = SPECTRA.get("method", "dct") From 38c9197ca9a8a7171e14dc7b3d3608daa465e0dc Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:18:31 +0100 Subject: [PATCH 03/21] fix xr.dataset typo in compute_source_spectra docstring --- src/spectra/compute.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spectra/compute.py b/src/spectra/compute.py index 336788fa..d48e3f05 100644 --- a/src/spectra/compute.py +++ b/src/spectra/compute.py @@ -17,7 +17,7 @@ def compute_source_spectra(ds: xr.Dataset, variables, lead_times, method, label): """Compute spectra for one source (already-loaded native dataset). - Returns an x.Dataset with dims (variable, leadtime, wavenumber) and a shared + Returns an xr.Dataset with dims (variable, leadtime, wavenumber) and a shared `wavelength` coordinate. The native grid is detected from the field length. """ if not variables or not lead_times: From 069f51d76b65c33a2f60143f36f738c85d522a43 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:19:09 +0100 Subject: [PATCH 04/21] move spectra plotting to spectra.plot --- src/spectra/__init__.py | 2 +- src/spectra/compute.py | 52 +---------- src/spectra/core.py | 94 ------------------- src/spectra/plot.py | 149 +++++++++++++++++++++++++++++++ workflow/scripts/spectra_plot.py | 2 +- 5 files changed, 153 insertions(+), 146 deletions(-) create mode 100644 src/spectra/plot.py diff --git a/src/spectra/__init__.py b/src/spectra/__init__.py index 7599fd4a..fa47a74d 100644 --- a/src/spectra/__init__.py +++ b/src/spectra/__init__.py @@ -4,4 +4,4 @@ directly. The spectral `core` is numpy-only and reusable for any model. """ -from spectra import compute, core, io, regrid # noqa: F401 +from spectra import compute, core, io, plot, regrid # noqa: F401 diff --git a/src/spectra/compute.py b/src/spectra/compute.py index d48e3f05..4f7b04fe 100644 --- a/src/spectra/compute.py +++ b/src/spectra/compute.py @@ -1,18 +1,13 @@ -"""Orchestration: per-source spectra computation, over-init aggregation, and the -experiment overlay plot. Defines the spectra.nc schema.""" +"""Orchestration: per-source spectra computation and over-init aggregation. +Defines the spectra.nc schema.""" from __future__ import annotations -from itertools import cycle -from pathlib import Path - import numpy as np import xarray as xr from spectra import core, io, regrid -_PALETTE = ["#1f3b73", "#c1272d", "#1b7837", "#6a3d9a", "#e08214", "#4d4d4d"] - def compute_source_spectra(ds: xr.Dataset, variables, lead_times, method, label): """Compute spectra for one source (already-loaded native dataset). @@ -74,46 +69,3 @@ def aggregate_spectra(spectra_files) -> xr.Dataset: for ds in datasets: ds.close() return agg - - -def plot_experiment_spectra( - truth_file, participant_files, out_dir, variables, lead_times -): - """One overlay figure per (variable, lead time): truth + all participants.""" - out_dir = Path(out_dir) - out_dir.mkdir(parents=True, exist_ok=True) - participants = [xr.open_dataset(f) for f in participant_files] - try: - with xr.open_dataset(truth_file) as truth: - for var in variables: - for step in lead_times: - spectra = {} - t_wl = truth["wavelength"].values - t_p = truth["power"].sel(variable=var, leadtime=step).values - spectra[f"{truth.attrs['label']} (truth)"] = ( - t_wl, - t_p, - "k", - "-", - 1.0, - 2.0, - ) - color = cycle(_PALETTE) - for ds in participants: - wl = ds["wavelength"].values - p = ds["power"].sel(variable=var, leadtime=step).values - spectra[ds.attrs["label"]] = (wl, p, next(color), "-", 0.9, 1.6) - out = out_dir / f"spectrum_{var}_{step:03d}.png" - # Reference grid/eff-res/Nyquist lines use the TRUTH grid's dx (labelled - # "truth"); each participant curve keeps its own wavelength axis. - core.plot_power_spectra( - spectra, - out, - f"Power spectrum — {var} @ +{step}h", - grid_dx_km=float(truth.attrs["dx_km"]), - model_short="truth", - model_color="k", - ) - finally: - for ds in participants: - ds.close() diff --git a/src/spectra/core.py b/src/spectra/core.py index 62ce95e3..32b3fe77 100644 --- a/src/spectra/core.py +++ b/src/spectra/core.py @@ -9,23 +9,11 @@ from __future__ import annotations -from itertools import cycle - -import matplotlib.pyplot as plt import numpy as np import scipy.fft KM_PER_DEG = 111.2 -REFERENCE_GRIDS = { - "O96 (~104 km)": 90.0 / 96.0 * KM_PER_DEG, - "N320 (~31 km)": 0.28125 * KM_PER_DEG, - "CERRA (~9 km)": 0.10 * KM_PER_DEG, -} - -EFF_RES_FACTOR = 6.0 -TRANSITION_KM = 400.0 - def _hann2d(ny: int, nx: int) -> np.ndarray: win = np.hanning(ny)[:, None] * np.hanning(nx)[None, :] @@ -117,85 +105,3 @@ def combined_spectrum( wl = w total = np.nan_to_num(p) if total is None else total + np.nan_to_num(p) return wl, factor * total - - -def _add_slope_guide(ax, n, label, wl_range, anchor): - wl = np.array(wl_range, dtype=float) - wl0, p0 = anchor - power = p0 * (wl / wl0) ** n - ax.plot(wl, power, ls=":", color="0.35", lw=1.3) - ax.annotate( - label, - xy=(wl.min(), power[np.argmin(wl)]), - xytext=(3, 3), - textcoords="offset points", - color="0.35", - fontsize=11, - style="italic", - ) - - -def plot_power_spectra( - spectra, - out_path, - title, - grid_dx_km=None, - model_short="model", - model_color="#1f3b73", - extra_lines=None, -): - """Plot one or more spectra vs wavelength with the standard annotations. - - spectra : {label: (wavelength_km, power, color, linestyle, alpha, lw)} - """ - fig, ax = plt.subplots(figsize=(10, 6)) - for label, (wl, power, color, ls, alpha, lw) in spectra.items(): - ax.loglog(wl, power, lw=lw, label=label, color=color, ls=ls, alpha=alpha) - - lines = [(lbl, wl, "0.55") for lbl, wl in REFERENCE_GRIDS.items()] - if grid_dx_km is not None: - lines += [ - (f"{model_short} grid", grid_dx_km, model_color), - (f"{model_short} eff.", EFF_RES_FACTOR * grid_dx_km, model_color), - (f"{model_short} Nyq.", 2.0 * grid_dx_km, model_color), - ] - lines += list(extra_lines or []) - _, ymax = ax.get_ylim() - label_y = cycle([1.0, 0.6, 0.35]) - for label, wl, color in sorted(lines, key=lambda t: t[1]): - ax.axvline(wl, color=color, ls="--", lw=1.0, alpha=0.8) - ax.text( - wl, - ymax * next(label_y), - " " + label, - rotation=90, - va="top", - ha="right", - color=color, - fontsize=8.5, - ) - - ref_wl, ref_p = next(iter(spectra.values()))[:2] - ref_wl = np.asarray(ref_wl) - ref_p = np.asarray(ref_p) - finite = np.isfinite(ref_wl) & np.isfinite(ref_p) - ref_wl, ref_p = ref_wl[finite], ref_p[finite] - meso = (ref_wl >= 5) & (ref_wl <= TRANSITION_KM) - if meso.any(): - _add_slope_guide( - ax, - 5.0 / 3.0, - r"$k^{-5/3}$", - (ref_wl[meso].min(), ref_wl[meso].max()), - (np.median(ref_wl[meso]), 5 * np.median(ref_p[meso])), - ) - - ax.set_xlabel("Wavelength [km]") - ax.set_ylabel("Power (variance density)") - ax.set_title(title) - ax.invert_xaxis() - ax.grid(True, which="both", ls=":", alpha=0.3) - ax.legend(loc="lower left", frameon=True) - fig.tight_layout() - fig.savefig(out_path, dpi=200, bbox_inches="tight") - plt.close(fig) diff --git a/src/spectra/plot.py b/src/spectra/plot.py new file mode 100644 index 00000000..316de825 --- /dev/null +++ b/src/spectra/plot.py @@ -0,0 +1,149 @@ +"""Plotting of power spectra: single-figure helper with reference annotations, +and the experiment overlay (truth + all participants).""" + +from __future__ import annotations + +from itertools import cycle +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +from spectra.core import KM_PER_DEG + +REFERENCE_GRIDS = { + "O96 (~104 km)": 90.0 / 96.0 * KM_PER_DEG, + "N320 (~31 km)": 0.28125 * KM_PER_DEG, + "CERRA (~9 km)": 0.10 * KM_PER_DEG, +} + +EFF_RES_FACTOR = 6.0 +TRANSITION_KM = 400.0 + +_PALETTE = ["#1f3b73", "#c1272d", "#1b7837", "#6a3d9a", "#e08214", "#4d4d4d"] + + +def _add_slope_guide(ax, n, label, wl_range, anchor): + wl = np.array(wl_range, dtype=float) + wl0, p0 = anchor + power = p0 * (wl / wl0) ** n + ax.plot(wl, power, ls=":", color="0.35", lw=1.3) + ax.annotate( + label, + xy=(wl.min(), power[np.argmin(wl)]), + xytext=(3, 3), + textcoords="offset points", + color="0.35", + fontsize=11, + style="italic", + ) + + +def plot_power_spectra( + spectra, + out_path, + title, + grid_dx_km=None, + model_short="model", + model_color="#1f3b73", + extra_lines=None, +): + """Plot one or more spectra vs wavelength with the standard annotations. + + spectra : {label: (wavelength_km, power, color, linestyle, alpha, lw)} + """ + fig, ax = plt.subplots(figsize=(10, 6)) + for label, (wl, power, color, ls, alpha, lw) in spectra.items(): + ax.loglog(wl, power, lw=lw, label=label, color=color, ls=ls, alpha=alpha) + + lines = [(lbl, wl, "0.55") for lbl, wl in REFERENCE_GRIDS.items()] + if grid_dx_km is not None: + lines += [ + (f"{model_short} grid", grid_dx_km, model_color), + (f"{model_short} eff.", EFF_RES_FACTOR * grid_dx_km, model_color), + (f"{model_short} Nyq.", 2.0 * grid_dx_km, model_color), + ] + lines += list(extra_lines or []) + _, ymax = ax.get_ylim() + label_y = cycle([1.0, 0.6, 0.35]) + for label, wl, color in sorted(lines, key=lambda t: t[1]): + ax.axvline(wl, color=color, ls="--", lw=1.0, alpha=0.8) + ax.text( + wl, + ymax * next(label_y), + " " + label, + rotation=90, + va="top", + ha="right", + color=color, + fontsize=8.5, + ) + + ref_wl, ref_p = next(iter(spectra.values()))[:2] + ref_wl = np.asarray(ref_wl) + ref_p = np.asarray(ref_p) + finite = np.isfinite(ref_wl) & np.isfinite(ref_p) + ref_wl, ref_p = ref_wl[finite], ref_p[finite] + meso = (ref_wl >= 5) & (ref_wl <= TRANSITION_KM) + if meso.any(): + _add_slope_guide( + ax, + 5.0 / 3.0, + r"$k^{-5/3}$", + (ref_wl[meso].min(), ref_wl[meso].max()), + (np.median(ref_wl[meso]), 5 * np.median(ref_p[meso])), + ) + + ax.set_xlabel("Wavelength [km]") + ax.set_ylabel("Power (variance density)") + ax.set_title(title) + ax.invert_xaxis() + ax.grid(True, which="both", ls=":", alpha=0.3) + ax.legend(loc="lower left", frameon=True) + fig.tight_layout() + fig.savefig(out_path, dpi=200, bbox_inches="tight") + plt.close(fig) + + +def plot_experiment_spectra( + truth_file, participant_files, out_dir, variables, lead_times +): + """One overlay figure per (variable, lead time): truth + all participants.""" + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + participants = [xr.open_dataset(f) for f in participant_files] + try: + with xr.open_dataset(truth_file) as truth: + for var in variables: + for step in lead_times: + spectra = {} + t_wl = truth["wavelength"].values + t_p = truth["power"].sel(variable=var, leadtime=step).values + spectra[f"{truth.attrs['label']} (truth)"] = ( + t_wl, + t_p, + "k", + "-", + 1.0, + 2.0, + ) + color = cycle(_PALETTE) + for ds in participants: + wl = ds["wavelength"].values + p = ds["power"].sel(variable=var, leadtime=step).values + spectra[ds.attrs["label"]] = (wl, p, next(color), "-", 0.9, 1.6) + out = out_dir / f"spectrum_{var}_{step:03d}.png" + # Reference grid/eff-res/Nyquist lines use the TRUTH grid's dx (labelled + # "truth"); each participant curve keeps its own wavelength axis. + plot_power_spectra( + spectra, + out, + f"Power spectrum — {var} @ +{step}h", + grid_dx_km=float(truth.attrs["dx_km"]), + model_short="truth", + model_color="k", + ) + finally: + for ds in participants: + ds.close() diff --git a/workflow/scripts/spectra_plot.py b/workflow/scripts/spectra_plot.py index b81cc1a8..6b40f5c0 100644 --- a/workflow/scripts/spectra_plot.py +++ b/workflow/scripts/spectra_plot.py @@ -3,7 +3,7 @@ import logging from argparse import ArgumentParser -from spectra.compute import plot_experiment_spectra +from spectra.plot import plot_experiment_spectra LOG = logging.getLogger(__name__) logging.basicConfig( From 3e2989c0392739e6fe813f05685ca70305f15d27 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:20:19 +0100 Subject: [PATCH 05/21] allow any loadable param as spectra variable, default to tot_prec1 --- src/evalml/config.py | 14 +++++------- src/spectra/io.py | 30 +++++++++++-------------- tests/unit/test_config.py | 8 +++---- tests/unit/test_spectra_io.py | 37 +++++++++++++++++++++++++------ workflow/rules/spectra.smk | 2 +- workflow/tools/config.schema.json | 4 ++-- 6 files changed, 56 insertions(+), 39 deletions(-) diff --git a/src/evalml/config.py b/src/evalml/config.py index d089c158..0f2571e7 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -485,8 +485,12 @@ class SpectraConfig(BaseModel): description="Representative lead times (hours) at which to compute spectra.", ) variables: List[str] = Field( - default=["T_2M", "WIND_KE", "TOT_PREC"], - description="Spectra variables. Supported: T_2M, WIND_KE (from U/V_10M), TOT_PREC.", + default=["T_2M", "WIND_KE", "TOT_PREC1"], + description=( + "Spectra variables: any parameter the forecast and truth loaders provide " + "(e.g. T_2M, SP_10M, TOT_PREC1), plus WIND_KE, the kinetic energy from " + "U_10M and V_10M." + ), ) init_hours: Optional[List[int]] = Field( default=None, @@ -500,12 +504,6 @@ def validate_variables(cls, v: List[str]) -> List[str]: raise ValueError( "`variables` must list at least one variable when spectra is configured." ) - allowed = {"T_2M", "WIND_KE", "TOT_PREC"} - invalid = set(v) - allowed - if invalid: - raise ValueError( - f"Unsupported spectra variable(s) {invalid!r}. Must be subset of {allowed}." - ) return v @model_validator(mode="after") diff --git a/src/spectra/io.py b/src/spectra/io.py index a289eb59..8e4bc088 100644 --- a/src/spectra/io.py +++ b/src/spectra/io.py @@ -6,25 +6,26 @@ import numpy as np import xarray as xr -# Logical spectra variable -> (component GRIB params, scale factor). -# WIND_KE combines U/V as kinetic energy E = 0.5*(|u_hat|^2 + |v_hat|^2). +# Spectra variables built from several GRIB params: name -> (component params, +# scale factor). WIND_KE combines U/V as kinetic energy +# E = 0.5*(|u_hat|^2 + |v_hat|^2). Any other variable is a single param, loaded +# as such by the data_input loaders (including derived params such as SP_10M +# and period params such as TOT_PREC1/TOT_PREC6). VARIABLE_COMPONENTS: dict[str, tuple[list[str], float]] = { - "T_2M": (["T_2M"], 1.0), "WIND_KE": (["U_10M", "V_10M"], 0.5), - "TOT_PREC": (["TOT_PREC"], 1.0), } +def variable_components(variable: str) -> tuple[list[str], float]: + """Component params and scale factor of a spectra variable.""" + return VARIABLE_COMPONENTS.get(variable, ([variable], 1.0)) + + def required_params(variables: list[str]) -> list[str]: - """All GRIB params needed to compute the requested spectra variables.""" + """All params to load to compute the requested spectra variables.""" params: list[str] = [] for var in variables: - if var not in VARIABLE_COMPONENTS: - raise KeyError( - f"Unknown spectra variable {var!r}. " - f"Known: {sorted(VARIABLE_COMPONENTS)}." - ) - for p in VARIABLE_COMPONENTS[var][0]: + for p in variable_components(var)[0]: if p not in params: params.append(p) return params @@ -60,10 +61,5 @@ def native_components( ds: xr.Dataset, variable: str, step: int ) -> tuple[list[np.ndarray], float]: """Native component arrays + scale factor for a spectra variable at `step`.""" - if variable not in VARIABLE_COMPONENTS: - raise KeyError( - f"Unknown spectra variable {variable!r}. " - f"Known: {sorted(VARIABLE_COMPONENTS)}." - ) - params, factor = VARIABLE_COMPONENTS[variable] + params, factor = variable_components(variable) return [native_field(ds, p, step) for p in params], factor diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 16cd0118..c1880c77 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -20,7 +20,7 @@ def test_spectra_config_defaults(): cfg = SpectraConfig() assert cfg.enabled is False assert cfg.method == "dct" - assert cfg.variables == ["T_2M", "WIND_KE", "TOT_PREC"] + assert cfg.variables == ["T_2M", "WIND_KE", "TOT_PREC1"] assert cfg.lead_times == [] @@ -41,9 +41,9 @@ def test_spectra_config_rejects_bad_method(): SpectraConfig(method="wavelet") -def test_spectra_config_rejects_unknown_variable(): - with pytest.raises(ValidationError): - SpectraConfig(variables=["GEOPOTENTIAL"]) +def test_spectra_config_accepts_any_variable(): + cfg = SpectraConfig(variables=["SP_10M", "TOT_PREC6", "WIND_KE"]) + assert cfg.variables == ["SP_10M", "TOT_PREC6", "WIND_KE"] def test_spectra_config_rejects_empty_variables(): diff --git a/tests/unit/test_spectra_io.py b/tests/unit/test_spectra_io.py index a9359b68..d1536994 100644 --- a/tests/unit/test_spectra_io.py +++ b/tests/unit/test_spectra_io.py @@ -18,10 +18,10 @@ def _toy_dataset(): ) -def test_variable_components_known(): - assert io.VARIABLE_COMPONENTS["T_2M"] == (["T_2M"], 1.0) - assert io.VARIABLE_COMPONENTS["WIND_KE"] == (["U_10M", "V_10M"], 0.5) - assert io.VARIABLE_COMPONENTS["TOT_PREC"] == (["TOT_PREC"], 1.0) +def test_variable_components(): + assert io.variable_components("WIND_KE") == (["U_10M", "V_10M"], 0.5) + assert io.variable_components("T_2M") == (["T_2M"], 1.0) + assert io.variable_components("TOT_PREC1") == (["TOT_PREC1"], 1.0) def test_native_field_squeezes_singleton_reftime_dim(): @@ -56,8 +56,15 @@ def test_native_components_returns_list_and_factor(): assert comps[0].shape == (6,) -def test_unknown_variable_raises(): - with pytest.raises(KeyError, match="Unknown spectra variable"): +def test_native_components_single_param(): + comps, factor = io.native_components(_toy_dataset(), "T_2M", step=6) + assert factor == 1.0 + assert len(comps) == 1 + np.testing.assert_array_equal(comps[0], np.arange(6, 12)) + + +def test_missing_variable_raises(): + with pytest.raises(KeyError): io.native_components(_toy_dataset(), "NOPE", step=0) @@ -80,8 +87,24 @@ def test_native_field_rejects_non_1d(): def test_required_params_deduplicates(): - assert io.required_params(["WIND_KE", "T_2M", "WIND_KE"]) == [ + assert io.required_params(["WIND_KE", "T_2M", "WIND_KE", "TOT_PREC1"]) == [ "U_10M", "V_10M", "T_2M", + "TOT_PREC1", ] + + +def test_period_param_from_loader_disaggregation(): + # TOT_PREC1 goes through the same disaggregation as in verification. + from data_input import _disaggregated_and_derived_params + + cumulative = np.repeat(np.arange(4.0)[:, None], 5, axis=1) + ds = xr.Dataset( + {"TOT_PREC": (("step", "values"), cumulative)}, + coords={"step": pd.to_timedelta(np.arange(4), unit="h")}, + ) + ds = _disaggregated_and_derived_params(ds, [1, 3], ["TOT_PREC1"]) + comps, factor = io.native_components(ds, "TOT_PREC1", step=3) + assert factor == 1.0 + np.testing.assert_array_equal(comps[0], np.ones(5)) diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index 08d40e5e..edf02b57 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -8,7 +8,7 @@ include: "common.smk" SPECTRA = EXPERIMENT_CONFIG.get("spectra") or {} SPECTRA_LEAD_TIMES = ",".join(str(s) for s in SPECTRA.get("lead_times", [])) -SPECTRA_VARIABLES = ",".join(SPECTRA.get("variables", ["T_2M", "WIND_KE", "TOT_PREC"])) +SPECTRA_VARIABLES = ",".join(SPECTRA.get("variables", ["T_2M", "WIND_KE", "TOT_PREC1"])) SPECTRA_METHOD = SPECTRA.get("method", "dct") diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 21d21929..a51c7939 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -1005,9 +1005,9 @@ "default": [ "T_2M", "WIND_KE", - "TOT_PREC" + "TOT_PREC1" ], - "description": "Spectra variables. Supported: T_2M, WIND_KE (from U/V_10M), TOT_PREC.", + "description": "Spectra variables: any parameter the forecast and truth loaders provide (e.g. T_2M, SP_10M, TOT_PREC1), plus WIND_KE, the kinetic energy from U_10M and V_10M.", "items": { "type": "string" }, From 0b48cf0a45a36d3b8b3d47d5491987c154e792ab Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:20:34 +0100 Subject: [PATCH 06/21] propagate nan when aggregating spectra over init times --- src/spectra/compute.py | 7 +++++-- tests/unit/test_spectra_core.py | 4 ++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/spectra/compute.py b/src/spectra/compute.py index 4f7b04fe..a4ae5217 100644 --- a/src/spectra/compute.py +++ b/src/spectra/compute.py @@ -57,11 +57,14 @@ def compute_source_spectra(ds: xr.Dataset, variables, lead_times, method, label) def aggregate_spectra(spectra_files) -> xr.Dataset: - """Average power over init times (nanmean). All inputs share one grid.""" + """Average power over init times. All inputs share one grid. + + NaN propagates: a value missing in any input is missing in the average. + """ datasets = [xr.open_dataset(f) for f in spectra_files] try: stacked = xr.concat(datasets, dim="init") - agg = stacked.mean(dim="init", skipna=True) + agg = stacked.mean(dim="init", skipna=False) agg.attrs = datasets[0].attrs agg["wavelength"] = datasets[0]["wavelength"].copy() agg.load() diff --git a/tests/unit/test_spectra_core.py b/tests/unit/test_spectra_core.py index 19f9deb9..54be8c54 100644 --- a/tests/unit/test_spectra_core.py +++ b/tests/unit/test_spectra_core.py @@ -77,7 +77,7 @@ def mk(power, init): np.testing.assert_array_equal(agg["wavelength"].values, wl) -def test_aggregate_spectra_nanmean_ignores_missing(tmp_path): +def test_aggregate_spectra_propagates_missing(tmp_path): wl = np.array([100.0, 50.0]) def mk(power, init): @@ -99,5 +99,5 @@ def mk(power, init): b = mk([4.0, 4.0], 1) agg = compute.aggregate_spectra([a, b]) np.testing.assert_allclose( - agg["power"].sel(variable="T_2M", leadtime=6).values, [4.0, 3.0] + agg["power"].sel(variable="T_2M", leadtime=6).values, [np.nan, 3.0] ) From a4fd7c54745f0e1330bda1dd15bb3169d3635a2b Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:20:52 +0100 Subject: [PATCH 07/21] use shared experiment participants for spectra --- workflow/rules/spectra.smk | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index edf02b57..f3289a92 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -29,18 +29,10 @@ def _spectra_truth_steps(): return "0/120/6" -def spectra_participants(): - """participant key -> aggregated spectra path (runs + baselines).""" - out = {} - for base in BASELINE_CONFIGS: - out[base] = OUT_ROOT / f"data/baselines/{base}/spectra_aggregated.nc" - for run_id, cfg in RUN_CONFIGS.items(): - if cfg.get("_is_candidate", False): - out[run_id] = OUT_ROOT / f"data/runs/{run_id}/spectra_aggregated.nc" - return out - - -SPECTRA_PARTICIPANTS = spectra_participants() +SPECTRA_PARTICIPANTS = { + pid: participant_data_dir(pid) / "spectra_aggregated.nc" + for pid in experiment_participant_ids() +} rule spectra_compute: From 65469ac02c3b515102777ed410516211140ad33f Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:37:18 +0100 Subject: [PATCH 08/21] make spectrum nan when an input field has missing values --- src/spectra/core.py | 7 ++++++- tests/unit/test_spectra_core.py | 19 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/spectra/core.py b/src/spectra/core.py index 32b3fe77..b917eef4 100644 --- a/src/spectra/core.py +++ b/src/spectra/core.py @@ -95,13 +95,18 @@ def combined_spectrum( """Sum the spectra of one or more component fields, scaled by `factor`. For wind kinetic energy pass [u, v] and factor=0.5 -> 0.5*(|u_hat|^2+|v_hat|^2). - NaN bins are treated as zero in the sum so the common coordinate is preserved. + Empty and zero-power bins count as zero, so a constant (e.g. dry) field gives a + valid all-zero spectrum. If any input field contains NaN (missing data), the + whole spectrum is NaN, so the gap carries through to the aggregation. """ if not fields_2d: raise ValueError("combined_spectrum requires at least one field") + missing = any(np.isnan(f).any() for f in fields_2d) wl = total = None for f in fields_2d: w, p = spectrum_fn(f, dx_km, nbins, oro_coeff=oro_coeff) wl = w total = np.nan_to_num(p) if total is None else total + np.nan_to_num(p) + if missing: + return wl, np.full_like(total, np.nan) return wl, factor * total diff --git a/tests/unit/test_spectra_core.py b/tests/unit/test_spectra_core.py index 54be8c54..ec1f4c7d 100644 --- a/tests/unit/test_spectra_core.py +++ b/tests/unit/test_spectra_core.py @@ -49,6 +49,25 @@ def test_combined_spectrum_scales_and_sums(): np.testing.assert_allclose(np.nan_to_num(p2), np.nan_to_num(p1), rtol=1e-6) +def test_combined_spectrum_constant_field_is_valid_zero(): + # An entirely dry precipitation field: zero power everywhere, not missing. + dry = np.zeros((64, 64)) + wl, p = core.combined_spectrum(core.dct_power_spectrum, [dry], 1.0) + assert np.all(np.isfinite(p)) + np.testing.assert_array_equal(p, 0.0) + + +def test_combined_spectrum_nan_input_gives_nan(): + field = _cosine_field(64, 8) + gappy = field.copy() + gappy[3, 5] = np.nan + wl, p = core.combined_spectrum( + core.dct_power_spectrum, [field, gappy], 1.0, factor=0.5 + ) + assert np.all(np.isfinite(wl)) + assert np.all(np.isnan(p)) + + def test_aggregate_spectra_averages_over_init(tmp_path): wl = np.array([100.0, 50.0, 25.0]) From b999f2cf8954e971f7b85b9ba4b340f35b714d1a Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:37:31 +0100 Subject: [PATCH 09/21] read sal config from optional experiment block --- workflow/rules/verification.smk | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index 1a00be99..4dc25bfa 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -374,7 +374,7 @@ rule verification_scoremaps_baseline: # the validated config, which fills in the SalConfig defaults. Absent when there # is no sal block at all, in which case these rules never fire and the script's # own defaults would apply anyway. -_SAL_CONFIG = config.get("experiment", {}).get("sal") or {} +_SAL_CONFIG = EXPERIMENT_CONFIG.get("sal") or {} _SAL_ARGS = ( f"--grid-extent {' '.join(str(x) for x in _SAL_CONFIG['grid_extent'])} " f"--grid-step-lat {_SAL_CONFIG['grid_step_lat']} " From 2bec0251a37f8f45c735c33331957ea70f9a4a48 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 19:37:40 +0100 Subject: [PATCH 10/21] document that tot_prec1 spectra need hourly steps --- src/evalml/config.py | 3 ++- workflow/tools/config.schema.json | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/evalml/config.py b/src/evalml/config.py index 0f2571e7..8db1a7eb 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -489,7 +489,8 @@ class SpectraConfig(BaseModel): description=( "Spectra variables: any parameter the forecast and truth loaders provide " "(e.g. T_2M, SP_10M, TOT_PREC1), plus WIND_KE, the kinetic energy from " - "U_10M and V_10M." + "U_10M and V_10M. TOT_PREC1 needs hourly output steps: runs with coarser " + "steps (e.g. 6h) must use TOT_PREC6 instead." ), ) init_hours: Optional[List[int]] = Field( diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index a51c7939..557ea336 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -1007,7 +1007,7 @@ "WIND_KE", "TOT_PREC1" ], - "description": "Spectra variables: any parameter the forecast and truth loaders provide (e.g. T_2M, SP_10M, TOT_PREC1), plus WIND_KE, the kinetic energy from U_10M and V_10M.", + "description": "Spectra variables: any parameter the forecast and truth loaders provide (e.g. T_2M, SP_10M, TOT_PREC1), plus WIND_KE, the kinetic energy from U_10M and V_10M. TOT_PREC1 needs hourly output steps: runs with coarser steps (e.g. 6h) must use TOT_PREC6 instead.", "items": { "type": "string" }, From 940bb1173eacde684e7f4e8cd094398dd7394e50 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 21:18:12 +0100 Subject: [PATCH 11/21] track spectra settings as rule params --- workflow/rules/spectra.smk | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index f3289a92..43d8e0e9 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -58,13 +58,16 @@ rule spectra_compute: ).resolve(), steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], label=lambda wc: RUN_CONFIGS[wc.run_id].get("label", wc.run_id), + lead_times=SPECTRA_LEAD_TIMES, + variables=SPECTRA_VARIABLES, + method=SPECTRA_METHOD, shell: """ export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) uv run {input.script} \ --forecast {params.grib} --reftime {wildcards.init_time} \ - --steps "{params.steps}" --lead_times "{SPECTRA_LEAD_TIMES}" \ - --variables "{SPECTRA_VARIABLES}" --method "{SPECTRA_METHOD}" \ + --steps "{params.steps}" --lead_times "{params.lead_times}" \ + --variables "{params.variables}" --method "{params.method}" \ --label "{params.label}" --output {output} >{log} 2>&1 """ @@ -86,13 +89,16 @@ rule spectra_compute_baseline: params: steps=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["steps"], label=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("label", wc.baseline_id), + lead_times=SPECTRA_LEAD_TIMES, + variables=SPECTRA_VARIABLES, + method=SPECTRA_METHOD, shell: """ export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) uv run {input.script} \ --forecast {input.forecast} --reftime {wildcards.init_time} \ - --steps "{params.steps}" --lead_times "{SPECTRA_LEAD_TIMES}" \ - --variables "{SPECTRA_VARIABLES}" --method "{SPECTRA_METHOD}" \ + --steps "{params.steps}" --lead_times "{params.lead_times}" \ + --variables "{params.variables}" --method "{params.method}" \ --label "{params.label}" --output {output} >{log} 2>&1 """ @@ -114,13 +120,16 @@ rule spectra_compute_truth: params: steps=_spectra_truth_steps(), label=config["truth"]["label"], + lead_times=SPECTRA_LEAD_TIMES, + variables=SPECTRA_VARIABLES, + method=SPECTRA_METHOD, shell: """ export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) uv run {input.script} \ --truth {input.truth} --reftime {wildcards.init_time} \ - --steps "{params.steps}" --lead_times "{SPECTRA_LEAD_TIMES}" \ - --variables "{SPECTRA_VARIABLES}" --method "{SPECTRA_METHOD}" \ + --steps "{params.steps}" --lead_times "{params.lead_times}" \ + --variables "{params.variables}" --method "{params.method}" \ --label "{params.label}" --output {output} >{log} 2>&1 """ From 5395bec320320ce05d223b5dd30acab27a63847a Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 21:19:07 +0100 Subject: [PATCH 12/21] pass the baseline member to the spectra loader --- tests/unit/test_spectra_compute_script.py | 59 +++++++++++++++++++++++ workflow/rules/spectra.smk | 2 + workflow/scripts/spectra_compute.py | 9 +++- 3 files changed, 69 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_spectra_compute_script.py diff --git a/tests/unit/test_spectra_compute_script.py b/tests/unit/test_spectra_compute_script.py new file mode 100644 index 00000000..342faca0 --- /dev/null +++ b/tests/unit/test_spectra_compute_script.py @@ -0,0 +1,59 @@ +import importlib.util +import sys +from pathlib import Path + +import xarray as xr + +SCRIPT = Path(__file__).resolve().parents[2] / "workflow/scripts/spectra_compute.py" +# spectra_compute.py loaded by path (workflow/scripts is not a package). +_spec = importlib.util.spec_from_file_location("spectra_compute", SCRIPT) +_script = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_script) + + +def _run(monkeypatch, tmp_path, extra_args): + calls = {} + + def fake_load_forecast_data(root, reftime, steps, params, member="000"): + calls["steps"] = steps + calls["member"] = member + return xr.Dataset() + + monkeypatch.setattr(_script, "load_forecast_data", fake_load_forecast_data) + monkeypatch.setattr( + _script, "compute_source_spectra", lambda *args, **kwargs: xr.Dataset() + ) + monkeypatch.setattr( + sys, + "argv", + [ + "spectra_compute.py", + "--forecast", + str(tmp_path), + "--reftime", + "202406150000", + "--steps", + "0/24/6", + "--lead_times", + "6,24", + "--variables", + "T_2M", + "--label", + "baseline", + "--output", + str(tmp_path / "spectra.nc"), + *extra_args, + ], + ) + _script.main() + return calls + + +def test_member_is_passed_to_loader(monkeypatch, tmp_path): + calls = _run(monkeypatch, tmp_path, ["--member", "median"]) + assert calls["member"] == "median" + + +def test_member_defaults_to_control(monkeypatch, tmp_path): + calls = _run(monkeypatch, tmp_path, []) + assert calls["member"] == "000" diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index 43d8e0e9..b9ba34fe 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -89,6 +89,7 @@ rule spectra_compute_baseline: params: steps=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["steps"], label=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("label", wc.baseline_id), + member=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("member", "000"), lead_times=SPECTRA_LEAD_TIMES, variables=SPECTRA_VARIABLES, method=SPECTRA_METHOD, @@ -97,6 +98,7 @@ rule spectra_compute_baseline: export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) uv run {input.script} \ --forecast {input.forecast} --reftime {wildcards.init_time} \ + --member "{params.member}" \ --steps "{params.steps}" --lead_times "{params.lead_times}" \ --variables "{params.variables}" --method "{params.method}" \ --label "{params.label}" --output {output} >{log} 2>&1 diff --git a/workflow/scripts/spectra_compute.py b/workflow/scripts/spectra_compute.py index 327ab59f..6f017a56 100644 --- a/workflow/scripts/spectra_compute.py +++ b/workflow/scripts/spectra_compute.py @@ -23,6 +23,11 @@ def main(): ap.add_argument("--forecast", help="Forecast root (grib dir or baseline archive).") ap.add_argument("--truth", help="Truth zarr root (use instead of --forecast).") ap.add_argument("--reftime", required=True) + ap.add_argument( + "--member", + default="000", + help="Baseline ensemble member: '000' for control, 'median', 'mean', or a 3-digit ID.", + ) ap.add_argument("--steps", required=True, help="start/stop/step in hours.") ap.add_argument("--lead_times", required=True, help="comma-separated hours.") ap.add_argument("--variables", required=True, help="comma-separated names.") @@ -66,7 +71,9 @@ def main(): ) ds = ds.assign_coords(step=("time", steps_h)).swap_dims({"time": "step"}) else: - ds = load_forecast_data(Path(args.forecast), reftime, all_steps, params) + ds = load_forecast_data( + Path(args.forecast), reftime, all_steps, params, member=args.member + ) result = compute_source_spectra(ds, variables, lead_times, args.method, args.label) out = Path(args.output) From de45a6c64f8a50f8c0433bd4af4ab3ece7392bae Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 21:19:35 +0100 Subject: [PATCH 13/21] load only the requested spectra lead times --- tests/unit/test_spectra_compute_script.py | 5 +++++ workflow/scripts/spectra_compute.py | 6 ++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_spectra_compute_script.py b/tests/unit/test_spectra_compute_script.py index 342faca0..92c73460 100644 --- a/tests/unit/test_spectra_compute_script.py +++ b/tests/unit/test_spectra_compute_script.py @@ -57,3 +57,8 @@ def test_member_is_passed_to_loader(monkeypatch, tmp_path): def test_member_defaults_to_control(monkeypatch, tmp_path): calls = _run(monkeypatch, tmp_path, []) assert calls["member"] == "000" + + +def test_loads_only_requested_lead_times(monkeypatch, tmp_path): + calls = _run(monkeypatch, tmp_path, []) + assert calls["steps"] == [6, 24] diff --git a/workflow/scripts/spectra_compute.py b/workflow/scripts/spectra_compute.py index 6f017a56..5a1dc5eb 100644 --- a/workflow/scripts/spectra_compute.py +++ b/workflow/scripts/spectra_compute.py @@ -55,8 +55,10 @@ def main(): variables = args.variables.split(",") params = io.required_params(variables) + # Load only the requested lead times: the loaders add the preceding steps + # that accumulation periods (e.g. TOT_PREC6) need. if args.truth: - ds = load_truth_data(Path(args.truth), reftime, all_steps, params) + ds = load_truth_data(Path(args.truth), reftime, lead_times, params) if "time" in ds.dims and "step" not in ds.dims: valid = pd.to_datetime(ds["time"].values) steps_h = ( @@ -72,7 +74,7 @@ def main(): ds = ds.assign_coords(step=("time", steps_h)).swap_dims({"time": "step"}) else: ds = load_forecast_data( - Path(args.forecast), reftime, all_steps, params, member=args.member + Path(args.forecast), reftime, lead_times, params, member=args.member ) result = compute_source_spectra(ds, variables, lead_times, args.method, args.label) From 27f7a57284c459d19e91b27f6af3d0dfb928225b Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 21:20:29 +0100 Subject: [PATCH 14/21] check spectra lead times against every participant --- src/evalml/config.py | 15 +++++++++------ tests/unit/test_config.py | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/src/evalml/config.py b/src/evalml/config.py index 8db1a7eb..04943631 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -888,22 +888,25 @@ class ConfigModel(BaseModel): @model_validator(mode="after") def validate_leadtimes_producible(self) -> "ConfigModel": - """Reject scoremaps/sal lead times not produced by every participant.""" + """Reject scoremaps/sal/spectra lead times not produced by every participant.""" if self.experiment is None: return self - for label, cfg in ( - ("scoremaps", self.experiment.scoremaps), - ("sal", self.experiment.sal), + # (section, config, name of its lead-time field) + for label, cfg, field in ( + ("scoremaps", self.experiment.scoremaps, "leadtimes"), + ("sal", self.experiment.sal, "leadtimes"), + ("spectra", self.experiment.spectra, "lead_times"), ): if cfg is None or not cfg.enabled: continue for item in self.runs: steps = getattr(item, next(iter(type(item).model_fields))).steps start, end, step = map(int, steps.split("/")) - unsupported = set(cfg.leadtimes) - set(range(start, end + 1, step)) + leadtimes = getattr(cfg, field) + unsupported = set(leadtimes) - set(range(start, end + 1, step)) if unsupported: raise ValueError( - f"{label}.leadtimes contains {sorted(unsupported)} h which are not " + f"{label}.{field} contains {sorted(unsupported)} h which are not " f"produced by participant with steps '{steps}'." ) return self diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index c1880c77..ec4db1c9 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -135,3 +135,18 @@ def test_scorecard_reference_must_exist(example_config): ConfigModel.model_validate( _with_scorecard_reference(example_config, "no-such-label") ) + + +def test_spectra_lead_times_must_be_produced_by_every_participant(example_config): + # the INCA baseline of the example config only covers 0/6/1 + example_config["experiment"]["spectra"] = {"enabled": True, "lead_times": [3, 6]} + ConfigModel.model_validate(example_config) + + example_config["experiment"]["spectra"]["lead_times"] = [6, 24] + with pytest.raises(ValueError, match=r"spectra.lead_times contains \[24\]"): + ConfigModel.model_validate(example_config) + + +def test_spectra_lead_times_unchecked_when_disabled(example_config): + example_config["experiment"]["spectra"] = {"enabled": False, "lead_times": [999]} + ConfigModel.model_validate(example_config) From b9795df472f6149e689edd79fafad063379e6545 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 21:21:44 +0100 Subject: [PATCH 15/21] require unique labels across runs and baselines --- src/evalml/config.py | 32 ++++++++++++++++++++++++++++++++ tests/unit/test_config.py | 21 +++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/evalml/config.py b/src/evalml/config.py index 04943631..7e5b77ab 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -911,6 +911,38 @@ def validate_leadtimes_producible(self) -> "ConfigModel": ) return self + @model_validator(mode="after") + def validate_unique_labels(self) -> "ConfigModel": + """Reject a label shared by different runs or baselines. + + Results are keyed by label, so a repeated label would merge or overwrite + participants. The inline forecaster of a temporal downscaler is included; + an entry repeated with the same configuration is the same run and is + allowed. + """ + configs_by_label: dict[str, list[dict]] = {} + for item in self.runs: + cfg = getattr(item, next(iter(type(item).model_fields))) + entries = [cfg] + if isinstance(cfg, TemporalDownscalerConfig) and cfg.forecaster: + entries.append(cfg.forecaster) + for entry in entries: + if entry.label is None: + continue + configs = configs_by_label.setdefault(entry.label, []) + dump = entry.model_dump() + if dump not in configs: + configs.append(dump) + duplicates = sorted( + label for label, configs in configs_by_label.items() if len(configs) > 1 + ) + if duplicates: + raise ValueError( + f"Labels must be unique across `runs`; used by several runs or " + f"baselines: {duplicates}" + ) + return self + @model_validator(mode="after") def validate_scorecard_references(self) -> "ConfigModel": sc = self.experiment.scorecards if self.experiment else None diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index ec4db1c9..f74fb7b6 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -150,3 +150,24 @@ def test_spectra_lead_times_must_be_produced_by_every_participant(example_config def test_spectra_lead_times_unchecked_when_disabled(example_config): example_config["experiment"]["spectra"] = {"enabled": False, "lead_times": [999]} ConfigModel.model_validate(example_config) + + +def test_labels_must_be_unique(example_config): + example_config["runs"][2]["baseline"]["label"] = "INCA" # same as runs[1] + with pytest.raises(ValueError, match=r"Labels must be unique.*\['INCA'\]"): + ConfigModel.model_validate(example_config) + + +def test_inline_forecaster_label_must_be_unique(example_config): + td = example_config["runs"][0]["temporal_downscaler"] + td["forecaster"]["label"] = td["label"] + with pytest.raises(ValueError, match="Labels must be unique"): + ConfigModel.model_validate(example_config) + + +def test_inline_forecaster_may_repeat_a_listed_forecaster(example_config): + """The same forecaster listed and used inline is one run: its label may repeat.""" + td = example_config["runs"][0]["temporal_downscaler"] + td["forecaster"]["label"] = "Varda forecaster" + example_config["runs"].append({"forecaster": dict(td["forecaster"])}) + ConfigModel.model_validate(example_config) From da06001b5c5528e7b0a1d200f0d408a78f8e3d03 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 21:24:28 +0100 Subject: [PATCH 16/21] fall back to run id for unlabelled spectra runs --- workflow/rules/spectra.smk | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index b9ba34fe..272c5365 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -57,7 +57,7 @@ rule spectra_compute: Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" ).resolve(), steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], - label=lambda wc: RUN_CONFIGS[wc.run_id].get("label", wc.run_id), + label=lambda wc: RUN_CONFIGS[wc.run_id].get("label") or wc.run_id, lead_times=SPECTRA_LEAD_TIMES, variables=SPECTRA_VARIABLES, method=SPECTRA_METHOD, From d3d606a292f979e406bdd92263f9558f86820736 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Fri, 6 Nov 2026 09:25:04 +0100 Subject: [PATCH 17/21] reject spectra sources that are not on an icon native grid --- src/evalml/config.py | 41 ++++++++++++++++++++- src/spectra/regrid.py | 5 +-- tests/unit/test_config.py | 60 +++++++++++++++++++++++++++---- tests/unit/test_spectra_regrid.py | 2 +- workflow/tools/config.schema.json | 2 +- 5 files changed, 99 insertions(+), 11 deletions(-) diff --git a/src/evalml/config.py b/src/evalml/config.py index 7e5b77ab..9d9d36a1 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -649,7 +649,11 @@ class ExperimentConfig(BaseModel): ) spectra: SpectraConfig = Field( default_factory=SpectraConfig, - description="Power-spectra QC configuration. Disabled by default.", + description=( + "Power-spectra QC configuration. Disabled by default. Spectra need every " + "participant and the truth on the ICON-CH1 or ICON-CH2 native grid: INCA " + "baselines, INCA truth and station (jretrievedwh) truth are not supported." + ), ) sal: Optional[SalConfig] = Field( default=None, @@ -911,6 +915,41 @@ def validate_leadtimes_producible(self) -> "ConfigModel": ) return self + @model_validator(mode="after") + def validate_spectra_sources(self) -> "ConfigModel": + """Reject spectra participants or truth that are not on an ICON native grid. + + The detection mirrors ``data_input.is_inca_archive`` and + ``data_input.truth_type``, which are not imported here to keep the config + module light. + """ + if self.experiment is None or not self.experiment.spectra.enabled: + return self + unsupported = [] + for item in self.runs: + if ( + isinstance(item, BaselineItem) + and "INCA" in Path(item.baseline.root).parts + ): + unsupported.append(f"baseline {item.baseline.label!r} (INCA)") + if self.truth is not None: + root = self.truth.root + if Path(root).suffix != ".zarr": + if "jretrieve" in root: + unsupported.append( + f"truth {self.truth.label!r} (station observations)" + ) + elif "INCA" in Path(root).parts: + unsupported.append(f"truth {self.truth.label!r} (INCA)") + if unsupported: + raise ValueError( + "experiment.spectra needs every participant and the truth on the " + "ICON-CH1 or ICON-CH2 native grid; not supported: " + + ", ".join(unsupported) + + ". Disable spectra or remove these sources." + ) + return self + @model_validator(mode="after") def validate_unique_labels(self) -> "ConfigModel": """Reject a label shared by different runs or baselines. diff --git a/src/spectra/regrid.py b/src/spectra/regrid.py index 3aa9ce39..c8106225 100644 --- a/src/spectra/regrid.py +++ b/src/spectra/regrid.py @@ -32,9 +32,10 @@ def weights_path_for_npoints(npoints: int) -> Path: try: return ICON_WEIGHTS[npoints] except KeyError as exc: - raise KeyError( + raise ValueError( f"No iconremap weights registered for native grid with {npoints} " - f"cells. Known grids: {sorted(ICON_WEIGHTS)}." + f"cells. Power spectra need fields on the ICON-CH1 ({CH1_NPOINTS} cells) " + f"or ICON-CH2 ({CH2_NPOINTS} cells) native grid." ) from exc diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index f74fb7b6..b36fb99f 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -138,13 +138,13 @@ def test_scorecard_reference_must_exist(example_config): def test_spectra_lead_times_must_be_produced_by_every_participant(example_config): - # the INCA baseline of the example config only covers 0/6/1 - example_config["experiment"]["spectra"] = {"enabled": True, "lead_times": [3, 6]} - ConfigModel.model_validate(example_config) + # the ICON-CH1-CTRL baseline of the example config only covers 0/33/1 + config = _spectra_ready(example_config, lead_times=[6, 24]) + ConfigModel.model_validate(config) - example_config["experiment"]["spectra"]["lead_times"] = [6, 24] - with pytest.raises(ValueError, match=r"spectra.lead_times contains \[24\]"): - ConfigModel.model_validate(example_config) + config["experiment"]["spectra"]["lead_times"] = [6, 48] + with pytest.raises(ValueError, match=r"spectra.lead_times contains \[48\]"): + ConfigModel.model_validate(config) def test_spectra_lead_times_unchecked_when_disabled(example_config): @@ -171,3 +171,51 @@ def test_inline_forecaster_may_repeat_a_listed_forecaster(example_config): td["forecaster"]["label"] = "Varda forecaster" example_config["runs"].append({"forecaster": dict(td["forecaster"])}) ConfigModel.model_validate(example_config) + + +def _spectra_ready(config, lead_times=(6,)): + """Example config with spectra on: ICON participants and a zarr truth only.""" + config["runs"] = [ + r for r in config["runs"] if r.get("baseline", {}).get("label") != "INCA" + ] + config["experiment"].pop("scorecards", None) # references the INCA baseline + config["experiment"]["spectra"] = {"enabled": True, "lead_times": list(lead_times)} + config["truth"] = {"label": "REA-L-CH1", "root": "/store/realch1.zarr"} + return config + + +def test_spectra_accepts_icon_sources_and_zarr_truth(example_config): + ConfigModel.model_validate(_spectra_ready(example_config)) + + +def test_spectra_rejects_inca_baseline(example_config): + config = _spectra_ready(example_config) + config["runs"].append( + { + "baseline": { + "label": "INCA", + "root": "/store_new/mch/msclim/INCA", + "steps": "0/6/1", + } + } + ) + with pytest.raises(ValueError, match=r"native grid.*baseline 'INCA' \(INCA\)"): + ConfigModel.model_validate(config) + + +@pytest.mark.parametrize( + "root, kind", + [ + ("jretrievedwh:SwissMetNet", "station observations"), + ("/store_new/mch/msclim/INCA", "INCA"), + ], +) +def test_spectra_rejects_station_and_inca_truth(example_config, root, kind): + config = _spectra_ready(example_config) + config["truth"]["root"] = root + with pytest.raises(ValueError, match=rf"native grid.*truth 'REA-L-CH1' \({kind}\)"): + ConfigModel.model_validate(config) + + +def test_inca_allowed_when_spectra_disabled(example_config): + ConfigModel.model_validate(example_config) # INCA baseline + station truth diff --git a/tests/unit/test_spectra_regrid.py b/tests/unit/test_spectra_regrid.py index bf439796..7f2e457d 100644 --- a/tests/unit/test_spectra_regrid.py +++ b/tests/unit/test_spectra_regrid.py @@ -42,7 +42,7 @@ def test_grid_geometry_reads_dx_in_km(tmp_path): def test_weights_path_for_npoints_unknown_raises(): - with pytest.raises(KeyError, match="No iconremap weights"): + with pytest.raises(ValueError, match="No iconremap weights.*ICON-CH1"): regrid.weights_path_for_npoints(999_999) diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 557ea336..69a2b5fd 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -346,7 +346,7 @@ }, "spectra": { "$ref": "#/$defs/SpectraConfig", - "description": "Power-spectra QC configuration. Disabled by default." + "description": "Power-spectra QC configuration. Disabled by default. Spectra need every participant and the truth on the ICON-CH1 or ICON-CH2 native grid: INCA baselines, INCA truth and station (jretrievedwh) truth are not supported." }, "sal": { "anyOf": [ From 0aa25923820717c72e01c0cbd3a12cb31d023a3c Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Fri, 6 Nov 2026 09:25:22 +0100 Subject: [PATCH 18/21] fail early when spectra init_hours match no date --- src/evalml/config.py | 1 + tests/unit/test_config.py | 6 ++++++ workflow/rules/spectra.smk | 8 ++++++++ workflow/tools/config.schema.json | 1 + 4 files changed, 16 insertions(+) diff --git a/src/evalml/config.py b/src/evalml/config.py index 9d9d36a1..823ad09f 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -495,6 +495,7 @@ class SpectraConfig(BaseModel): ) init_hours: Optional[List[int]] = Field( default=None, + min_length=1, description="Optional subset of init hours to average over. None = all.", ) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index b36fb99f..00befe23 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -219,3 +219,9 @@ def test_spectra_rejects_station_and_inca_truth(example_config, root, kind): def test_inca_allowed_when_spectra_disabled(example_config): ConfigModel.model_validate(example_config) # INCA baseline + station truth + + +def test_spectra_init_hours_must_not_be_empty(): + with pytest.raises(ValidationError): + SpectraConfig(init_hours=[]) + assert SpectraConfig(init_hours=[0, 12]).init_hours == [0, 12] diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index 272c5365..76b6e442 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -19,6 +19,14 @@ def _spectra_reftimes(): return [t.strftime("%Y%m%d%H%M") for t in REFTIMES if t.hour in hours] +if SPECTRA.get("enabled", False) and not _spectra_reftimes(): + raise ValueError( + f"experiment.spectra.init_hours {SPECTRA.get('init_hours')} matches none of " + f"the experiment dates (init hours {sorted({t.hour for t in REFTIMES})}); " + "adjust init_hours or remove it to use all dates." + ) + + def _spectra_truth_steps(): candidates = list(collect_all_candidates()) if candidates: diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 69a2b5fd..1228579e 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -1020,6 +1020,7 @@ "items": { "type": "integer" }, + "minItems": 1, "type": "array" }, { From 333d15ef1ffc765b370783b9774e4124605d7bf1 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Fri, 6 Nov 2026 09:25:32 +0100 Subject: [PATCH 19/21] reserve memory-proportional cores for spectra jobs --- workflow/rules/spectra.smk | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/workflow/rules/spectra.smk b/workflow/rules/spectra.smk index 76b6e442..d4afbec6 100644 --- a/workflow/rules/spectra.smk +++ b/workflow/rules/spectra.smk @@ -57,7 +57,7 @@ rule spectra_compute: run_id=r"[^/]+/[^/]+", init_time=r"\d+", resources: - cpus_per_task=8, + cpus_per_task=balfrin_cpus(50_000, actual_cpus=8), mem_mb=50_000, runtime="30m", params: @@ -91,7 +91,7 @@ rule spectra_compute_baseline: log: OUT_ROOT / "logs/spectra_compute_baseline/{baseline_id}-{init_time}.log", resources: - cpus_per_task=8, + cpus_per_task=balfrin_cpus(50_000, actual_cpus=8), mem_mb=50_000, runtime="30m", params: @@ -124,7 +124,7 @@ rule spectra_compute_truth: log: OUT_ROOT / "logs/spectra_truth/{init_time}.log", resources: - cpus_per_task=8, + cpus_per_task=balfrin_cpus(50_000, actual_cpus=8), mem_mb=50_000, runtime="30m", params: @@ -209,7 +209,7 @@ rule spectra_plot: log: OUT_ROOT / "logs/spectra_plot/{experiment}.log", resources: - cpus_per_task=2, + cpus_per_task=balfrin_cpus(20_000, actual_cpus=2), mem_mb=20_000, runtime="20m", shell: From 845b9784858298fbb6c37df780bd2bada326759c Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Fri, 6 Nov 2026 12:44:24 +0100 Subject: [PATCH 20/21] move regridding out of spectra into a top-level regrid package --- pyproject.toml | 3 ++- src/{spectra/regrid.py => regrid/__init__.py} | 9 +++++---- src/spectra/__init__.py | 2 +- src/spectra/compute.py | 3 ++- src/spectra/core.py | 2 -- src/spectra/plot.py | 2 +- tests/unit/{test_spectra_regrid.py => test_regrid.py} | 5 ++--- 7 files changed, 13 insertions(+), 13 deletions(-) rename src/{spectra/regrid.py => regrid/__init__.py} (95%) rename tests/unit/{test_spectra_regrid.py => test_regrid.py} (96%) diff --git a/pyproject.toml b/pyproject.toml index 5943a199..a2587ec1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,8 @@ packages = [ "src/evalml", "src/verification", "src/data_input", - "src/spectra" + "src/spectra", + "src/regrid" ] [tool.uv.sources] diff --git a/src/spectra/regrid.py b/src/regrid/__init__.py similarity index 95% rename from src/spectra/regrid.py rename to src/regrid/__init__.py index c8106225..67fe7f4d 100644 --- a/src/spectra/regrid.py +++ b/src/regrid/__init__.py @@ -1,6 +1,7 @@ -# src/spectra/regrid.py """ICON native (unstructured) -> regular rotlatlon regridding via iconremap RBF -weights. The weights file is selected at runtime from the native point count.""" +weights. The weights file is selected at runtime from the native point count. + +Shared by the spectra diagnostic and usable elsewhere (e.g. plotting).""" from __future__ import annotations @@ -11,7 +12,7 @@ import scipy.sparse import xarray as xr -from spectra import core +KM_PER_DEG = 111.2 # Hardcoded iconremap RBF weight files on Balfrin, keyed by native cell count. # CH1 native count is 1,147,980 (per the icon-power-spectra skill). @@ -68,7 +69,7 @@ def grid_geometry(iconremap_nc): return ( int(ds.attrs["ny"]), int(ds.attrs["nx"]), - float(ds.attrs["dx"]) * core.KM_PER_DEG, + float(ds.attrs["dx"]) * KM_PER_DEG, ) diff --git a/src/spectra/__init__.py b/src/spectra/__init__.py index fa47a74d..42206476 100644 --- a/src/spectra/__init__.py +++ b/src/spectra/__init__.py @@ -4,4 +4,4 @@ directly. The spectral `core` is numpy-only and reusable for any model. """ -from spectra import compute, core, io, plot, regrid # noqa: F401 +from spectra import compute, core, io, plot # noqa: F401 diff --git a/src/spectra/compute.py b/src/spectra/compute.py index a4ae5217..b13f0833 100644 --- a/src/spectra/compute.py +++ b/src/spectra/compute.py @@ -6,7 +6,8 @@ import numpy as np import xarray as xr -from spectra import core, io, regrid +import regrid +from spectra import core, io def compute_source_spectra(ds: xr.Dataset, variables, lead_times, method, label): diff --git a/src/spectra/core.py b/src/spectra/core.py index b917eef4..4025dadc 100644 --- a/src/spectra/core.py +++ b/src/spectra/core.py @@ -12,8 +12,6 @@ import numpy as np import scipy.fft -KM_PER_DEG = 111.2 - def _hann2d(ny: int, nx: int) -> np.ndarray: win = np.hanning(ny)[:, None] * np.hanning(nx)[None, :] diff --git a/src/spectra/plot.py b/src/spectra/plot.py index 316de825..489d95f9 100644 --- a/src/spectra/plot.py +++ b/src/spectra/plot.py @@ -10,7 +10,7 @@ import numpy as np import xarray as xr -from spectra.core import KM_PER_DEG +from regrid import KM_PER_DEG REFERENCE_GRIDS = { "O96 (~104 km)": 90.0 / 96.0 * KM_PER_DEG, diff --git a/tests/unit/test_spectra_regrid.py b/tests/unit/test_regrid.py similarity index 96% rename from tests/unit/test_spectra_regrid.py rename to tests/unit/test_regrid.py index 7f2e457d..d5eececa 100644 --- a/tests/unit/test_spectra_regrid.py +++ b/tests/unit/test_regrid.py @@ -1,9 +1,8 @@ -# tests/unit/test_spectra_regrid.py import numpy as np import pytest import xarray as xr -from spectra import regrid +import regrid def _fake_weights(tmp_path, ny=4, nx=5, n_source=20, npoints=3): @@ -38,7 +37,7 @@ def test_grid_geometry_reads_dx_in_km(tmp_path): path, ny, nx, _ = _fake_weights(tmp_path) g_ny, g_nx, dx_km = regrid.grid_geometry(path) assert (g_ny, g_nx) == (ny, nx) - assert dx_km == pytest.approx(0.01 * regrid.core.KM_PER_DEG) + assert dx_km == pytest.approx(0.01 * regrid.KM_PER_DEG) def test_weights_path_for_npoints_unknown_raises(): From d075b18a2647aa8e856cbfbece0672ac22d635d7 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Fri, 6 Nov 2026 16:56:03 +0100 Subject: [PATCH 21/21] drop mislabelled cerra reference line from spectra plots --- src/spectra/plot.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/spectra/plot.py b/src/spectra/plot.py index 489d95f9..8246c34a 100644 --- a/src/spectra/plot.py +++ b/src/spectra/plot.py @@ -15,7 +15,6 @@ REFERENCE_GRIDS = { "O96 (~104 km)": 90.0 / 96.0 * KM_PER_DEG, "N320 (~31 km)": 0.28125 * KM_PER_DEG, - "CERRA (~9 km)": 0.10 * KM_PER_DEG, } EFF_RES_FACTOR = 6.0