File size: 4,156 Bytes
5c365c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """ACE 40-channel input, 44-channel output and autoregressive rollout."""
from __future__ import annotations
from dataclasses import asdict, dataclass
from typing import Callable
import torch
from torch import nn
from ACE.model.variables import (
FORCING_CHANNELS,
INPUT_CHANNELS,
OUTPUT_CHANNELS,
PROGNOSTIC_CHANNELS,
split_output,
validate_channels,
)
from ACE.model.sfno import SFNOAdapter, SFNOConfig
@dataclass
class ACEModelConfig:
nlat: int = 180
nlon: int = 360
input_channels: int = len(INPUT_CHANNELS)
output_channels: int = len(OUTPUT_CHANNELS)
prognostic_channels: int = len(PROGNOSTIC_CHANNELS)
forcing_channels: int = len(FORCING_CHANNELS)
embed_dim: int = 256
num_layers: int = 8
filter_type: str = "linear"
operator_type: str = "dhconv"
scale_factor: int = 1
spectral_layers: int = 3
grid: str = "legendre-gauss"
grid_internal: str = "legendre-gauss"
mlp_ratio: float = 2.0
fallback: bool = False
def to_dict(self) -> dict:
return asdict(self)
class ACEModel(nn.Module):
def __init__(self, config: ACEModelConfig | None = None) -> None:
super().__init__()
self.config = config or ACEModelConfig()
if self.config.input_channels != len(INPUT_CHANNELS) or self.config.output_channels != len(OUTPUT_CHANNELS):
raise ValueError("ACE channel contract must remain 40 input and 44 output channels")
sfno_config = SFNOConfig(
nlat=self.config.nlat,
nlon=self.config.nlon,
in_channels=self.config.input_channels,
out_channels=self.config.output_channels,
embed_dim=self.config.embed_dim,
num_layers=self.config.num_layers,
filter_type=self.config.filter_type,
operator_type=self.config.operator_type,
scale_factor=self.config.scale_factor,
spectral_layers=self.config.spectral_layers,
grid=self.config.grid,
grid_internal=self.config.grid_internal,
mlp_ratio=self.config.mlp_ratio,
fallback=self.config.fallback,
)
self.sfno = SFNOAdapter(sfno_config)
def forward(self, x: torch.Tensor) -> torch.Tensor:
validate_channels(x, self.config.input_channels, "ACE input")
return self.sfno(x)
def step(self, prognostic: torch.Tensor, forcing: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
validate_channels(prognostic, self.config.prognostic_channels, "prognostic state")
validate_channels(forcing, self.config.forcing_channels, "forcing")
predicted = self.forward(torch.cat([prognostic, forcing], dim=1))
return split_output(predicted)
@torch.no_grad()
def rollout(
self,
initial_prognostic: torch.Tensor,
forcings: torch.Tensor | Callable[[int, torch.Tensor], torch.Tensor],
steps: int | None = None,
) -> torch.Tensor:
"""Return predictions with shape ``[B,T,44,H,W]``.
``forcings`` is either `[B,T,6,H,W]` or a callable receiving
`(step, current_prognostic)` and returning `[B,6,H,W]`.
"""
validate_channels(initial_prognostic, self.config.prognostic_channels, "initial prognostic")
if callable(forcings):
if steps is None or steps < 1:
raise ValueError("steps is required for callable forcing")
forcing_steps = steps
else:
if forcings.ndim != 5 or forcings.shape[2] != self.config.forcing_channels:
raise ValueError("tensor forcings must have shape [B,T,6,H,W]")
forcing_steps = forcings.shape[1] if steps is None else min(steps, forcings.shape[1])
state = initial_prognostic
outputs = []
for step in range(forcing_steps):
forcing = forcings[:, step] if not callable(forcings) else forcings(step, state)
predicted_state, diagnostics = self.step(state, forcing)
outputs.append(torch.cat([predicted_state, diagnostics], dim=1))
state = predicted_state
return torch.stack(outputs, dim=1)
|