[priyDe (priyald)]

Follow-up to #196, rebased onto main now that it has landed. The checks from both PRs sit in one
Batch.__post_init__.

#196 validates that metadata.time and metadata.atmos_levels match the data. This adds the
remaining shape invariants that Batch documents but does not enforce.

What it checks

1. Every surface-level, static and atmospheric variable describes the same grid.
2. The surface-level and atmospheric variables agree on the history dimension.
3. The latitudes and longitudes match that grid, in both the vector and the matrix form.

Why these three

The latitudes and longitudes were already checked in Perceiver3DEncoder.forward, but with

assert lat.shape[0] == H and lon.shape[-1] == W

which carries no message, fires deep inside the forward pass rather than where the batch was
built, and is removed entirely under python -O. Since the batch now validates the grid at
construction, this PR removes that assert. The other two currently surface as shape errors from
inside torch, for example Sizes of tensors must match except in dimension 2, which does not tell
the caller which variable is wrong.

Decoder change

The checks assume the documented shapes. The one batch that did not follow them was the decoder
output, which lacked the history dimension until Aurora.forward inserted it. The decoder now
inserts it directly, and the block in Aurora.forward is gone. The _post_decoder_hook
implementations already ran after that block, so they see the same shapes as before.

Tests

Added to tests/test_batch.py:

• latitudes or longitudes that do not match the grid raise, in vector and matrix form;
• matrix latitudes and longitudes are accepted;
• variables that disagree on the grid raise;
• variables that disagree on the history dimension raise.

Full suite passes apart from test_aurora_small, which fails for me at 0.04224444 on unmodified
main as well (likely the machine-dependent numerics in #169). ruff check and
ruff format --check are clean.

[Wessel (wesselb)]

Thanks for the PR, priyDe (@priyald). :)

I've left some Claude-assisted comments and suggestions.

[priyDe (priyald)]

Sorry for the slow reply on this! Rebased onto main and applied all your suggestions, including the decoder change and removing the assert in the encoder. Thanks for the detailed review.

[Wessel (wesselb)]

Suggested change

	{v: surf_preds[:, i, None] for i, v in enumerate(surf_vars)},
	# Insert a history dimension for the surf. and atmos variables.
	{v: surf_preds[:, i, None] for i, v in enumerate(surf_vars)},

[Wessel (wesselb)]

Left two more comments. Happy to merge after those. Thanks, priyDe (@priyald)! This looks really great. :)

[Wessel (wesselb)]

Document the fixed grid and rename.

Suggested change

	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)})
	def _make_batch(n_lat: int = 17, n_lon: int = 32, matrix: bool = False) -> Batch:
	"""Construct a batch whose variables have spatial shape `(17, 32)`.
	The latitudes and longitudes describe an `n_lat` by `n_lon` grid, so values other than the
	defaults give latitudes and longitudes that do not match the variables. If `matrix` is `True`,
	they are matrices rather than vectors.
	"""
	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"):
	_make_batch(n_lat, n_lon, matrix)
	def test_matrix_lat_lon_is_accepted() -> None:
	_make_batch(matrix=True)
	def test_inconsistent_spatial_shape_between_variables() -> None:
	with pytest.raises(AssertionError, match="same spatial shape"):
	dataclasses.replace(_make_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(_make_batch(), atmos_vars={"z": torch.randn(1, 3, 4, 17, 32)})

A caller needs to know that only the latitudes and longitudes follow n_lat and n_lon while the variables stay on a 17 by 32 grid. _make_batch matches the helper of the same name in tests/v1p5/_helpers.py.

Comment generated in collaboration with Claude. I take full responsibility for the content of this comment.