model/ace2.py

import os
import random
from pathlib import Path

import numpy as np
import torch
import yaml
from torch import nn


LEVELS_HPA = (1000, 850, 700, 500, 300, 200, 100, 50)
CHANNELS = (
    [f"temperature_{p}" for p in LEVELS_HPA]
    + [f"specific_humidity_{p}" for p in LEVELS_HPA]
    + [f"u_wind_{p}" for p in LEVELS_HPA]
    + [f"v_wind_{p}" for p in LEVELS_HPA]
    + [f"geopotential_{p}" for p in LEVELS_HPA]
    + [
        "surface_pressure",
        "air_temperature_2m",
        "specific_humidity_2m",
        "eastward_wind_10m",
        "northward_wind_10m",
        "sea_surface_temperature",
        "total_precipitation_6h",
        "surface_downward_shortwave",
        "surface_downward_longwave",
        "toa_outgoing_longwave",
    ]
)
assert len(CHANNELS) == 50
Q_INDICES = tuple(range(8, 16)) + (42,)
SURFACE_PRESSURE = 40
PRECIPITATION = 46
RADIATION_INDICES = (47, 48, 49)


class SpectralConv2d(nn.Module):
    def __init__(self, width, modes_lat, modes_lon):
        super().__init__()
        self.modes_lat, self.modes_lon = modes_lat, modes_lon
        scale = 1.0 / width
        self.weight = nn.Parameter(
            scale * torch.randn(width, width, modes_lat, modes_lon, dtype=torch.cfloat)
        )

    def forward(self, x):
        spectrum = torch.fft.rfft2(x, norm="ortho")
        out = torch.zeros_like(spectrum)
        ml = min(self.modes_lat, spectrum.shape[-2])
        mn = min(self.modes_lon, spectrum.shape[-1])
        out[:, :, :ml, :mn] = torch.einsum(
            "bixy,ioxy->boxy", spectrum[:, :, :ml, :mn], self.weight[:, :, :ml, :mn]
        )
        return torch.fft.irfft2(out, s=x.shape[-2:], norm="ortho")


class SFNOBlock(nn.Module):
    def __init__(self, width, modes_lat, modes_lon):
        super().__init__()
        self.spectral = SpectralConv2d(width, modes_lat, modes_lon)
        self.mlp = nn.Sequential(
            nn.Conv2d(width, width * 2, 1), nn.GELU(), nn.Conv2d(width * 2, width, 1)
        )
        self.norm = nn.GroupNorm(1, width)

    def forward(self, x):
        return x + self.mlp(self.norm(self.spectral(x)))


class CompactSFNO(nn.Module):
    def __init__(self, channels=50, forcing_channels=4, width=4, depth=1,
                 modes_lat=4, modes_lon=4):
        super().__init__()
        self.lift = nn.Conv2d(channels + forcing_channels, width, 1)
        self.blocks = nn.Sequential(
            *[SFNOBlock(width, modes_lat, modes_lon) for _ in range(depth)]
        )
        self.project = nn.Sequential(nn.GELU(), nn.Conv2d(width, channels, 1))

    def forward(self, state, forcing):
        features = self.blocks(self.lift(torch.cat((state, forcing), dim=1)))
        return state + self.project(features)


def area_weights(height, device, dtype):
    lat = torch.linspace(-89.5, 89.5, height, device=device, dtype=dtype)
    return torch.cos(torch.deg2rad(lat)).view(1, 1, height, 1)


def weighted_mean(x, weights):
    return (x * weights).sum(dim=(-2, -1), keepdim=True) / (
        weights.sum(dim=(-2, -1), keepdim=True) * x.shape[-1]
    )


def hard_correct(previous, predicted):
    """Apply differentiable positivity, dry-mass, and global-water constraints."""
    out = predicted.clone()
    positive = list(Q_INDICES) + [PRECIPITATION] + list(RADIATION_INDICES)
    out[:, positive] = torch.clamp_min(out[:, positive], 0.0)
    weights = area_weights(out.shape[-2], out.device, out.dtype)

    q_prev = previous[:, Q_INDICES].sum(dim=1, keepdim=True)
    water_target = weighted_mean(q_prev, weights)
    precip = weighted_mean(out[:, PRECIPITATION:PRECIPITATION + 1], weights)
    precip_scale = torch.clamp(
        0.5 * water_target / torch.clamp_min(precip, 1e-8), max=1.0
    )
    out[:, PRECIPITATION:PRECIPITATION + 1] *= precip_scale
    precip = weighted_mean(out[:, PRECIPITATION:PRECIPITATION + 1], weights)
    q_target = torch.clamp_min(water_target - precip, 0.0)
    q_now = weighted_mean(out[:, Q_INDICES].sum(dim=1, keepdim=True), weights)
    out[:, Q_INDICES] *= q_target / torch.clamp_min(q_now, 1e-8)

    q_new = out[:, Q_INDICES].sum(dim=1, keepdim=True)
    dry_target = weighted_mean(
        previous[:, SURFACE_PRESSURE:SURFACE_PRESSURE + 1] - q_prev, weights
    )
    dry_now = weighted_mean(
        out[:, SURFACE_PRESSURE:SURFACE_PRESSURE + 1] - q_new, weights
    )
    out[:, SURFACE_PRESSURE:SURFACE_PRESSURE + 1] += dry_target - dry_now
    return out


def load_config(root=None):
    root = Path(root) if root is not None else Path(__file__).resolve().parents[1]
    with (root / "conf" / "config.yaml").open(encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def seed_all(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def forcing_for_hours(hours, height=180, width=360):
    hours = np.asarray(hours, dtype=np.float32)
    phase = 2 * np.pi * hours / (365.25 * 24)
    lat = np.deg2rad(np.linspace(-89.5, 89.5, height, dtype=np.float32))
    lon = np.deg2rad(np.linspace(0.5, 359.5, width, dtype=np.float32))
    solar = np.maximum(
        0,
        np.cos(lat)[None, :, None]
        * np.cos(lon[None, None, :] + phase[:, None, None]),
    )
    fields = np.empty((len(hours), 4, height, width), dtype=np.float32)
    fields[:, 0] = np.sin(phase)[:, None, None]
    fields[:, 1] = np.cos(phase)[:, None, None]
    fields[:, 2] = (400.0 + 0.01 * hours)[:, None, None] / 500.0
    fields[:, 3] = solar
    return fields


def init_distributed():
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    use_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= world_size
    if distributed:
        backend = "nccl" if use_cuda else "gloo"
        torch.distributed.init_process_group(backend=backend)
        rank = torch.distributed.get_rank()
        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    else:
        rank = local_rank = 0
    device = torch.device(
        f"cuda:{local_rank}" if use_cuda else "cpu"
    )
    if device.type == "cuda":
        torch.cuda.set_device(device)
    return distributed, rank, device


def build_model(config):
    return CompactSFNO(channels=config["data"]["channels"], **config["model"])