feat(metrics): add spread_squared ensemble variance metric

[Sir-Sloth-The-Lazy]

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 (..., S, N,
  d_state) convention used throughout the metrics module. T from the (B, T, N, F) deterministic
  shape ([RFC/Design] Standardize probabilistic vs deterministic return contract to unblock evaluation integrations  #335) is a batch dim folded into (...) before calling metrics to match how ar_model.py
  already calls all other metrics in practice.
• assert S > 1: Single-member variance is undefined as
  discussed in feat/add-ensemble-metrics-crps-spread #226 by @kshirajahere and @Panchadip-128
• 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.

[Sir-Sloth-The-Lazy]

@sadamov @joeloskarsson if this makes any good addition to #226 , I would like to keep contributing to that PR

[kshirajahere]

Thanks for summarizing the design choices so clearly. The mask_and_reduce_metric alignment and the explicit S > 1 check both make sense to me.

One thing I am curious about: do you expect spread_squared itself to be the primitive metric that gets aggregated/logged, with spread = sqrt(spread_squared) handled later as metric-specific post-processing, or do you want the square root applied directly in the metric implementation?

I am asking because it seems closely related to the per-metric aggregation / post-processing semantics discussed in #343 / #344, and I want to understand which level you want that distinction to live at.

[Sir-Sloth-The-Lazy]

I intentionally kept spread_squared as the raw variance (no sqrt) to match the convention in the prob_model_lam branch and in WeatherBench 2, where spread-skill analysis compares spread_squared directly against MSE (both in squared units). The sqrt would only be needed for display purposes (e.g. logging as "spread" in physical units alongside RMSE).
So yes, the intent is exactly what you described: spread_squared is the primitive metric, and spread = sqrt(spread_squared) would be handled as post-processing at the logging stage — the same pattern MSE uses today (computed as MSE, logged as RMSE after sqrt).
I'm aware of the class-based refactor in #343 / #344. Once that lands, converting spread_squared to a SpreadSquared(BaseMetric) subclass with post_process = sqrt and display_name = "spread" would be straightforward. For now this PR targets the current function-based API on main to keep things simple, but happy to rebase against the refactored structure if that's preferred.

[Panchadip-128]

The spread_squared as primitive with sqrt deferred to post-processing is the right call - mirrors the MSE/RMSE pattern and keeps things composable with the class-based refactor in #343/#344.

One system-level consideration worth raising: spread_squared and CRPS (being worked on in #226) will need to interoperate for spread-skill analysis - SSR requires both in consistent units and with aligned reduction paths. The ens_dim=-3 fixed convention here vs the explicit ens_dim parameter approach in #226 is a small divergence worth resolving before both land, otherwise downstream calibration metrics will need to special-case each one.

I've been thinking about this coordination across the evaluation layer as part of a broader probabilistic evaluation system , happy to help align with the two PRs so the full pipeline from CRPS to spread-skill works consistently end to end.

[sadamov]

This PR adds spread_squared which is also included in #226 (along with crps_ens). #226 is a superset of this PR, if that merges first, this one will conflict and be redundant. Please coordinate with the author of #226 or consider closing this in favour of the broader PR.