Power spectra plots

pyproject.toml

@@ -60,7 +60,9 @@ markers = [
 packages = [
   "src/evalml",
   "src/verification",
-  "src/data_input"
+  "src/data_input",
+  "src/spectra",
+  "src/regrid"
 ]
 
 [tool.uv.sources]

src/evalml/config.py

@@ -1,6 +1,6 @@
 from math import isfinite
 from pathlib import Path
-from typing import Dict, List, Any, ClassVar, FrozenSet, Optional, Union
+from typing import Dict, List, Any, ClassVar, FrozenSet, Literal, Optional, Union
 
 from pydantic import BaseModel, Field, RootModel, field_validator, model_validator
 
@@ -469,6 +469,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_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. TOT_PREC1 needs hourly output steps: runs with coarser "
+            "steps (e.g. 6h) must use TOT_PREC6 instead."
+        ),
+    )
+    init_hours: Optional[List[int]] = Field(
+        default=None,
+        min_length=1,
+        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."
+            )
+        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."""
 
@@ -597,6 +648,14 @@ class ExperimentConfig(BaseModel):
         default=None,
         description="Score map plot configuration. Omit or set enabled: false to disable.",
     )
+    spectra: SpectraConfig = Field(
+        default_factory=SpectraConfig,
+        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,
         description="SAL precipitation verification. Omit or set enabled: false to disable.",
@@ -834,26 +893,96 @@ 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
 
+    @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.
+
+        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

src/regrid/__init__.py

@@ -0,0 +1,92 @@
+"""ICON native (unstructured) -> regular rotlatlon regridding via iconremap RBF
+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
+
+from functools import lru_cache
+from pathlib import Path
+
+import numpy as np
+import scipy.sparse
+import xarray as xr
+
+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).
+# 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 ValueError(
+            f"No iconremap weights registered for native grid with {npoints} "
+            f"cells. Power spectra need fields on the ICON-CH1 ({CH1_NPOINTS} cells) "
+            f"or ICON-CH2 ({CH2_NPOINTS} cells) native grid."
+        ) 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"]) * 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

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, plot  # noqa: F401

src/spectra/compute.py

@@ -0,0 +1,75 @@
+"""Orchestration: per-source spectra computation and over-init aggregation.
+Defines the spectra.nc schema."""
+
+from __future__ import annotations
+
+import numpy as np
+import xarray as xr
+
+import regrid
+from spectra import core, io
+
+
+def compute_source_spectra(ds: xr.Dataset, variables, lead_times, method, label):
+    """Compute spectra for one source (already-loaded native dataset).
+
+    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:
+        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. 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=False)
+        agg.attrs = datasets[0].attrs
+        agg["wavelength"] = datasets[0]["wavelength"].copy()
+        agg.load()
+    finally:
+        for ds in datasets:
+            ds.close()
+    return agg