feat: Zarr export bridge for ensemble probabilistic forecasts

[Debadri-das]

Describe your changes

This PR adds comprehensive support for 4D ensemble tensors (S, T, N, F) to enable probabilistic weather forecast exports to Zarr format for mllam-verification integration.

Problem: Because the neural-lam codebase did not support 4D ensemble tensors and could only handle 2D/3D deterministic forecasts, probabilistic predictions could not be exported to Zarr format for external verification using mllam-verification.

Solution: Implemented full ensemble support across the data pipeline:

1. WeatherDataset (neural_lam/weather_dataset.py):

   • Added 4D tensor shape validation with elif len(tensor.shape) == 4: branch
   • Implemented ensemble_member dimension in dims list: ["ensemble_member", "time", "grid_index", f"{category}_feature"]
   • Added ensemble_member coordinate binding logic with fallback to np.arange()
   • Updated docstring documenting 2D/3D/4D tensor support
   • Maintains full backward compatibility with existing deterministic workflows

2. ARModel Documentation (neural_lam/models/ar_model.py):

   • Updated _create_dataarray_from_tensor() docstring with detailed ensemble support documentation
   • Clarified parameter expectations for 3D (deterministic) vs 4D (ensemble) tensors
   • Added Returns section documenting coordinate structure for Zarr export

3. ARModel Ensemble Handling (neural_lam/models/ar_model.py):

   • Modified unroll_prediction() to detect and handle 5D ensemble inputs (vs 4D deterministic)
   • Reshapes tensors for consistent processing: (B, S, T, N, F) → (B*S, T, N, F)
   • Reshapes predictions back to ensemble form after stacking: (B*S, T, N, F) → (B, S, T, N, F)
   • Maintains deterministic behavior for 4D inputs

4. Test Coverage (tests/test_datasets.py):

   • Added test_dataset_item_create_dataarray_from_tensor_4d_ensemble() test
   • Verifies 4D tensor conversion to xarray DataArray with correct dimensions
   • Validates ensemble_member coordinate exists with correct size
   • Checks time coordinate mapping and tensor value preservation

5. Documentation (CHANGELOG.md):

   • Added 4 entries documenting all changes linked to issue [Feature] Zarr Export Bridge: Addition of ensemble support to DataArray creation for mllam-verification #520

Motivation and Context

As stated in #62, the core evaluation pipeline is moving toward using mllam-verification (supported by the scores package). We require a bridge to export ensemble predictions to Zarr in order to support probabilistic models (such as Graph-EFM). Ensemble forecasts cannot be used with external verification tools.

Dependencies

No new dependencies are required. All code uses existing packages: PyTorch, xarray, and numpy.

Issue Link

#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
• 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 (not applicable - internal API change)
• I have added tests that prove my feature works
• I have given the PR a name that clearly describes the change, written in imperative form
• I have requested a reviewer and an assignee (no write access - requesting maintainer to add)

Commits

1. Add 4D ensemble tensor support to WeatherDataset - Implements shape validation, dimension mapping, and coordinate binding for 4D tensors
2. Document ensemble support in ARModel._create_dataarray_from_tensor() - Updates docstring with 4D tensor documentation
3. Add test for 4D ensemble DataArray creation - Comprehensive test coverage for ensemble DataArray conversion
4. Update CHANGELOG for ensemble Zarr export support - Documents all changes with issue [Feature] Zarr Export Bridge: Addition of ensemble support to DataArray creation for mllam-verification #520 reference

Related Issues

• Closes [Feature] Zarr Export Bridge: Addition of ensemble support to DataArray creation for mllam-verification #520 (Zarr Export Bridge)
• Related to Merge Graph-EFM model from prob_model_lam branch #62 (evaluation pipeline migration to mllam-verification)
• Related to Output predictions as zarr dataset #89 (probabilistic forecasting support)
• Related to Refactor: Decouple WeatherDataset from ARModel._create_dataarray_from_tensor #223 (ensemble forecasting)

[Debadri-das]

Hi @joeloskarsson @sadamov , due to an ongoing nan bounds check issue, I am observing a local failure in test_clamping.py::test_clamping. Should we create a different issue to monitor it even if it has nothing to do with these DataArray modifications?

[kshirajahere]

I reproduced a follow-up break in ARModel.plot_examples() on top of this branch: once the prediction carries an ensemble dimension, the current two-panel example plotting path no longer has a single deterministic field to pass into vis.plot_prediction().

I kept that as a separate narrow stacked PR here: Debadri-das#1

Scope is intentionally small: it preserves the existing example-plotting path by collapsing ensemble predictions to their ensemble mean before plotting, with a focused regression test.

[AswaniSahoo]

Nice work on the ensemble support. The (B, S, T, N, F) → (B*S, T, N, F) batch-fold in unroll_prediction() aligns with the ensemble tensor contract discussed in #335.

One thing I noticed while looking through: the current reshape sits inside unroll_prediction() and wraps around the existing predict_step() calls. This works cleanly for deterministic models running with replicated ensemble members.

However, for CVAE-like probabilistic models (the Graph-EFM direction in #62), each ensemble member would have its own latent variable z sampled during predict_step(). In that case, the reshape to (BS, ...) before predict_step() would need the Z sampling to happen per-member within that call - which it would, since each of the BS items in the batch gets an independent sample. So the batch-fold approach actually generalizes naturally to latent variable models too, as long as the stochasticity is internal to predict_step().

Just wanted to flag this since it connects to the broader #62 roadmap. The approach here looks solid for the current scope.

[kshirajahere]

@Debadri-das this PR adds 4D ensemble support to WeatherDataset.create_dataarray_from_tensor(), which is needed for probabilistic / ensemble exports. But ARModel.plot_examples() still assumes that the prediction side of the example plot is a single deterministic field.

With an ensemble prediction tensor shaped (B, S, T, N, F), the current plotting path reaches vis.plot_prediction() with an ensemble_member dimension still present on the prediction DataArray. In the existing two-panel example view (Ground Truth vs Prediction), that leaves no defined scalar field to plot and crashes during plotting / color-range handling.
Observed failure:
ValueError: can only convert an array of size 1 to a Python scalar

[sadamov]

Closing this in favour of #649, which consolidates the probabilistic metrics + ensemble plumbing track for v0.8.0. The 4D ensemble tensor support and zarr export bridge will inform the consolidated PR. Note that this PR touches ARModel which was removed in #208, so a rewrite against ForecasterModule would have been needed regardless. Thanks for the contribution @Debadri-das.