| """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 |
| |
| 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} |
|
|
