---

name: ensembles description: "Run Earth2Studio ensemble forecasts with controlled perturbations, bounded member batching, temporal interpolation, checkpoint restart, and ensemble-aware downstream analysis." disable-model-invocation: true metadata: disco-role: operating license: Apache 2.0

Ensemble forecasts

Use this skill when the task needs multiple prognostic rollouts from one or more initial conditions, controlled initial-condition perturbations, member-wise uncertainty, or ensemble verification. The primary entry point is earth2studio.run.ensemble; it returns the same IO object supplied by the caller, populated with an explicit ensemble coordinate.

This skill owns:

• nensemble, batch_size, device, output shape, and member bookkeeping.
• Built-in and custom perturbation selection, normalization, and coordinate requirements.
• Medium-range, temporally interpolated, lagged, HENS-style, and cyclone- analysis handoffs.
• Reducing or verifying ensemble outputs with diagnostics and statistics.

It does not replace deterministic model/data setup, the general IO backend catalog, model-package download, or a complete cyclone tracker recipe. Keep those concerns in the caller's model/data/IO skills. Support is not exhaustive: model-specific extras, accelerator support, licenses, and data-source access remain prerequisites for the selected components.

Fast route

1. Confirm that the prognostic model accepts a leading dynamic member axis and record its input_coords(), output step, variables, and spatial grid.
2. Choose a perturbation whose units and coordinate assumptions match the model. Use Zero() only for a deliberately identical control ensemble.
3. Choose nensemble and a GPU-fitting batch_size; do not use the full ensemble as a batch merely because it is the default.
4. Pre-create an IO store whose coordinates have ensemble first, followed by time, lead_time, and the model output dimensions. See workflows.
5. Run the contract in API reference, then validate member count, lead-time count, finite values, and perturbation spread before computing statistics.
6. For a large or asynchronous store, run the offline configuration check: python scripts/check_ensemble_config.py --help.

Minimal invocation

import numpy as np
from earth2studio.io import ZarrBackend
from earth2studio.perturbation import SphericalGaussian
from earth2studio.run import ensemble

io = ZarrBackend(
    "forecast.zarr",
    chunks={"ensemble": 1, "time": 1, "lead_time": 1},
    backend_kwargs={"overwrite": True},
)
result = ensemble(
    time=["2024-01-01"],
    nsteps=10,
    nensemble=8,
    prognostic=model,
    data=data_source,
    io=io,
    perturbation=SphericalGaussian(noise_amplitude=0.15),
    batch_size=2,
    output_coords={"variable": np.array(["t2m", "tcwv"])},
    device=device,
)

nsteps=10 produces nsteps + 1 lead-time slices, including the initial slice. batch_size=2 still writes all eight members with member coordinates 0..7; it only limits concurrent model state. output_coords is a coordinate selection, not a perturbation selection.

Decide before running

• Small/medium ensemble: use Gaussian for independent unstructured noise, SphericalGaussian for a spatially correlated spherical field, or a custom wrapper for variable-specific amplitudes. Set batch_size to the largest value that fits the model and perturbation working set.
• Analysis-time lags: use LaggedEnsemble only when the source can fetch each lag and len(lags) == nensemble. Negative lags are appropriate for forecasting; positive lags are for hindcast-like use, not future prediction.
• Medium-range finer cadence: pass a correctly configured prognostic interpolation wrapper (for example InterpModAFNO) as prognostic. The ensemble workflow detects its interp_method while fetching initial data; do not manually resize the initial tensor.
• HENS or cyclone uncertainty: use the selected model/checkpoint and its compatible HENS perturbation as a separate run. Keep batch_size small, write one store per checkpoint when needed, then combine outputs only after checking member coordinate semantics. The HENS perturbation needs model/data warmup and spherical-grid assumptions; it is not a generic noise shortcut.
• Huge ensemble: stream batches to a disk-backed, ensemble-aware store. Avoid in-memory output and avoid batch_size=nensemble unless the state is known to fit. For AsyncZarrBackend, ensemble and lead_time must be in parallel_coords because built-in ensemble writes partial slices on both axes; call close() for non-blocking writes. See the synthetic check and recovery notes in troubleshooting.

Output and downstream handoff

The built-in workflow writes one array per selected variable. The normal array layout is [ensemble, time, lead_time, ...model output dimensions...]; the coordinate dictionary is the authority when a model has extra dimensions. Check array.shape[0] == nensemble and array.shape[2] == nsteps + 1 for a single start time before analysis. A subset output_coords can remove variables or spatial points, so validate against the requested subset rather than the full model grid.

For member summaries, apply a statistic to tensors with an ensemble coordinate and reduce that named dimension, for example earth2studio.statistics.mean(["ensemble"]) or earth2studio.statistics.std(["ensemble"]). Use crps, energy_score, brier_score, or rank_histogram only with their metric contracts: the observation tensor must not contain ensemble, and expensive full-field energy scores should be avoided unless memory is budgeted. For a diagnostic model, map each forecast tensor to diagnostic.input_coords() and call the diagnostic on the member-batched tensor in a custom loop; run.diagnostic itself is not an ensemble runner. Concrete patterns are in workflows.

Verification gates

Before accepting a run, confirm:

• the selected extras and model/data prerequisites import;
• the perturbation returns the same tensor rank, member length, device, and compatible coordinates;
• all requested member IDs are present, with no missing or duplicate writes;
• initial perturbation spread is plausible in the model's physical units;
• all output arrays have the requested nsteps + 1 lead times and finite values;
• a tiny CPU fixture works before a long GPU run; and
• the chosen downstream diagnostic/statistic receives matching coordinates.

Use API reference for exact signatures and troubleshooting for predictable failures.

Bundled helper and evidence limits

• scripts/check_ensemble_config.py performs only local arithmetic and compatibility checks; it never downloads data, models, or credentials.
• The built-in list in the API reference is a useful selection map, not an exhaustive model or backend inventory.
• Native candidates for verification are the ensemble workflow tests, the perturbation contract tests, the getting-started ensemble example, and the medium-range interpolation/HENS examples. They require the relevant model, data, optional extras, and often GPU/network resources; this skill does not claim to have run them.
• Intentionally omitted: deterministic fetch/model construction, generic IO tuning beyond ensemble slicing, serving clients, distributed orchestration, and a full cyclone tracker implementation.