Skip to content

Add spread_squared ensemble variance metric for probabilistic evaluation - #430

Closed
Sir-Sloth-The-Lazy wants to merge 7 commits into
mllam:mainfrom
Sir-Sloth-The-Lazy:metrics/spread-squared-implementation
Closed

Add spread_squared ensemble variance metric for probabilistic evaluation#430
Sir-Sloth-The-Lazy wants to merge 7 commits into
mllam:mainfrom
Sir-Sloth-The-Lazy:metrics/spread-squared-implementation

Conversation

@Sir-Sloth-The-Lazy

Copy link
Copy Markdown
Contributor

Describe your changes

Add spread_squared (ensemble variance) metric to the metrics module.

This metric computes the unbiased sample variance across ensemble members, following the same API pattern as all existing metrics: (pred, target, pred_std, mask, average_grid, sum_vars). Entry-wise variance is computed first across the ensemble dimension (dim=-3), then passed through mask_and_reduce_metric for grid masking and optional reduction.

Key design decisions:

  • Matches existing API: Same signature as mse, mae, nll, etc. target and pred_std are accepted but unused, maintaining the uniform interface expected by get_metric() callers in ar_model.py.
  • ens_dim = -3: Ensemble members are at the third-from-last dim, consistent with the (..., N, d_state) convention used throughout the metrics module. For (B, S, T, N, F) input per [RFC/Design] Standardize probabilistic vs deterministic return contract to unblock evaluation integrations  #335, this correctly reduces over S.
  • assert S > 1: Single-member variance is undefined
  • Bessel's correction: Unbiased estimator (divides by S-1), consistent with the unbiased CRPS convention in @joeloskarsson's reference code shared in feat/add-ensemble-metrics-crps-spread #226.
  • Internal mask_and_reduce_metric call: Follows the same pattern as every other metric in the file.

This metric is a prerequisite for spread-skill analysis of ensemble and probabilistic models (Graph-EFM, Diffusion-LAM) tracked in #62.

Dependencies: None

Issue Link

Addresses parts of #62
Related discussion: #226, #335

Type of change

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

  • My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
  • I have performed a self-review of my code
  • For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
  • I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
  • I have updated the README to cover introduced code changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have given the PR a name that clearly describes the change, written in imperative form (context).
  • I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • the code is readable
  • the code is well tested
  • the code is documented (including return types and parameters)
  • the code is easy to maintain

Author checklist after completed review

  • I have added a line to the CHANGELOG describing this change, in a section
    reflecting type of change (add section where missing):
    • added: when you have added new functionality
    • changed: when default behaviour of the code has been changed
    • fixes: when your contribution fixes a bug
    • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

  • PR is up to date with the base branch
  • the tests pass
  • (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
  • author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  • Once the PR is ready to be merged, squash commits and merge the PR.

kshirajahere and others added 7 commits March 18, 2026 19:07
Refactor error map plotting function to heatmap and improve layout computation.
Added a HeatmapDatastore class for testing heatmap plots and updated tests to verify heatmap behavior.
Scale metric heatmap figure size and improve readability for larger outputs.
…cale control

- expose optional vmin/vmax args so callers can fix a shared colour
  scale across multiple heatmaps (e.g. val vs test comparison)
- propagate params through the plot_error_map deprecation wrapper
- add tests: explicit vmin/vmax respected, cross-run scale consistency
- update docstring to document new parameters

Follow-up to mllam#376
Implements spread_squared following the existing metric API contract:
- Signature: (pred, target, pred_std, mask, average_grid, sum_vars)
- Accepts (..., S, N, d_state) ensemble predictions with S at dim=-3
- Uses Bessel's correction (S/(S-1)) for unbiased sample variance
- Calls mask_and_reduce_metric internally, consistent with mse/mae/nll
- Registered in DEFINED_METRICS as "spread_squared"
- Asserts S > 1 (single-member variance is undefined)

Also fixes pre-existing vis.py breakage from the heatmap refactor:
- Restores _TITLE_SIZE, _LABEL_SIZE, _TICK_SIZE, _tex_safe, plot_on_axis
  (used by plot_prediction and plot_spatial_error but accidentally removed)
- Removes unused top-level cartopy.crs import (moved inside plot_on_axis)
- Removes unused Iterator and patch imports from tests/test_plotting.py

Adds tests/test_metrics.py with 9 tests covering mathematical
correctness, Bessel's correction, S=1 guard, shape contract,
mask filtering, and consistency with torch.var(unbiased=True).
@Sir-Sloth-The-Lazy
Sir-Sloth-The-Lazy deleted the metrics/spread-squared-implementation branch March 18, 2026 15:09
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.

2 participants