Internal ensembling in Aurora models

aurora/__init__.py

@@ -14,7 +14,7 @@
     AuroraV1p5Ensemble,
     AuroraWave,
 )
-from aurora.rollout import rollout
+from aurora.rollout import rollout, rollout_ensemble
 from aurora.tracker import Tracker
 
 __all__ = [
@@ -32,5 +32,6 @@
     "Metadata",
     "insolation",
     "rollout",
+    "rollout_ensemble",
     "Tracker",
 ]

aurora/batch.py

@@ -311,6 +311,37 @@ def from_netcdf(cls, path: str | Path) -> "Batch":
         )
 
 
+def _tile_batch(batch: Batch, n: int) -> Batch:
+    """Tile `batch` along the batch dimension `n` times.
+
+    Not part of the public `Batch` API. Used only by `aurora.Aurora.forward` and
+    `aurora.rollout.rollout` to run `n` ensemble members as a single fused computation. The
+    tiled batch dimension is an internal implementation detail and must be undone with
+    `_split_batch` before any result derived from it is returned to a caller.
+    """
+    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]:
+    """Undo `_tile_batch`, splitting a tiled batch back into `n` standard-shaped batches."""
+    b = next(iter(batch.surf_vars.values())).shape[0] // 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/model/aurora.py

@@ -14,7 +14,7 @@
     apply_activation_checkpointing,
 )
 
-from aurora.batch import Batch
+from aurora.batch import Batch, _split_batch, _tile_batch
 from aurora.insolation import insolation
 from aurora.model.compat import (
     _adapt_checkpoint_air_pollution,
@@ -103,6 +103,7 @@ def __init__(
         clamp_at_first_step: bool = False,
         simulate_indexing_bug: bool = False,
         stochastic: bool = False,
+        num_ensemble_members: int = 1,
         use_updated_lead_time_embedding: bool = False,
         variable_lead_time: bool = False,
         rollout_input_clipping: Optional[dict[str, dict[str, Optional[float]]]] = None,
@@ -200,6 +201,17 @@ def __init__(
                 to the original implementation. Defaults to `False`.
             stochastic (bool, optional): If `True`, enable stochastic mode with noise injection.
                 Defaults to `False`.
+            num_ensemble_members (int, optional): Number of ensemble members to produce
+                *internally* on every call to :meth:`forward_ensemble`, as an alternative to
+                looping over separate :meth:`forward` calls and combining the results externally
+                yourself (which remains perfectly valid, e.g. if you need more control over how
+                members are seeded or combined). When set to a value greater than `1`, the batch
+                is tiled `num_ensemble_members` times internally and run through the model in a
+                single, fully-batched pass, which is far more efficient on a GPU than looping.
+                This is most useful in combination with `stochastic=True`, since every tiled copy
+                then receives independent noise. When greater than `1`, plain :meth:`forward`
+                raises, since it can only ever return a single `Batch`; use
+                :meth:`forward_ensemble` instead. Defaults to `1`, i.e. no internal ensembling.
             use_updated_lead_time_embedding (bool, optional): Whether to use the updated lead time
                 embedding with a minimum wavelength of 2 hours. Defaults to `False`.
             variable_lead_time (bool, optional): If `True`, use per-sample lead times passed
@@ -236,6 +248,10 @@ def __init__(
         self.output_only_surf_vars = output_only_surf_vars
         self.output_only_atmos_vars = output_only_atmos_vars
 
+        if num_ensemble_members < 1:
+            raise ValueError("`num_ensemble_members` must be at least `1`.")
+        self.num_ensemble_members = num_ensemble_members
+
         if self.surf_stats:
             warnings.warn(
                 f"The normalisation statics for the following surface-level variables are manually "
@@ -284,6 +300,14 @@ def __init__(
             use_updated_lead_time_embedding=use_updated_lead_time_embedding,
         )
 
+        if num_ensemble_members > 1 and not self.backbone.stochastic:
+            warnings.warn(
+                f"`num_ensemble_members={num_ensemble_members}` was requested, but `stochastic="
+                f"False`, so the model has no source of randomness. All ensemble members will be "
+                f"identical.",
+                stacklevel=2,
+            )
+
         self.decoder = Perceiver3DDecoder(
             surf_vars=surf_vars,
             atmos_vars=atmos_vars,
@@ -349,7 +373,47 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba
                 `variable_lead_time=True`. Ignored otherwise.
 
         Returns:
-            :class:`Batch`: Prediction for the batch.
+            :class:`Batch`: Prediction for `batch`.
+        """
+        if self.num_ensemble_members > 1:
+            raise RuntimeError(
+                f"This model was constructed with `num_ensemble_members="
+                f"{self.num_ensemble_members}`. Use `forward_ensemble` instead of `forward` to "
+                f"obtain all ensemble members."
+            )
+        return self._forward_impl(batch, lead_times)
+
+    def forward_ensemble(
+        self, batch: Batch, lead_times: Optional[torch.Tensor] = None
+    ) -> list[Batch]:
+        """Forward pass producing all `self.num_ensemble_members` ensemble members internally.
+
+        All members are computed internally as a single fused pass through the
+        encoder/backbone/decoder, rather than looping over separate `forward` calls and combining
+        the results externally yourself (which remains equally valid and unaffected). This is most
+        useful in combination with `stochastic=True`, since every internally-tiled copy then
+        receives independent noise.
+
+        Args:
+            batch (:class:`aurora.Batch`): Batch to run the model on.
+            lead_times (:class:`torch.Tensor`, optional): Per-sample lead times of shape
+                `(batch,)` in hours. Required when the model was configured with
+                `variable_lead_time=True`. Ignored otherwise.
+
+        Returns:
+            list[:class:`Batch`]: A list of `self.num_ensemble_members` standard-shaped `Batch`\\
+                s, one per ensemble member, each with the same batch dimension as `batch`.
+        """
+        pred = self._forward_impl(batch, lead_times)
+        return _split_batch(pred, self.num_ensemble_members)
+
+    def _forward_impl(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Batch:
+        """Shared implementation for `forward` and `forward_ensemble`.
+
+        Internally tiles `batch` by `self.num_ensemble_members` before running it through the
+        encoder/backbone/decoder as a single fused batch, when greater than `1`. The tiled batch
+        dimension is a private implementation detail: `forward` forbids it (see above) and
+        `forward_ensemble` splits it back apart before returning.
         """
         batch = self.batch_transform_hook(batch)
 
@@ -361,6 +425,11 @@ def forward(self, batch: Batch, lead_times: Optional[torch.Tensor] = None) -> Ba
         batch = batch.crop(patch_size=self.patch_size)
         batch = batch.to(p.device)
 
+        if self.num_ensemble_members > 1:
+            if lead_times is not None:
+                lead_times = lead_times.repeat(self.num_ensemble_members)
+            batch = _tile_batch(batch, self.num_ensemble_members)
+
         H, W = batch.spatial_shape
         patch_res = (
             self.encoder.latent_levels,

aurora/rollout.py

@@ -6,10 +6,10 @@
 
 import torch
 
-from aurora.batch import Batch
+from aurora.batch import Batch, _split_batch, _tile_batch
 from aurora.model.aurora import Aurora
 
-__all__ = ["rollout"]
+__all__ = ["rollout", "rollout_ensemble"]
 
 
 def _make_lead_time_tensor(batch: Batch, lead_time_hours: float) -> torch.Tensor:
@@ -148,3 +148,44 @@ 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,
+    fine_lead_times: Optional[Sequence[float]] = None,
+    use_noise_accumulation: bool = True,
+    apply_rollout_input_clipping: bool = True,
+) -> Generator[list[Batch], None, None]:
+    """Like `rollout`, but produces `model.num_ensemble_members` ensemble members internally on
+    every step, as a single fused pass, instead of `rollout` yielding one `Batch` per step.
+
+    All arguments are identical to `rollout`; see there for details.
+
+    Yields:
+        list[:class:`aurora.Batch`]: A list of `model.num_ensemble_members` standard-shaped
+            `Batch`\\ s after every (sub-)step, one per ensemble member; see
+            :meth:`aurora.Aurora.forward_ensemble`.
+    """
+    num_ensemble_members = model.num_ensemble_members
+
+    # Tile the batch once up front, then temporarily disable further expansion so it isn't
+    # repeated by `_forward_impl` on every step. The tiled representation is kept purely internal
+    # to this loop: every `pred` yielded to the caller is split back into standard-shaped batches.
+    batch = _tile_batch(batch, num_ensemble_members)
+    model.num_ensemble_members = 1
+    try:
+        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,
+        ):
+            yield _split_batch(pred, num_ensemble_members)
+    finally:
+        # Restore the model's ensemble configuration, whether the roll-out ran to completion or
+        # was abandoned early.
+        model.num_ensemble_members = num_ensemble_members