diff --git a/aurora/batch.py b/aurora/batch.py index 2b3b491c..f8bf952d 100644 --- a/aurora/batch.py +++ b/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() diff --git a/aurora/model/aurora.py b/aurora/model/aurora.py index 4dcb6877..baa58267 100644 --- a/aurora/model/aurora.py +++ b/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. diff --git a/aurora/model/decoder.py b/aurora/model/decoder.py index e63508fc..4f9d58bc 100644 --- a/aurora/model/decoder.py +++ b/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, diff --git a/aurora/model/encoder.py b/aurora/model/encoder.py index 197f471c..fe382587 100644 --- a/aurora/model/encoder.py +++ b/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") diff --git a/tests/test_batch.py b/tests/test_batch.py index be43d058..2829944f 100644 --- a/tests/test_batch.py +++ b/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)})