model/w_mae.py

from __future__ import annotations

from functools import partial
from typing import NamedTuple

import numpy as np
import torch
from torch import nn

from .afno import Block


class WMAEOutput(NamedTuple):
    loss: torch.Tensor
    prediction: torch.Tensor
    mask: torch.Tensor


def _sincos_1d(embed_dim: int, positions: np.ndarray) -> np.ndarray:
    if embed_dim % 2 != 0:
        raise ValueError("The 1D sine-cosine embedding dimension must be even.")
    omega = np.arange(embed_dim // 2, dtype=np.float64)
    omega = 1.0 / (10000 ** (omega / (embed_dim / 2.0)))
    values = np.einsum("m,d->md", positions.reshape(-1), omega)
    return np.concatenate((np.sin(values), np.cos(values)), axis=1)


def build_2d_sincos_position_embedding(
    embed_dim: int, grid_size: tuple[int, int], include_cls_token: bool
) -> torch.Tensor:
    if embed_dim % 4 != 0:
        raise ValueError("The 2D sine-cosine embedding dimension must be divisible by four.")
    grid_h = np.arange(grid_size[0], dtype=np.float32)
    grid_w = np.arange(grid_size[1], dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)
    embedding = np.concatenate(
        (_sincos_1d(embed_dim // 2, grid[0]), _sincos_1d(embed_dim // 2, grid[1])), axis=1
    )
    if include_cls_token:
        embedding = np.concatenate((np.zeros((1, embed_dim)), embedding), axis=0)
    return torch.from_numpy(embedding).float().unsqueeze(0)


class PatchEmbed(nn.Module):
    def __init__(
        self,
        img_size: tuple[int, int],
        patch_size: tuple[int, int],
        in_chans: int,
        embed_dim: int,
    ) -> None:
        super().__init__()
        if img_size[0] % patch_size[0] or img_size[1] % patch_size[1]:
            raise ValueError(f"img_size={img_size} must be divisible by patch_size={patch_size}.")
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.ndim != 4 or tuple(x.shape[-2:]) != self.img_size:
            raise ValueError(f"PatchEmbed expects [B,C,{self.img_size[0]},{self.img_size[1]}], got {tuple(x.shape)}.")
        return self.proj(x).flatten(2).transpose(1, 2)


class MaskedAutoencoderAFNO(nn.Module):
    """W-MAE pretraining model reconstructed from the official AFNO source."""

    def __init__(
        self,
        img_size: tuple[int, int] = (720, 1440),
        patch_size: tuple[int, int] = (8, 8),
        in_chans: int = 20,
        embed_dim: int = 768,
        depth: int = 12,
        decoder_embed_dim: int = 512,
        decoder_depth: int = 6,
        mlp_ratio: float = 4.0,
        norm_layer: type[nn.Module] = nn.LayerNorm,
        norm_pix_loss: bool = False,
        num_blocks: int = 8,
        sparsity_threshold: float = 0.01,
        hard_thresholding_fraction: float = 1.0,
    ) -> None:
        super().__init__()
        self.img_size = tuple(img_size)
        self.patch_size = tuple(patch_size)
        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.decoder_embed_dim = decoder_embed_dim
        self.norm_pix_loss = norm_pix_loss

        self.patch_embed = PatchEmbed(self.img_size, self.patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches
        # The official AFNO path keeps this checkpoint key but does not prepend
        # a class token in forward_encoder, so it must not enter DDP reduction.
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim), requires_grad=False)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)
        block_args = dict(
            mlp_ratio=mlp_ratio,
            norm_layer=norm_layer,
            num_blocks=num_blocks,
            sparsity_threshold=sparsity_threshold,
            hard_thresholding_fraction=hard_thresholding_fraction,
        )
        self.blocks = nn.ModuleList([Block(dim=embed_dim, **block_args) for _ in range(depth)])
        self.norm = norm_layer(embed_dim)

        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
        self.decoder_pos_embed = nn.Parameter(
            torch.zeros(1, num_patches, decoder_embed_dim), requires_grad=False
        )
        decoder_args = dict(block_args)
        decoder_args["norm_layer"] = norm_layer
        self.decoder_blocks = nn.ModuleList(
            [Block(dim=decoder_embed_dim, **decoder_args) for _ in range(decoder_depth)]
        )
        self.decoder_norm = norm_layer(decoder_embed_dim)
        patch_area = self.patch_size[0] * self.patch_size[1]
        self.decoder_pred = nn.Linear(decoder_embed_dim, in_chans * patch_area)
        self.initialize_weights()

    def initialize_weights(self) -> None:
        grid_size = self.patch_embed.grid_size
        self.pos_embed.data.copy_(build_2d_sincos_position_embedding(self.embed_dim, grid_size, True))
        self.decoder_pos_embed.data.copy_(
            build_2d_sincos_position_embedding(self.decoder_embed_dim, grid_size, False)
        )
        nn.init.xavier_uniform_(self.patch_embed.proj.weight.data.flatten(1))
        nn.init.normal_(self.cls_token, std=0.02)
        nn.init.normal_(self.mask_token, std=0.02)
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.LayerNorm):
            nn.init.zeros_(module.bias)
            nn.init.ones_(module.weight)

    def patchify(self, images: torch.Tensor) -> torch.Tensor:
        if images.ndim != 4 or images.shape[1] != self.in_chans or tuple(images.shape[-2:]) != self.img_size:
            raise ValueError(
                f"patchify expects [B,{self.in_chans},{self.img_size[0]},{self.img_size[1]}], "
                f"got {tuple(images.shape)}."
            )
        ph, pw = self.patch_size
        gh, gw = self.patch_embed.grid_size
        x = images.reshape(images.shape[0], self.in_chans, gh, ph, gw, pw)
        x = torch.einsum("nchpwq->nhwpqc", x)
        return x.reshape(images.shape[0], gh * gw, ph * pw * self.in_chans)

    def unpatchify(self, patches: torch.Tensor) -> torch.Tensor:
        ph, pw = self.patch_size
        gh, gw = self.patch_embed.grid_size
        expected_dim = ph * pw * self.in_chans
        if patches.ndim != 3 or patches.shape[1:] != (gh * gw, expected_dim):
            raise ValueError(f"unpatchify expects [B,{gh * gw},{expected_dim}], got {tuple(patches.shape)}.")
        x = patches.reshape(patches.shape[0], gh, gw, ph, pw, self.in_chans)
        x = torch.einsum("nhwpqc->nchpwq", x)
        return x.reshape(patches.shape[0], self.in_chans, gh * ph, gw * pw)

    @staticmethod
    def random_masking(
        x: torch.Tensor, mask_ratio: float
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if mask_ratio not in {0.0, 0.75}:
            raise ValueError("Official W-MAE AFNO grid reshaping supports mask_ratio 0.0 or 0.75 only.")
        batch, length, channels = x.shape
        len_keep = int(length * (1.0 - mask_ratio))
        noise = torch.rand(batch, length, device=x.device)
        ids_shuffle = torch.argsort(noise, dim=1)
        ids_restore = torch.argsort(ids_shuffle, dim=1)
        ids_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(x, 1, ids_keep.unsqueeze(-1).expand(-1, -1, channels))
        mask = torch.ones(batch, length, device=x.device)
        mask[:, :len_keep] = 0
        return x_masked, torch.gather(mask, 1, ids_restore), ids_restore

    def forward_encoder(
        self, images: torch.Tensor, mask_ratio: float
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        x = self.patch_embed(images) + self.pos_embed[:, 1:]
        x, mask, ids_restore = self.random_masking(x, mask_ratio)
        grid_h, grid_w = self.patch_embed.grid_size
        divisor = 2 if mask_ratio == 0.75 else 1
        expected_tokens = (grid_h // divisor) * (grid_w // divisor)
        if x.shape[1] != expected_tokens:
            raise ValueError("The selected mask ratio does not form the rectangular AFNO grid expected by W-MAE.")
        x = x.reshape(x.shape[0], grid_h // divisor, grid_w // divisor, self.embed_dim)
        for block in self.blocks:
            x = block(x)
        x = self.norm(x)
        return x.flatten(1, 2), mask, ids_restore

    def forward_decoder(self, latent: torch.Tensor, ids_restore: torch.Tensor) -> torch.Tensor:
        x = self.decoder_embed(latent)
        missing_tokens = ids_restore.shape[1] - x.shape[1]
        if missing_tokens < 0:
            raise ValueError("Latent token count exceeds the decoder target token count.")
        mask_tokens = self.mask_token.expand(x.shape[0], missing_tokens, -1)
        x = torch.cat((x, mask_tokens), dim=1)
        x = torch.gather(x, 1, ids_restore.unsqueeze(-1).expand(-1, -1, x.shape[-1]))
        x = x + self.decoder_pos_embed
        grid_h, grid_w = self.patch_embed.grid_size
        x = x.reshape(x.shape[0], grid_h, grid_w, self.decoder_embed_dim)
        for block in self.decoder_blocks:
            x = block(x)
        x = self.decoder_pred(self.decoder_norm(x))
        return x.flatten(1, 2)

    def forward_loss(
        self, images: torch.Tensor, prediction: torch.Tensor, mask: torch.Tensor, mask_ratio: float
    ) -> torch.Tensor:
        target = self.patchify(images)
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            variance = target.var(dim=-1, keepdim=True)
            target = (target - mean) / torch.sqrt(variance + 1e-6)
        loss = (prediction - target).pow(2).mean(dim=-1)
        if mask_ratio == 0.0:
            return loss.mean()
        masked_count = mask.sum()
        if masked_count.item() == 0:
            raise ValueError("Masked reconstruction loss requires at least one masked patch.")
        return (loss * mask).sum() / masked_count

    def forward(self, images: torch.Tensor, mask_ratio: float = 0.75) -> WMAEOutput:
        latent, mask, ids_restore = self.forward_encoder(images, mask_ratio)
        prediction = self.forward_decoder(latent, ids_restore)
        loss = self.forward_loss(images, prediction, mask, mask_ratio)
        return WMAEOutput(loss, prediction, mask)


def w_mae_base(
    embed_dim: int = 768,
    depth: int = 12,
    decoder_embed_dim: int = 512,
    decoder_depth: int = 6,
    mlp_ratio: float = 4.0,
    norm_layer: type[nn.Module] = partial(nn.LayerNorm, eps=1e-6),
    **kwargs: object,
) -> MaskedAutoencoderAFNO:
    return MaskedAutoencoderAFNO(
        embed_dim=embed_dim,
        depth=depth,
        decoder_embed_dim=decoder_embed_dim,
        decoder_depth=decoder_depth,
        mlp_ratio=mlp_ratio,
        norm_layer=norm_layer,
        **kwargs,
    )


mae_vit_base_patch16 = w_mae_base