diff --git a/src/ocean_emulators/config.py b/src/ocean_emulators/config.py index bee3c053..c5e7a6d7 100644 --- a/src/ocean_emulators/config.py +++ b/src/ocean_emulators/config.py @@ -386,48 +386,43 @@ class PerceiverConfig(BaseConfig): def build( self, - in_channels: int, - out_channels: int, - max_patch_size: tuple[int, int], + token_dim: int, implementation: PerceiverImpl, ) -> nn.Module: - """Build a regular Perceiver (used by the encoder).""" - # This is not really a "frequency" but a maximum of the width appears to be reasonable from looking at the code. - max_freq = max(*max_patch_size) + """Build a Perceiver that maps ``(B, N, token_dim)`` to ``(B, latent_dim)``. + Position encoding is the caller's responsibility — the encoder applies + a custom 2-D Fourier scheme before concatenating prognostic and + boundary tokens into the sequence. + """ if _use_flash(implementation): try: from flash_perceiver import Perceiver as FlashPerceiver # type: ignore except ImportError as e: raise _flash_import_error() from e - from einops.layers.torch import Rearrange - - # Flash perceiver expects (batch, seq_len, dim); naive handles - # (batch, ph, pw, dim) natively via input_axis=2. Bake the - # spatial-flatten into the module so callers don't need to care. - perceiver: nn.Module = nn.Sequential( - Rearrange("b ph pw v -> b (ph pw) v"), - FlashPerceiver( - latent_rotary_emb_dim=max_freq, - depth=self.depth, - input_dim=in_channels, - output_dim=out_channels, - output_mode="average", - latent_dim=self.latent_dim, - num_latents=self.num_latents, - use_flash_attn=True, - weight_tie_layers=True, - self_per_cross_attn=2, - ), + + return FlashPerceiver( + depth=self.depth, + input_dim=token_dim, + output_dim=self.latent_dim, + output_mode="average", + latent_dim=self.latent_dim, + num_latents=self.num_latents, + use_flash_attn=True, + weight_tie_layers=True, + self_per_cross_attn=2, ) elif _use_naive(implementation): - perceiver = NaivePerceiver( - num_freq_bands=4, - max_freq=max_freq, + # ``num_freq_bands`` / ``max_freq`` are required positional kwargs + # but unused when ``fourier_encode_data=False``. + return NaivePerceiver( depth=self.depth, - input_axis=2, - input_channels=in_channels, - num_classes=out_channels, + input_axis=1, + input_channels=token_dim, + fourier_encode_data=False, + num_freq_bands=0, + max_freq=1.0, + num_classes=self.latent_dim, latent_dim=self.latent_dim, num_latents=self.num_latents, weight_tie_layers=True, @@ -436,8 +431,6 @@ def build( else: raise ValueError(f"Unknown perceiver implementation: {implementation}.") - return perceiver - def build_io( self, in_channels: int, @@ -504,24 +497,32 @@ def _flash_import_error() -> ValueError: class EncoderConfig(BaseConfig): perceiver: PerceiverConfig = PerceiverConfig() + token_dim: int = Field( + default=256, + description=( + "Per-token dimension fed into the Perceiver after projecting " + "prognostic and boundary pixels to a common space." + ), + ) def build( self, - in_channels: int, + prog_channels: int, + boundary_channels: int, out_channels: int, patch_extent: tuple[float, float], - max_lat_size: int, - max_lon_size: int, + max_patch_size: tuple[int, int], implementation: PerceiverImpl, ) -> PerceiverEncoder: - max_patch_size = patch_from(patch_extent, max_lat_size, max_lon_size) return PerceiverEncoder( - in_channels=in_channels, + prog_channels=prog_channels, + boundary_channels=boundary_channels, out_channels=out_channels, + token_dim=self.token_dim, + latent_dim=self.perceiver.latent_dim, patch_extent=patch_extent, - perceiver=self.perceiver.build( - in_channels, out_channels, max_patch_size, implementation - ), + max_patch_size=max_patch_size, + perceiver=self.perceiver.build(self.token_dim, implementation), ) @@ -769,12 +770,6 @@ def build( assert len(self.patch_extent) == 2, "patch_extent must be a pair of floats." extent = self.patch_extent[0], self.patch_extent[1] - all_grid_sizes = [s.grid_size for s in srcs] - max_lat_size, max_lon_size = ( - max(g[0] for g in all_grid_sizes), - max(g[1] for g in all_grid_sizes), - ) - impl = self.perceiver_implementation if _use_flash(impl) and not self.use_bfloat16: raise ValueError( @@ -782,15 +777,25 @@ def build( "Please set `use_bfloat16=True` or `perceiver_implementation='naive'`." ) - in_channels = prog_channels + boundary_channels - total_in_channels = in_channels + (3 if self.add_3d_coordinates else 0) + # 3D coordinates are appended to the prognostic stream only. + encoder_prog_channels = prog_channels + (3 if self.add_3d_coordinates else 0) + + # Compute the largest patch size (in pixels) across all data source + # resolutions. Used to set the Perceiver's Fourier frequency range. + all_grid_sizes = [s.grid_size for s in srcs] + max_ph, max_pw = 0, 0 + for grid_h, grid_w in all_grid_sizes: + ph, pw = patch_from(extent, grid_h, grid_w) + max_ph = max(max_ph, ph) + max_pw = max(max_pw, pw) + max_patch_size = (max_ph, max_pw) encoder = self.encoder.build( - total_in_channels, + encoder_prog_channels, + boundary_channels, self.embedding_dim, extent, - max_lat_size, - max_lon_size, + max_patch_size, impl, ) processor = self.processor.build( @@ -806,8 +811,9 @@ def build( ) add_3d_coordinates = Concat3dCoordinates() if self.add_3d_coordinates else None + # TODO(alxmrs): `in_channels` isn't used anywhere :/ Consider removing from base model. return FOMO( - in_channels=total_in_channels, + in_channels=encoder_prog_channels + boundary_channels, out_channels=out_channels, pred_residuals=self.pred_residuals, last_kernel_size=self.last_kernel_size, diff --git a/src/ocean_emulators/models/fomo.py b/src/ocean_emulators/models/fomo.py index dd2570dd..1a4c1d2b 100644 --- a/src/ocean_emulators/models/fomo.py +++ b/src/ocean_emulators/models/fomo.py @@ -89,15 +89,14 @@ def __init__( def forward_once( self, prognostic: Prognostic, boundary: Boundary, ctx: GridContext ) -> Prognostic: - # Prognostic and boundary are carried as separate tensors through the - # data pipeline, but this encoder still expects a single concatenated - # input. The dual-perceiver encoder that fuses them at the token level - # (enabling cross-resolution) lands in a follow-up PR. - fts = torch.cat((prognostic, boundary), dim=1) with autocast(enabled=self.use_bfloat16, dtype=torch.bfloat16): + # 3D coordinates describe the grid the prognostic tensor sits on. if self.maybe_add_3d_coordinates is not None: - fts = self.maybe_add_3d_coordinates(fts, ctx.input_resolution_cpu) - fts = self.encoder(fts, ctx.input_resolution_cpu) + prognostic = self.maybe_add_3d_coordinates( + prognostic, ctx.input_resolution_cpu + ) + + fts = self.encoder(prognostic, boundary, ctx) fts = self.processor(fts) # TODO(alxmrs): When the output resolution differs from the input (i.e. in a "mix" schedule), we cannot use diff --git a/src/ocean_emulators/models/modules/encoder.py b/src/ocean_emulators/models/modules/encoder.py index c40c2cf0..f4a9aea2 100644 --- a/src/ocean_emulators/models/modules/encoder.py +++ b/src/ocean_emulators/models/modules/encoder.py @@ -2,15 +2,23 @@ # - https://github.com/microsoft/aurora/blob/main/aurora/model/patchembed.py # - https://github.com/microsoft/aurora/blob/main/aurora/model/encoder.py # - https://github.com/lucidrains/vit-pytorch +# - https://github.com/lucidrains/perceiver-pytorch (intra-patch Fourier encoding) + import torch from aurora.model.fourier import pos_expansion, scale_expansion from aurora.model.posencoding import pos_scale_enc from einops import rearrange from jaxtyping import Float +from perceiver_pytorch.perceiver_pytorch import fourier_encode from torch import nn -from ocean_emulators.constants import Input, Lat, Lon +from ocean_emulators.constants import Boundary, Prognostic +from ocean_emulators.utils.ctx import GridContext + +# Stream-type indices for the learned stream embedding. +_PROG_STREAM = 0 +_BOUNDARY_STREAM = 1 def patch_from( @@ -27,27 +35,62 @@ def patch_from( return patch_h, patch_w -class PerceiverEncoder(nn.Module): - """A perceiver-based encoder for Samudra's flattened data (a whole column of the ocean, with history). +def _intra_patch_fourier( + ph: int, + pw: int, + num_bands: int, + max_freq: float, + device: torch.device, + dtype: torch.dtype, +) -> Float[torch.Tensor, "n d"]: + """2-D Fourier positional encoding for one ``ph x pw`` patch. + + Mirrors what ``perceiver_pytorch`` does internally when ``input_axis=2`` + and ``fourier_encode_data=True``: each pixel's normalized position in + ``[-1, 1]^2`` is encoded with ``num_bands`` log-spaced frequencies. - We adopt some of Aurora's positional encodings[1], which uses log-spaced fourier features with geometry-informed - wavelengths. These encode 2d positions (the average latitude and longitude of each patch) as well as grid cell area - (measured in km^2) for each token before it enters the processor. + Returns ``(ph * pw, 2 * (2 * num_bands + 1))``: row-major flatten over + the patch, then concatenated ``[sin, cos, raw]`` features per axis. + """ + y = torch.linspace(-1.0, 1.0, ph, device=device, dtype=dtype) + x = torch.linspace(-1.0, 1.0, pw, device=device, dtype=dtype) + pos = torch.stack(torch.meshgrid(y, x, indexing="ij"), dim=-1) + enc = fourier_encode(pos, max_freq, num_bands) + return rearrange(enc, "ph pw a d -> (ph pw) (a d)") - > Note: We assume that data along the lat/lon coordinates are positioned at the center of each grid point! Please - > ensure this is the case at the data processing time. - This encoder is designed to make the same number of patches with the same spatial extents across different scales - of input data (input data may vary in resolution of lat/lng grid). To accomplish this with a single perceiver model, - our `forward` call requires supplementary information: the resolution (a pair of Lat/Lon tensors), which is used to - make consistent positional encodings for patches across different scales. While higher resolution scales will - contain more data per patch, the patch will refer to the same physical area on Earth as all other scales. +class PerceiverEncoder(nn.Module): + """Single-perceiver encoder over a concatenated prog + boundary sequence. + + Each input stream is patchified in 2-D using the same ``patch_extent`` + in degrees, so both produce the same patch grid regardless of native + resolution. Within each patch: + + * Prog and boundary pixels are flattened to 1-D token sequences and + linearly projected to a common ``token_dim``. + * A 2-D Fourier positional encoding (in normalized patch coordinates + ``[-1, 1]^2``) is added to every token; the same projection is + applied to both streams, so geometrically equivalent positions + share the same feature. + * A learned stream-type embedding is added so the Perceiver can tell + prog tokens apart from boundary tokens. + + The two token sequences are concatenated and fed to a single Perceiver, + which pools each patch's sequence to a latent vector. We project that + to ``out_channels`` and add Aurora-style patch-level positional and + scale encodings [1] (computed on the prognostic grid). Args: - in_channels (int): the number of input channels (roughly: time x variable x (surface + depths)). - out_channels (int): size of the latent dimension (aka, the embedding dimension). - patch_extent (tuple[float, float]): spatial extent of each patch measured in degrees of lat/lon. - perceiver (nn.Module): the perceiver module implementation to use. + prog_channels: Number of prognostic input channels. + boundary_channels: Number of boundary input channels. + out_channels: Final embedding dimension produced by this encoder. + token_dim: Per-token dimension fed into the Perceiver. + latent_dim: Pooled output dimension of the Perceiver. + patch_extent: Spatial extent of each patch in degrees ``(lat, lon)``. + max_patch_size: Largest ``(ph, pw)`` across input sources, used to + set the intra-patch Fourier max frequency (Nyquist-style). + perceiver: Perceiver mapping ``(B, N, token_dim) -> (B, latent_dim)``. + num_freq_bands: Fourier frequency bands for intra-patch encoding. References: [1]: https://ar5iv.labs.arxiv.org/html/2405.13063#A2.SS4 @@ -56,59 +99,142 @@ class PerceiverEncoder(nn.Module): # TODO(alxmrs): Implement gradient checkpointing def __init__( self, - in_channels: int, + prog_channels: int, + boundary_channels: int, out_channels: int, + token_dim: int, + latent_dim: int, patch_extent: tuple[float, float], + max_patch_size: tuple[int, int], perceiver: nn.Module, + num_freq_bands: int = 4, ) -> None: super().__init__() - self.in_channels = in_channels - self.out_channels: int = out_channels # aka, `embed_dim`. + self.prog_channels = prog_channels + self.boundary_channels = boundary_channels + self.out_channels: int = out_channels + self.token_dim = token_dim self.patch_extent = patch_extent + self.num_freq_bands = num_freq_bands + # Nyquist-style cap on the intra-patch Fourier features, mirroring + # the previous behavior when the Perceiver did its own Fourier. + self.max_intra_patch_freq = float(max(*max_patch_size)) + + self.prog_proj = nn.Linear(prog_channels, token_dim) + self.boundary_proj = nn.Linear(boundary_channels, token_dim) + + fourier_channels = 2 * (2 * num_freq_bands + 1) + self.pos_proj = nn.Linear(fourier_channels, token_dim) + + # 0 -> prog, 1 -> boundary. + self.stream_embed = nn.Embedding(2, token_dim) + self.perceiver = perceiver + + self.fusion_proj = nn.Linear(latent_dim, out_channels) + # TODO(#451): The input to these position and scale linear units could be a hparam. self.pos_embed = nn.Linear(self.out_channels, self.out_channels) self.scale_embed = nn.Linear(self.out_channels, self.out_channels) - def forward( - self, x: Input, resolution: tuple[Lat, Lon] - ) -> Float[torch.Tensor, "batch {self.embed_dim} h w"]: - _, V, H, W = x.shape - lat, lon = resolution - patch_h, patch_w = patch_from(self.patch_extent, H, W) - # V is a cross product of variable, level (encoded in vars), and time (has history). - assert V == self.in_channels - # Ensure patch_size is appropriate for the data. - assert H % patch_h == 0, f"{H} % {patch_h} != 0." - assert W % patch_w == 0, f"{W} % {patch_w} != 0." - - # Perceiver experiment ideas: - # 1. leave it as it is: treating each pixel as a token -- i.e. all channels (includes depths) per pixel - # 2. change to original plan, where each float is its own token - # 3. Add a third dim -- ph pw d v -- so each spatial position is a token - x = rearrange( - x, - "b v (h ph) (w pw) -> (b h w) ph pw v", - ph=patch_h, - pw=patch_w, + def _patchify_params( + self, shape: torch.Size, expected_channels: int + ) -> tuple[int, int, int, int]: + """Validate channels and compute patch / latent-grid dims. + + Returns ``(ph, pw, lat_h, lat_w)``. + """ + _, v, h, w = shape + assert v == expected_channels, ( + f"Expected {expected_channels} channels, got {v}." ) - # NB(alxmrs): This is includes a mean and LayerNorm before linear projection! - x = self.perceiver(x) # (B_H_W, ..., V) -> (B_H_W, out_channels) + ph, pw = patch_from(self.patch_extent, h, w) + assert h % ph == 0, f"{h} % {ph} != 0." + assert w % pw == 0, f"{w} % {pw} != 0." + return ph, pw, h // ph, w // pw - # Make `x` amenable to adding position + scale encoding - x = rearrange( + def _tokenize( + self, + x: torch.Tensor, + proj: nn.Linear, + ph: int, + pw: int, + stream_id: int, + ) -> Float[torch.Tensor, "bp n d"]: + """Flatten patches to 1-D, project, add intra-patch Fourier + stream embed. + + Returns ``(B * lat_h * lat_w, ph * pw, token_dim)``. + """ + tokens = rearrange( x, - "(b h w) l -> b (h w) l ", - h=(H // patch_h), - w=(W // patch_w), + "b v (h ph) (w pw) -> (b h w) (ph pw) v", + ph=ph, + pw=pw, + ) + # Per-pixel channel mix into a common token_dim — equivalent to a 1x1 + # conv on each patch, and analogous to how the lucidrains Perceiver + # lifts raw channels before adding its positional features. Lets us + # concat prog and boundary (different V) and add the additive Fourier + # / stream signals at a meaningful scale. + tokens = proj(tokens) + + pos_enc = _intra_patch_fourier( + ph, + pw, + self.num_freq_bands, + self.max_intra_patch_freq, + tokens.device, + tokens.dtype, + ) + # TODO(alxmrs): Should we make a pos_proj for each stream, or use a shared projection? + tokens = tokens + self.pos_proj(pos_enc).unsqueeze(0) + + stream_idx = torch.tensor(stream_id, device=tokens.device) + tokens = tokens + self.stream_embed(stream_idx) + + return tokens + + def forward( + self, + prog: Prognostic, + boundary: Boundary, + ctx: GridContext, + ) -> Float[torch.Tensor, "batch {self.out_channels} h w"]: + prog_ph, prog_pw, lat_h, lat_w = self._patchify_params( + prog.shape, self.prog_channels + ) + b_ph, b_pw, b_lat_h, b_lat_w = self._patchify_params( + boundary.shape, self.boundary_channels + ) + assert lat_h == b_lat_h and lat_w == b_lat_w, ( + f"Latent grid mismatch: prog ({lat_h}, {lat_w}) vs " + f"boundary ({b_lat_h}, {b_lat_w}). Check that patch_extent " + f"divides both grids evenly." + ) + + prog_tokens = self._tokenize( + prog, self.prog_proj, prog_ph, prog_pw, _PROG_STREAM + ) + boundary_tokens = self._tokenize( + boundary, self.boundary_proj, b_ph, b_pw, _BOUNDARY_STREAM ) - # Calculate and add positional + scale encoding + # Concat prog + boundary tokens along the sequence dim. The Perceiver + # cross-attends over the unified sequence, so prog and boundary mix + # inside the latent set rather than being fused after the fact. + seq = torch.cat([prog_tokens, boundary_tokens], dim=1) + + pooled = self.perceiver(seq) # (B*lat_h*lat_w, latent_dim) + x = self.fusion_proj(pooled) # (B*lat_h*lat_w, out_channels) + + # --- Patch-level positional + scale encoding (Aurora) --- + x = rearrange(x, "(b h w) l -> b (h w) l", h=lat_h, w=lat_w) + lat, lon = ctx.input_resolution_cpu pos_encode, scale_encode = pos_scale_enc( - self.out_channels, # aka "embed_dim" + self.out_channels, lat, lon, - (patch_h, patch_w), + (prog_ph, prog_pw), # TODO(#452): Pos and scale wavelengths range all the way to the whole Earth by default; we could probably # better tune these for our Oceans modeling use case. pos_expansion=pos_expansion, @@ -122,12 +248,6 @@ def forward( ).unsqueeze(0) x = x + pos_encoding + scale_encoding - # Unpack spatial channels, move channel dimension to correct location. - x = rearrange( - x, - "b (h w) l -> b l h w", - h=(H // patch_h), - w=(W // patch_w), - ) + x = rearrange(x, "b (h w) l -> b l h w", h=lat_h, w=lat_w) return x diff --git a/tests/test_decoder.py b/tests/test_decoder.py index 24251a8b..7184fae7 100644 --- a/tests/test_decoder.py +++ b/tests/test_decoder.py @@ -1,9 +1,9 @@ import pytest import torch from perceiver_pytorch.perceiver_io import PerceiverIO -from test_encoder import make_resolution # type: ignore +from test_encoder import make_ctx, make_encoder, make_resolution # type: ignore -from ocean_emulators.models.modules import PerceiverDecoder, PerceiverEncoder +from ocean_emulators.models.modules import PerceiverDecoder # Small values for fast tests. LATENT_DIM = 8 @@ -21,24 +21,6 @@ H, W = 8, 16 -def make_perceiver_encoder(in_channels, out_channels, *, num_latents=2): - """Build a regular Perceiver for the encoder (uses mean-pooling).""" - from perceiver_pytorch import Perceiver - - return Perceiver( - num_freq_bands=4, - max_freq=1.0, - depth=2, - input_axis=2, - input_channels=in_channels, - latent_dim=3, - num_latents=num_latents, - num_classes=out_channels, - weight_tie_layers=True, - self_per_cross_attn=2, - ) - - def make_decoder_perceiver_io(in_channels=IN_CHANNELS, out_channels=OUT_CHANNELS): """Build a PerceiverIO for the decoder.""" return PerceiverIO( @@ -119,29 +101,30 @@ def make_decoder_with_shared_weights( def test_roundtrip(): H_rt, W_rt = 4, 8 - x = torch.randn(3, 10, H_rt, W_rt) - - patch_embed = PerceiverEncoder( - in_channels=10, - out_channels=4, + embed_dim = 4 + prog = torch.randn(3, 7, H_rt, W_rt) + boundary = torch.randn(3, 3, H_rt, W_rt) + + patch_embed = make_encoder( + prog_channels=7, + boundary_channels=3, + out_channels=embed_dim, patch_extent=(180, 180), - perceiver=make_perceiver_encoder(10, 4), ) - patches = patch_embed(x, make_resolution(x)) - res = make_resolution(x) + patches = patch_embed(prog, boundary, make_ctx(prog)) decode = PerceiverDecoder( - in_channels=4, + in_channels=embed_dim, out_channels=10, patch_extent=(180, 180), queries_dim=QUERIES_DIM, - perceiver_io=make_decoder_perceiver_io(4, 10), + perceiver_io=make_decoder_perceiver_io(embed_dim, 10), window_patches=None, context_patches=None, ) - y_hat = decode(patches, res) + y_hat = decode(patches, make_resolution(prog)) assert y_hat.shape == (3, 10, H_rt, W_rt), ( f"Decoder should produce full-resolution output, got {y_hat.shape}." diff --git a/tests/test_encoder.py b/tests/test_encoder.py index 1b983cb7..2a4bbb30 100644 --- a/tests/test_encoder.py +++ b/tests/test_encoder.py @@ -3,92 +3,123 @@ from ocean_emulators.constants import Lat, Lon from ocean_emulators.models.modules.encoder import PerceiverEncoder, patch_from +from ocean_emulators.utils.ctx import GridContext +LATENT_DIM = 4 +TOKEN_DIM = 8 -def make_perceiver(in_channels, out_channels, *, num_latents=2, input_axis=2): + +def make_perceiver(*, num_latents=2): + """Build a 1-D Perceiver matching the encoder's expected sequence shape.""" return Perceiver( - num_freq_bands=4, + # num_freq_bands / max_freq are required positional kwargs but unused + # when fourier_encode_data=False. + num_freq_bands=0, max_freq=1.0, depth=2, - input_axis=input_axis, - input_channels=in_channels, - latent_dim=3, + input_axis=1, + input_channels=TOKEN_DIM, + fourier_encode_data=False, + latent_dim=LATENT_DIM, num_latents=num_latents, - num_classes=out_channels, + num_classes=LATENT_DIM, weight_tie_layers=True, self_per_cross_attn=2, ) +def make_encoder( + prog_channels, + boundary_channels, + out_channels, + patch_extent, + max_patch_size=(8, 8), +): + return PerceiverEncoder( + prog_channels=prog_channels, + boundary_channels=boundary_channels, + out_channels=out_channels, + token_dim=TOKEN_DIM, + latent_dim=LATENT_DIM, + patch_extent=patch_extent, + max_patch_size=max_patch_size, + perceiver=make_perceiver(), + ) + + def make_resolution(x: torch.Tensor) -> tuple[Lat, Lon]: lat = torch.linspace(start=-90, end=90, steps=x.shape[-2]) lon = torch.linspace(start=0, end=360, steps=x.shape[-1]) return lat, lon -def test_makes_patches(): - x = torch.randn(3, 10, 4, 8) - - patch_embed = PerceiverEncoder( - in_channels=10, - out_channels=4, - patch_extent=(180, 180), - perceiver=make_perceiver(10, 4), +def make_ctx(prog: torch.Tensor) -> GridContext: + """A minimal GridContext for encoder/FOMO tests; label_mask is a placeholder.""" + res = make_resolution(prog) + mask = torch.ones(prog.shape[1], prog.shape[2], prog.shape[3], dtype=torch.bool) + return GridContext( + label_mask=mask, + input_resolution_cpu=res, + output_resolution_cpu=res, ) - patches = patch_embed(x, make_resolution(x)) - assert patches.shape == (3, 4, 1, 2) +def test_makes_patches(): + prog = torch.randn(3, 7, 4, 8) + boundary = torch.randn(3, 3, 4, 8) + embed_dim = 4 + encoder = make_encoder(7, 3, embed_dim, (180, 180)) + patches = encoder(prog, boundary, make_ctx(prog)) -def test_makes_rectangular_patches(): - x = torch.randn(1, 10, 4, 8) + assert patches.shape == (3, embed_dim, 1, 2) - patch_embed = PerceiverEncoder( - in_channels=10, - out_channels=4, - patch_extent=(180, 90), - perceiver=make_perceiver(10, 4), - ) - patches = patch_embed(x, make_resolution(x)) +def test_makes_rectangular_patches(): + prog = torch.randn(1, 7, 4, 8) + boundary = torch.randn(1, 3, 4, 8) + embed_dim = 4 - assert patches.shape == ( - 1, - 4, - 1, - 4, - ) + encoder = make_encoder(7, 3, embed_dim, (180, 90)) + patches = encoder(prog, boundary, make_ctx(prog)) + assert patches.shape == (1, embed_dim, 1, 4) -def test_makes_patches__high_res(): - x = torch.randn(1, 10, 14, 21) - patch_embed = PerceiverEncoder( - in_channels=10, - out_channels=4, - patch_extent=(90.0, 120.0), - perceiver=make_perceiver(10, 4), - ) +def test_makes_patches__high_res(): + prog = torch.randn(1, 7, 14, 21) + boundary = torch.randn(1, 3, 14, 21) + embed_dim = 4 - patches = patch_embed(x, make_resolution(x)) + encoder = make_encoder(7, 3, embed_dim, (90.0, 120.0)) + patches = encoder(prog, boundary, make_ctx(prog)) - assert patches.shape == (1, 4, 2, 3) + assert patches.shape == (1, embed_dim, 2, 3) def test_makes_patches__more_variables(): - x = torch.randn(1, 20, 4, 8) + prog = torch.randn(1, 17, 4, 8) + boundary = torch.randn(1, 3, 4, 8) + embed_dim = 4 - patch_embed = PerceiverEncoder( - in_channels=20, - out_channels=4, - patch_extent=(180, 180), - perceiver=make_perceiver(20, 4), - ) + encoder = make_encoder(17, 3, embed_dim, (180, 180)) + patches = encoder(prog, boundary, make_ctx(prog)) + + assert patches.shape == (1, embed_dim, 1, 2) + + +def test_cross_resolution_patches(): + """Prog and boundary at different resolutions still produce the same patch grid.""" + # Prog at 1/2 deg (8x16), boundary at 1 deg (4x8), patches of 90 degrees. + prog = torch.randn(1, 7, 8, 16) + boundary = torch.randn(1, 3, 4, 8) + embed_dim = 4 - patches = patch_embed(x, make_resolution(x)) + encoder = make_encoder(7, 3, embed_dim, (90.0, 90.0)) + patches = encoder(prog, boundary, make_ctx(prog)) - assert patches.shape == (1, 4, 1, 2) + # 8 / patch_h(=4) = 2, 16 / patch_w(=4) = 4 + assert patches.shape == (1, embed_dim, 2, 4) def test_patch_from__full_globe(): diff --git a/tests/test_fomo_cross_resolution.py b/tests/test_fomo_cross_resolution.py new file mode 100644 index 00000000..ff87e664 --- /dev/null +++ b/tests/test_fomo_cross_resolution.py @@ -0,0 +1,189 @@ +"""End-to-end test: FOMO.forward_once with mismatched prog/boundary resolutions.""" + +import pytest +import torch +from perceiver_pytorch import Perceiver +from perceiver_pytorch.perceiver_io import PerceiverIO + +from ocean_emulators.models.fomo import FOMO +from ocean_emulators.models.modules import ( + PerceiverDecoder, + PerceiverEncoder, + UNetBackbone, +) +from ocean_emulators.models.modules.augment_input import Concat3dCoordinates +from ocean_emulators.models.modules.blocks import ( + BilinearUpsample, + ConvNeXtBlock, + CoreBlock, +) +from ocean_emulators.utils.ctx import GridContext + +LATENT_DIM = 4 +TOKEN_DIM = 8 +EMBED_DIM = 8 +QUERIES_DIM = 16 +PATCH_EXTENT = (90.0, 90.0) + + +def _create_block( + in_channels: int, + out_channels: int, + dilation: int, + n_layers: int, + pad: str, + checkpoint_simple: bool, +) -> CoreBlock: + return ConvNeXtBlock( + in_channels=in_channels, + out_channels=out_channels, + dilation=dilation, + n_layers=n_layers, + pad=pad, + checkpoint_simple=checkpoint_simple, + ) + + +def _create_upsample(in_channels: int, out_channels: int) -> BilinearUpsample: + return BilinearUpsample() + + +def _make_perceiver() -> Perceiver: + return Perceiver( + # num_freq_bands / max_freq are required positional kwargs but unused + # when fourier_encode_data=False. + num_freq_bands=0, + max_freq=1.0, + depth=2, + input_axis=1, + input_channels=TOKEN_DIM, + fourier_encode_data=False, + latent_dim=LATENT_DIM, + num_latents=2, + num_classes=LATENT_DIM, + weight_tie_layers=True, + self_per_cross_attn=2, + ) + + +def _make_fomo( + prog_channels: int, + boundary_channels: int, + out_channels: int, + add_3d_coordinates: bool = False, +) -> FOMO: + encoder_prog_channels = prog_channels + (3 if add_3d_coordinates else 0) + encoder = PerceiverEncoder( + prog_channels=encoder_prog_channels, + boundary_channels=boundary_channels, + out_channels=EMBED_DIM, + token_dim=TOKEN_DIM, + latent_dim=LATENT_DIM, + patch_extent=PATCH_EXTENT, + max_patch_size=(4, 4), + perceiver=_make_perceiver(), + ) + processor = UNetBackbone( + in_channels=EMBED_DIM, + ch_width=[EMBED_DIM], + dilation=[1], + n_layers=[1], + pad="circular", + create_block=_create_block, + downsampling_block=torch.nn.Identity(), + create_upsampling_block=_create_upsample, + checkpointing=None, + ) + perceiver_io = PerceiverIO( + depth=2, + dim=EMBED_DIM, + queries_dim=QUERIES_DIM, + logits_dim=out_channels, + num_latents=4, + latent_dim=LATENT_DIM, + weight_tie_layers=True, + decoder_ff=True, + ) + decoder = PerceiverDecoder( + in_channels=EMBED_DIM, + out_channels=out_channels, + patch_extent=PATCH_EXTENT, + queries_dim=QUERIES_DIM, + perceiver_io=perceiver_io, + window_patches=None, + context_patches=None, + ) + return FOMO( + in_channels=prog_channels + boundary_channels, + out_channels=out_channels, + pred_residuals=False, + last_kernel_size=1, + pad="circular", + add_3d_coordinates=Concat3dCoordinates() if add_3d_coordinates else None, + encoder=encoder, + processor=processor, + decoder=decoder, + hist=0, + checkpointing=None, + gradient_detach_interval=0, + use_bfloat16=False, + ) + + +def _make_resolution(h: int, w: int) -> tuple[torch.Tensor, torch.Tensor]: + return torch.linspace(-90, 90, h), torch.linspace(0, 360, w) + + +def test_fomo_forward_once_cross_resolution(): + """FOMO produces output at prognostic resolution when boundary resolution differs.""" + prog_channels, boundary_channels, out_channels = 7, 3, 7 + # Prognostic: 1/4 degree (8x16), Boundary: 1 degree (2x4) + prog_h, prog_w = 8, 16 + boundary_h, boundary_w = 2, 4 + + model = _make_fomo(prog_channels, boundary_channels, out_channels) + prog = torch.randn(2, prog_channels, prog_h, prog_w) + boundary = torch.randn(2, boundary_channels, boundary_h, boundary_w) + + prog_res = _make_resolution(prog_h, prog_w) + ctx = GridContext( + label_mask=torch.ones(out_channels, prog_h, prog_w, dtype=torch.bool), + input_resolution_cpu=prog_res, + output_resolution_cpu=prog_res, + ) + + output = model.forward_once(prog, boundary, ctx) + + assert output.shape == (2, out_channels, prog_h, prog_w) + assert torch.isfinite(output).all(), "Output contains NaN or Inf." + + +@pytest.mark.parametrize("add_3d_coordinates", [False, True]) +def test_fomo_forward_once_mix_schedule(add_3d_coordinates: bool): + """Mix schedule: prog at 1° input grid, output at ¼° grid.""" + prog_channels, boundary_channels, out_channels = 7, 3, 7 + # Input (prognostic) at 1 degree (4x8), output at 1/4 degree (16x32) + input_h, input_w = 4, 8 + output_h, output_w = 16, 32 + + model = _make_fomo( + prog_channels, + boundary_channels, + out_channels, + add_3d_coordinates=add_3d_coordinates, + ) + prog = torch.randn(1, prog_channels, input_h, input_w) + boundary = torch.randn(1, boundary_channels, input_h, input_w) + + input_res = _make_resolution(input_h, input_w) + output_res = _make_resolution(output_h, output_w) + ctx = GridContext( + label_mask=torch.ones(out_channels, output_h, output_w, dtype=torch.bool), + input_resolution_cpu=input_res, + output_resolution_cpu=output_res, + ) + + output = model.forward_once(prog, boundary, ctx) + + assert output.shape == (1, out_channels, output_h, output_w) + assert torch.isfinite(output).all(), "Output contains NaN or Inf."