model/surya.py

"""Self-contained, paper-aligned Surya forecasting model."""

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


CHANNEL_NAMES = ("AIA_94", "AIA_131", "AIA_171", "AIA_193", "AIA_211", "AIA_304",
                 "AIA_335", "AIA_1600", "HMI_magnetogram", "HMI_continuum",
                 "HMI_doppler", "HMI_vector_x", "HMI_vector_y")


def signum_log(x):
    """Compress signed solar products while retaining their sign."""
    return torch.sign(x) * torch.log1p(torch.abs(x))


class SpectralGating(nn.Module):
    def __init__(self, grid, dim):
        super().__init__()
        self.grid = grid
        self.weight = nn.Parameter(torch.randn(grid, grid // 2 + 1, dim, 2) * 0.02)
        self.norm = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim))

    def forward(self, tokens):
        residual = tokens
        values = self.norm(tokens).reshape(tokens.shape[0], self.grid, self.grid, -1).float()
        spectrum = torch.fft.rfft2(values, dim=(1, 2), norm="ortho")
        spectrum = spectrum * torch.view_as_complex(self.weight.float())
        values = torch.fft.irfft2(spectrum, s=(self.grid, self.grid), dim=(1, 2), norm="ortho")
        return residual + self.mlp(values.reshape_as(tokens).to(tokens.dtype))


class LongShortAttention(nn.Module):
    def __init__(self, grid, dim, heads, window, global_tokens):
        super().__init__()
        self.grid, self.window = grid, window
        self.norm = nn.LayerNorm(dim)
        self.local = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.global_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.projection = nn.Linear(dim * 2, dim)
        self.score = nn.Linear(dim, global_tokens)
        self.global_norm = nn.LayerNorm(dim)
        self.output_norm = nn.LayerNorm(dim)
        self.mlp_norm = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim))

    def forward(self, tokens):
        residual = tokens
        values = self.norm(tokens)
        b, _, d = values.shape
        grid = values.reshape(b, self.grid, self.grid, d)
        w = self.window
        # Unfold with padding creates overlapping neighborhoods at every token.
        padded = F.pad(grid.permute(0, 3, 1, 2), (w // 2, w // 2, w // 2, w // 2), mode="replicate")
        neighborhoods = F.unfold(padded, kernel_size=w, padding=0).transpose(1, 2)
        windows = neighborhoods.reshape(b, self.grid * self.grid, d, w * w)
        windows = windows.permute(0, 1, 3, 2).reshape(-1, w * w, d)
        local = self.local(windows, windows, windows, need_weights=False)[0]
        local = local[:, (w * w) // 2].reshape(b, self.grid * self.grid, d)
        mixing = self.score(values).transpose(1, 2).softmax(-1)
        compressed = self.global_norm(mixing @ values)
        global_context = self.global_attn(values, compressed, compressed, need_weights=False)[0]
        tokens = residual + self.projection(torch.cat((local, global_context), dim=-1))
        tokens = self.output_norm(tokens)
        return tokens + self.mlp(self.mlp_norm(tokens))


class Surya(nn.Module):
    def __init__(self, image_size=32, patch_size=4, channels=13, input_steps=2,
                 embed_dim=64, depth=4, spectral_blocks=1, num_heads=4,
                 window_size=2, global_tokens=4):
        super().__init__()
        if image_size % patch_size:
            raise ValueError("image_size must be divisible by patch_size")
        if spectral_blocks >= depth:
            raise ValueError("spectral_blocks must be smaller than depth")
        if embed_dim % num_heads:
            raise ValueError("embed_dim must be divisible by num_heads")
        self.image_size, self.patch_size = image_size, patch_size
        self.channels, self.input_steps = channels, input_steps
        self.grid = image_size // patch_size
        if channels != 13:
            raise ValueError("Surya requires the 13 SDO AIA/HMI channels")
        patch_dim = channels * patch_size ** 2
        self.patch_embed = nn.Linear(patch_dim * 2, embed_dim)
        self.temporal_embed = nn.Parameter(torch.zeros(1, input_steps, 1, embed_dim))
        self.position = nn.Parameter(torch.zeros(1, self.grid ** 2, embed_dim))
        blocks = [SpectralGating(self.grid, embed_dim) for _ in range(spectral_blocks)]
        blocks += [LongShortAttention(self.grid, embed_dim, num_heads, window_size, global_tokens)
                   for _ in range(depth - spectral_blocks)]
        self.blocks = nn.ModuleList(blocks)
        self.norm = nn.LayerNorm(embed_dim)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(embed_dim, embed_dim // 2, patch_size, stride=patch_size),
            nn.GELU(), nn.Conv2d(embed_dim // 2, channels, 3, padding=1))
        nn.init.normal_(self.position, std=0.02)
        nn.init.zeros_(self.decoder[-1].weight)
        nn.init.zeros_(self.decoder[-1].bias)

    def tokenize(self, frames):
        b, t, c, h, w = frames.shape
        if (t, c, h, w) != (self.input_steps, self.channels, self.image_size, self.image_size):
            raise ValueError("Expected BTCHW input matching configured dimensions")
        p = self.patch_size
        current = frames[:, -1]
        delta = current - frames[:, -2]
        values = torch.stack((current, delta), dim=1).reshape(b, 2 * c, h // p, p, w // p, p)
        values = values.permute(0, 2, 4, 1, 3, 5).reshape(b, self.grid ** 2, -1)
        return self.patch_embed(values) + self.position + self.temporal_embed[:, -1]

    def _predict(self, frames):
        tokens = self.tokenize(frames)
        for block in self.blocks:
            tokens = block(tokens)
        tokens = self.norm(tokens)
        features = tokens.mean(1)
        grid = tokens.transpose(1, 2).reshape(tokens.shape[0], -1, self.grid, self.grid)
        prediction = frames[:, -1] + self.decoder(grid)
        if not torch.isfinite(prediction).all():
            raise FloatingPointError("Surya produced a non-finite prediction")
        return {"prediction": prediction, "features": features}

    def forward(self, frames, steps=None):
        if steps is None:
            return self._predict(frames)
        if steps < 1:
            raise ValueError("steps must be positive")
        history, predictions = frames, []
        for _ in range(steps):
            prediction = self._predict(history)["prediction"]
            predictions.append(prediction)
            history = torch.cat((history[:, 1:], prediction[:, None]), dim=1)
        return torch.stack(predictions, dim=1)

    def rollout(self, frames, steps):
        return self(frames, steps=steps)

    def rollout_loss(self, frames, targets):
        predictions = self.rollout(frames, targets.shape[1])
        step_losses = (predictions - targets).square().mean(dim=(0, 2, 3, 4))
        return {"loss": step_losses.mean(), "step_losses": step_losses,
                "predictions": predictions}