fix(batch): validate Batch metadata dimensions in __post_init__

aurora/batch.py

@@ -310,6 +310,21 @@ 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"`Metadata.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"`Metadata.atmos_levels` has length {len(self.metadata.atmos_levels)}, but the "
+                f"atmospheric variables have {c} pressure levels. These must be equal."
+            )
+
 
 def _np(x: torch.Tensor) -> np.ndarray:
     return x.detach().cpu().numpy()

tests/test_batch.py

@@ -1,8 +1,10 @@
 """Copyright (c) Microsoft Corporation. Licensed under the MIT license."""
 
+import dataclasses
 from pathlib import Path
 
 import numpy as np
+import pytest
 
 from tests.conftest import SavedBatch
 
@@ -57,3 +59,21 @@ def test_save_load(test_input_output: tuple[Batch, SavedBatch], tmp_path: Path)
     assert batch.metadata.time == batch_loaded.metadata.time
     assert batch.metadata.atmos_levels == batch_loaded.metadata.atmos_levels
     assert batch.metadata.rollout_step == batch_loaded.metadata.rollout_step
+
+
+def test_batch_rejects_metadata_time_length_mismatch(
+    test_input_output: tuple[Batch, SavedBatch],
+) -> None:
+    batch, _ = test_input_output
+    bad_metadata = dataclasses.replace(batch.metadata, time=batch.metadata.time * 2)
+    with pytest.raises(ValueError, match=r"`Metadata\.time`"):
+        dataclasses.replace(batch, metadata=bad_metadata)
+
+
+def test_batch_rejects_atmos_levels_mismatch(test_input_output: tuple[Batch, SavedBatch]) -> None:
+    batch, _ = test_input_output
+    bad_metadata = dataclasses.replace(
+        batch.metadata, atmos_levels=batch.metadata.atmos_levels[:-1]
+    )
+    with pytest.raises(ValueError, match=r"`Metadata\.atmos_levels`"):
+        dataclasses.replace(batch, metadata=bad_metadata)

tests/test_model.py

@@ -10,7 +10,7 @@
 import torch
 import torch.distributed as dist
 
-from tests.conftest import SavedBatch
+from tests.conftest import SavedBatch, SavedMetadata
 
 from aurora import Aurora, AuroraSmallPretrained, Batch
 
@@ -28,16 +28,21 @@ def test_aurora_small(aurora_small: Aurora, test_input_output: tuple[Batch, Save
     batch, test_output = test_input_output
 
     # Run the test with batch size two.
+    batch_metadata = dataclasses.replace(batch.metadata, time=batch.metadata.time * 2)
     batch = dataclasses.replace(
         batch,
         surf_vars={k: v.repeat(2, 1, 1, 1) for k, v in batch.surf_vars.items()},
         atmos_vars={k: v.repeat(2, 1, 1, 1, 1) for k, v in batch.atmos_vars.items()},
+        metadata=batch_metadata,
     )
     test_output = cast(SavedBatch, dict(test_output))  #  Copy before mutating.
     test_output["surf_vars"] = {k: v.repeat(2, axis=0) for k, v in test_output["surf_vars"].items()}
     test_output["atmos_vars"] = {
         k: v.repeat(2, axis=0) for k, v in test_output["atmos_vars"].items()
     }
+    test_metadata = cast(SavedMetadata, dict(test_output["metadata"]))
+    test_metadata["time"] = test_metadata["time"] * 2  # list * 2, not tuple
+    test_output["metadata"] = test_metadata
 
     with torch.inference_mode():
         pred = aurora_small.forward(batch)