code/training/src/model/simvp_bt_aux.py

from __future__ import annotations

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

from ..openstl.models.simvp_model import Encoder, MidIncepNet, MidMetaNet, sampling_generator
from ..openstl.modules.simvp_modules import ConvSC


class BTAuxDecoder(nn.Module):
    """SimVP decoder with a CI readout and a shallow future-BT auxiliary head."""

    def __init__(
        self,
        c_hid: int,
        ci_out_channels: int,
        bt_out_channels: int,
        n_s: int,
        spatio_kernel: int,
        act_inplace: bool = True,
    ):
        super().__init__()
        samplings = sampling_generator(n_s, reverse=True)
        self.dec = nn.Sequential(
            *[
                ConvSC(c_hid, c_hid, spatio_kernel, upsampling=s, act_inplace=act_inplace)
                for s in samplings[:-1]
            ],
            ConvSC(c_hid * 2, c_hid, spatio_kernel, upsampling=samplings[-1], act_inplace=act_inplace),
        )

        first_c = max(c_hid // 3 * 2, 1)
        second_c = max(c_hid // 3, 1)
        self.ci_readout = nn.Sequential(
            nn.Conv2d(c_hid, first_c, 3, padding="same"),
            nn.SiLU(),
            nn.Conv2d(first_c, second_c, 3, padding="same"),
            nn.SiLU(),
            nn.Conv2d(second_c, ci_out_channels, 3, padding="same"),
        )

        bt_hidden = max(c_hid // 4, 1)
        self.bt_head = nn.Sequential(
            nn.Conv2d(c_hid, bt_hidden, 3, padding="same"),
            nn.SiLU(),
            nn.Conv2d(bt_hidden, bt_out_channels, 1),
        )

    def forward(self, hid: torch.Tensor, enc1: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        for i in range(0, len(self.dec) - 1):
            hid = self.dec[i](hid)
        if hid.shape[-2:] != enc1.shape[-2:]:
            hid = F.interpolate(hid, size=(enc1.shape[2], enc1.shape[3]), mode="bilinear", align_corners=False)
        shared_feat = self.dec[-1](torch.cat([hid, enc1], dim=1))
        return self.ci_readout(shared_feat), self.bt_head(shared_feat)


class SimVPBTAux(nn.Module):
    """SimVP CI model with a shallow future-BT auxiliary prediction branch.

    Returns:
        {
            "ci": sigmoid CI map with shape (B, H, W),
            "bt": normalized BT prediction with shape (B, T_bt, C_bt, H, W),
        }
    """

    def __init__(
        self,
        in_shape,
        hid_S: int = 16,
        hid_T: int = 256,
        N_S: int = 4,
        N_T: int = 4,
        model_type: str = "gSTA",
        mlp_ratio: float = 8.0,
        drop: float = 0.0,
        drop_path: float = 0.0,
        spatio_kernel_enc: int = 3,
        spatio_kernel_dec: int = 3,
        act_inplace: bool = True,
        bt_out_shape=(6, 1, None, None),
        **kwargs,
    ):
        super().__init__()
        t_in, c_in, h, w = in_shape
        self.out_shape = tuple(kwargs.get("out_shape", (1, 1, h, w)))
        self.bt_out_shape = self._normalize_bt_out_shape(bt_out_shape, h, w)
        self._validate_shapes()

        kwargs["enc_C"] = hid_S
        hid_h = int(h / 2 ** (N_S / 2))
        hid_w = int(w / 2 ** (N_S / 2))
        act_inplace = False

        self.enc = Encoder(c_in, hid_S, N_S, spatio_kernel_enc, act_inplace=act_inplace)
        self.dec = BTAuxDecoder(
            c_hid=hid_S,
            ci_out_channels=self.out_shape[1],
            bt_out_channels=self.bt_out_shape[0] * self.bt_out_shape[1],
            n_s=N_S,
            spatio_kernel=spatio_kernel_dec,
            act_inplace=act_inplace,
        )

        model_type = "gsta" if model_type is None else str(model_type).lower()
        if model_type == "incepu":
            self.hid = MidIncepNet(t_in * hid_S, hid_T, N_T)
        else:
            self.hid = MidMetaNet(
                t_in * hid_S,
                hid_T,
                N_T,
                input_resolution=(hid_h, hid_w),
                model_type=model_type,
                mlp_ratio=mlp_ratio,
                drop=drop,
                drop_path=drop_path,
                **kwargs,
            )

        if t_in != self.out_shape[0]:
            self.skipconv = nn.Sequential(
                nn.Conv2d(t_in * hid_S, self.out_shape[0] * hid_S, 7, padding="same"),
                nn.SiLU(),
            )

    @staticmethod
    def _normalize_bt_out_shape(bt_out_shape, h: int, w: int) -> tuple[int, int, int, int]:
        values = list(bt_out_shape)
        if len(values) != 4:
            raise ValueError(f"bt_out_shape must be [T, C, H, W], got {bt_out_shape}")
        values[2] = h if values[2] is None else values[2]
        values[3] = w if values[3] is None else values[3]
        return tuple(int(v) for v in values)

    def _validate_shapes(self) -> None:
        if len(self.out_shape) != 4:
            raise ValueError(f"out_shape must be [T, C, H, W], got {self.out_shape}")
        if self.out_shape[0] != 1 or self.out_shape[1] != 1:
            raise ValueError(f"SimVPBTAux expects CI out_shape [1, 1, H, W], got {self.out_shape}")
        if self.bt_out_shape[2:] != self.out_shape[2:]:
            raise ValueError(f"bt_out_shape H,W {self.bt_out_shape[2:]} must match out_shape H,W {self.out_shape[2:]}")

    def forward(self, x_raw: torch.Tensor, **kwargs) -> dict[str, torch.Tensor]:
        b, t, c, h, w = x_raw.shape
        x = x_raw.view(b * t, c, h, w)

        embed, skip = self.enc(x)
        _, c_hid, h_hid, w_hid = embed.shape

        z = embed.view(b, t, c_hid, h_hid, w_hid)
        hid = self.hid(z)
        hid = hid.reshape(b * self.out_shape[0], c_hid, h_hid, w_hid)

        if hid.shape[0] != skip.shape[0]:
            skip = skip.view(b, t, c_hid, h, w)
            skip = skip.view(b, t * c_hid, h, w)
            skip = self.skipconv(skip)
            skip = skip.view(b * self.out_shape[0], c_hid, h, w)

        ci_logits, bt_logits = self.dec(hid, skip)
        ci = torch.sigmoid(ci_logits.reshape(b, self.out_shape[0], self.out_shape[1], h, w)).reshape(b, h, w)
        bt = bt_logits.reshape(b, self.bt_out_shape[0], self.bt_out_shape[1], h, w)
        return {"ci": ci, "bt": bt}