model/earthformer.py

from __future__ import annotations

from dataclasses import dataclass

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


@dataclass
class CuboidMeta:
    batch: int
    shape: tuple[int, int, int]
    padded: tuple[int, int, int]
    cuboid: tuple[int, int, int]


def cuboid_partition(x: torch.Tensor, cuboid: tuple[int, int, int]) -> tuple[torch.Tensor, CuboidMeta]:
    b, t, h, w, c = x.shape
    bt, bh, bw = (min(size, dim) for size, dim in zip(cuboid, (t, h, w)))
    pt, ph, pw = (-t) % bt, (-h) % bh, (-w) % bw
    padded = F.pad(x.permute(0, 4, 1, 2, 3), (0, pw, 0, ph, 0, pt)).permute(0, 2, 3, 4, 1)
    tp, hp, wp = padded.shape[1:4]
    windows = padded.reshape(b, tp // bt, bt, hp // bh, bh, wp // bw, bw, c)
    windows = windows.permute(0, 1, 3, 5, 2, 4, 6, 7).reshape(-1, bt * bh * bw, c)
    return windows, CuboidMeta(b, (t, h, w), (tp, hp, wp), (bt, bh, bw))


def cuboid_merge(windows: torch.Tensor, meta: CuboidMeta) -> torch.Tensor:
    b, (t, h, w), (tp, hp, wp), (bt, bh, bw) = meta.batch, meta.shape, meta.padded, meta.cuboid
    c = windows.shape[-1]
    x = windows.reshape(b, tp // bt, hp // bh, wp // bw, bt, bh, bw, c)
    x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).reshape(b, tp, hp, wp, c)
    return x[:, :t, :h, :w]


class FeedForward(nn.Module):
    def __init__(self, dim: int, ratio: float, dropout: float):
        super().__init__()
        hidden = int(dim * ratio)
        self.net = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class CuboidAttentionLayer(nn.Module):
    def __init__(self, dim: int, heads: int, cuboid: tuple[int, int, int], ff_ratio: float, dropout: float, use_global: bool):
        super().__init__()
        self.cuboid = cuboid
        self.use_global = use_global
        self.local_norm = nn.LayerNorm(dim)
        self.global_norm = nn.LayerNorm(dim) if use_global else None
        self.local_attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        self.global_attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) if use_global else None
        self.local_ff_norm = nn.LayerNorm(dim)
        self.local_ff = FeedForward(dim, ff_ratio, dropout)
        self.global_ff_norm = nn.LayerNorm(dim) if use_global else None
        self.global_ff = FeedForward(dim, ff_ratio, dropout) if use_global else None

    def forward(self, x: torch.Tensor, global_vectors: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor | None]:
        normalized = self.local_norm(x)
        windows, meta = cuboid_partition(normalized, self.cuboid)
        if self.use_global:
            if global_vectors is None:
                raise ValueError("global vectors are required when use_global=True")
            windows_per_batch = windows.shape[0] // x.shape[0]
            repeated_global = self.global_norm(global_vectors).repeat_interleave(windows_per_batch, dim=0)
            key_value = torch.cat((windows, repeated_global), dim=1)
        else:
            key_value = windows
        attended = self.local_attention(windows, key_value, key_value, need_weights=False)[0]
        x = x + cuboid_merge(attended, meta)
        x = x + self.local_ff(self.local_ff_norm(x))
        if self.use_global:
            global_query = self.global_norm(global_vectors)
            all_tokens = self.local_norm(x).reshape(x.shape[0], -1, x.shape[-1])
            global_kv = torch.cat((global_query, all_tokens), dim=1)
            global_vectors = global_vectors + self.global_attention(global_query, global_kv, global_kv, need_weights=False)[0]
            global_vectors = global_vectors + self.global_ff(self.global_ff_norm(global_vectors))
        return x, global_vectors


def resolve_pattern(pattern: str | list[list[int]], shape: tuple[int, int, int]) -> list[tuple[int, int, int]]:
    if pattern == "axial":
        t, h, w = shape
        return [(t, 1, 1), (1, h, 1), (1, 1, w)]
    return [tuple(int(value) for value in item) for item in pattern]


class CuboidBlock(nn.Module):
    def __init__(self, dim: int, heads: int, pattern: str | list[list[int]], shape: tuple[int, int, int], ff_ratio: float, dropout: float, use_global: bool):
        super().__init__()
        self.layers = nn.ModuleList(
            CuboidAttentionLayer(dim, heads, cuboid, ff_ratio, dropout, use_global)
            for cuboid in resolve_pattern(pattern, shape)
        )

    def forward(self, x: torch.Tensor, global_vectors: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor | None]:
        for layer in self.layers:
            x, global_vectors = layer(x, global_vectors)
        return x, global_vectors


class CuboidCrossAttention(nn.Module):
    """CuboidCross(T,1,1): future queries attend to history at each spatial site."""

    def __init__(self, dim: int, heads: int, ff_ratio: float, dropout: float):
        super().__init__()
        self.query_norm = nn.LayerNorm(dim)
        self.memory_norm = nn.LayerNorm(dim)
        self.attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        self.ff_norm = nn.LayerNorm(dim)
        self.ff = FeedForward(dim, ff_ratio, dropout)

    def forward(self, query: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
        b, k, h, w, c = query.shape
        if memory.shape[0] != b or memory.shape[2:4] != (h, w):
            raise ValueError("cross-attention memory must match batch and spatial dimensions")
        q = self.query_norm(query).permute(0, 2, 3, 1, 4).reshape(b * h * w, k, c)
        m = self.memory_norm(memory).permute(0, 2, 3, 1, 4).reshape(b * h * w, memory.shape[1], c)
        attended = self.attention(q, m, m, need_weights=False)[0]
        attended = attended.reshape(b, h, w, k, c).permute(0, 3, 1, 2, 4)
        query = query + attended
        return query + self.ff(self.ff_norm(query))


class DecoderBlock(nn.Module):
    def __init__(self, dim: int, heads: int, pattern: str | list[list[int]], shape: tuple[int, int, int], ff_ratio: float, dropout: float, use_global: bool):
        super().__init__()
        self.self_block = CuboidBlock(dim, heads, pattern, shape, ff_ratio, dropout, use_global)
        self.cross = CuboidCrossAttention(dim, heads, ff_ratio, dropout)

    def forward(self, x: torch.Tensor, memory: torch.Tensor, global_vectors: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor | None]:
        x, global_vectors = self.self_block(x, global_vectors)
        return self.cross(x, memory), global_vectors


class Earthformer(nn.Module):
    """Two-level Cuboid Attention Earthformer with BTHWC input and output."""

    def __init__(self, config: dict):
        super().__init__()
        data, model = config["data"], config["model"]
        self.input_length = int(data["input_length"])
        self.output_length = int(data["output_length"])
        self.height, self.width = int(data["height"]), int(data["width"])
        channels = int(data["channels"])
        d0, d1 = (int(value) for value in model["dims"])
        depths = model["depths"]
        heads = int(model["heads"])
        pattern = model.get("pattern", "axial")
        ff_ratio, dropout = float(model.get("ff_ratio", 2.0)), float(model.get("dropout", 0.0))
        self.num_global = int(model.get("num_global_vectors", 0))
        use_global = self.num_global > 0
        h0, w0, h1, w1 = self.height // 2, self.width // 2, self.height // 4, self.width // 4
        self.stem = nn.Conv2d(channels, d0, 3, stride=2, padding=1)
        self.downsample = nn.Conv2d(d0, d1, 3, stride=2, padding=1)
        self.encoder_pos0 = nn.Parameter(torch.zeros(1, self.input_length, h0, w0, d0))
        self.encoder_pos1 = nn.Parameter(torch.zeros(1, self.input_length, h1, w1, d1))
        self.future_query = nn.Parameter(torch.empty(1, self.output_length, h1, w1, d1))
        nn.init.trunc_normal_(self.future_query, std=0.02)
        self.encoder0 = nn.ModuleList(CuboidBlock(d0, heads, pattern, (self.input_length, h0, w0), ff_ratio, dropout, use_global) for _ in range(depths[0]))
        self.encoder1 = nn.ModuleList(CuboidBlock(d1, heads, pattern, (self.input_length, h1, w1), ff_ratio, dropout, use_global) for _ in range(depths[1]))
        self.decoder1 = nn.ModuleList(DecoderBlock(d1, heads, "axial", (self.output_length, h1, w1), ff_ratio, dropout, use_global) for _ in range(depths[1]))
        self.decoder0 = nn.ModuleList(DecoderBlock(d0, heads, "axial", (self.output_length, h0, w0), ff_ratio, dropout, use_global) for _ in range(depths[0]))
        self.up_project = nn.Conv2d(d1, d0, 3, padding=1)
        self.skip_project = nn.Linear(d0, d0)
        self.head = nn.Conv2d(d0, channels, 3, padding=1)
        if use_global:
            self.encoder_global0 = nn.Parameter(torch.zeros(1, self.num_global, d0))
            self.encoder_global1 = nn.Parameter(torch.zeros(1, self.num_global, d1))
            self.decoder_global1 = nn.Parameter(torch.zeros(1, self.num_global, d1))
            self.decoder_global0 = nn.Parameter(torch.zeros(1, self.num_global, d0))

    @staticmethod
    def _frames(module: nn.Module, x: torch.Tensor) -> torch.Tensor:
        b, t, h, w, c = x.shape
        result = module(x.permute(0, 1, 4, 2, 3).reshape(b * t, c, h, w))
        return result.reshape(b, t, result.shape[1], result.shape[2], result.shape[3]).permute(0, 1, 3, 4, 2)

    def _global(self, name: str, batch: int) -> torch.Tensor | None:
        value = getattr(self, name, None)
        return value.expand(batch, -1, -1) if value is not None else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        expected = (self.input_length, self.height, self.width)
        if x.ndim != 5 or x.shape[1:4] != expected:
            raise ValueError(f"expected input [B,{expected[0]},{expected[1]},{expected[2]},C], got {tuple(x.shape)}")
        batch = x.shape[0]
        e0 = self._frames(self.stem, x) + self.encoder_pos0
        g0 = self._global("encoder_global0", batch)
        for block in self.encoder0:
            e0, g0 = block(e0, g0)
        e1 = self._frames(self.downsample, e0) + self.encoder_pos1
        g1 = self._global("encoder_global1", batch)
        for block in self.encoder1:
            e1, g1 = block(e1, g1)
        d1 = self.future_query.expand(batch, -1, -1, -1, -1)
        gd1 = self._global("decoder_global1", batch)
        for block in self.decoder1:
            d1, gd1 = block(d1, e1, gd1)
        b, k, h, w, c = d1.shape
        up = F.interpolate(d1.permute(0, 1, 4, 2, 3).reshape(b * k, c, h, w), scale_factor=2, mode="nearest")
        d0 = self.up_project(up).reshape(b, k, -1, h * 2, w * 2).permute(0, 1, 3, 4, 2)
        d0 = d0 + self.skip_project(e0.mean(dim=1, keepdim=True)).expand(-1, k, -1, -1, -1)
        gd0 = self._global("decoder_global0", batch)
        for block in self.decoder0:
            d0, gd0 = block(d0, e0, gd0)
        b, k, h, w, c = d0.shape
        full = F.interpolate(d0.permute(0, 1, 4, 2, 3).reshape(b * k, c, h, w), scale_factor=2, mode="nearest")
        output = self.head(full)
        return output.reshape(b, k, -1, self.height, self.width).permute(0, 1, 3, 4, 2)