Generation cleanout

[maritsandstad]

Refactor and consistency fix for timeseries and gridded generation

• Tests added
• Documentation added
• Example added (in the documentation, to an existing notebook, or in a new notebook)
• Description in CHANGELOG.rst added (single line such as: (`#XX <https://github.com/benmsanderson/METEOR/pull/XX>`_) Added feature which does something)

[maritsandstad]

Based on analysis in collaboration with Claude (Sonnet and then Opus):

_generate_timeseries vs _generate_gridded — Refactoring Analysis

Commonalities and Differences

Shared steps

Step	_generate_timeseries	_generate_gridded
Pattern scaling	_get_or_compute_pattern_scaling → monthly_prediction, monthly_warming	identical call
Noise model access	self.noise_models[variable]	identical
include_noise guard	generates PCs or falls back to pattern_agg[np.newaxis,:]	forces n_realizations=1
Month-index slicing	monthly_prediction.isel(month=slice(…))	identical
Realization stacking	done inside the noise model call	xr.concat / expand_dims(realization=[0]) — verbatim 3×

Key differences

1. Temporal structure — continuous vs. snapshot

• _generate_timeseries produces a contiguous (n_realizations, n_months) array from start_year to end_year.
• _generate_gridded produces discrete slices: one field per year (annual), one field per month of a specific year (monthly), or one field per multi-year period (climatology).

2. Spatial treatment — aggregated vs. full field

• _generate_timeseries spatially collapses the field (global mean / AR6 region / point), then adds noise via generate_regional_mean_realizations.
• _generate_gridded preserves the full (lat, lon) field and uses generate_realization, a completely different noise-model entry point.

3. Noise coherence strategy
This is the deepest conceptual difference. _generate_timeseries calls generate_stochastic_pcs once before the aggregation loop, so all spatial scales (global, EAS, point:59.9,10.8) are driven by the same underlying PC realisations — spatial consistency is guaranteed. _generate_gridded generates a fresh generate_realization call per year-slice per realization, so there is no enforced temporal continuity across year-slices in an ensemble member.

4. Effect of averaging on spread

• _generate_timeseries: all monthly noise is preserved; ensemble spread ≈ full internal variability.
• _generate_gridded "annual": 12-month mean reduces spread by ~1/√12 relative to monthly.
• _generate_gridded "climatology" over N years: spread collapses further by ~1/√(12N).
• _generate_gridded "monthly": the only gridded spec that keeps full monthly variability — but only for individual years, not connected across years.

5. Variable transform (Gamma for pr)
Only _generate_timeseries applies the distribution transform. _generate_gridded has no transform path at all.

6. CMIP6 reference data loading
Only _generate_timeseries loads ssp_data and picontrol_data (for transform fitting and PR baseline). _generate_gridded never touches them.

---

Code blocks that lend themselves to extraction

A. Named return from _get_or_compute_pattern_scaling

Both methods unpack pattern_result[0] / pattern_result[1] from the tuple return. The result tuple should become a small dataclass or namedtuple (PatternScalingResult) with .monthly_prediction, .monthly_warming, .em_data, .conc_data fields. This removes positional indexing across the whole file.

B. _stack_realizations(realizations) — one-liner helper

The pattern:

if len(lst) > 1:
    xr.concat(lst, dim="realization")
else:
    lst[0].expand_dims(realization=[0])

    ### C. `_generate_gridded_slice(...)` — unify the three gridded loops

The bodies of the `annual`, `monthly`, and `climatology` loops inside `_generate_gridded`
are structurally identical. Only two things vary: the month-index range and whether the
time axis is averaged away after generation.

```python
def _generate_gridded_slice(
    self,
    monthly_prediction,
    monthly_warming,
    noise_model,
    start_idx,
    end_idx,
    n_realizations,
    include_noise,
    reduce_time,          # True → .mean(dim="month"), False → keep months
):
    realizations = []
    for _ in range(n_realizations):
        if include_noise:
            r = noise_model.generate_realization(
                monthly_warming[start_idx:end_idx],
                n_realizations=1,
                noise_only=True,
                add_base=monthly_prediction.isel(month=slice(start_idx, end_idx)),
            )
        else:
            r = monthly_prediction.isel(month=slice(start_idx, end_idx))
        if reduce_time:
            r = r.mean(dim="month")
        realizations.append(r)
    return _stack_realizations(realizations)

---

Review and expansion (verified against the code)

I read through _generate_timeseries, _generate_gridded,
_get_or_compute_pattern_scaling, and the three noise-model entry points
(generate_stochastic_pcs, generate_realization,
generate_regional_mean_realizations). Overall the first-version analysis is
sound, but there is one factual error, one missing (and dominant)
variability mechanism, and a broken helper that the refactoring section
should account for.

Verdict on each original conclusion

#	Claim	Verdict
1	Continuous vs. snapshot temporal structure	✅ Correct
2	Aggregated vs. full-field spatial treatment	✅ Correct
3	Noise coherence: timeseries shares PCs, gridded does not	✅ Correct direction, but see correction below
4	Averaging reduces spread (1/√12, 1/√(12N))	⚠️ True but not the main cause of the spread gap
5	Gamma transform only on timeseries path	✅ Correct
6	CMIP6 reference data only loaded on timeseries path	✅ Correct

Correction to conclusion #3 — sharing PCs is not currently possible on the gridded path

The document implies the two paths could already share PCs and simply choose
not to. That is not the case. generate_regional_mean_realizations accepts a
stochastic_pcs= argument, but generate_realization (the gridded entry point)
has no such parameter — its signature is:

def generate_realization(self, global_temp_trajectory, n_realizations=1,
                         random_seed=None, noise_only=False, add_base=None):

It always calls self._generate_stochastic_pcs(X_exog, n_time) inside its own
realization loop. So consistent noise across timeseries and gridded is a
feature that must be built, not merely wired up. This is the key prerequisite
for the "consistent realisations" goal.

New finding — the dominant variability gap is AR spin-up from zero initial conditions

This is the most important addition and it is not in the original analysis.

_generate_stochastic_pcs seeds the VAR(p) process with zeros:

synthetic_pcs = np.zeros((n_time, self.n_modes))
synthetic_pcs[: self.lag_order] = 0
for t in range(self.lag_order, n_time):
    ...
    synthetic_pcs[t] = forecast + all_shocks[t]

The variance of an autoregressive process started from zero grows from ~0 toward
its stationary value over a spin-up window. Consequences:

• Timeseries path: PCs are generated once over the full sliced trajectory
  (monthly_warming for the whole start_year…end_year span, often hundreds of
  months). By the time you reach any year of interest the process is fully spun
  up, so its noise has stationary variance.
• Gridded path: PCs are regenerated per year-slice from monthly_warming[start_idx:end_idx]
  — i.e. from a 12-month window (or the climatology window) that restarts at
  zero each time. Those 12 months sit squarely inside the spin-up transient, so
  the realised variance is systematically too low, independent of any
  averaging.

So the gridded spread is suppressed by two stacked effects:

1. AR spin-up from zero ICs (the dominant one, present even for the monthly
   spec that keeps all 12 months), then
2. the 1/√12 (annual) / 1/√(12N) (climatology) averaging the document already
   identified.

The monthly gridded spec is the tell: the document claims it "keeps full
monthly variability", but because each year is generated from a zero start it
still under-disperses relative to the same months pulled from the timeseries
path. This is worth verifying empirically (compare the ensemble std of the
gridded monthly field global-mean against the timeseries global std for the
same years) — it should expose the gap cleanly.

Unified strategy — one change fixes both consistency and variability

Both goals collapse to a single design: generate the stochastic PCs once, over
the full contiguous trajectory, and have every consumer slice into that same PC
array.

1. In generate_ensemble_outputs, after pattern scaling, call
   generate_stochastic_pcs(monthly_warming, n_realizations=...) once per
   variable over the whole start_year…end_year span.
2. Pass that array into both _generate_timeseries (already supported via
   stochastic_pcs=) and _generate_gridded.
3. Add a stochastic_pcs= parameter to generate_realization mirroring the one
   in generate_regional_mean_realizations; when supplied, skip internal PC
   generation and reconstruct from pcs[:, start_idx:end_idx] @ self.pca.components_.
4. In _generate_gridded, slice the shared PC array by month index
   (pcs[..., start_idx:end_idx, :]) instead of regenerating per slice.

Why this fixes both problems at once:

• Consistency: a gridded field reconstructed as pcs @ pca.components_ and a
  regional mean reconstructed as pcs @ eof_projections are then driven by
  identical PC loadings, so the area-mean of the gridded field equals the
  regional timeseries (up to the EOF projection algebra). Realization k means
  the same physical draw everywhere.
• Variability: slices now come from a fully-spun-up long trajectory, so each
  year-slice carries stationary variance. The zero-IC transient is incurred once
  at the trajectory start (outside any output window) instead of at every
  requested year.

Caveats to flag before implementing

• Seeding: to truly reproduce realization k across calls, either reuse the
  generated PC array directly (cleaner) or thread random_seed. Array reuse
  avoids coupling to global RNG state.
• Time-axis alignment: the shared PC array is indexed off base_year inside
  _get_or_compute_pattern_scaling; the gridded year_to_month_idx is indexed
  off start_year. These must use a single, consistent origin or the slices will
  silently misalign. Worth a unit test asserting
  global_mean(gridded_annual[year]) == timeseries["global"].sel(year=year).mean().
• Transform mismatch (conclusion Make notebook to simulate 1pctCO2 response #5): even with shared PCs, the pr Gamma
  transform is applied on the timeseries path only, so gridded pr will not be
  distributionally consistent with timeseries pr. If consistency for pr
  matters, the transform must move into (or be shared with) the gridded path —
  otherwise scope the "consistent noise" promise to tas.
• Memory: the full-trajectory PC array is (n_realizations, n_time, n_modes)
  — modest, far smaller than gridded fields — so caching it on the run is cheap
  and removes most redundant PC generation.

Correction to refactoring block B — the existing _stack_realizations is broken

The nested helper currently in _generate_gridded is dead, buggy code:

def _stack_realizations(lst):
    if len(lst) > 1:
        return xr.concat(realizations, dim="realization")   # wrong var: `realizations`
    else:
        return realizatio[0].expand_dims(realization=[0])    # typo: `realizatio`

It references realizations/realizatio instead of its own lst argument and
is never actually called (the three loops inline their own xr.concat). So
proposal B is not just a tidy-up — it fixes a latent NameError waiting to fire
if anyone wires the helper in. The corrected version:

def _stack_realizations(lst):
    if len(lst) > 1:
        return xr.concat(lst, dim="realization")
    return lst[0].expand_dims(realization=[0])

Refactoring proposals — assessment

• A (named/dataclass return from _get_or_compute_pattern_scaling): Good.
  Both callers use monthly_prediction, monthly_warming, _, _ = ...; a small
  result object removes positional unpacking and the awkward throwaway tuple.
• B (_stack_realizations): Good, but fix the bug above as part of it.
• C (_generate_gridded_slice): Good — the three loops differ only in the
  month range and the reduce_time flag. Extend the proposed signature to
  also accept stochastic_pcs (or per-slice PC views) so the extraction helper
  becomes the single place that consumes shared PCs, tying the refactor and the
  consistency fix together rather than doing them as two passes over the same
  code.

Suggested ordering

1. Fix _stack_realizations (B) — trivial, removes a latent bug.
2. Introduce the PatternScalingResult return (A) — pure refactor, no behaviour
   change.
3. Add stochastic_pcs= to generate_realization and extract
   _generate_gridded_slice (C) together.
4. Generate PCs once in generate_ensemble_outputs and thread them through both
   paths — this is the step that actually changes (and aligns) the science, so
   it should land last with the alignment unit test described above.

[maritsandstad]

Superseded by #80 which is a clean version directly on base