Fix: plotting bounds calculation for multidimensional metrics dataarr…

[Debadri-das]

Describe your changes

This PR fixes an ensemble-handling bug in the plotting path of ARModel.plot_examples() that becomes visible once probabilistic / ensemble predictions are passed through the model evaluation pipeline.

Summary of the problem

plot_examples() currently computes per-variable plotting bounds using:

var_vmin = torch.minimum(
    pred_slice.flatten(0, 1).min(dim=0)[0],
    target_slice.flatten(0, 1).min(dim=0)[0],
).cpu().numpy()

var_vmax = torch.maximum(
    pred_slice.flatten(0, 1).max(dim=0)[0],
    target_slice.flatten(0, 1).max(dim=0)[0],
).cpu().numpy()

This logic is valid only when each example slice has deterministic shape:

• pred_slice: (pred_steps, num_grid_nodes, d_f)
• target_slice: (pred_steps, num_grid_nodes, d_f)

In that deterministic case, flatten(0, 1) collapses (time, grid) into one axis, and the reduction over dim=0 correctly returns a (d_f,) vector of per-variable extrema.

However, after ensemble/probabilistic support is introduced, the prediction tensor can carry an additional ensemble axis. In that case the reduction logic is no longer shape-correct.

For example, if pred_slice has shape:

• (ensemble_member, pred_steps, num_grid_nodes, d_f)

then:

pred_slice.flatten(0, 1)

produces shape:

• (ensemble_member * pred_steps, num_grid_nodes, d_f)

and:

.min(dim=0)[0]

returns shape:

• (num_grid_nodes, d_f)

instead of the expected:

• (d_f,)

This breaks the contract expected by the downstream plotting code, where var_vmin / var_vmax are zipped as one (vmin, vmax) pair per state variable. At best this produces semantically incorrect bounds; at worst it causes shape mismatch failures in the plotting routine.

What this PR changes

This PR updates the resampling / plotting-range calculation in ARModel.plot_examples() so that per-variable extrema are reduced across all non-feature dimensions, not only across the deterministic (time, grid) case.

Concretely, the fix makes the vrange computation robust for both:

• deterministic example slices, and
• ensemble/probabilistic example slices with an additional ensemble dimension.

The intended invariant is:

• regardless of whether the example is deterministic or ensemble-based,
• var_vmin.shape == (d_f,)
• var_vmax.shape == (d_f,)

This preserves the existing plotting interface and ensures vis.plot_prediction(...) continues to receive one valid color range per state variable.

Why this fix matters

This is a small but important follow-up for the ensemble-support work around issue #520.

The codebase is moving toward probabilistic verification workflows and ensemble-aware tensor/DataArray conversion. Once that support lands, the plotting path must also remain stable during validation/test-time visualization. Otherwise, users can successfully produce ensemble predictions but still hit failures when generating example plots or GIFs.

Dependencies

No new dependencies are required.

This change only uses existing project dependencies:

• torch
• numpy
• matplotlib (indirectly via the existing plotting path)

Issue Link

closes #520

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.

[Debadri-das]

It fixes only the per-variable plotting-bound reduction in ARModel.plot_examples() (var_vmin / var_vmax), where the current code assumes deterministic (pred_steps, grid_index, feature) slices and can produce incorrectly shaped extrema once an ensemble dimension is present.

[Debadri-das]

@sadamov @joeloskarsson would you please review this?

[sadamov]

This will indeed become an issue once we move to probabilistic modelling. But this is currently not a bug in main, as main does not support prob. modelling at all.
The right time to land both the fix and its test is alongside your #521 PR, which updates unroll_prediction() to handle 4D init_states -- at that point the test is motivated by real code, and the fix can be verified end-to-end.