model/samudra.py

"""Paper-version Samudra ConvNeXt U-Net implementation.

The layer layout follows the official ``samudra_om4_v1/model.yaml`` and the
official v1 ``blocks.py`` implementation. Inputs are channel-stacked tensors
with shape ``[batch, 158, lat, lon]`` for the full thermo-dynamic variant.
"""

from __future__ import annotations

from pathlib import Path
from typing import Iterable

import torch
from torch import nn
from torch.nn import functional as F


def circular_pad_width(x: torch.Tensor, padding: int) -> torch.Tensor:
    """Apply periodic longitude padding, including widths smaller than padding."""
    if padding == 0:
        return x
    width = x.shape[-1]
    repeats = (padding + width - 1) // width
    tiled = x.repeat(1, 1, 1, repeats)
    left = tiled[..., -padding:]
    right = tiled[..., :padding]
    return torch.cat((left, x, right), dim=-1)


def globe_pad(x: torch.Tensor, padding: int) -> torch.Tensor:
    x = circular_pad_width(x, padding)
    return F.pad(x, (0, 0, padding, padding), mode="constant")


class CappedGELU(nn.Module):
    def __init__(self, cap_value: float = 10.0):
        super().__init__()
        self.gelu = nn.GELU()
        # Keep this constant out of DDP buffers. DDP synchronizes buffers at
        # every forward, which would mutate the autograd version counter during
        # a recurrent multi-step rollout.
        self.cap = float(cap_value)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.clamp(self.gelu(x), max=self.cap)


class ConvNeXtBlock(nn.Module):
    """Official v1 residual block with dilated 3x3 convolutions."""

    def __init__(self, in_channels: int, out_channels: int, dilation: int):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.dilation = dilation
        self.padding = dilation
        self.skip = nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, 1)
        hidden = in_channels * 4
        self.layers = nn.ModuleList(
            [
                nn.Conv2d(in_channels, hidden, 3, dilation=dilation),
                nn.BatchNorm2d(hidden),
                CappedGELU(),
                nn.Conv2d(hidden, hidden, 3, dilation=dilation),
                nn.BatchNorm2d(hidden),
                CappedGELU(),
                nn.Conv2d(hidden, out_channels, 1),
            ]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        skip = self.skip(x)
        for layer in self.layers:
            if isinstance(layer, nn.Conv2d) and layer.kernel_size == (3, 3):
                x = globe_pad(x, self.padding)
            x = layer(x)
        return skip + x


class PeriodicBilinearUpsample(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        width = x.shape[-1]
        padded = F.pad(x, (1, 1, 0, 0), mode="circular")
        result = F.interpolate(padded, scale_factor=2, mode="bilinear", align_corners=False)
        return result[..., 2 : 2 + width * 2]


class SamudraUNet(nn.Module):
    def __init__(self, in_channels: int, widths: Iterable[int], dilations: Iterable[int]):
        super().__init__()
        widths = list(widths)
        dilations = list(dilations)
        if len(widths) != 4 or len(dilations) != 4:
            raise ValueError("paper v1 requires four widths and four dilation rates")
        channels = [in_channels, *widths]
        self.down_blocks = nn.ModuleList(
            ConvNeXtBlock(channels[i], channels[i + 1], dilations[i]) for i in range(4)
        )
        self.pools = nn.ModuleList(nn.AvgPool2d(2) for _ in range(4))
        self.middle = ConvNeXtBlock(widths[-1], widths[-1], dilations[-1])
        self.first_up = PeriodicBilinearUpsample()
        reversed_widths = list(reversed(widths))
        reversed_dilations = list(reversed(dilations))
        self.up_blocks = nn.ModuleList()
        self.upsamples = nn.ModuleList()
        current = widths[-1]
        for index in range(3):
            target = reversed_widths[index + 1]
            self.up_blocks.append(ConvNeXtBlock(current, target, reversed_dilations[index]))
            self.upsamples.append(PeriodicBilinearUpsample())
            current = target
        self.final = ConvNeXtBlock(current, widths[0], reversed_dilations[-1])
        self.out_channels = widths[0]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        skips = []
        for block, pool in zip(self.down_blocks, self.pools):
            x = block(x)
            skips.append(x)
            x = pool(x)
        x = self.middle(x)
        x = self.first_up(x)
        x = self._merge(x, skips[-1])
        for block, up, skip in zip(self.up_blocks, self.upsamples, reversed(skips[:-1])):
            x = block(x)
            x = up(x)
            x = self._merge(x, skip)
        return self.final(x)

    @staticmethod
    def _merge(x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
        height = min(x.shape[-2], skip.shape[-2])
        width = min(x.shape[-1], skip.shape[-1])
        x = x[..., :height, :width]
        skip = skip[..., :height, :width]
        if x.shape[1] != skip.shape[1]:
            raise RuntimeError(f"skip channel mismatch: {x.shape[1]} != {skip.shape[1]}")
        return x + skip


class Samudra(nn.Module):
    """Samudra v1 single-scale emulator.

    ``variant='thermo_dynamic'`` uses the paper's 158-to-154 interface.
    ``variant='thermo'`` uses 78 input channels and 154 output channels.
    """

    def __init__(
        self,
        variant: str = "thermo_dynamic",
        input_channels: int | None = None,
        output_channels: int | None = None,
        widths: Iterable[int] = (200, 250, 300, 400),
        dilations: Iterable[int] = (1, 2, 4, 8),
    ):
        super().__init__()
        if variant not in {"thermo_dynamic", "thermo"}:
            raise ValueError("variant must be thermo_dynamic or thermo")
        default_input = 158 if variant == "thermo_dynamic" else 78
        default_output = 154
        self.variant = variant
        self.input_channels = default_input if input_channels is None else input_channels
        self.output_channels = default_output if output_channels is None else output_channels
        self.unet = SamudraUNet(self.input_channels, widths, dilations)
        self.decoder = nn.Conv2d(self.unet.out_channels, self.output_channels, 3)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.ndim != 4 or x.shape[1] != self.input_channels:
            raise ValueError(f"expected [batch, {self.input_channels}, lat, lon], got {tuple(x.shape)}")
        features = self.unet(x)
        features = globe_pad(features, 1)
        return self.decoder(features)

    def load_official_checkpoint(self, path: str | Path, strict: bool = True) -> None:
        checkpoint = torch.load(path, map_location="cpu")
        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
            checkpoint = checkpoint["state_dict"]
        if not isinstance(checkpoint, dict):
            raise TypeError("official checkpoint must contain a state dictionary")
        # Older local checkpoints stored the fixed GELU cap as a buffer.
        checkpoint = {key: value for key, value in checkpoint.items() if not key.endswith(".cap")}
        self.load_state_dict(checkpoint, strict=strict)


def build_model(config: dict) -> Samudra:
    """Build a model from the project's YAML model section."""
    model_config = config.get("model", config)
    model = Samudra(
        variant=model_config.get("variant", "thermo_dynamic"),
        input_channels=model_config.get("input_channels"),
        output_channels=model_config.get("output_channels"),
        widths=model_config.get("widths", (200, 250, 300, 400)),
        dilations=model_config.get("dilations", (1, 2, 4, 8)),
    )
    checkpoint = model_config.get("checkpoint")
    if checkpoint:
        model.load_official_checkpoint(checkpoint)
    return model