samudra_model.py

# SPDX-FileCopyrightText: 2023 Matthias Karlbauer, Nathaniel Cresswell-Clay, Thorsten Kurth
# SPDX-FileCopyrightText: 2026 Samudra Authors
#
# SPDX-License-Identifier: Apache-2.0
# SPDX-License-Identifier: MIT
"""Minimal, dependency-light copy of the Samudra 2 architecture.

Vendored (and trimmed to the inference path) from
https://github.com/m2lines/Samudra — `samudra.models.samudra`,
`samudra.models.modules.{blocks,activations,unet_backbone}` — so this Space can
build the released `M2LInES/Samudra2` checkpoints without pulling in the full
training stack (dask / wandb / xesmf / skypilot).

The layer construction order is byte-for-byte compatible with the upstream
`state_dict`, which is what makes `load_state_dict(..., strict=True)` work.
"""

from __future__ import annotations

from itertools import tee

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


def pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


class CappedGELU(nn.Module):
    """GELU with a capped maximum value."""

    def __init__(self, cap_value: float = 10.0, **kwargs):
        super().__init__()
        self.gelu = nn.GELU(**kwargs)
        self.cap = nn.Buffer(torch.tensor(cap_value, dtype=torch.float32))

    def forward(self, inputs):
        return torch.clamp(self.gelu(inputs), max=self.cap)


class AvgPool(nn.Module):
    def __init__(self, pooling: int = 2):
        super().__init__()
        self.avgpool = nn.AvgPool2d(pooling)

    def forward(self, x):
        return self.avgpool(x)


class ZonallyPeriodicBilinearUpsample(nn.Module):
    """Bilinear upsampling that enforces periodicity along the x/longitude axis."""

    def __init__(self, upsampling: int = 2):
        super().__init__()
        self.scale_h, self.scale_w = upsampling, upsampling

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        width = x.shape[-1]
        padded = F.pad(x, (1, 1, 0, 0), mode="circular")
        upsampled = F.interpolate(
            padded,
            scale_factor=(self.scale_h, self.scale_w),
            mode="bilinear",
            align_corners=False,
        )
        start = self.scale_w
        end = start + width * self.scale_w
        return upsampled[..., start:end]


class DropPath(nn.Module):
    def __init__(self, drop_prob: float = 0.0):
        super().__init__()
        self.dropout = nn.Dropout(p=drop_prob)

    def forward(self, skip_conn):
        if not self.training or self.dropout.p == 0.0:
            return skip_conn
        mask = self.dropout(
            torch.ones(
                skip_conn.shape[0], 1, 1, 1,
                device=skip_conn.device, dtype=skip_conn.dtype,
            )
        )
        return skip_conn * mask


class CoreBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation, pad):
        super().__init__()
        assert kernel_size % 2 != 0, "Cannot use even kernel sizes!"
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.N_pad = int((kernel_size + (kernel_size - 1) * (dilation - 1) - 1) / 2)
        self.pad = pad


class ConvNeXtBlock(CoreBlock):
    """ConvNeXt-style block with dilations (as used by Samudra / DLWP-HPX)."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int = 3,
        dilation: int = 1,
        n_layers: int = 1,
        activation=CappedGELU,
        pad: str = "circular",
        upscale_factor: int = 4,
        norm: str = "batch",
    ):
        super().__init__(in_channels, out_channels, kernel_size, dilation, pad)
        assert n_layers == 1, "Can only use a single layer here!"

        if in_channels == out_channels:
            self.skip_module = lambda x: x
        else:
            self.skip_module = nn.Conv2d(
                in_channels, out_channels, kernel_size=1, padding="same"
            )

        hidden = int(in_channels * upscale_factor)
        convblock: list[nn.Module] = []
        convblock.append(
            nn.Conv2d(in_channels, hidden, kernel_size, dilation=dilation)
        )
        convblock.append(self._norm(norm, hidden))
        convblock.append(activation())
        convblock.append(nn.Conv2d(hidden, hidden, kernel_size, dilation=dilation))
        convblock.append(self._norm(norm, hidden))
        convblock.append(activation())
        convblock.append(nn.Conv2d(hidden, out_channels, kernel_size=1, padding="same"))
        self.convblock = nn.Sequential(*[m for m in convblock if m is not None])

    @staticmethod
    def _norm(norm: str, channels: int):
        if norm == "batch":
            return nn.BatchNorm2d(channels)
        if norm == "instance":
            return nn.InstanceNorm2d(channels)
        if norm == "nonorm":
            return None
        raise NotImplementedError(norm)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        skip = self.skip_module(x)
        for layer in self.convblock:
            if isinstance(layer, nn.Conv2d) and layer.kernel_size[0] != 1:
                x = F.pad(x, (self.N_pad, self.N_pad, 0, 0), mode=self.pad)
                x = F.pad(x, (0, 0, self.N_pad, self.N_pad), mode="constant")
            x = layer(x)
        return skip + x


class UNetBackbone(nn.Module):
    """Configurable ConvNeXt U-Net backbone (inference path only)."""

    def __init__(
        self,
        in_channels: int,
        ch_width: list[int],
        dilation: list[int],
        n_layers: list[int],
        pad: str,
        create_block,
        downsampling_block: nn.Module,
        create_upsampling_block,
        drop_path_rate: float = 0.0,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = ch_width[0]

        ch_width = [in_channels] + list(ch_width)
        dilation = list(dilation)
        n_layers = list(n_layers)
        self.pad = pad

        layers: list[nn.Module] = []
        for i, (a, b) in enumerate(pairwise(ch_width)):
            layers.append(create_block(a, b, dilation[i], n_layers[i], pad))
            layers.append(downsampling_block)

        layers.append(create_block(b, b, dilation[i], n_layers[i], pad))
        layers.append(create_upsampling_block(b, b))

        ch_width.reverse()
        dilation.reverse()
        n_layers.reverse()

        for i, (a, b) in enumerate(pairwise(ch_width[:-1])):
            layers.append(create_block(a, b, dilation[i], n_layers[i], pad))
            layers.append(create_upsampling_block(b, b))

        layers.append(create_block(b, b, dilation[i], n_layers[i], pad))

        first_block = layers[0]
        assert isinstance(first_block, CoreBlock)
        self.N_pad = first_block.N_pad
        self.layers = nn.ModuleList(layers)
        self.num_steps = int(len(ch_width) - 1)
        self.drop_path = DropPath(drop_path_rate)

    def forward(self, fts: torch.Tensor) -> torch.Tensor:
        skip_inputs: list[torch.Tensor] = [
            torch.zeros_like(fts) for _ in range(self.num_steps)
        ]
        count = 0
        for layer in self.layers:
            if isinstance(layer, nn.Conv2d):
                fts = F.pad(fts, (self.N_pad, self.N_pad, 0, 0), mode=self.pad)
                fts = F.pad(fts, (0, 0, self.N_pad, self.N_pad), mode="constant")

            fts = layer(fts)

            if count < self.num_steps:
                if isinstance(layer, CoreBlock):
                    skip_inputs[count] = fts
                    count += 1
            else:
                if isinstance(layer, ZonallyPeriodicBilinearUpsample):
                    crop = np.array(fts.shape[2:])
                    shape = np.array(
                        skip_inputs[int(2 * self.num_steps - count - 1)].shape[2:]
                    )
                    pads = shape - crop
                    pads = [
                        pads[1] // 2,
                        pads[1] - pads[1] // 2,
                        pads[0] // 2,
                        pads[0] - pads[0] // 2,
                    ]
                    fts = F.pad(fts, pads)
                    fts = fts + self.drop_path(
                        skip_inputs[int(2 * self.num_steps - count - 1)]
                    )
                    count += 1
        return fts


class Samudra(nn.Module):
    """Samudra 2 ocean emulator (single-scale, ConvNeXt U-Net backbone)."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        pred_residuals: bool,
        last_kernel_size: int,
        pad: str,
        unet: UNetBackbone,
        use_bfloat16: bool = False,
    ):
        super().__init__()
        assert last_kernel_size % 2 != 0, "Cannot use even kernel sizes!"
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.N_pad = int((last_kernel_size - 1) / 2)
        self.pad = pad
        self.pred_residuals = pred_residuals
        self.unet = unet
        self.decoder = nn.Conv2d(unet.out_channels, out_channels, last_kernel_size)
        self.use_bfloat16 = use_bfloat16

    def forward_once(self, prognostic, boundary, label_mask):
        fts = torch.cat((prognostic, boundary), dim=1)
        device_type = "cuda" if fts.is_cuda else "cpu"
        with torch.autocast(device_type, dtype=torch.bfloat16, enabled=self.use_bfloat16):
            fts = self.unet(fts)
            fts = F.pad(fts, (self.N_pad, self.N_pad, 0, 0), mode=self.pad)
            fts = F.pad(fts, (0, 0, self.N_pad, self.N_pad), mode="constant")
        fts = fts.to(torch.float32)
        fts = self.decoder(fts)
        return torch.where(label_mask, fts, 0.0)


def build_samudra(
    in_channels: int = 162,
    out_channels: int = 154,
    ch_width: tuple[int, ...] = (280, 380, 480, 520),
    dilation: tuple[int, ...] = (1, 2, 4, 8),
    n_layers: tuple[int, ...] = (1, 1, 1, 1),
    pad: str = "circular",
    norm: str = "batch",
    upscale_factor: int = 2,
    kernel_size: int = 3,
    last_kernel_size: int = 3,
    pred_residuals: bool = False,
    use_bfloat16: bool = False,
) -> Samudra:
    """Build the Samudra 2 model matching `configs/samudra_om4_v2/model.yaml`."""

    def create_block(in_ch, out_ch, dil, n_lay, pad_):
        return ConvNeXtBlock(
            in_channels=in_ch,
            out_channels=out_ch,
            kernel_size=kernel_size,
            dilation=dil,
            n_layers=n_lay,
            activation=CappedGELU,
            pad=pad_,
            upscale_factor=upscale_factor,
            norm=norm,
        )

    def create_upsampling_block(in_ch, out_ch):
        return ZonallyPeriodicBilinearUpsample(2)

    unet = UNetBackbone(
        in_channels=in_channels,
        ch_width=list(ch_width),
        dilation=list(dilation),
        n_layers=list(n_layers),
        pad=pad,
        create_block=create_block,
        downsampling_block=AvgPool(2),
        create_upsampling_block=create_upsampling_block,
    )
    return Samudra(
        in_channels=in_channels,
        out_channels=out_channels,
        pred_residuals=pred_residuals,
        last_kernel_size=last_kernel_size,
        pad=pad,
        unet=unet,
        use_bfloat16=use_bfloat16,
    )