model/seeds.py

"""PyTorch implementation of the SEEDS conditional diffusion core."""

from __future__ import annotations

import math
from typing import Optional

import torch
from torch import Tensor, nn


def _fourier_embedding(value: Tensor, dim: int, max_period: float = 10000.0) -> Tensor:
    """Return a deterministic sinusoidal embedding for diffusion time."""
    half = dim // 2
    frequencies = torch.exp(
        -math.log(max_period) * torch.arange(half, device=value.device, dtype=value.dtype) / max(half, 1)
    )
    angles = value[..., None] * frequencies
    embedding = torch.cat((angles.sin(), angles.cos()), dim=-1)
    if dim % 2:
        embedding = torch.nn.functional.pad(embedding, (0, 1))
    return embedding


class _AxialBlock(nn.Module):
    def __init__(self, dim: int, heads: int, mlp_ratio: int, dropout: float) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        hidden = dim * mlp_ratio
        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim), nn.Dropout(dropout))

    def forward(self, sequence: Tensor) -> Tensor:
        normalized = self.norm1(sequence)
        attended, _ = self.attention(normalized, normalized, normalized, need_weights=False)
        sequence = sequence + attended
        return sequence + self.mlp(self.norm2(sequence))


class SEEDS(nn.Module):
    """Conditional score network for cubed-sphere atmospheric fields.

    Inputs use ``[batch, channels, faces, height, width]`` for one snapshot and
    ``[batch, seeds, channels, faces, height, width]`` for seed forecasts.
    The output is the normalized noise prediction with the target snapshot shape.
    """

    def __init__(
        self,
        channels: int = 8,
        faces: int = 6,
        height: int = 48,
        width: int = 48,
        patch_size: int = 12,
        embed_dim: int = 768,
        spatial_layers: int = 6,
        field_layers: int = 4,
        sequence_layers: int = 6,
        mlp_ratio: int = 4,
        dropout: float = 0.0,
        seed_count: int = 2,
        sigma_min: float = 0.01,
        sigma_max: float = 100.0,
    ) -> None:
        super().__init__()
        if height % patch_size or width % patch_size:
            raise ValueError("height and width must be divisible by patch_size")
        if embed_dim % 2:
            raise ValueError("embed_dim must be even")
        self.channels, self.faces = channels, faces
        self.height, self.width = height, width
        self.patch_size, self.seed_count = patch_size, seed_count
        self.patch_rows, self.patch_cols = height // patch_size, width // patch_size
        self.patch_count = faces * self.patch_rows * self.patch_cols
        self.sigma_min, self.sigma_max = sigma_min, sigma_max
        heads = max(1, min(12, embed_dim // 64))
        while embed_dim % heads:
            heads -= 1

        self.patch_embedding = nn.Conv2d(channels, embed_dim, patch_size, patch_size)
        self.output_projection = nn.Linear(embed_dim, patch_size * patch_size)
        self.position_embedding = nn.Parameter(torch.zeros(1, 1, 1, self.patch_count, embed_dim))
        self.field_embedding = nn.Parameter(torch.zeros(1, 1, channels, 1, embed_dim))
        self.snapshot_embedding = nn.Parameter(torch.zeros(1, seed_count + 2, 1, 1, embed_dim))
        self.time_projection = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim))
        self.spatial_transformer = nn.ModuleList(
            [_AxialBlock(embed_dim, heads, mlp_ratio, dropout) for _ in range(spatial_layers)]
        )
        self.field_transformer = nn.ModuleList(
            [_AxialBlock(embed_dim, heads, mlp_ratio, dropout) for _ in range(field_layers)]
        )
        self.sequence_transformer = nn.ModuleList(
            [_AxialBlock(embed_dim, heads, mlp_ratio, dropout) for _ in range(sequence_layers)]
        )
        nn.init.normal_(self.position_embedding, std=0.02)
        nn.init.normal_(self.field_embedding, std=0.02)
        nn.init.normal_(self.snapshot_embedding, std=0.02)

    def _check_inputs(self, noisy: Tensor, seeds: Tensor, climate: Optional[Tensor]) -> None:
        expected = (self.channels, self.faces, self.height, self.width)
        if noisy.ndim != 5 or tuple(noisy.shape[1:]) != expected:
            raise ValueError(f"noisy must have shape [B, {expected}], got {tuple(noisy.shape)}")
        if seeds.ndim != 6 or tuple(seeds.shape[2:]) != expected or seeds.shape[1] != self.seed_count:
            raise ValueError(f"seeds must have shape [B, {self.seed_count}, {expected}], got {tuple(seeds.shape)}")
        if climate is not None and (climate.ndim != 5 or tuple(climate.shape[1:]) != expected):
            raise ValueError(f"climate must have shape [B, {expected}], got {tuple(climate.shape)}")

    def _embed_snapshot(self, snapshot: Tensor) -> Tensor:
        batch, channels, faces, _, _ = snapshot.shape
        embedded = self.patch_embedding(snapshot.permute(0, 2, 1, 3, 4).reshape(batch * faces, channels, self.height, self.width))
        embedded = embedded.flatten(2).transpose(1, 2).reshape(batch, faces * self.patch_rows * self.patch_cols, -1)
        return embedded

    def forward(self, noisy: Tensor, seeds: Tensor, climate: Optional[Tensor] = None, diffusion_time: Optional[Tensor] = None) -> Tensor:
        self._check_inputs(noisy, seeds, climate)
        batch = noisy.shape[0]
        if climate is None:
            climate = torch.zeros_like(noisy)
        snapshots = torch.cat((noisy[:, None], seeds, climate[:, None]), dim=1)
        sequence = torch.stack([self._embed_snapshot(snapshots[:, index]) for index in range(snapshots.shape[1])], dim=1)
        sequence = sequence[:, :, None] + self.position_embedding + self.field_embedding
        if diffusion_time is None:
            diffusion_time = torch.zeros(batch, device=noisy.device, dtype=noisy.dtype)
        time = self.time_projection(_fourier_embedding(diffusion_time, sequence.shape[-1])).to(sequence.dtype)
        sequence[:, 0] = sequence[:, 0] + time[:, None, None]
        sequence = sequence.expand(-1, -1, self.channels, -1, -1) + self.snapshot_embedding[:, : sequence.shape[1]]

        shape = sequence.shape
        sequence = sequence.reshape(batch * shape[1] * shape[2], shape[3], shape[4])
        for block in self.spatial_transformer:
            sequence = block(sequence)
        sequence = sequence.reshape(batch * shape[1] * shape[3], shape[2], shape[4])
        for block in self.field_transformer:
            sequence = block(sequence)
        sequence = sequence.reshape(batch * shape[2] * shape[3], shape[1], shape[4])
        for block in self.sequence_transformer:
            sequence = block(sequence)
        sequence = sequence.reshape(batch, shape[1], shape[2], shape[3], shape[4])[:, 0]
        patches = self.output_projection(sequence).reshape(batch, self.channels, self.faces, self.patch_rows, self.patch_cols, self.patch_size, self.patch_size)
        return patches.permute(0, 1, 2, 3, 5, 4, 6).reshape(batch, self.channels, self.faces, self.height, self.width)

    def sigma(self, diffusion_time: Tensor) -> Tensor:
        return self.sigma_min * (self.sigma_max / self.sigma_min) ** diffusion_time

    def denoising_loss(
        self,
        clean: Tensor,
        seeds: Tensor,
        climate: Optional[Tensor] = None,
        diffusion_time: Optional[Tensor] = None,
        noise: Optional[Tensor] = None,
    ) -> Tensor:
        if diffusion_time is None:
            diffusion_time = torch.rand(clean.shape[0], device=clean.device, dtype=clean.dtype)
        if noise is None:
            noise = torch.randn_like(clean)
        sigma = self.sigma(diffusion_time).view(-1, 1, 1, 1, 1)
        noisy = clean + sigma * noise
        model_input = noisy / torch.sqrt(1.0 + sigma.square())
        prediction = self(model_input, seeds, climate, diffusion_time)
        return ((prediction - noise) ** 2).flatten(1).mean()

    @torch.no_grad()
    def sample(
        self,
        seeds: Tensor,
        climate: Optional[Tensor] = None,
        members: int = 1,
        steps: int = 64,
        member_batch_size: Optional[int] = None,
    ) -> Tensor:
        if members < 1 or steps < 1:
            raise ValueError("members and steps must be positive")
        chunk_size = min(member_batch_size or members, members)
        generated = []
        schedule = torch.linspace(1.0, 0.0, steps + 1, device=seeds.device, dtype=seeds.dtype)
        sigma_schedule = self.sigma(schedule)
        for start in range(0, members, chunk_size):
            current_members = min(chunk_size, members - start)
            expanded_seeds = seeds.repeat_interleave(current_members, dim=0)
            expanded_climate = None if climate is None else climate.repeat_interleave(current_members, dim=0)
            sample = torch.randn_like(expanded_seeds[:, 0]) * sigma_schedule[0]
            for index, current in enumerate(schedule[:-1]):
                current_time = torch.full((sample.shape[0],), current, device=sample.device, dtype=sample.dtype)
                sigma = sigma_schedule[index]
                model_input = sample / torch.sqrt(1.0 + sigma.square())
                predicted_noise = self(model_input, expanded_seeds, expanded_climate, current_time)
                sample = sample + (sigma_schedule[index + 1] - sigma) * predicted_noise
            generated.append(
                sample.reshape(seeds.shape[0], current_members, self.channels, self.faces, self.height, self.width)
            )
        return torch.cat(generated, dim=1)


SEEDSModel = SEEDS