Skip to content

Add per-sample global-mean drift inference aggregator - #1617

Draft
yyexela wants to merge 4 commits into
mainfrom
feature/global-mean-drift-aggregator
Draft

yyexela wants to merge 4 commits into
mainfrom
feature/global-mean-drift-aggregator

Conversation

@yyexela

@yyexela yyexela commented Nov 2, 2026 •

Copy link
Copy Markdown
Contributor

Adds a default-on inference metric that measures slow drift of the global-mean bias away from the target, one slope per sample (initial condition or ensemble member) per variable. The per-sample slopes are exposed via InferenceEvaluatorAggregator.get_global_mean_drift_slopes() so a downstream trainer can test whether the drift is significant using the spread across initial conditions. That accessor has no in-repo caller yet; it is a hook for a downstream trainer.

The aggregator accumulates running OLS sums (n, Σt, Σt², Σy, Σty) in O(n_samples × n_vars) memory and tracks the recorded step count as a scalar (identical on every rank). Slopes are NaN when n_time * timestep < 730 days, matching the annual-metric minimum. Samples are gathered across data-parallel ranks at finalize time.

record_batch(data)                        # per sample s, per variable v
  y[s, t] = area_weighted_mean(gen - target)[v]
  t       = years since sample's first recorded step
  accumulate n, Σt, Σt², Σy, Σt·y

finalize
  data_parallel_gather(stats)
  slope[s, v] = (nΣty − ΣtΣy) / (nΣt² − (Σt)²)
  slope = NaN where n_time * timestep < 730 days

Logs (under the metric name global_mean_drift):

  • global_mean_drift/slope/{v}: sample-mean slope (units of v per year)
  • global_mean_drift/t/{v}: mean / (std(ddof=1) / sqrt(n_samples)); NaN when n_samples < 2
  • global_mean_drift/max_abs_t: max over variables of |t|, ignoring NaN

global_mean_drift_diagnostics.nc holds slope(sample, variable) and slope_units(variable).

The metric is added only to InferenceEvaluatorAggregatorConfig — it is not added to LegacyFlagInferenceEvaluatorAggregatorConfig or the coupled InferenceEvaluatorAggregatorConfig (fme/coupled/aggregator.py).

Changes:

  • fme.ace.aggregator.inference.global_mean_drift: new GlobalMeanDriftAggregator and GlobalMeanDriftConfig (default-on, skipped for rollouts ≤ 730 days). GlobalMeanDriftConfig is exported from fme.ace.

  • fme.ace.aggregator.inference.data: new PerSampleSlopes protocol and MetricBuildResult.per_sample_slopes slot, so the build loop hands the aggregator to the top level without isinstance checks.

  • fme.ace.aggregator.inference.main: registers the config in the MetricConfig union, adds global_mean_drift field to InferenceEvaluatorAggregatorConfig, and adds InferenceEvaluatorAggregator.get_global_mean_drift_slopes().

  • fme.core.testing.MockDistributed: adds data_parallel_gather.

  • Tests added

  • If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated (no dependency changes)

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().
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.
Drop the zero-drift test (covered by per-sample slope recovery) and merge the
two accessor-returns-None tests into one parametrized test.
yyexela added a commit that referenced this pull request Nov 3, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant