diff --git a/aurora/__init__.py b/aurora/__init__.py index 9403fb7f..5ca4f317 100644 --- a/aurora/__init__.py +++ b/aurora/__init__.py @@ -1,6 +1,6 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" -from aurora.batch import Batch, Metadata +from aurora.batch import Batch, Metadata, split_batch, tile_batch from aurora.insolation import insolation from aurora.model.aurora import ( Aurora, @@ -14,7 +14,7 @@ AuroraV1p5Ensemble, AuroraWave, ) -from aurora.rollout import rollout +from aurora.rollout import rollout, rollout_ensemble from aurora.tracker import Tracker __all__ = [ @@ -30,7 +30,10 @@ "AuroraV1p5Ensemble", "Batch", "Metadata", + "tile_batch", + "split_batch", "insolation", "rollout", + "rollout_ensemble", "Tracker", ] diff --git a/aurora/batch.py b/aurora/batch.py index 33f3792c..2b3b491c 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -17,7 +17,7 @@ unnormalise_surf_var, ) -__all__ = ["Metadata", "Batch"] +__all__ = ["Metadata", "Batch", "tile_batch", "split_batch"] @dataclasses.dataclass @@ -326,6 +326,57 @@ def __post_init__(self): ) +def tile_batch(batch: Batch, n: int) -> Batch: + """Tile `batch` along the batch dimension `n` times. + + Use `split_batch` to split the tiled batch, or a prediction for it, back into `n` batches. + + Args: + batch (:class:`aurora.Batch`): The batch to tile. + n (int): Number of times to tile. + + Returns: + :class:`aurora.Batch`: `batch` tiled `n` times along the batch dimension. + """ + return dataclasses.replace( + batch, + surf_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.surf_vars.items()}, + atmos_vars={k: v.repeat(n, *([1] * (v.dim() - 1))) for k, v in batch.atmos_vars.items()}, + metadata=dataclasses.replace(batch.metadata, time=batch.metadata.time * n), + ) + + +def split_batch(batch: Batch, n: int) -> list[Batch]: + """Split `batch` into `n` equally sized batches along the batch dimension. + + This undoes `tile_batch`. + + Args: + batch (:class:`aurora.Batch`): The batch to split. + n (int): Number of batches to split into. + + Returns: + list[:class:`aurora.Batch`]: `batch` split into `n` equally sized batches. + + Raises: + ValueError: If the batch size of `batch` is not divisible by `n`. + """ + batch_size = next(iter(batch.surf_vars.values())).shape[0] + b, remainder = divmod(batch_size, n) + if remainder != 0: + raise ValueError(f"Batch size `{batch_size}` must be divisible by `n = {n}`.") + time = batch.metadata.time + return [ + dataclasses.replace( + batch, + surf_vars={k: v[m * b : (m + 1) * b] for k, v in batch.surf_vars.items()}, + atmos_vars={k: v[m * b : (m + 1) * b] for k, v in batch.atmos_vars.items()}, + metadata=dataclasses.replace(batch.metadata, time=time[m * b : (m + 1) * b]), + ) + for m in range(n) + ] + + def _np(x: torch.Tensor) -> np.ndarray: return x.detach().cpu().numpy() diff --git a/aurora/rollout.py b/aurora/rollout.py index bcd133d4..e9861e69 100644 --- a/aurora/rollout.py +++ b/aurora/rollout.py @@ -6,11 +6,11 @@ import torch -from aurora.batch import Batch +from aurora.batch import Batch, split_batch, tile_batch from aurora.model.aurora import Aurora from aurora.model.swin3d import NoiseGenerator -__all__ = ["rollout"] +__all__ = ["rollout", "rollout_ensemble"] def _make_lead_time_tensor(batch: Batch, lead_time_hours: float) -> torch.Tensor: @@ -153,3 +153,52 @@ def rollout( # Disable noise accumulation after roll-out is complete, in case the model will be used for # normal inference or training afterwards. model.set_noise_accumulation(n=0) + + +def rollout_ensemble( + model: Aurora, + batch: Batch, + steps: int, + num_ensemble_members: int, + fine_lead_times: Optional[Sequence[float]] = None, + use_noise_accumulation: bool = True, + apply_rollout_input_clipping: bool = True, + generator: NoiseGenerator = None, +) -> Generator[list[Batch], None, None]: + """Perform a roll-out for `num_ensemble_members` ensemble members simultaneously. + + The members are computed in one pass by repeating `batch` along the batch dimension, so every + member receives independent noise. All other arguments are as for :func:`rollout`. + + Args: + num_ensemble_members (int): Number of ensemble members. + generator (:class:`torch.Generator` or tuple of :class:`torch.Generator` or `None`, + optional): Generator for the noise, as for :func:`rollout`. A tuple gives one + generator per element of the tiled batch, which is `batch` repeated + `num_ensemble_members` times. Defaults to `None`. + + Yields: + list[:class:`aurora.Batch`]: After every (sub-)step, one prediction per ensemble member, + each with the batch size of `batch`. + + Raises: + ValueError: If `num_ensemble_members` is less than `1` or if `model` is not stochastic. + """ + if num_ensemble_members < 1: + raise ValueError( + f"`num_ensemble_members` must be at least `1`, but is `{num_ensemble_members}`." + ) + if not model.backbone.stochastic: + raise ValueError("`rollout_ensemble` requires a stochastic model.") + + batch = tile_batch(batch, num_ensemble_members) + for pred in rollout( + model, + batch, + steps, + fine_lead_times=fine_lead_times, + use_noise_accumulation=use_noise_accumulation, + apply_rollout_input_clipping=apply_rollout_input_clipping, + generator=generator, + ): + yield split_batch(pred, num_ensemble_members) diff --git a/docs/api.rst b/docs/api.rst index 7540ae0d..236fabcf 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -9,10 +9,15 @@ Batch .. autoclass:: aurora.Metadata :members: +.. autofunction:: aurora.tile_batch + +.. autofunction:: aurora.split_batch + Roll-Outs --------- -.. autoclass:: aurora.rollout - :members: +.. autofunction:: aurora.rollout + +.. autofunction:: aurora.rollout_ensemble Tropical Cyclone Tracking ------------------------- diff --git a/docs/models.md b/docs/models.md index a2c3e258..e4d62c1d 100644 --- a/docs/models.md +++ b/docs/models.md @@ -450,6 +450,28 @@ smoother intra-step transitions while using independent effective noise between matching the training regimen. Set `use_noise_accumulation=False` to draw independent noise at each sub-step instead, though this is not recommended. +(internal-ensembling)= +### Internal Ensembling + +For a stochastic model, such as `AuroraV1p5Ensemble`, you can use `rollout_ensemble` to +generate an autoregressive forecast with multiple ensemble members at once: + +```python +from aurora import rollout_ensemble + +with torch.inference_mode(): + preds = [ + [member.to("cpu") for member in members] + for members in rollout_ensemble(model, batch, steps=10, num_ensemble_members=4) + ] +``` + +Alternatively, you can produce an ensemble forecast by calling `rollout` multiple times, once for +every ensemble member. `rollout_ensemble` can be faster, at the cost of requiring more memory. + +Internally, `rollout_ensemble` uses `aurora.tile_batch` and `aurora.split_batch`. You can also use +these functions to run multiple ensemble members in a single forward pass. + ### Reproducible Noise By default, the noise is drawn from the global RNG. To control the noise, pass a `torch.Generator` diff --git a/tests/v1p5/test_ensemble.py b/tests/v1p5/test_ensemble.py new file mode 100644 index 00000000..7427283c --- /dev/null +++ b/tests/v1p5/test_ensemble.py @@ -0,0 +1,107 @@ +"""Copyright (c) Microsoft Corporation. Licensed under the MIT license. + +Tests for ensemble rollout and tiling utility functions. +""" + +import itertools + +import pytest +import torch + +from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, BATCH, _make_batch, _make_small_v1p5 +from aurora import Aurora, rollout_ensemble, split_batch, tile_batch +from aurora.model.film import AdaptiveLayerNorm + + +def _unzero_adaptive_layer_norms(model: Aurora, std: float = 0.1) -> None: + """Nudge every `AdaptiveLayerNorm`'s modulation away from its zero initialisation. + + A freshly constructed model is exactly insensitive to its conditioning signal, and hence to + the injected noise, because `AdaptiveLayerNorm.ln_modulation` is zero-initialised. + """ + for m in model.modules(): + if isinstance(m, AdaptiveLayerNorm): + with torch.no_grad(): + m.ln_modulation[-1].weight.normal_(std=std) + m.ln_modulation[-1].bias.normal_(std=std) + + +def test_tile_and_split_batch_roundtrip(): + b, n = 2, 3 + batch = _make_batch(batch_size=b) + + tiled = tile_batch(batch, n) + assert len(tiled.metadata.time) == n * b + for v in (*tiled.surf_vars.values(), *tiled.atmos_vars.values()): + assert v.shape[0] == n * b + + members = split_batch(tiled, n) + assert len(members) == n + for member in members: + assert member.metadata.time == batch.metadata.time + for k, v in member.surf_vars.items(): + torch.testing.assert_close(v, batch.surf_vars[k]) + for k, v in member.atmos_vars.items(): + torch.testing.assert_close(v, batch.atmos_vars[k]) + + # `batch` has batch size `b`, which is not divisible by `n`. + with pytest.raises(ValueError): + split_batch(batch, n) + + +@pytest.mark.parametrize("stochastic", [True, False]) +def test_forward_ensemble(stochastic: bool): + torch.manual_seed(0) + model = _make_small_v1p5(stochastic=stochastic) + _unzero_adaptive_layer_norms(model) # Otherwise, the noise has no effect on the output. + model.eval() + surf_vars = tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF) + batch = _make_batch(surf_vars=surf_vars) + n = 3 + batch = tile_batch(batch, n) + b = next(iter(batch.surf_vars.values())).shape[0] + + with torch.inference_mode(): + pred = model.forward(batch, lead_times=torch.full((b,), 6.0)) + members = split_batch(pred, n) + + if stochastic: + # Every member receives independent noise, so members must differ. + for member1, member2 in itertools.combinations(members, 2): + assert (member1.surf_vars["2t"] - member2.surf_vars["2t"]).abs().mean() > 1e-2 + else: + # Check that all are equal. + for m in range(1, n): + torch.testing.assert_close( + members[0].surf_vars["2t"], members[m].surf_vars["2t"], atol=1e-3, rtol=1e-3 + ) + + +def test_rollout_ensemble(): + num_ensemble_members = 3 + torch.manual_seed(0) + model = _make_small_v1p5(stochastic=True) + _unzero_adaptive_layer_norms(model) # Otherwise, the noise has no effect on the output. + model.eval() + batch = _make_batch(surf_vars=tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF)) + steps = 2 + + with torch.inference_mode(): + preds = list(rollout_ensemble(model, batch, steps, num_ensemble_members)) + + assert len(preds) == steps + for members in preds: + assert len(members) == num_ensemble_members + for member in members: + for v in member.surf_vars.values(): + assert v.shape[0] == BATCH + # Every member receives independent noise, so members must differ. + for member1, member2 in itertools.combinations(members, 2): + assert (member1.surf_vars["2t"] - member2.surf_vars["2t"]).abs().mean() > 1e-2 + + +@pytest.mark.parametrize("stochastic, num_ensemble_members", [(True, 0), (False, 3)]) +def test_rollout_ensemble_raises(stochastic: bool, num_ensemble_members: int): + model = _make_small_v1p5(stochastic=stochastic) + with pytest.raises(ValueError): + next(rollout_ensemble(model, _make_batch(), 2, num_ensemble_members))