Stream gridded transform: bound memory to O(chunk_size), unlock N > 8

[benmsanderson]

Summary

The gridded distribution-transform path used to build the entire (N, T, lat, lon) transformed ensemble in memory (~3× peak-copy overhead at intermediate steps), which capped n_realizations at ~8 for NorESM2-MM full grid and silently OOMed above that.

This PR refactors that path to stream realizations in chunks. Peak memory is now bounded by the chunk size, not by N, which unlocks arbitrarily large gridded ensembles.

Design

Two-pass approach in _build_ensemble_slices_streaming:

• Pass 1: iterate chunks, reconstruct each, accumulate per-month-of-year, per-gridpoint sum + sum-of-squares. Compute Gaussian params (mean, std) from moments once all chunks are seen.
• Pass 2: iterate chunks again, reconstruct, apply the seasonal quantile transform, and extract only the requested time slices into per-slice output arrays. Chunk memory is freed between iterations.

The Gaussian half of the quantile map is still fit on the FULL window (correctness-equivalent to the previous path; the climate-change trend in per-month sigma is preserved). The transform itself factors cleanly across chunks because the target-distribution params only depend on the reference data.

_generate_gridded builds a slice_specs list from gridded_spec upfront and passes it to the streaming builder; the transform result comes back as a dict[(s_idx, e_idx, reduce_time)] of per-slice DataArrays. This replaces the pattern where a 4-D pre-transformed ensemble was sliced lazily by a lambda.

Chunk size defaults to ~5 GB per chunk (measured from the per-realization reconstruction footprint) and can be overridden with the METEOR_GRIDDED_CHUNK_SIZE env variable. Malformed values silently fall back to the auto-sized default so a typo in a batch script cannot crash METEOR at import time.

Removed the now-unused _build_full_window_transformed_ensemble.

Correctness

• Verified output matches the pre-refactor full-window path to ≤1e-16 relative on N=3 pr with the real NorESM2-MM noise model.
• New test_resolve_gridded_chunk_size_env_var pins env-var handling and the auto-size formula.
• All 130 existing unit tests still pass.

Wins

Peak memory (NorESM2-MM full grid, gen_gridded workload, measured with the profiling harness):

N	Before (peak GB)	After (peak GB)	Wall (s)
1	11	9	55
5	33	17	108
10	OOM (~65 est.)	22	168
20	OOM (~130 est.)	22	356
50	OOM (~325 est.)	22	1090

Peak memory is essentially constant in N. Wall time scales linearly.

Notes for reviewers

• The equivalence proof relied on running both the old and new paths side-by-side during development. Since the old path is removed here, that check is not runnable at head; the profiling harness introduced in the sibling PR bundles a memory/timing script (verify_streaming.py) for future memory-scaling investigations.
• chunk_size >= n_realizations reproduces the pre-refactor behavior (one pass per phase), so the code is a strict superset of the old semantics.
• The env-var name mirrors the METEOR_GAMMA_PPF_THREADS pattern introduced in the sibling precip-transform PR (safe default, batch-script friendly).

Test plan

• pytest tests/unit/test_meteor_interface.py::test_resolve_gridded_chunk_size_env_var
• pytest tests/unit/test_meteor_interface.py
• Manual: run generate_ensemble_outputs with gridded={...} at n_realizations=10 (previously OOM) and confirm it completes with reasonable memory.

🤖 Generated with Claude Code

[maritsandstad]

Ok, this also looks ok, just minor comment todo with coding style preference

[maritsandstad]

Is there a clear advantage to having this function be internal, or could it be pulled out as a more general help function?