Validate the grid, the history dimension and the latitudes and longitudes

aurora/batch.py

@@ -4,7 +4,7 @@
 from datetime import datetime
 from functools import partial
 from pathlib import Path
-from typing import Callable, List
+from typing import Callable, Iterable, List
 
 import numpy as np
 import torch
@@ -86,6 +86,53 @@ class Batch:
     atmos_vars: dict[str, torch.Tensor]
     metadata: Metadata
 
+    def __post_init__(self):
+        b = next(iter(self.surf_vars.values())).shape[0]  # First dim. of (b, t, h, w)
+        c = next(iter(self.atmos_vars.values())).shape[2]  # Pressure-level dim. of (b, t, c, h, w)
+
+        if len(self.metadata.time) != b:
+            raise ValueError(
+                f"`time` has length {len(self.metadata.time)}, but the batch size "
+                f"is {b}. `time` must contain exactly one entry per batch element."
+            )
+        if len(self.metadata.atmos_levels) != c:
+            raise ValueError(
+                f"`atmos_levels` has length {len(self.metadata.atmos_levels)}, but the "
+                f"atmospheric variables have {c} pressure levels. These must be equal."
+            )
+
+        all_vars = [*self.surf_vars.values(), *self.static_vars.values(), *self.atmos_vars.values()]
+
+        # The model expands the static variables to `(b, t, h, w)`, so compare only the last two
+        # dimensions.
+        _assert_all_equal((tuple(v.shape[-2:]) for v in all_vars), "spatial shape", "variables")
+        _assert_all_equal(
+            (v.shape[1] for v in [*self.surf_vars.values(), *self.atmos_vars.values()]),
+            "history size",
+            "surface-level and atmospheric variables",
+        )
+
+        self._check_lat_lon(*self.spatial_shape)
+
+    def _check_lat_lon(self, h: int, w: int) -> None:
+        """Check that the latitudes and longitudes describe an `h` by `w` grid."""
+        lat, lon = self.metadata.lat, self.metadata.lon
+
+        # `Metadata.__post_init__` has already checked that the latitudes and longitudes are either
+        # both vectors or both matrices.
+        if lat.dim() == 1:
+            if lat.shape[0] != h or lon.shape[0] != w:
+                raise ValueError(
+                    f"The latitudes and longitudes describe a grid of shape "
+                    f"`{(lat.shape[0], lon.shape[0])}`, but the variables have spatial shape "
+                    f"`{(h, w)}`."
+                )
+        elif tuple(lat.shape) != (h, w) or tuple(lon.shape) != (h, w):
+            raise ValueError(
+                f"The latitude and longitude matrices have shapes `{tuple(lat.shape)}` and "
+                f"`{tuple(lon.shape)}`, but the variables have spatial shape `{(h, w)}`."
+            )
+
     @property
     def spatial_shape(self) -> tuple[int, int]:
         """Get the spatial shape from an arbitrary surface-level variable."""
@@ -310,21 +357,6 @@ def from_netcdf(cls, path: str | Path) -> "Batch":
             ),
         )
 
-    def __post_init__(self):
-        b = next(iter(self.surf_vars.values())).shape[0]  # First dim. of (b, t, h, w)
-        c = next(iter(self.atmos_vars.values())).shape[-3]  # Pressure-level dim. of (b, t, c, h, w)
-
-        if len(self.metadata.time) != b:
-            raise ValueError(
-                f"`time` has length {len(self.metadata.time)}, but the batch size "
-                f"is {b}. `time` must contain exactly one entry per batch element."
-            )
-        if len(self.metadata.atmos_levels) != c:
-            raise ValueError(
-                f"`atmos_levels` has length {len(self.metadata.atmos_levels)}, but the "
-                f"atmospheric variables have {c} pressure levels. These must be equal."
-            )
-
 
 def tile_batch(batch: Batch, n: int) -> Batch:
     """Tile `batch` along the batch dimension `n` times.
@@ -377,6 +409,15 @@ def split_batch(batch: Batch, n: int) -> list[Batch]:
     ]
 
 
+def _assert_all_equal(values: Iterable, what: str, of_what: str) -> None:
+    """Assert that all `values` are equal."""
+    unique = set(values)
+    if len(unique) > 1:
+        raise AssertionError(
+            f"The {of_what} must all have the same {what}, but found `{sorted(unique)}`."
+        )
+
+
 def _np(x: torch.Tensor) -> np.ndarray:
     return x.detach().cpu().numpy()
 

aurora/model/aurora.py

@@ -462,13 +462,6 @@ def forward(
             static_vars={k: v[0, 0] for k, v in batch.static_vars.items()},
         )
 
-        # Insert history dimension in prediction. The time should already be right.
-        pred = dataclasses.replace(
-            pred,
-            surf_vars={k: v[:, None] for k, v in pred.surf_vars.items()},
-            atmos_vars={k: v[:, None] for k, v in pred.atmos_vars.items()},
-        )
-
         pred = self._post_decoder_hook(batch, pred)
 
         # Clamp positive variables.

aurora/model/decoder.py

@@ -268,9 +268,9 @@ def forward(
         )
 
         return Batch(
-            {v: surf_preds[:, i] for i, v in enumerate(surf_vars)},
+            {v: surf_preds[:, i, None] for i, v in enumerate(surf_vars)},
             batch.static_vars,
-            {v: atmos_preds[:, i] for i, v in enumerate(atmos_vars)},
+            {v: atmos_preds[:, i, None] for i, v in enumerate(atmos_vars)},
             Metadata(
                 lat=lat,
                 lon=lon,

aurora/model/encoder.py

@@ -289,7 +289,6 @@ def forward(
         lat, lon = batch.metadata.lat, batch.metadata.lon
         check_lat_lon_dtype(lat, lon)
         lat, lon = lat.to(dtype=torch.float32), lon.to(dtype=torch.float32)
-        assert lat.shape[0] == H and lon.shape[-1] == W
 
         # Patch embed the surface level.
         x_surf = rearrange(x_surf, "b t v h w -> b v t h w")

tests/test_batch.py

@@ -1,14 +1,16 @@
 """Copyright (c) Microsoft Corporation. Licensed under the MIT license."""
 
 import dataclasses
+from datetime import datetime
 from pathlib import Path
 
 import numpy as np
 import pytest
+import torch
 
 from tests.conftest import SavedBatch
 
-from aurora import Batch
+from aurora import Batch, Metadata
 
 
 def test_interpolation(test_input_output: tuple[Batch, SavedBatch]) -> None:
@@ -77,3 +79,41 @@ def test_batch_rejects_atmos_levels_mismatch(test_input_output: tuple[Batch, Sav
     )
     with pytest.raises(ValueError, match=r"`atmos_levels`"):
         dataclasses.replace(batch, metadata=bad_metadata)
+
+
+def _batch(n_lat: int = 17, n_lon: int = 32, matrix: bool = False) -> Batch:
+    lat = torch.linspace(90, -90, n_lat)
+    lon = torch.linspace(0, 360, n_lon + 1)[:-1]
+    if matrix:
+        lat, lon = lat[:, None].expand(n_lat, n_lon), lon[None, :].expand(n_lat, n_lon)
+    return Batch(
+        surf_vars={"2t": torch.randn(1, 2, 17, 32)},
+        static_vars={"lsm": torch.randn(17, 32)},
+        atmos_vars={"z": torch.randn(1, 2, 4, 17, 32)},
+        metadata=Metadata(
+            lat=lat,
+            lon=lon,
+            time=(datetime(2020, 6, 1, 12, 0),),
+            atmos_levels=(100, 250, 500, 850),
+        ),
+    )
+
+
+@pytest.mark.parametrize("n_lat, n_lon, matrix", [(18, 32, False), (17, 33, False), (18, 32, True)])
+def test_mismatching_lat_lon(n_lat: int, n_lon: int, matrix: bool) -> None:
+    with pytest.raises(ValueError, match="but the variables have spatial shape"):
+        _batch(n_lat, n_lon, matrix)
+
+
+def test_matrix_lat_lon_is_accepted() -> None:
+    _batch(matrix=True)
+
+
+def test_inconsistent_spatial_shape_between_variables() -> None:
+    with pytest.raises(AssertionError, match="same spatial shape"):
+        dataclasses.replace(_batch(), static_vars={"lsm": torch.randn(9, 32)})
+
+
+def test_inconsistent_history_between_variables() -> None:
+    with pytest.raises(AssertionError, match="same history size"):
+        dataclasses.replace(_batch(), atmos_vars={"z": torch.randn(1, 3, 4, 17, 32)})