model/aurora_decoder.py

"""Copyright (c) Microsoft Corporation. Licensed under the MIT license."""

from datetime import timedelta
from typing import Optional

import torch
from einops import rearrange
from torch import nn

from .aurora_batch import Batch, Metadata
from .aurora_fourier import levels_expansion
from .aurora_levelcond import LevelConditioned
from .aurora_perceiver import PerceiverResampler
from .aurora_util import (
    check_lat_lon_dtype,
    init_weights,
    unpatchify,
)

__all__ = ["Perceiver3DDecoder"]


class LinearPatchReconstruction(nn.Linear):
    """Linear layer for patch reconstruction."""


class Perceiver3DDecoder(nn.Module):
    """Multi-scale multi-source multi-variable decoder based on the Perceiver architecture."""

    def __init__(
        self,
        surf_vars: tuple[str, ...],
        atmos_vars: tuple[str, ...],
        patch_size: int = 4,
        embed_dim: int = 1024,
        depth: int = 1,
        head_dim: int = 64,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
        drop_rate: float = 0.0,
        perceiver_ln_eps: float = 1e-5,
        level_condition: Optional[tuple[int | float, ...]] = None,
        separate_perceiver: tuple[str, ...] = (),
        modulation_heads: tuple[str, ...] = (),
    ) -> None:
        """Initialise.

        Args:
            surf_vars (tuple[str, ...]): All supported surface-level variables.
            atmos_vars (tuple[str, ...]): All supported atmospheric variables.
            patch_size (int, optional): Patch size. Defaults to `4`.
            embed_dim (int, optional): Embedding dim.. Defaults to `1024`.
            depth (int, optional): Number of Perceiver cross-attention and feed-forward blocks.
                Defaults to `1`.
            head_dim (int, optional): Dimension of the attention heads used in the aggregation
                blocks. Defaults to `64`.
            num_heads (int, optional): Number of attention heads used in the aggregation blocks.
                Defaults to `8`.
            mlp_ratio (float, optional): Ratio of MLP hidden dimension to embedding dimensionality.
                Defaults to `4.0`.
            drop_rate (float, optional): Drop-out rate for input patches. Defaults to `0.0`.
            perceiver_ln_eps (float, optional): Layer norm. epsilon for the Perceiver blocks.
                Defaults to `1e-5`.
            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.
            separate_perceiver (tuple[str, ...], optional): In the decoder, use a separate Perceiver
                for specific atmospheric variables. This can be helpful at fine-tuning time to deal
                with variables that have a significantly different behaviour. If you want to enable
                this features, set this to the collection of variables that should be run on a
                separate Perceiver.
            modulation_heads (tuple[str, ...], optional): Names of every variable for which to
                enable an additional head, the so-called modulation head, that can be used to
                predict the difference.
        """
        super().__init__()

        # If additional modulation heads are required, simulate them as different variables with
        # the suffix `_mod`.
        surf_vars += tuple(f"{name}_mod" for name in surf_vars if name in modulation_heads)
        atmos_vars += tuple(f"{name}_mod" for name in atmos_vars if name in modulation_heads)
        if modulation_heads:
            separate_perceiver += tuple(f"{name}_mod" for name in separate_perceiver)

        self.patch_size = patch_size
        self.surf_vars = surf_vars
        self.atmos_vars = atmos_vars
        self.embed_dim = embed_dim
        self.level_condition = level_condition
        self.separate_perceiver = separate_perceiver
        self.modulation_heads = modulation_heads

        self.level_decoder = PerceiverResampler(
            latent_dim=embed_dim,
            context_dim=embed_dim,
            depth=depth,
            head_dim=head_dim,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            drop=drop_rate,
            residual_latent=True,
            ln_eps=perceiver_ln_eps,
        )
        if self.separate_perceiver:
            self.level_decoder_alternate = PerceiverResampler(
                latent_dim=embed_dim,
                context_dim=embed_dim,
                depth=depth,
                head_dim=head_dim,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                drop=drop_rate,
                residual_latent=True,
                ln_eps=perceiver_ln_eps,
            )

        self.surf_heads = nn.ParameterDict(
            {name: LinearPatchReconstruction(embed_dim, patch_size**2) for name in surf_vars}
        )
        if not self.level_condition:
            self.atmos_heads = nn.ParameterDict(
                {name: LinearPatchReconstruction(embed_dim, patch_size**2) for name in atmos_vars}
            )
        else:
            self.atmos_heads = nn.ParameterDict(
                {
                    name: LevelConditioned(
                        lambda: LinearPatchReconstruction(embed_dim, patch_size**2),
                        levels=self.level_condition,
                        levels_dim=-2,
                    )
                    for name in atmos_vars
                }
            )

        self.atmos_levels_embed = nn.Linear(embed_dim, embed_dim)

        self.apply(init_weights)

    def deaggregate_levels(
        self,
        level_embed: torch.Tensor,
        x: torch.Tensor,
        level_decoder: nn.Module,
    ) -> torch.Tensor:
        """Deaggregate pressure level information.

        Args:
            level_embed (torch.Tensor): Level embedding of shape `(B, L, C, D)`.
            x (torch.Tensor): Aggregated input of shape `(B, L, C', D)`.
            level_decoder (nn.Module): Pressure level decoder.

        Returns:
            torch.Tensor: Deaggregate output of shape `(B, L, C, D)`.
        """
        B, L, C, D = level_embed.shape
        level_embed = level_embed.flatten(0, 1)  # (BxL, C, D)
        x = x.flatten(0, 1)  # (BxL, C', D)
        _msg = f"Batch size mismatch. Found {level_embed.size(0)} and {x.size(0)}."
        assert level_embed.size(0) == x.size(0), _msg
        assert len(level_embed.shape) == 3, f"Expected 3 dims, found {level_embed.dims()}."
        assert x.dim() == 3, f"Expected 3 dims, found {x.dim()}."

        x = level_decoder(level_embed, x)  # (BxL, C, D)
        x = x.reshape(B, L, C, D)
        return x

    def forward(
        self,
        x: torch.Tensor,
        batch: Batch,
        patch_res: tuple[int, int, int],
        lead_times: torch.Tensor,
    ) -> Batch:
        """Forward pass.

        Args:
            x (torch.Tensor): Backbone output of shape `(B, L, D)`.
            batch (:class:`aurora.Batch`): Batch to make predictions for.
            patch_res (tuple[int, int, int]): Patch resolution
            lead_times (torch.Tensor): Lead times of shape `(batch,)` in hours.

        Returns:
            :class:`aurora.batch.Batch`: Prediction for `batch`.
        """
        surf_vars = tuple(batch.surf_vars.keys())
        atmos_vars = tuple(batch.atmos_vars.keys())
        atmos_levels = batch.metadata.atmos_levels

        # If additional modulation heads are required, simulate them as different variables with
        # the suffix `_mod`.
        surf_vars += tuple(f"{name}_mod" for name in surf_vars if name in self.modulation_heads)
        atmos_vars += tuple(f"{name}_mod" for name in atmos_vars if name in self.modulation_heads)

        # Compress the latent dimension from the U-net skip concatenation.
        B, L, D = x.shape

        # Extract the lat, lon and convert to float32.
        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)
        H, W = lat.shape[0], lon.shape[-1]

        # Unwrap the latent level dimension.
        x = rearrange(
            x,
            "B (C H W) D -> B (H W) C D",
            C=patch_res[0],
            H=patch_res[1],
            W=patch_res[2],
        )

        # Decode surface vars. Run the head for every surface-level variable.
        x_surf = torch.stack([self.surf_heads[name](x[..., :1, :]) for name in surf_vars], dim=-1)
        x_surf = x_surf.reshape(*x_surf.shape[:3], -1)  # (B, L, 1, V_S*p*p)
        surf_preds = unpatchify(x_surf, len(surf_vars), H, W, self.patch_size)
        surf_preds = surf_preds.squeeze(2)  # (B, V_S, H, W)

        # Embed the atmospheric levels.
        atmos_levels_encode = levels_expansion(
            torch.tensor(atmos_levels, device=x.device), self.embed_dim
        ).to(dtype=x.dtype)
        levels_embed = self.atmos_levels_embed(atmos_levels_encode)  # (C_A, D)

        # De-aggregate the hidden levels into the physical levels.
        levels_embed = levels_embed.expand(B, x.size(1), -1, -1)
        x_atmos = self.deaggregate_levels(
            levels_embed,
            x[..., 1:, :],
            self.level_decoder,
        )  # (B, L, C_A, D)
        if self.separate_perceiver:
            x_atmos_alternate = self.deaggregate_levels(
                levels_embed,
                x[..., 1:, :],
                self.level_decoder_alternate,
            )
        else:
            # `x_atmos_alternate` won't be used, but we define the variable anyway for type
            # stability.
            x_atmos_alternate = x_atmos

        # Decode the atmospheric vars. Per variable, first determine whether the main or alternate
        # Perceiver pressure level decoder should be used.
        head_inputs = [
            x_atmos if name not in self.separate_perceiver else x_atmos_alternate
            for name in atmos_vars
        ]
        if not self.level_condition:
            x_atmos = torch.stack(
                [self.atmos_heads[name](x) for name, x in zip(atmos_vars, head_inputs)],
                dim=-1,
            )
        else:
            x_atmos = torch.stack(
                [
                    self.atmos_heads[name](x, levels=atmos_levels)
                    for name, x in zip(atmos_vars, head_inputs)
                ],
                dim=-1,
            )
        x_atmos = x_atmos.reshape(*x_atmos.shape[:3], -1)  # (B, L, C_A, V_A*p*p)
        atmos_preds = unpatchify(x_atmos, len(atmos_vars), H, W, self.patch_size)

        # Compute output times from per-sample lead times.
        pred_time = tuple(
            t + timedelta(hours=float(lead_times[i])) for i, t in enumerate(batch.metadata.time)
        )

        return Batch(
            {v: surf_preds[:, i] for i, v in enumerate(surf_vars)},
            batch.static_vars,
            {v: atmos_preds[:, i] for i, v in enumerate(atmos_vars)},
            Metadata(
                lat=lat,
                lon=lon,
                time=pred_time,
                atmos_levels=atmos_levels,
                rollout_step=batch.metadata.rollout_step + 1,
            ),
        )