diff --git a/aurora/batch.py b/aurora/batch.py index d0616d7..3c82613 100644 --- a/aurora/batch.py +++ b/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() diff --git a/tests/test_batch.py b/tests/test_batch.py index c1b1f0c..7099b88 100644 --- a/tests/test_batch.py +++ b/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) diff --git a/tests/test_model.py b/tests/test_model.py index 09671c6..a6180ad 100644 --- a/tests/test_model.py +++ b/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)