Internal ensembling in Aurora models

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",
 ]

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()
 

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)

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
 -------------------------

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`

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))