From 0f492cc4e0bd6a82b60cb3171b1179d892cacd1f Mon Sep 17 00:00:00 2001 From: priyald Date: Wed, 9 Sep 2026 16:51:51 -0500 Subject: [PATCH 1/5] Validate the grid, the history dimension and the latitudes and longitudes `Batch` documents the shapes of its variables but does not enforce them. Add a `__post_init__` that checks the three invariants not covered by #196: * every surface-level, static and atmospheric variable describes the same grid; * the surface-level and atmospheric variables agree on the history dimension; * the latitudes and longitudes match that grid, in both the vector and the matrix form. The latitudes and longitudes are checked in the encoder today, but with a bare `assert` that carries no message and is removed under `python -O`. The other two surface as shape errors from inside torch, which do not say which variable is wrong. Raising at construction is consistent with the existing latitude and longitude validation in `Metadata.__post_init__`. The spatial shape is read from the last two dimensions because the static variables are expanded to `(b, t, h, w)` inside the model and the decoder builds a batch before the history dimension is inserted. The history dimension is checked only on variables that still have one, for the same reason. An empty `surf_vars` or `atmos_vars` is left alone rather than treated as an error, since running without atmospheric variables is suggested in #176. --- aurora/batch.py | 96 +++++++++++++++++++++++++++++++++++++-------- tests/test_batch.py | 65 +++++++++++++++++++++++++++++- 2 files changed, 144 insertions(+), 17 deletions(-) diff --git a/aurora/batch.py b/aurora/batch.py index 2b3b491c..2ae07caf 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,66 @@ 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[-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." + ) + + surf = tuple(self.surf_vars.values()) + static = tuple(self.static_vars.values()) + atmos = tuple(self.atmos_vars.values()) + + # Every variable must describe the same grid. Compare the last two dimensions rather than + # the full shape: the static variables are `(h, w)` here but are expanded to `(b, t, h, w)` + # inside the model, and the decoder constructs a batch before the history dimension is + # inserted. + spatial_shape = _consistent( + (tuple(v.shape[-2:]) for v in surf + static + atmos if v.dim() >= 2), + "spatial shape", + "surface-level, static and atmospheric variables", + ) + + # The history dimension is only present on fully shaped variables, so check it only where + # it exists. + _consistent( + [v.shape[1] for v in surf if v.dim() == 4] + + [v.shape[1] for v in atmos if v.dim() == 5], + "history size", + "surface-level and atmospheric variables", + ) + + if spatial_shape is not None: + self._check_lat_lon(*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 + + # :meth:`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 +370,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 +422,25 @@ def split_batch(batch: Batch, n: int) -> list[Batch]: ] +def _consistent(values: Iterable, what: str, of_what: str) -> object | None: + """Check that `values` are all equal. + + Args: + values (iterable): Values to compare. + what (str): Description of the values, used in the error message. + of_what (str): Description of where the values came from, used in the error message. + + Returns: + object or None: The common value, or `None` if `values` is empty. + """ + unique = set(values) + if len(unique) > 1: + raise ValueError( + f"The {of_what} must all have the same {what}, but found {sorted(unique)}." + ) + return unique.pop() if unique else None + + 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 be43d058..173df463 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,64 @@ 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) + + +_LEVELS = (100, 250, 500, 850) + + +def _metadata(n_lat: int = 17, n_lon: int = 32, n_levels: int = 4) -> Metadata: + return Metadata( + lat=torch.linspace(90, -90, n_lat), + lon=torch.linspace(0, 360, n_lon + 1)[:-1], + time=(datetime(2020, 6, 1, 12, 0),), + atmos_levels=_LEVELS[:n_levels], + ) + + +def _batch(metadata: Metadata | None = None) -> Batch: + 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 if metadata is not None else _metadata(), + ) + + +@pytest.mark.parametrize("n_lat, n_lon", [(9, 32), (17, 16), (18, 32), (17, 33)]) +def test_mismatching_lat_lon_and_spatial_shape(n_lat: int, n_lon: int) -> None: + # The encoder only asserts this. That assertion carries no message and is removed under + # `python -O`, so the mismatch reaches the caller as a bare `AssertionError` at best. + with pytest.raises(ValueError, match="describe a grid of shape"): + _batch(_metadata(n_lat=n_lat, n_lon=n_lon)) + + +def test_matrix_lat_lon_is_accepted() -> None: + metadata = _metadata() + metadata.lat = metadata.lat[:, None].expand(17, 32) + metadata.lon = metadata.lon[None, :].expand(17, 32) + assert _batch(metadata).spatial_shape == (17, 32) + + +def test_mismatching_matrix_lat_lon() -> None: + metadata = _metadata(n_lat=18) + metadata.lat = metadata.lat[:, None].expand(18, 32) + metadata.lon = metadata.lon[None, :].expand(18, 32) + with pytest.raises(ValueError, match="latitude and longitude matrices"): + _batch(metadata) + + +def test_inconsistent_spatial_shape_between_variables() -> None: + with pytest.raises(ValueError, match="same spatial shape"): + dataclasses.replace(_batch(), static_vars={"lsm": torch.randn(9, 32)}) + + +def test_inconsistent_history_between_variables() -> None: + with pytest.raises(ValueError, match="same history size"): + dataclasses.replace(_batch(), atmos_vars={"z": torch.randn(1, 3, 4, 17, 32)}) + + +def test_derived_batches_stay_valid(test_input_output: tuple[Batch, SavedBatch]) -> None: + # The operations that construct new batches must not trip the validation. + batch, _ = test_input_output + for derived in (batch.to("cpu"), batch.crop(4), batch.regrid(0.45)): + assert derived.spatial_shape == next(iter(derived.surf_vars.values())).shape[-2:] From 2dea1aa6c640c02810cd46a6a4a803648372daae Mon Sep 17 00:00:00 2001 From: priyald Date: Wed, 28 Oct 2026 04:54:23 -0500 Subject: [PATCH 2/5] Assume the documented shapes in Batch validation --- aurora/batch.py | 47 +++++++++++------------------------------ aurora/model/aurora.py | 7 ------ aurora/model/decoder.py | 4 ++-- tests/test_batch.py | 4 ++-- 4 files changed, 16 insertions(+), 46 deletions(-) diff --git a/aurora/batch.py b/aurora/batch.py index 2ae07caf..3676ae15 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -88,7 +88,7 @@ class 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) + 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( @@ -101,31 +101,18 @@ def __post_init__(self): f"atmospheric variables have {c} pressure levels. These must be equal." ) - surf = tuple(self.surf_vars.values()) - static = tuple(self.static_vars.values()) - atmos = tuple(self.atmos_vars.values()) - - # Every variable must describe the same grid. Compare the last two dimensions rather than - # the full shape: the static variables are `(h, w)` here but are expanded to `(b, t, h, w)` - # inside the model, and the decoder constructs a batch before the history dimension is - # inserted. - spatial_shape = _consistent( - (tuple(v.shape[-2:]) for v in surf + static + atmos if v.dim() >= 2), - "spatial shape", - "surface-level, static and atmospheric variables", - ) + all_vars = [*self.surf_vars.values(), *self.static_vars.values(), *self.atmos_vars.values()] - # The history dimension is only present on fully shaped variables, so check it only where - # it exists. - _consistent( - [v.shape[1] for v in surf if v.dim() == 4] - + [v.shape[1] for v in atmos if v.dim() == 5], + # 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", ) - if spatial_shape is not None: - self._check_lat_lon(*spatial_shape) + 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.""" @@ -422,23 +409,13 @@ def split_batch(batch: Batch, n: int) -> list[Batch]: ] -def _consistent(values: Iterable, what: str, of_what: str) -> object | None: - """Check that `values` are all equal. - - Args: - values (iterable): Values to compare. - what (str): Description of the values, used in the error message. - of_what (str): Description of where the values came from, used in the error message. - - Returns: - object or None: The common value, or `None` if `values` is empty. - """ +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 ValueError( - f"The {of_what} must all have the same {what}, but found {sorted(unique)}." + raise AssertionError( + f"The {of_what} must all have the same {what}, but found `{sorted(unique)}`." ) - return unique.pop() if unique else None def _np(x: torch.Tensor) -> np.ndarray: 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/tests/test_batch.py b/tests/test_batch.py index 173df463..d57c1f4d 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -126,12 +126,12 @@ def test_mismatching_matrix_lat_lon() -> None: def test_inconsistent_spatial_shape_between_variables() -> None: - with pytest.raises(ValueError, match="same spatial shape"): + 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(ValueError, match="same history size"): + with pytest.raises(AssertionError, match="same history size"): dataclasses.replace(_batch(), atmos_vars={"z": torch.randn(1, 3, 4, 17, 32)}) From 8ad38a3615359616c1ae68984fc10739add011d0 Mon Sep 17 00:00:00 2001 From: priyald Date: Wed, 28 Oct 2026 04:54:34 -0500 Subject: [PATCH 3/5] Tidy the validation messages and comments --- aurora/batch.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/aurora/batch.py b/aurora/batch.py index 3676ae15..f8bf952d 100644 --- a/aurora/batch.py +++ b/aurora/batch.py @@ -118,19 +118,19 @@ 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 - # :meth:`Metadata.__post_init__` has already checked that the latitudes and longitudes are - # either both vectors or both matrices. + # `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)}." + 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)}." + 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 From 402155d7d135ccb93f72459440491e5ca32f3faa Mon Sep 17 00:00:00 2001 From: priyald Date: Wed, 28 Oct 2026 04:54:39 -0500 Subject: [PATCH 4/5] Remove the latitude and longitude assertion from the encoder --- aurora/model/encoder.py | 1 - 1 file changed, 1 deletion(-) 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") From 16202d9ad51da0a620684171c569f6c3787882b5 Mon Sep 17 00:00:00 2001 From: priyald Date: Wed, 28 Oct 2026 04:54:54 -0500 Subject: [PATCH 5/5] Trim the batch validation tests --- tests/test_batch.py | 55 +++++++++++++-------------------------------- 1 file changed, 16 insertions(+), 39 deletions(-) diff --git a/tests/test_batch.py b/tests/test_batch.py index d57c1f4d..2829944f 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -81,48 +81,32 @@ def test_batch_rejects_atmos_levels_mismatch(test_input_output: tuple[Batch, Sav dataclasses.replace(batch, metadata=bad_metadata) -_LEVELS = (100, 250, 500, 850) - - -def _metadata(n_lat: int = 17, n_lon: int = 32, n_levels: int = 4) -> Metadata: - return Metadata( - lat=torch.linspace(90, -90, n_lat), - lon=torch.linspace(0, 360, n_lon + 1)[:-1], - time=(datetime(2020, 6, 1, 12, 0),), - atmos_levels=_LEVELS[:n_levels], - ) - - -def _batch(metadata: Metadata | None = None) -> Batch: +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 if metadata is not None else _metadata(), + 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", [(9, 32), (17, 16), (18, 32), (17, 33)]) -def test_mismatching_lat_lon_and_spatial_shape(n_lat: int, n_lon: int) -> None: - # The encoder only asserts this. That assertion carries no message and is removed under - # `python -O`, so the mismatch reaches the caller as a bare `AssertionError` at best. - with pytest.raises(ValueError, match="describe a grid of shape"): - _batch(_metadata(n_lat=n_lat, n_lon=n_lon)) +@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: - metadata = _metadata() - metadata.lat = metadata.lat[:, None].expand(17, 32) - metadata.lon = metadata.lon[None, :].expand(17, 32) - assert _batch(metadata).spatial_shape == (17, 32) - - -def test_mismatching_matrix_lat_lon() -> None: - metadata = _metadata(n_lat=18) - metadata.lat = metadata.lat[:, None].expand(18, 32) - metadata.lon = metadata.lon[None, :].expand(18, 32) - with pytest.raises(ValueError, match="latitude and longitude matrices"): - _batch(metadata) + _batch(matrix=True) def test_inconsistent_spatial_shape_between_variables() -> None: @@ -133,10 +117,3 @@ def test_inconsistent_spatial_shape_between_variables() -> None: 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)}) - - -def test_derived_batches_stay_valid(test_input_output: tuple[Batch, SavedBatch]) -> None: - # The operations that construct new batches must not trip the validation. - batch, _ = test_input_output - for derived in (batch.to("cpu"), batch.crop(4), batch.regrid(0.45)): - assert derived.spatial_shape == next(iter(derived.surf_vars.values())).shape[-2:]