From f78f9dd731505c6e26c6be4f212b3ed709d0c40a Mon Sep 17 00:00:00 2001 From: Alexey Yermakov Date: Mon, 2 Nov 2026 15:59:22 -0800 Subject: [PATCH 1/4] Add per-sample global-mean drift inference aggregator Add GlobalMeanDriftAggregator, which streams per-sample least-squares slopes (per year) of the area-weighted global-mean bias of each variable over an inference rollout. It logs the sample-mean slope, its t-statistic across samples and the max |t| over variables, and saves the per-sample slopes as a (sample, variable) array in the diagnostics. Register it on InferenceEvaluatorAggregatorConfig as global_mean_drift (enabled by default, skipped for rollouts of ~2 years or less like the annual metric) and expose the per-sample slopes via InferenceEvaluatorAggregator.get_global_mean_drift_slopes(). --- fme/ace/__init__.py | 1 + fme/ace/aggregator/inference/__init__.py | 1 + fme/ace/aggregator/inference/data.py | 5 + .../aggregator/inference/global_mean_drift.py | 275 ++++++++++++++++++ fme/ace/aggregator/inference/main.py | 41 ++- .../inference/test_global_mean_drift.py | 257 ++++++++++++++++ fme/core/testing/distributed.py | 3 + 7 files changed, 582 insertions(+), 1 deletion(-) create mode 100644 fme/ace/aggregator/inference/global_mean_drift.py create mode 100644 fme/ace/aggregator/inference/test_global_mean_drift.py diff --git a/fme/ace/__init__.py b/fme/ace/__init__.py index 1ca547f1cf..885665c08e 100644 --- a/fme/ace/__init__.py +++ b/fme/ace/__init__.py @@ -5,6 +5,7 @@ from fme.ace.aggregator.inference.enso.enso_coefficient import ( EnsoCoefficientMetricConfig, ) +from fme.ace.aggregator.inference.global_mean_drift import GlobalMeanDriftConfig from fme.ace.aggregator.inference.histogram import HistogramMetricConfig from fme.ace.aggregator.inference.ipo.ipo_index import IpoIndexMetricConfig from fme.ace.aggregator.inference.main import ( diff --git a/fme/ace/aggregator/inference/__init__.py b/fme/ace/aggregator/inference/__init__.py index 3fe3da6985..c0529114dd 100644 --- a/fme/ace/aggregator/inference/__init__.py +++ b/fme/ace/aggregator/inference/__init__.py @@ -3,6 +3,7 @@ from .annual import AnnualMetricConfig from .enso.dynamic_index import EnsoIndexMetricConfig from .enso.enso_coefficient import EnsoCoefficientMetricConfig +from .global_mean_drift import GlobalMeanDriftConfig from .histogram import HistogramMetricConfig from .ipo.ipo_index import IpoIndexMetricConfig from .main import ( diff --git a/fme/ace/aggregator/inference/data.py b/fme/ace/aggregator/inference/data.py index 9a7b00dcd4..6ec741374a 100644 --- a/fme/ace/aggregator/inference/data.py +++ b/fme/ace/aggregator/inference/data.py @@ -157,6 +157,10 @@ def get_logs(self, label: str) -> dict[str, Any]: ... def get_dataset(self) -> xr.Dataset: ... +class PerSampleSlopes(Protocol): + def get_per_sample_slopes(self) -> dict[str, torch.Tensor]: ... + + class TimeSeriesLogs(Protocol): def get_logs(self, label: str, step_slice: slice = ...) -> dict[str, Any]: ... @@ -172,3 +176,4 @@ class MetricBuildResult: aggregator: SubAggregator | None = None time_series: TimeSeriesLogs | None = None ensemble: Any = None + per_sample_slopes: PerSampleSlopes | None = None diff --git a/fme/ace/aggregator/inference/global_mean_drift.py b/fme/ace/aggregator/inference/global_mean_drift.py new file mode 100644 index 0000000000..adad099c13 --- /dev/null +++ b/fme/ace/aggregator/inference/global_mean_drift.py @@ -0,0 +1,275 @@ +import dataclasses +import datetime +from collections.abc import Mapping +from typing import Any + +import numpy as np +import torch +import xarray as xr + +from fme.core.dataset.data_typing import VariableMetadata +from fme.core.device import get_device +from fme.core.distributed import Distributed +from fme.core.gridded_ops import GriddedOperations + +from .build_context import MetricBuildContext, MetricNotSupportedError, maybe_filter +from .data import InferenceBatchData, MetricBuildResult, SubAggregator +from .trend import SECONDS_PER_YEAR, _years_since_reference + +# Minimum rollout duration for drift, matching the ~2 year minimum used by the +# annual metric: a slow drift cannot be separated from interannual variability +# over a shorter span. The metric is not built for configured rollouts at or +# below this duration, and slopes are reported as NaN for samples whose +# recorded data covers less than it. +MIN_DRIFT_DURATION = datetime.timedelta(days=730) + +# Index of each per-sample sufficient statistic of the time regressor. +_N, _SUM_T, _SUM_TT, _T_MIN, _T_MAX = range(5) + + +class GlobalMeanDriftAggregator: + """Per-sample linear drift of the global-mean bias over a rollout. + + For each variable and each sample (initial condition or ensemble member), + the drift is the ordinary least-squares slope, against time in years, of + the area-weighted global mean of ``prediction - target``. As in + :class:`TrendEvaluatorAggregator`, it is computed from running sums of + ``t``, ``t**2``, ``y``, ``t*y`` and ``n`` accumulated over every + ``record_batch`` call, here kept separately for each sample, so memory is + ``O(n_samples * n_variables)`` regardless of the rollout length. Time is + measured from each sample's first recorded timestep. + + Logged scalars are the sample-mean slope per variable, its t-statistic + ``mean / (std / sqrt(n_samples))`` across samples (NaN with fewer than two + samples), and the maximum absolute t-statistic over variables. Slopes are + NaN for samples whose rollout covers less than ``MIN_DRIFT_DURATION``. + """ + + def __init__( + self, + gridded_operations: GriddedOperations, + timestep: datetime.timedelta, + variable_metadata: Mapping[str, VariableMetadata] | None = None, + ): + """ + Args: + gridded_operations: Computes the area-weighted global means. + timestep: Timestep of the data, used to compute the duration + covered by the recorded timesteps. + variable_metadata: Mapping of variable names to metadata used in + dataset attributes. + """ + self._ops = gridded_operations + self._timestep_years = timestep.total_seconds() / SECONDS_PER_YEAR + self._variable_metadata = variable_metadata or {} + # per-sample start time [years since reference], set on first record + self._t0: torch.Tensor | None = None + # [n_sample, 5] time-regressor statistics, indexed by _N, _SUM_T, ... + self._time_stats: torch.Tensor | None = None + # per-variable [n_sample] running sums of y and t*y + self._sum_y: dict[str, torch.Tensor] = {} + self._sum_ty: dict[str, torch.Tensor] = {} + + @torch.no_grad() + def record_batch(self, data: InferenceBatchData) -> None: + # Mirror TrendEvaluatorAggregator: when the initial condition is the + # first timestep of the first batch, skip it so it is not counted. + time_slice = slice(1, None) if data.i_time_start == 0 else slice(None) + device = get_device() + years = torch.tensor( + _years_since_reference(data.time)[:, time_slice], + dtype=torch.float64, + device=device, + ) + if years.shape[1] == 0: + return + if self._t0 is None: + self._t0 = years[:, 0].clone() + self._time_stats = torch.zeros( + years.shape[0], 5, dtype=torch.float64, device=device + ) + self._time_stats[:, _T_MIN] = torch.inf + self._time_stats[:, _T_MAX] = -torch.inf + assert self._time_stats is not None + t = years - self._t0[:, None] + stats = self._time_stats.clone() + stats[:, _N] += t.shape[1] + stats[:, _SUM_T] += t.sum(dim=1) + stats[:, _SUM_TT] += (t * t).sum(dim=1) + stats[:, _T_MIN] = torch.minimum(stats[:, _T_MIN], t.min(dim=1).values) + stats[:, _T_MAX] = torch.maximum(stats[:, _T_MAX], t.max(dim=1).values) + self._time_stats = stats + for name, prediction in data.prediction.items(): + bias = self._ops.area_weighted_mean_bias( + truth=data.target[name][:, time_slice], + predicted=prediction[:, time_slice], + name=name, + ).to(torch.float64) # [n_sample, n_time] + contrib_y = bias.sum(dim=1) + contrib_ty = (t * bias).sum(dim=1) + if name in self._sum_y: + self._sum_y[name] = self._sum_y[name] + contrib_y + self._sum_ty[name] = self._sum_ty[name] + contrib_ty + else: + self._sum_y[name] = contrib_y + self._sum_ty[name] = contrib_ty + + def _get_slopes(self) -> tuple[list[str], torch.Tensor] | None: + """Gather the per-sample statistics across data-parallel ranks and + finalize the slopes. + + Returns: + The sorted variable names and a ``[n_sample, n_variable]`` tensor + of slopes on the root rank, or ``None`` on other ranks. + """ + if self._time_stats is None: + raise ValueError("No data has been recorded yet.") + names = sorted(self._sum_y) + # [n_sample, 5 + 2 * n_variable], one collective for all statistics + local = torch.cat( + [ + self._time_stats, + torch.stack([self._sum_y[name] for name in names], dim=1), + torch.stack([self._sum_ty[name] for name in names], dim=1), + ], + dim=1, + ) + # Samples are split across data-parallel ranks; spatial co-ranks hold + # identical global means, so gather within the data-parallel group. + dist = Distributed.get_instance() + gathered = dist.data_parallel_gather(local) + if not dist.is_root() or gathered is None: + return None + stats = torch.cat(gathered, dim=0) + n_var = len(names) + n = stats[:, _N, None] + sum_t = stats[:, _SUM_T, None] + sum_tt = stats[:, _SUM_TT, None] + sum_y = stats[:, 5 : 5 + n_var] + sum_ty = stats[:, 5 + n_var :] + slopes = (n * sum_ty - sum_t * sum_y) / (n * sum_tt - sum_t * sum_t) + duration = stats[:, _T_MAX] - stats[:, _T_MIN] + self._timestep_years + min_duration = MIN_DRIFT_DURATION.total_seconds() / SECONDS_PER_YEAR + slopes[duration < min_duration] = torch.nan + return names, slopes + + def get_per_sample_slopes(self) -> dict[str, torch.Tensor]: + """Per-sample drift slopes of the global-mean bias. + + This is a collective operation and must be called on every rank. + + Returns: + Mapping from variable name to a ``[n_sample]`` float64 tensor of + slopes in units of the variable per year, with samples from all + data-parallel ranks. Empty on non-root ranks. + """ + result = self._get_slopes() + if result is None: + return {} + names, slopes = result + return {name: slopes[:, i] for i, name in enumerate(names)} + + @torch.no_grad() + def get_logs(self, label: str) -> dict[str, Any]: + per_sample = self.get_per_sample_slopes() + if len(per_sample) == 0: + return {} # only the root process returns logs + if len(label) > 0: + label = label + "/" + logs: dict[str, float] = {} + abs_t_values = [] + for name, slopes in per_sample.items(): + slopes_np = slopes.cpu().numpy() + logs[f"{label}slope/{name}"] = float(np.mean(slopes_np)) + t = _t_statistic(slopes_np) + logs[f"{label}t/{name}"] = t + if not np.isnan(t): + abs_t_values.append(abs(t)) + logs[f"{label}max_abs_t"] = max(abs_t_values) if abs_t_values else np.nan + return logs + + def _units(self, name: str) -> str: + if name in self._variable_metadata: + units = self._variable_metadata[name].display_units("unknown units") + else: + units = "unknown units" + return f"{units} / year" + + def get_dataset(self) -> xr.Dataset: + result = self._get_slopes() + if result is None: + return xr.Dataset() + names, slopes = result + slope = xr.DataArray( + slopes.cpu().numpy(), + dims=["sample", "variable"], + coords={"variable": names}, + attrs={ + "long_name": ( + "least-squares slope vs. time of the global-mean bias " + "(prediction - target)" + ), + }, + ) + units = xr.DataArray( + [self._units(name) for name in names], + dims=["variable"], + coords={"variable": names}, + ) + return xr.Dataset({"slope": slope, "slope_units": units}) + + +def _t_statistic(slopes: np.ndarray) -> float: + """One-sample t-statistic of the slopes against zero drift, using the + spread across samples; NaN for fewer than two samples. + """ + n_sample = slopes.shape[0] + if n_sample < 2: + return np.nan + std = np.std(slopes, ddof=1) + with np.errstate(divide="ignore", invalid="ignore"): + return float(np.mean(slopes) / (std / np.sqrt(n_sample))) + + +@dataclasses.dataclass +class GlobalMeanDriftConfig: + """Configuration for the per-sample global-mean drift inference metric. + + Logs the sample-mean least-squares slope (per year) of the global-mean + bias of each variable, its t-statistic across samples, and the maximum + absolute t-statistic over variables. Per-sample slopes are saved in the + diagnostics dataset. Like the annual metric, it is skipped for rollouts + of ~2 years or less, and slopes are NaN for samples whose recorded data + covers less than ~2 years. + + Attributes: + variables: Variables to compute drift for. If ``None``, all available + variables are included. + name: Name used to label the metric's logs and diagnostics. + enabled: Whether the metric is computed. + strict: If True, raise rather than skip when the metric is not + supported for the current configuration. + """ + + variables: list[str] | None = None + name: str = "global_mean_drift" + enabled: bool = True + strict: bool = False + + def get_name(self) -> str: + return self.name + + def build(self, ctx: MetricBuildContext) -> MetricBuildResult: + total_duration = ctx.n_timesteps * ctx.timestep + if total_duration <= MIN_DRIFT_DURATION: + raise MetricNotSupportedError( + "global_mean_drift metric requires > ~2 years of data, " + f"got {total_duration.days} days" + ) + agg = GlobalMeanDriftAggregator( + ctx.ops, + timestep=ctx.timestep, + variable_metadata=ctx.variable_metadata, + ) + filtered: SubAggregator = maybe_filter(agg, self.variables) + return MetricBuildResult(aggregator=filtered, per_sample_slopes=agg) diff --git a/fme/ace/aggregator/inference/main.py b/fme/ace/aggregator/inference/main.py index e998263968..e9bc6d4b17 100644 --- a/fme/ace/aggregator/inference/main.py +++ b/fme/ace/aggregator/inference/main.py @@ -29,10 +29,17 @@ from ..one_step.reduced import StepMeanMetricConfig from .annual import AnnualMetricConfig, GlobalMeanAnnualAggregator from .build_context import MetricBuildContext, MetricNotSupportedError -from .data import InferenceBatchData, MetricBuildResult, SubAggregator, TimeSeriesLogs +from .data import ( + InferenceBatchData, + MetricBuildResult, + PerSampleSlopes, + SubAggregator, + TimeSeriesLogs, +) from .enso import RegionalIndexAggregator from .enso.dynamic_index import EnsoIndexMetricConfig from .enso.enso_coefficient import EnsoCoefficientMetricConfig +from .global_mean_drift import GlobalMeanDriftConfig from .histogram import HistogramMetricConfig from .ipo.ipo_index import MIN_YEARS_FOR_FILTERED_TPI, IpoIndexMetricConfig from .near_zero_fraction import NearZeroFractionMetricConfig @@ -68,6 +75,7 @@ | EnsembleMetricConfig | IpoIndexMetricConfig | TrendMetricConfig + | GlobalMeanDriftConfig | NearZeroFractionMetricConfig ) @@ -139,6 +147,7 @@ def build_inference_evaluator_aggregator( aggregators: dict[str, SubAggregator] = {} time_series_aggregators: dict[str, TimeSeriesLogs] = {} ensemble_aggregators: dict[str, SelectStepEnsembleAggregator] = {} + global_mean_drift: PerSampleSlopes | None = None for metric in metrics: name = metric.get_name() @@ -158,6 +167,10 @@ def build_inference_evaluator_aggregator( time_series_aggregators[name] = result.time_series if result.ensemble is not None: ensemble_aggregators[name] = result.ensemble + if result.per_sample_slopes is not None: + if global_mean_drift is not None: + raise ValueError("Only one global-mean drift metric is supported.") + global_mean_drift = result.per_sample_slopes step_diagnostics_aggregator = ( step_diagnostics or StepDiagnosticsMetricConfig() @@ -180,6 +193,7 @@ def build_inference_evaluator_aggregator( n_ensemble_per_ic=n_ensemble_per_ic, ensemble_aggregators=ensemble_aggregators, step_diagnostics_aggregator=step_diagnostics_aggregator, + global_mean_drift=global_mean_drift, ) @@ -219,6 +233,9 @@ class InferenceEvaluatorAggregatorConfig: ipo_index: Interdecadal Pacific Oscillation index metrics. trend: Per-grid-cell linear trend (slope vs. time) map metrics. Disabled by default. + global_mean_drift: Per-sample linear drift (slope vs. time) of the + global-mean bias, with its t-statistic across samples. Skipped + for rollouts of ~2 years or less. near_zero_fraction: Area-weighted fraction of cells at or below a small non-negative ``eps`` per variable, reported for prediction and its difference from the target. Optionally (``include_maps``) also @@ -278,6 +295,9 @@ class InferenceEvaluatorAggregatorConfig: default_factory=IpoIndexMetricConfig ) trend: TrendMetricConfig = dataclasses.field(default_factory=TrendMetricConfig) + global_mean_drift: GlobalMeanDriftConfig = dataclasses.field( + default_factory=GlobalMeanDriftConfig + ) near_zero_fraction: NearZeroFractionMetricConfig = dataclasses.field( default_factory=NearZeroFractionMetricConfig ) @@ -325,6 +345,7 @@ def _get_metrics(self) -> list[MetricConfig]: self.enso_coefficient, self.ipo_index, self.trend, + self.global_mean_drift, self.near_zero_fraction, ] return [m for m in all_metrics if m.enabled] @@ -545,11 +566,13 @@ def __init__( n_ensemble_per_ic: int = 1, ensemble_aggregators: dict[str, SelectStepEnsembleAggregator] | None = None, step_diagnostics_aggregator: StepDiagnosticsAggregator | None = None, + global_mean_drift: PerSampleSlopes | None = None, ): if save_diagnostics and output_dir is None: raise ValueError("Output directory must be set to save diagnostics") self._aggregators = aggregators self._time_series_aggregators = time_series_aggregators + self._global_mean_drift = global_mean_drift self.n_ensemble_per_ic = n_ensemble_per_ic self._ensemble_aggregators = ensemble_aggregators or {} summary_aggregators: dict[str, SubAggregator | SelectStepEnsembleAggregator] = { @@ -678,6 +701,22 @@ def get_summary(self) -> InferenceSummary: def get_summary_logs(self) -> InferenceLog: return self.get_summary().logs + @torch.no_grad() + def get_global_mean_drift_slopes(self) -> dict[str, torch.Tensor] | None: + """Per-sample drift slopes of the global-mean bias. + + This is a collective operation and must be called on every rank. + + Returns: + ``None`` if the global-mean drift metric is not configured. + Otherwise a mapping from variable name to a ``[n_sample]`` tensor + of slopes (variable units per year) over all samples, empty on + non-root ranks. + """ + if self._global_mean_drift is None: + return None + return self._global_mean_drift.get_per_sample_slopes() + @torch.no_grad() def _get_logs(self): """Returns logs as can be reported to WandB.""" diff --git a/fme/ace/aggregator/inference/test_global_mean_drift.py b/fme/ace/aggregator/inference/test_global_mean_drift.py new file mode 100644 index 0000000000..df3821ca54 --- /dev/null +++ b/fme/ace/aggregator/inference/test_global_mean_drift.py @@ -0,0 +1,257 @@ +import datetime +from collections.abc import Sequence + +import cftime +import numpy as np +import pytest +import torch +import xarray as xr + +from fme.ace.data_loading.batch_data import BatchData, PairedData +from fme.core.coordinates import LatLonCoordinates +from fme.core.dataset_info import DatasetInfo +from fme.core.device import get_device +from fme.core.gridded_ops import LatLonOperations +from fme.core.testing import mock_distributed + +from .data import InferenceBatchData +from .global_mean_drift import GlobalMeanDriftAggregator, GlobalMeanDriftConfig +from .main import ( + InferenceEvaluatorAggregator, + InferenceEvaluatorAggregatorConfig, + MetricConfig, + build_inference_evaluator_aggregator, +) +from .trend import SECONDS_PER_YEAR + +N_LAT, N_LON = 3, 4 +TIMESTEP = datetime.timedelta(days=1) + + +def _daily_time(n_sample: int, n_time: int, start_day: int = 0) -> xr.DataArray: + base = cftime.DatetimeNoLeap(2000, 1, 1) + times = [base + datetime.timedelta(days=start_day + i) for i in range(n_time)] + return xr.DataArray([times for _ in range(n_sample)], dims=["sample", "time"]) + + +def _drift_fields( + drift_per_year: Sequence[float], + time: xr.DataArray, + names: Sequence[str] = ("a",), +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Target is random noise; prediction is target plus ``drift * t`` for each + sample, with ``t`` in years since 2000-01-01.""" + n_sample, n_time = time.shape + base = cftime.DatetimeNoLeap(2000, 1, 1) + days = np.array([[(t - base).days for t in row] for row in time.values]) + t = torch.tensor(days * 86400.0 / SECONDS_PER_YEAR, dtype=torch.float32) + drift = torch.tensor(drift_per_year, dtype=torch.float32)[:, None] + offset = (drift * t)[:, :, None, None].to(get_device()) + target = { + name: torch.randn(n_sample, n_time, N_LAT, N_LON, device=get_device()) + for name in names + } + gen = {name: target[name] + offset for name in names} + return target, gen + + +def _batch( + drift_per_year: Sequence[float], + time: xr.DataArray, + i_time_start: int = 1, + names: Sequence[str] = ("a",), +) -> InferenceBatchData: + target, gen = _drift_fields(drift_per_year, time, names) + return InferenceBatchData( + prediction=gen, target=target, time=time, i_time_start=i_time_start + ) + + +def _aggregator() -> GlobalMeanDriftAggregator: + ops = LatLonOperations(torch.ones(N_LAT, N_LON).to(get_device())) + return GlobalMeanDriftAggregator(ops, timestep=TIMESTEP) + + +def _record_drift( + drift_per_year: Sequence[float], n_days: int = 3 * 365, n_batches: int = 3 +) -> GlobalMeanDriftAggregator: + agg = _aggregator() + n_sample = len(drift_per_year) + per_batch = n_days // n_batches + for i in range(n_batches): + time = _daily_time(n_sample, per_batch, start_day=i * per_batch) + agg.record_batch(_batch(drift_per_year, time, i_time_start=1 + i * per_batch)) + return agg + + +def test_recovers_per_sample_slope(): + torch.manual_seed(0) + drift = [0.5, -1.0, 2.0] + agg = _record_drift(drift) + slopes = agg.get_per_sample_slopes()["a"] + assert slopes.shape == (3,) + np.testing.assert_allclose(slopes.cpu().numpy(), drift, atol=1e-4) + logs = agg.get_logs(label="drift") + assert logs["drift/slope/a"] == pytest.approx(np.mean(drift), abs=1e-4) + + +def test_streaming_matches_single_batch(): + torch.manual_seed(0) + time = _daily_time(2, 3 * 365) + batch = _batch([1.0, 2.0], time) + single = _aggregator() + single.record_batch(batch) + streamed = _aggregator() + for start, stop in [(0, 400), (400, 3 * 365)]: + streamed.record_batch( + InferenceBatchData( + prediction={k: v[:, start:stop] for k, v in batch.prediction.items()}, + target={k: v[:, start:stop] for k, v in batch.target.items()}, + time=time.isel(time=slice(start, stop)), + i_time_start=1 + start, + ) + ) + xr.testing.assert_allclose(single.get_dataset(), streamed.get_dataset()) + + +def test_zero_drift_gives_zero_slope(): + torch.manual_seed(0) + agg = _record_drift([0.0, 0.0, 0.0, 0.0]) + slopes = agg.get_per_sample_slopes()["a"].cpu().numpy() + np.testing.assert_allclose(slopes, 0.0, atol=1e-5) + + +def test_t_statistic_large_for_consistent_drift(): + torch.manual_seed(0) + drift = [1.0, 1.05, 0.95, 1.02, 0.98, 1.01, 0.99, 1.03] + logs = _record_drift(drift).get_logs(label="") + assert logs["t/a"] > 50.0 + assert logs["max_abs_t"] == pytest.approx(logs["t/a"]) + + +def test_t_statistic_small_for_random_sign_drift(): + torch.manual_seed(0) + drift = [1.0, -1.2, 0.8, -0.9, 1.1, -0.7, 0.9, -1.0] + logs = _record_drift(drift).get_logs(label="") + assert abs(logs["t/a"]) < 1.0 + + +def test_max_abs_t_over_variables(): + torch.manual_seed(0) + agg = _aggregator() + time = _daily_time(4, 3 * 365) + target_a, gen_a = _drift_fields([1.0, 1.1, 0.9, 1.0], time, names=("a",)) + target_b, gen_b = _drift_fields([-1.0, -1.01, -0.99, -1.0], time, names=("b",)) + agg.record_batch( + InferenceBatchData( + prediction={**gen_a, **gen_b}, + target={**target_a, **target_b}, + time=time, + i_time_start=1, + ) + ) + logs = agg.get_logs(label="") + assert logs["t/b"] < 0 + assert logs["max_abs_t"] == pytest.approx(max(abs(logs["t/a"]), abs(logs["t/b"]))) + + +def test_short_rollout_gives_nan(): + torch.manual_seed(0) + agg = _record_drift([1.0, 2.0], n_days=365, n_batches=1) + assert torch.isnan(agg.get_per_sample_slopes()["a"]).all() + logs = agg.get_logs(label="") + assert np.isnan(logs["slope/a"]) + assert np.isnan(logs["t/a"]) + assert np.isnan(logs["max_abs_t"]) + + +def test_single_sample_gives_nan_t_statistic(): + torch.manual_seed(0) + logs = _record_drift([1.5]).get_logs(label="") + assert logs["slope/a"] == pytest.approx(1.5, abs=1e-4) + assert np.isnan(logs["t/a"]) + assert np.isnan(logs["max_abs_t"]) + + +def test_dataset_has_sample_variable_slopes(): + torch.manual_seed(0) + agg = _aggregator() + time = _daily_time(3, 3 * 365) + agg.record_batch(_batch([1.0, 2.0, 3.0], time, names=("b", "a"))) + ds = agg.get_dataset() + assert ds["slope"].dims == ("sample", "variable") + assert list(ds["variable"].values) == ["a", "b"] + np.testing.assert_allclose( + ds["slope"].sel(variable="a").values, [1.0, 2.0, 3.0], atol=1e-4 + ) + + +def test_samples_gathered_across_data_parallel_ranks(): + torch.manual_seed(0) + with mock_distributed(world_size=2): + agg = _record_drift([1.0, 2.0]) + slopes = agg.get_per_sample_slopes()["a"].cpu().numpy() + np.testing.assert_allclose(slopes, [1.0, 2.0, 1.0, 2.0], atol=1e-4) + + +def test_enabled_in_default_metrics(): + names = [m.get_name() for m in InferenceEvaluatorAggregatorConfig()._get_metrics()] + assert "global_mean_drift" in names + + +def _build_top_level( + metrics: list[MetricConfig], n_forward_steps: int = 3 * 365 +) -> InferenceEvaluatorAggregator: + n_sample = 4 + return build_inference_evaluator_aggregator( + metrics=metrics, + dataset_info=DatasetInfo( + horizontal_coordinates=LatLonCoordinates( + lat=torch.linspace(-60, 60, N_LAT), lon=torch.arange(N_LON) * 90.0 + ), + timestep=TIMESTEP, + ), + n_ic_steps=1, + n_forward_steps=n_forward_steps, + initial_time=xr.DataArray(np.zeros(n_sample), dims=["sample"]), + normalize=lambda x, apply_mean=True: dict(x), + save_diagnostics=False, + raise_on_unsupported=False, + ) + + +def test_top_level_aggregator_logs_and_slopes(): + torch.manual_seed(0) + drift = [1.0, 1.1, 0.9, 1.0] + agg = _build_top_level([GlobalMeanDriftConfig()]) + time = _daily_time(len(drift), 3 * 365, start_day=1) + target, gen = _drift_fields(drift, time, names=("a", "b")) + agg.record_batch( + PairedData.from_batch_data( + prediction=BatchData(data=gen, time=time), + reference=BatchData(data=target, time=time), + ) + ) + logs = agg.get_summary().logs + for key in [ + "global_mean_drift/slope/a", + "global_mean_drift/slope/b", + "global_mean_drift/t/a", + "global_mean_drift/t/b", + "global_mean_drift/max_abs_t", + ]: + assert key in logs + assert logs["global_mean_drift/slope/a"] == pytest.approx(1.0, abs=1e-3) + slopes = agg.get_global_mean_drift_slopes() + assert slopes is not None + np.testing.assert_allclose(slopes["a"].cpu().numpy(), drift, atol=1e-3) + + +def test_top_level_slopes_none_when_not_configured(): + agg = _build_top_level([]) + assert agg.get_global_mean_drift_slopes() is None + + +def test_top_level_skips_metric_for_short_rollout(): + agg = _build_top_level([GlobalMeanDriftConfig()], n_forward_steps=365) + assert agg.get_global_mean_drift_slopes() is None diff --git a/fme/core/testing/distributed.py b/fme/core/testing/distributed.py index 867e556fb8..98fa19a7a2 100644 --- a/fme/core/testing/distributed.py +++ b/fme/core/testing/distributed.py @@ -37,6 +37,9 @@ def is_distributed(self) -> bool: def gather(self, tensor: torch.Tensor) -> list[torch.Tensor]: return [tensor for i in range(self.world_size)] + def data_parallel_gather(self, tensor: torch.Tensor) -> list[torch.Tensor]: + return self.gather(tensor) + def gather_irregular(self, tensor: torch.Tensor) -> list[torch.Tensor] | None: """ Note this uses the actual implementation but mocks the underlying From ac7aa3bc3e9f283be85582f6bf422200aab8f6f0 Mon Sep 17 00:00:00 2001 From: Alexey Yermakov Date: Mon, 2 Nov 2026 16:18:01 -0800 Subject: [PATCH 2/4] Simplify global-mean drift aggregator Track the recorded step count as a scalar and check the minimum duration as n_time * timestep, replacing the packed per-sample [n, sum_t, sum_tt, t_min, t_max] tensor and its index constants. Trim docstrings. --- .../aggregator/inference/global_mean_drift.py | 209 ++++++------------ fme/ace/aggregator/inference/main.py | 9 +- 2 files changed, 75 insertions(+), 143 deletions(-) diff --git a/fme/ace/aggregator/inference/global_mean_drift.py b/fme/ace/aggregator/inference/global_mean_drift.py index adad099c13..05bcef24c4 100644 --- a/fme/ace/aggregator/inference/global_mean_drift.py +++ b/fme/ace/aggregator/inference/global_mean_drift.py @@ -1,7 +1,6 @@ import dataclasses import datetime from collections.abc import Mapping -from typing import Any import numpy as np import torch @@ -13,36 +12,21 @@ from fme.core.gridded_ops import GriddedOperations from .build_context import MetricBuildContext, MetricNotSupportedError, maybe_filter -from .data import InferenceBatchData, MetricBuildResult, SubAggregator -from .trend import SECONDS_PER_YEAR, _years_since_reference +from .data import InferenceBatchData, MetricBuildResult +from .trend import _years_since_reference -# Minimum rollout duration for drift, matching the ~2 year minimum used by the -# annual metric: a slow drift cannot be separated from interannual variability -# over a shorter span. The metric is not built for configured rollouts at or -# below this duration, and slopes are reported as NaN for samples whose -# recorded data covers less than it. +# Same ~2 year minimum as the annual metric: shorter spans cannot separate a +# slow drift from interannual variability. MIN_DRIFT_DURATION = datetime.timedelta(days=730) -# Index of each per-sample sufficient statistic of the time regressor. -_N, _SUM_T, _SUM_TT, _T_MIN, _T_MAX = range(5) - class GlobalMeanDriftAggregator: - """Per-sample linear drift of the global-mean bias over a rollout. - - For each variable and each sample (initial condition or ensemble member), - the drift is the ordinary least-squares slope, against time in years, of - the area-weighted global mean of ``prediction - target``. As in - :class:`TrendEvaluatorAggregator`, it is computed from running sums of - ``t``, ``t**2``, ``y``, ``t*y`` and ``n`` accumulated over every - ``record_batch`` call, here kept separately for each sample, so memory is - ``O(n_samples * n_variables)`` regardless of the rollout length. Time is - measured from each sample's first recorded timestep. + """Per-sample least-squares slope vs. time (per year) of the area-weighted + global-mean bias (prediction - target) of each variable. - Logged scalars are the sample-mean slope per variable, its t-statistic - ``mean / (std / sqrt(n_samples))`` across samples (NaN with fewer than two - samples), and the maximum absolute t-statistic over variables. Slopes are - NaN for samples whose rollout covers less than ``MIN_DRIFT_DURATION``. + Slopes are computed from running sums, so memory does not grow with the + rollout length. They are NaN if the recorded data covers less than + ``MIN_DRIFT_DURATION``. """ def __init__( @@ -51,117 +35,84 @@ def __init__( timestep: datetime.timedelta, variable_metadata: Mapping[str, VariableMetadata] | None = None, ): - """ - Args: - gridded_operations: Computes the area-weighted global means. - timestep: Timestep of the data, used to compute the duration - covered by the recorded timesteps. - variable_metadata: Mapping of variable names to metadata used in - dataset attributes. - """ self._ops = gridded_operations - self._timestep_years = timestep.total_seconds() / SECONDS_PER_YEAR + self._timestep = timestep self._variable_metadata = variable_metadata or {} - # per-sample start time [years since reference], set on first record + self._n_time = 0 + # Per-sample time of the first recorded step, so the regressor stays + # small and the float64 sums well conditioned. self._t0: torch.Tensor | None = None - # [n_sample, 5] time-regressor statistics, indexed by _N, _SUM_T, ... - self._time_stats: torch.Tensor | None = None - # per-variable [n_sample] running sums of y and t*y + zero = torch.zeros((), dtype=torch.float64, device=get_device()) + # [n_sample] running sums; y and t*y are per variable. + self._sum_t = zero + self._sum_tt = zero self._sum_y: dict[str, torch.Tensor] = {} self._sum_ty: dict[str, torch.Tensor] = {} @torch.no_grad() def record_batch(self, data: InferenceBatchData) -> None: - # Mirror TrendEvaluatorAggregator: when the initial condition is the - # first timestep of the first batch, skip it so it is not counted. + # As in TrendEvaluatorAggregator, skip the initial condition. time_slice = slice(1, None) if data.i_time_start == 0 else slice(None) - device = get_device() years = torch.tensor( _years_since_reference(data.time)[:, time_slice], dtype=torch.float64, - device=device, + device=get_device(), ) if years.shape[1] == 0: return if self._t0 is None: self._t0 = years[:, 0].clone() - self._time_stats = torch.zeros( - years.shape[0], 5, dtype=torch.float64, device=device - ) - self._time_stats[:, _T_MIN] = torch.inf - self._time_stats[:, _T_MAX] = -torch.inf - assert self._time_stats is not None t = years - self._t0[:, None] - stats = self._time_stats.clone() - stats[:, _N] += t.shape[1] - stats[:, _SUM_T] += t.sum(dim=1) - stats[:, _SUM_TT] += (t * t).sum(dim=1) - stats[:, _T_MIN] = torch.minimum(stats[:, _T_MIN], t.min(dim=1).values) - stats[:, _T_MAX] = torch.maximum(stats[:, _T_MAX], t.max(dim=1).values) - self._time_stats = stats + self._n_time += t.shape[1] + self._sum_t = self._sum_t + t.sum(dim=1) + self._sum_tt = self._sum_tt + (t * t).sum(dim=1) for name, prediction in data.prediction.items(): bias = self._ops.area_weighted_mean_bias( truth=data.target[name][:, time_slice], predicted=prediction[:, time_slice], name=name, ).to(torch.float64) # [n_sample, n_time] - contrib_y = bias.sum(dim=1) - contrib_ty = (t * bias).sum(dim=1) - if name in self._sum_y: - self._sum_y[name] = self._sum_y[name] + contrib_y - self._sum_ty[name] = self._sum_ty[name] + contrib_ty - else: - self._sum_y[name] = contrib_y - self._sum_ty[name] = contrib_ty + self._sum_y[name] = self._sum_y.get(name, 0.0) + bias.sum(dim=1) + self._sum_ty[name] = self._sum_ty.get(name, 0.0) + (t * bias).sum(dim=1) def _get_slopes(self) -> tuple[list[str], torch.Tensor] | None: - """Gather the per-sample statistics across data-parallel ranks and - finalize the slopes. - - Returns: - The sorted variable names and a ``[n_sample, n_variable]`` tensor - of slopes on the root rank, or ``None`` on other ranks. + """Sorted variable names and ``[n_sample, n_variable]`` slopes over all + data-parallel ranks on the root rank, ``None`` on other ranks. """ - if self._time_stats is None: + if self._n_time == 0: raise ValueError("No data has been recorded yet.") names = sorted(self._sum_y) - # [n_sample, 5 + 2 * n_variable], one collective for all statistics - local = torch.cat( + local = torch.stack( [ - self._time_stats, - torch.stack([self._sum_y[name] for name in names], dim=1), - torch.stack([self._sum_ty[name] for name in names], dim=1), + self._sum_t, + self._sum_tt, + *(self._sum_y[name] for name in names), + *(self._sum_ty[name] for name in names), ], dim=1, ) # Samples are split across data-parallel ranks; spatial co-ranks hold - # identical global means, so gather within the data-parallel group. + # identical global means. dist = Distributed.get_instance() gathered = dist.data_parallel_gather(local) if not dist.is_root() or gathered is None: return None - stats = torch.cat(gathered, dim=0) n_var = len(names) - n = stats[:, _N, None] - sum_t = stats[:, _SUM_T, None] - sum_tt = stats[:, _SUM_TT, None] - sum_y = stats[:, 5 : 5 + n_var] - sum_ty = stats[:, 5 + n_var :] + sum_t, sum_tt, sum_y, sum_ty = torch.cat(gathered).split( + [1, 1, n_var, n_var], dim=1 + ) + n = self._n_time slopes = (n * sum_ty - sum_t * sum_y) / (n * sum_tt - sum_t * sum_t) - duration = stats[:, _T_MAX] - stats[:, _T_MIN] + self._timestep_years - min_duration = MIN_DRIFT_DURATION.total_seconds() / SECONDS_PER_YEAR - slopes[duration < min_duration] = torch.nan + if n * self._timestep < MIN_DRIFT_DURATION: + slopes = torch.full_like(slopes, torch.nan) return names, slopes def get_per_sample_slopes(self) -> dict[str, torch.Tensor]: - """Per-sample drift slopes of the global-mean bias. - - This is a collective operation and must be called on every rank. + """Collective; must be called on every rank. Returns: - Mapping from variable name to a ``[n_sample]`` float64 tensor of - slopes in units of the variable per year, with samples from all - data-parallel ranks. Empty on non-root ranks. + Variable name to ``[n_sample]`` slopes (variable units per year) + over all data-parallel ranks. Empty on non-root ranks. """ result = self._get_slopes() if result is None: @@ -170,77 +121,64 @@ def get_per_sample_slopes(self) -> dict[str, torch.Tensor]: return {name: slopes[:, i] for i, name in enumerate(names)} @torch.no_grad() - def get_logs(self, label: str) -> dict[str, Any]: + def get_logs(self, label: str) -> dict[str, float]: per_sample = self.get_per_sample_slopes() if len(per_sample) == 0: - return {} # only the root process returns logs - if len(label) > 0: - label = label + "/" + return {} + prefix = f"{label}/" if label else "" logs: dict[str, float] = {} - abs_t_values = [] + abs_t = [] for name, slopes in per_sample.items(): slopes_np = slopes.cpu().numpy() - logs[f"{label}slope/{name}"] = float(np.mean(slopes_np)) + logs[f"{prefix}slope/{name}"] = float(np.mean(slopes_np)) t = _t_statistic(slopes_np) - logs[f"{label}t/{name}"] = t + logs[f"{prefix}t/{name}"] = t if not np.isnan(t): - abs_t_values.append(abs(t)) - logs[f"{label}max_abs_t"] = max(abs_t_values) if abs_t_values else np.nan + abs_t.append(abs(t)) + logs[f"{prefix}max_abs_t"] = max(abs_t) if abs_t else np.nan return logs - def _units(self, name: str) -> str: - if name in self._variable_metadata: - units = self._variable_metadata[name].display_units("unknown units") - else: - units = "unknown units" - return f"{units} / year" - def get_dataset(self) -> xr.Dataset: result = self._get_slopes() if result is None: return xr.Dataset() names, slopes = result - slope = xr.DataArray( - slopes.cpu().numpy(), - dims=["sample", "variable"], - coords={"variable": names}, - attrs={ - "long_name": ( - "least-squares slope vs. time of the global-mean bias " - "(prediction - target)" + units = [ + self._variable_metadata[name].display_units("unknown units") + if name in self._variable_metadata + else "unknown units" + for name in names + ] + return xr.Dataset( + { + "slope": ( + ["sample", "variable"], + slopes.cpu().numpy(), + {"long_name": "slope vs. time of global-mean bias"}, ), + "slope_units": (["variable"], [f"{u} / year" for u in units]), }, - ) - units = xr.DataArray( - [self._units(name) for name in names], - dims=["variable"], coords={"variable": names}, ) - return xr.Dataset({"slope": slope, "slope_units": units}) def _t_statistic(slopes: np.ndarray) -> float: - """One-sample t-statistic of the slopes against zero drift, using the - spread across samples; NaN for fewer than two samples. - """ + """t-statistic of the mean slope against zero; NaN for < 2 samples.""" n_sample = slopes.shape[0] if n_sample < 2: return np.nan - std = np.std(slopes, ddof=1) with np.errstate(divide="ignore", invalid="ignore"): - return float(np.mean(slopes) / (std / np.sqrt(n_sample))) + return float(np.mean(slopes) / (np.std(slopes, ddof=1) / np.sqrt(n_sample))) @dataclasses.dataclass class GlobalMeanDriftConfig: - """Configuration for the per-sample global-mean drift inference metric. + """Per-sample drift of the global-mean bias. - Logs the sample-mean least-squares slope (per year) of the global-mean - bias of each variable, its t-statistic across samples, and the maximum + Logs, per variable, the sample-mean slope (units per year) of the + global-mean bias and its t-statistic across samples, plus the maximum absolute t-statistic over variables. Per-sample slopes are saved in the - diagnostics dataset. Like the annual metric, it is skipped for rollouts - of ~2 years or less, and slopes are NaN for samples whose recorded data - covers less than ~2 years. + diagnostics dataset. Skipped for rollouts of ~2 years or less. Attributes: variables: Variables to compute drift for. If ``None``, all available @@ -267,9 +205,8 @@ def build(self, ctx: MetricBuildContext) -> MetricBuildResult: f"got {total_duration.days} days" ) agg = GlobalMeanDriftAggregator( - ctx.ops, - timestep=ctx.timestep, - variable_metadata=ctx.variable_metadata, + ctx.ops, timestep=ctx.timestep, variable_metadata=ctx.variable_metadata + ) + return MetricBuildResult( + aggregator=maybe_filter(agg, self.variables), per_sample_slopes=agg ) - filtered: SubAggregator = maybe_filter(agg, self.variables) - return MetricBuildResult(aggregator=filtered, per_sample_slopes=agg) diff --git a/fme/ace/aggregator/inference/main.py b/fme/ace/aggregator/inference/main.py index e9bc6d4b17..47f95d4f0e 100644 --- a/fme/ace/aggregator/inference/main.py +++ b/fme/ace/aggregator/inference/main.py @@ -705,13 +705,8 @@ def get_summary_logs(self) -> InferenceLog: def get_global_mean_drift_slopes(self) -> dict[str, torch.Tensor] | None: """Per-sample drift slopes of the global-mean bias. - This is a collective operation and must be called on every rank. - - Returns: - ``None`` if the global-mean drift metric is not configured. - Otherwise a mapping from variable name to a ``[n_sample]`` tensor - of slopes (variable units per year) over all samples, empty on - non-root ranks. + Collective; must be called on every rank. ``None`` if the metric was + not built, otherwise as ``GlobalMeanDriftAggregator.get_per_sample_slopes``. """ if self._global_mean_drift is None: return None From 66a9165c3f59ece4baedcc74f3713858be0e4a9e Mon Sep 17 00:00:00 2001 From: Alexey Yermakov Date: Mon, 2 Nov 2026 16:18:05 -0800 Subject: [PATCH 3/4] Trim redundant global-mean drift tests Drop the zero-drift test (covered by per-sample slope recovery) and merge the two accessor-returns-None tests into one parametrized test. --- .../inference/test_global_mean_drift.py | 28 ++++++++----------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/fme/ace/aggregator/inference/test_global_mean_drift.py b/fme/ace/aggregator/inference/test_global_mean_drift.py index df3821ca54..e9807c951f 100644 --- a/fme/ace/aggregator/inference/test_global_mean_drift.py +++ b/fme/ace/aggregator/inference/test_global_mean_drift.py @@ -39,8 +39,7 @@ def _drift_fields( time: xr.DataArray, names: Sequence[str] = ("a",), ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - """Target is random noise; prediction is target plus ``drift * t`` for each - sample, with ``t`` in years since 2000-01-01.""" + """Random target; prediction is target + drift * (years since 2000).""" n_sample, n_time = time.shape base = cftime.DatetimeNoLeap(2000, 1, 1) days = np.array([[(t - base).days for t in row] for row in time.values]) @@ -114,13 +113,6 @@ def test_streaming_matches_single_batch(): xr.testing.assert_allclose(single.get_dataset(), streamed.get_dataset()) -def test_zero_drift_gives_zero_slope(): - torch.manual_seed(0) - agg = _record_drift([0.0, 0.0, 0.0, 0.0]) - slopes = agg.get_per_sample_slopes()["a"].cpu().numpy() - np.testing.assert_allclose(slopes, 0.0, atol=1e-5) - - def test_t_statistic_large_for_consistent_drift(): torch.manual_seed(0) drift = [1.0, 1.05, 0.95, 1.02, 0.98, 1.01, 0.99, 1.03] @@ -247,11 +239,15 @@ def test_top_level_aggregator_logs_and_slopes(): np.testing.assert_allclose(slopes["a"].cpu().numpy(), drift, atol=1e-3) -def test_top_level_slopes_none_when_not_configured(): - agg = _build_top_level([]) - assert agg.get_global_mean_drift_slopes() is None - - -def test_top_level_skips_metric_for_short_rollout(): - agg = _build_top_level([GlobalMeanDriftConfig()], n_forward_steps=365) +@pytest.mark.parametrize( + "metrics, n_forward_steps", + [ + pytest.param([], 3 * 365, id="not_configured"), + pytest.param([GlobalMeanDriftConfig()], 365, id="short_rollout"), + ], +) +def test_top_level_slopes_none_when_metric_not_built( + metrics: list[MetricConfig], n_forward_steps: int +): + agg = _build_top_level(metrics, n_forward_steps=n_forward_steps) assert agg.get_global_mean_drift_slopes() is None From 08221a14da4a292230b28ad0b1589754b3fc2b59 Mon Sep 17 00:00:00 2001 From: Alexey Yermakov Date: Mon, 2 Nov 2026 16:45:57 -0800 Subject: [PATCH 4/4] Document GlobalMeanDriftConfig and regenerate default aggregator config --- docs/default-aggregator-config.yaml | 5 +++++ docs/evaluator_config.rst | 3 +++ 2 files changed, 8 insertions(+) diff --git a/docs/default-aggregator-config.yaml b/docs/default-aggregator-config.yaml index c695c85c80..7d433ac4b6 100644 --- a/docs/default-aggregator-config.yaml +++ b/docs/default-aggregator-config.yaml @@ -105,6 +105,11 @@ aggregator: name: trend enabled: false strict: false + global_mean_drift: + variables: null + name: global_mean_drift + enabled: true + strict: false near_zero_fraction: variables: [] eps: 0.0 diff --git a/docs/evaluator_config.rst b/docs/evaluator_config.rst index caaffcbbc6..2460cd0b6e 100644 --- a/docs/evaluator_config.rst +++ b/docs/evaluator_config.rst @@ -154,6 +154,9 @@ in the default aggregator configuration above. .. autoclass:: fme.ace.TrendMetricConfig :noindex: +.. autoclass:: fme.ace.GlobalMeanDriftConfig + :noindex: + .. autoclass:: fme.ace.EnsembleMetricConfig :noindex: