"""Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" from typing import Optional import numpy as np import torch from einops import rearrange from torch import nn from .aurora_batch import Batch from .aurora_fourier import ( absolute_time_expansion, lead_time_expansion, lead_time_expansion_v3, levels_expansion, pos_expansion, scale_expansion, ) from .aurora_levelcond import LevelConditioned from .aurora_patchembed import LevelPatchEmbed from .aurora_perceiver import MLP, PerceiverResampler from .aurora_posencoding import pos_scale_enc from .aurora_util import ( check_lat_lon_dtype, init_weights, ) __all__ = ["Perceiver3DEncoder"] class Perceiver3DEncoder(nn.Module): """Multi-scale multi-source multi-variable encoder based on the Perceiver architecture.""" def __init__( self, surf_vars: tuple[str, ...], static_vars: tuple[str, ...] | None, atmos_vars: tuple[str, ...], patch_size: int = 4, latent_levels: int = 8, embed_dim: int = 1024, num_heads: int = 16, head_dim: int = 64, drop_rate: float = 0.0, depth: int = 2, mlp_ratio: float = 4.0, max_history_size: int = 2, perceiver_ln_eps: float = 1e-5, stabilise_level_agg: bool = False, level_condition: Optional[tuple[int | float, ...]] = None, dynamic_vars: bool = False, atmos_static_vars: bool = False, simulate_indexing_bug: bool = False, use_updated_lead_time_embedding: bool = False, ) -> None: """Initialise. Args: surf_vars (tuple[str, ...]): All supported surface-level variables. static_vars (tuple[str, ...], optional): All supported static variables. atmos_vars (tuple[str, ...]): All supported atmospheric variables. patch_size (int, optional): Patch size. Defaults to `4`. latent_levels (int): Number of latent pressure levels. Defaults to `8`. embed_dim (int, optional): Embedding dim. used in the aggregation blocks. Defaults to `1024`. num_heads (int, optional): Number of attention heads used in aggregation blocks. Defaults to `16`. head_dim (int, optional): Dimension of attention heads used in aggregation blocks. Defaults to `64`. drop_rate (float, optional): Drop out rate for input patches. Defaults to `0.0`. depth (int, optional): Number of Perceiver cross-attention and feed-forward blocks. Defaults to `2`. mlp_ratio (float, optional): Ratio of hidden dimensionality to embedding dimensionality for MLPs. Defaults to `4.0`. max_history_size (int, optional): Maximum number of history steps to consider. Defaults to `2`. perceiver_ln_eps (float, optional): Epsilon value for layer normalisation in the Perceiver. Defaults to `1e-5`. stabilise_level_agg (bool, optional): Stabilise the level aggregation by inserting an additional layer normalisation. Defaults to `False`. level_condition (tuple[int | float, ...], optional): Make the patch embeddings dependent on pressure level. If you want to enable this feature, provide a tuple of all possible pressure levels. dynamic_vars (bool, optional): Use dynamically generated static variables, like time of day. Defaults to `False`. atmos_static_vars (bool, optional): Also concatenate the static variables to the atmospheric variables. Defaults to `False`. simulate_indexing_bug (bool, optional): Simulate an indexing bug that's present for the air pollution version of Aurora. This is necessary to obtain numerical equivalence to the original implementation. Defaults to `False`. use_updated_lead_time_embedding (bool, optional): Whether to use the updated lead time embedding with a minimum wavelength of 6 hours. Defaults to `False`. """ super().__init__() self.drop_rate = drop_rate self.embed_dim = embed_dim self.patch_size = patch_size self.level_condition = level_condition self.dynamic_vars = dynamic_vars self.atmos_static_vars = atmos_static_vars self.simulate_indexing_bug = simulate_indexing_bug # Add in the dynamic variables first. if self.dynamic_vars: if static_vars is None: static_vars = () static_vars += ("tod_cos", "tod_sin", "dow_cos", "dow_sin", "doy_cos", "doy_sin") # We treat the static variables as surface variables in the model (and possibly even as # atmospheric variables!). if static_vars: surf_vars += static_vars if self.atmos_static_vars: # In this case, we prefix the static variables to avoid name clashes. E.g., `z` is # both a static variable and an atmospheric variable. atmos_vars += tuple(f"static_{v}" for v in static_vars) # Latent tokens assert latent_levels > 1, "At least two latent levels are required." self.latent_levels = latent_levels # One latent level will be used by the surface level. self.atmos_latents = nn.Parameter(torch.randn(latent_levels - 1, embed_dim)) # Learnable embedding to encode the surface level. self.surf_level_encoding = nn.Parameter(torch.randn(embed_dim)) self.surf_mlp = MLP(embed_dim, int(embed_dim * mlp_ratio), dropout=drop_rate) self.surf_norm = nn.LayerNorm(embed_dim) # Position, scale, and time embeddings self.pos_embed = nn.Linear(embed_dim, embed_dim) self.scale_embed = nn.Linear(embed_dim, embed_dim) self.lead_time_embed = nn.Linear(embed_dim, embed_dim) self.absolute_time_embed = nn.Linear(embed_dim, embed_dim) self.atmos_levels_embed = nn.Linear(embed_dim, embed_dim) self.use_updated_lead_time_embedding = use_updated_lead_time_embedding # Patch embeddings: assert max_history_size > 0, "At least one history step is required." self.surf_token_embeds = LevelPatchEmbed(surf_vars, patch_size, embed_dim, max_history_size) if not self.level_condition: self.atmos_token_embeds = LevelPatchEmbed( atmos_vars, patch_size, embed_dim, max_history_size ) else: self.atmos_token_embeds = LevelConditioned( lambda: LevelPatchEmbed(atmos_vars, patch_size, embed_dim, max_history_size), levels=self.level_condition, levels_dim=-5, ) # Learnable pressure level aggregation: self.level_agg = PerceiverResampler( latent_dim=embed_dim, context_dim=embed_dim, depth=depth, head_dim=head_dim, num_heads=num_heads, drop=drop_rate, mlp_ratio=mlp_ratio, ln_eps=perceiver_ln_eps, ln_k_q=stabilise_level_agg, ) # Drop patches after encoding. self.pos_drop = nn.Dropout(p=drop_rate) self.apply(init_weights) # Initialize the latents like in the Huggingface implementation of the Perceiver: # # https://github.com/huggingface/transformers/blob/v4.36.1/src/transformers/models/perceiver/modeling_perceiver.py#L628 # torch.nn.init.trunc_normal_(self.atmos_latents, std=0.02) torch.nn.init.trunc_normal_(self.surf_level_encoding, std=0.02) def aggregate_levels(self, x: torch.Tensor) -> torch.Tensor: """Aggregate pressure level information. Args: x (torch.Tensor): Tensor of shape `(B, C_A, L, D)` where `C_A` refers to the number of pressure levels. Returns: torch.Tensor: Tensor of shape `(B, C, L, D)` where `C` is the number of aggregated pressure levels. """ B, _, L, _ = x.shape latents = self.atmos_latents.to(dtype=x.dtype) latents = latents.unsqueeze(1).expand(B, -1, L, -1) # (C_A, D) to (B, C_A, L, D) x = torch.einsum("bcld->blcd", x) x = x.flatten(0, 1) # (B * L, C_A, D) latents = torch.einsum("bcld->blcd", latents) latents = latents.flatten(0, 1) # (B * L, C_A, D) x = self.level_agg(latents, x) # (B * L, C, D) x = x.unflatten(dim=0, sizes=(B, L)) # (B, L, C, D) x = torch.einsum("blcd->bcld", x) # (B, C, L, D) return x def forward( self, batch: Batch, lead_times: torch.Tensor, ) -> torch.Tensor: """Peform encoding. Args: batch (:class:`aurora.Batch`): Batch to encode. lead_times (:class:`torch.Tensor`): Lead times of shape `(batch,)` in hours. Returns: torch.Tensor: Encoding of shape `(B, L, D)`. """ surf_vars = tuple(batch.surf_vars.keys()) static_vars = tuple(batch.static_vars.keys()) atmos_vars = tuple(batch.atmos_vars.keys()) atmos_levels = batch.metadata.atmos_levels x_surf = torch.stack(tuple(batch.surf_vars.values()), dim=2) x_static = torch.stack(tuple(batch.static_vars.values()), dim=2) x_atmos = torch.stack(tuple(batch.atmos_vars.values()), dim=2) B, T, _, C, H, W = x_atmos.size() assert x_surf.shape[:2] == (B, T), f"Expected shape {(B, T)}, got {x_surf.shape[:2]}." if static_vars is None: assert x_static is None, "Static variables given, but not configured." else: assert x_static is not None, "Static variables not given." x_static = x_static.expand((B, T, -1, -1, -1)) if self.dynamic_vars: ones = torch.ones((1, T, 1, H, W), device=x_static.device, dtype=x_static.dtype) time = batch.metadata.time x_dynamic = torch.cat( [ torch.cat( ( ones * np.cos(2 * np.pi * time[b].hour / 24), ones * np.sin(2 * np.pi * time[b].hour / 24), ones * np.cos(2 * np.pi * time[b].weekday() / 7), ones * np.sin(2 * np.pi * time[b].weekday() / 7), ones * np.cos(2 * np.pi * time[b].day / 365.25), ones * np.sin(2 * np.pi * time[b].day / 365.25), ), dim=-3, ) for b in range(B) ], dim=0, ) dynamic_vars = ("tod_cos", "tod_sin", "dow_cos", "dow_sin", "doy_cos", "doy_sin") x_surf = torch.cat((x_surf, x_static, x_dynamic), dim=2) surf_vars = surf_vars + static_vars + dynamic_vars # Add to atmospheric variables too. if self.atmos_static_vars: # in this case, we prefix the static variables to avoid name clashes. e.g., `z` # is both a static variable and an atmospheric variable. atmos_vars += tuple(f"static_{v}" for v in static_vars + dynamic_vars) inds = (-1, -1, -1, len(atmos_levels), -1, -1) x_atmos = torch.cat( ( x_atmos, # Repeat for every pressure level. x_static[..., None, :, :].expand(*inds), x_dynamic[..., None, :, :].expand(*inds), ), dim=2, ) else: x_surf = torch.cat((x_surf, x_static), dim=2) # (B, T, V_S + V_Static, H, W) surf_vars = surf_vars + static_vars # Add to atmospheric variables too. if self.atmos_static_vars: atmos_vars = atmos_vars + static_vars x_atmos = torch.cat( ( x_atmos, # Repeat for every pressure level. x_static[..., None, :, :].expand(-1, -1, -1, len(atmos_levels), -1, -1), ), dim=2, ) 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") x_surf = self.surf_token_embeds(x_surf, surf_vars) # (B, L, D) dtype = x_surf.dtype # When using mixed precision, we need to keep track of the dtype. # In the original implementation, both `z` and `static_z` point towards the same index, # meaning that they select the same slice. Simulate this bug. if self.simulate_indexing_bug and "z" in atmos_vars: i_z = atmos_vars.index("z") i_static_z = atmos_vars.index("static_z") x_atmos = torch.cat( ( x_atmos[:, :, :i_static_z], x_atmos[:, :, i_z : i_z + 1], x_atmos[:, :, i_static_z + 1 :], ), dim=2, ) # Patch embed the atmospheric levels. if not self.level_condition: x_atmos = rearrange(x_atmos, "b t v c h w -> (b c) v t h w") x_atmos = self.atmos_token_embeds(x_atmos, atmos_vars) x_atmos = rearrange(x_atmos, "(b c) l d -> b c l d", b=B, c=C) else: # In this case we need to keep the levels dimension separate. x_atmos = rearrange(x_atmos, "b t v c h w -> b c v t h w") x_atmos = self.atmos_token_embeds(x_atmos, atmos_vars, levels=atmos_levels) # The levels dimension is now already in the right place. # Add surface level encoding. This helps the model distinguish between surface and # atmospheric levels. x_surf = x_surf + self.surf_level_encoding[None, None, :].to(dtype=dtype) # Since the surface level is not aggregated, we add a Perceiver-like MLP only. x_surf = x_surf + self.surf_norm(self.surf_mlp(x_surf)) # Add atmospheric pressure encoding of shape (C_A, D) and subsequent embedding. atmos_levels_tensor = torch.tensor(atmos_levels, device=x_atmos.device) atmos_levels_encode = levels_expansion(atmos_levels_tensor, self.embed_dim).to(dtype=dtype) atmos_levels_embed = self.atmos_levels_embed(atmos_levels_encode)[None, :, None, :] x_atmos = x_atmos + atmos_levels_embed # (B, C_A, L, D) # Aggregate over pressure levels. x_atmos = self.aggregate_levels(x_atmos) # (B, C_A, L, D) to (B, C, L, D) # Concatenate the surface level with the amospheric levels. x = torch.cat((x_surf.unsqueeze(1), x_atmos), dim=1) # Add position and scale embeddings to the 3D tensor. pos_encode, scale_encode = pos_scale_enc( self.embed_dim, lat, lon, self.patch_size, pos_expansion=pos_expansion, scale_expansion=scale_expansion, ) # Encodings are (L, D). pos_encode = self.pos_embed(pos_encode[None, None, :].to(dtype=dtype)) scale_encode = self.scale_embed(scale_encode[None, None, :].to(dtype=dtype)) x = x + pos_encode + scale_encode # Flatten the tokens. x = x.reshape(B, -1, self.embed_dim) # (B, C + 1, L, D) to (B, L', D) # Add lead time embedding. expansion = ( lead_time_expansion_v3 if self.use_updated_lead_time_embedding else lead_time_expansion ) lead_time_encode = expansion(lead_times, self.embed_dim).to(dtype=dtype) lead_time_emb = self.lead_time_embed(lead_time_encode) # (B, D) x = x + lead_time_emb.unsqueeze(1) # (B, L', D) + (B, 1, D) # Add absolute time embedding. absolute_times_list = [t.timestamp() / 3600 for t in batch.metadata.time] # Times in hours absolute_times = torch.tensor(absolute_times_list, dtype=torch.float32, device=x.device) absolute_time_encode = absolute_time_expansion(absolute_times, self.embed_dim) absolute_time_embed = self.absolute_time_embed(absolute_time_encode.to(dtype=dtype)) x = x + absolute_time_embed.unsqueeze(1) # (B, L, D) + (B, 1, D) x = self.pos_drop(x) return x