model/metnet_2.py

"""Shape-faithful, memory-bounded MetNet-2 engineering implementation."""
from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Iterable

import numpy as np
import torch
from torch import Tensor, nn
import torch.nn.functional as F
from torch.utils.data import Dataset
import yaml

CHANNEL_GROUPS = (
    ("mrms_radar_history", 33), ("hrrr_atmosphere_history", 484),
    ("goes_satellite_history", 96), ("static_geography", 24),
    ("time_coordinates", 4),
)
LOGICAL_SHAPE = (641, 512, 512)
CLASS_RATES = np.linspace(0.0, 102.4, 512, dtype=np.float32)
assert sum(size for _, size in CHANNEL_GROUPS) == LOGICAL_SHAPE[0]


def load_config(path: str | Path = "conf/config.yaml") -> dict:
    with Path(path).open(encoding="utf-8") as handle:
        return yaml.safe_load(handle)


class ProceduralField:
    """Generate crops of a logical [641, 512, 512] field without materializing it."""

    shape = LOGICAL_SHAPE

    def __init__(self, seed: int):
        self.seed = int(seed)

    def window(self, y: int, x: int, size: int, halo: int = 0) -> Tensor:
        if size <= 0 or halo < 0 or not (0 <= y < 512 and 0 <= x < 512):
            raise ValueError("invalid selected-window coordinates")
        yy = torch.arange(y - halo, y + size + halo).clamp(0, 511).float()
        xx = torch.arange(x - halo, x + size + halo).clamp(0, 511).float()
        channels = torch.arange(641).float()[:, None, None]
        return (torch.sin((channels + self.seed) * .017 + yy[None, :, None] * .031)
                + torch.cos((channels + 3 * self.seed) * .011 + xx[None, None, :] * .023)).float()

    def target_window(self, y: int, x: int, size: int, lead: int) -> Tensor:
        yy = torch.arange(y, y + size)[:, None]
        xx = torch.arange(x, x + size)[None, :]
        return ((yy * 7 + xx * 11 + self.seed + lead // 2) % 512).long()


class WindowDataset(Dataset):
    def __init__(self, data_path: str | Path, split: str = "train"):
        with np.load(data_path) as data:
            required = ("seed", "split", "y", "x", "size", "halo", "lead_minutes")
            missing = set(required).difference(data.files)
            if missing:
                raise ValueError(f"dataset is missing fields: {sorted(missing)}")
            indices = np.flatnonzero(data["split"].astype(str) == split)
            self.records = [{key: data[key][i].item() for key in required} for i in indices]

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> tuple[Tensor, Tensor, Tensor]:
        record = self.records[index]
        field = ProceduralField(record["seed"])
        args = record["y"], record["x"], record["size"]
        return (field.window(*args, record["halo"]),
                field.target_window(*args, record["lead_minutes"]),
                torch.tensor(record["lead_minutes"], dtype=torch.long))


def write_fake_data(path: str | Path, samples: int = 8, window: int = 32, halo: int = 8) -> Path:
    if samples < 3 or window != 32 or window + 2 * halo > 512:
        raise ValueError("fake data requires at least 3 selected 32x32 windows with a valid halo")
    records = []
    for i in range(samples):
        records.append({
            "id": f"sample-{i:04d}", "seed": 1000 + i,
            "split": "train" if i < samples - 2 else "test",
            "y": (i * 47) % (512 - window + 1), "x": (i * 83) % (512 - window + 1),
            "size": window, "halo": halo, "lead_minutes": 2 + 2 * (i % 360),
        })
    output = Path(path)
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, **{key: np.asarray([r[key] for r in records]) for key in records[0]})
    return output


class LeadFiLMConv(nn.Module):
    def __init__(self, cin: int, cout: int, dilation: int = 1):
        super().__init__()
        self.conv = nn.Conv2d(cin, cout, 3, padding=dilation, dilation=dilation)
        self.film = nn.Linear(cout, 2 * cout)

    def forward(self, x: Tensor, lead: Tensor) -> Tensor:
        result = self.conv(x)
        add, multiply = self.film(lead).chunk(2, dim=1)
        return result * (1.0 + torch.tanh(multiply)[:, :, None, None]) + add[:, :, None, None]


class ConvLSTMCell(nn.Module):
    def __init__(self, cin: int, hidden: int):
        super().__init__()
        self.hidden = hidden
        self.gates = nn.Conv2d(cin + hidden, 4 * hidden, 3, padding=1)

    def forward(self, x: Tensor, state: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, Tensor]:
        if state is None:
            shape = (x.shape[0], self.hidden, x.shape[-2], x.shape[-1])
            state = x.new_zeros(shape), x.new_zeros(shape)
        hidden, cell = state
        in_gate, forget, candidate, out_gate = self.gates(torch.cat((x, hidden), 1)).chunk(4, 1)
        cell = torch.sigmoid(forget) * cell + torch.sigmoid(in_gate) * torch.tanh(candidate)
        return torch.sigmoid(out_gate) * torch.tanh(cell), cell


class DilatedResidualBlock(nn.Module):
    def __init__(self, width: int, dilation: int):
        super().__init__()
        self.conv1 = LeadFiLMConv(width, width, dilation)
        self.conv2 = LeadFiLMConv(width, width, dilation)

    def forward(self, x: Tensor, lead: Tensor) -> Tensor:
        return x + self.conv2(F.relu(self.conv1(F.relu(x), lead)), lead)


class MetNet2(nn.Module):
    """MetNet-2 concept model retaining the 641-channel and 512-class contracts."""

    def __init__(self, input_channels: int = 641, classes: int = 512, width: int = 8,
                 stacks: int = 1, dilations: Iterable[int] = (1, 2, 4, 8, 16, 32, 64, 128),
                 lead_max_minutes: int = 720):
        super().__init__()
        if input_channels != 641 or classes != 512:
            raise ValueError("MetNet-2 requires 641 input channels and 512 output classes")
        self.input_channels, self.classes = input_channels, classes
        self.width, self.stacks = width, stacks
        self.dilations = tuple(dilations)
        self.lead_max_minutes, self.upscale = lead_max_minutes, 4
        self.lead_embedding = nn.Sequential(nn.Linear(1, width), nn.SiLU(), nn.Linear(width, width))
        self.input_projection = nn.Conv2d(input_channels, width, 1)
        self.temporal = ConvLSTMCell(width, width)
        self.blocks = nn.ModuleList(DilatedResidualBlock(width, dilation)
                                    for _ in range(stacks) for dilation in self.dilations)
        self.spatial = LeadFiLMConv(width, width)
        self.head = nn.Conv2d(width, classes, 1)

    def _lead(self, minutes: Tensor) -> Tensor:
        if torch.any((minutes < 2) | (minutes > self.lead_max_minutes) | (minutes % 2 != 0)):
            raise ValueError("lead time must be 2..720 minutes in 2-minute increments")
        return self.lead_embedding((minutes.float() / self.lead_max_minutes).unsqueeze(1))

    def _features(self, x: Tensor, lead_minutes: Tensor, output_size: int) -> Tensor:
        if x.ndim != 4 or x.shape[1] != 641:
            raise ValueError("x must have shape [B, 641, H, W]")
        if output_size <= 0 or output_size % self.upscale:
            raise ValueError("output_size must be positive and divisible by four")
        lead = self._lead(lead_minutes.to(x.device))
        features, _ = self.temporal(self.input_projection(x))
        for block in self.blocks:
            features = block(features, lead)
        features = self.spatial(F.relu(features), lead)
        crop = output_size // self.upscale
        if min(features.shape[-2:]) < crop:
            raise ValueError("input window is smaller than the requested output")
        top, left = (features.shape[-2] - crop) // 2, (features.shape[-1] - crop) // 2
        return F.interpolate(features[:, :, top:top + crop, left:left + crop], size=(output_size, output_size),
                             mode="bilinear", align_corners=False)

    def forward_window(self, x: Tensor, lead_minutes: Tensor, output_size: int = 32,
                       class_slice: tuple[int, int] | None = None) -> Tensor:
        features = self._features(x, lead_minutes, output_size)
        start, end = class_slice or (0, self.classes)
        if not (0 <= start < end <= self.classes):
            raise ValueError("invalid class slice")
        return F.conv2d(features, self.head.weight[start:end], self.head.bias[start:end])

    def forward(self, x: Tensor, lead_minutes: Tensor, output_size: int = 32) -> Tensor:
        return self.forward_window(x, lead_minutes, output_size)

    @torch.no_grad()
    def assemble_full(self, source: ProceduralField, lead_minutes: int, output_path: str | Path,
                      tile: int = 32, halo: int = 8, class_chunk: int = 64,
                      output: str = "probability", device: str | torch.device = "cpu") -> Path:
        """Stream a complete [512, 512, 512] probability or CDF array to disk."""
        if output not in {"probability", "cdf"}:
            raise ValueError("output must be probability or cdf")
        path = Path(output_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        array = np.lib.format.open_memmap(path, mode="w+", dtype=np.float16, shape=(512, 512, 512))
        self.eval().to(device)
        lead = torch.tensor([lead_minutes], device=device)
        for y in range(0, 512, tile):
            for x0 in range(0, 512, tile):
                size = min(tile, 512 - y, 512 - x0)
                features = self._features(source.window(y, x0, size, halo).unsqueeze(0).to(device), lead, size)[0]
                maximum = None
                for start in range(0, 512, class_chunk):
                    logits = F.conv2d(features.unsqueeze(0), self.head.weight[start:start + class_chunk],
                                      self.head.bias[start:start + class_chunk])[0]
                    value = logits.amax(0)
                    maximum = value if maximum is None else torch.maximum(maximum, value)
                denominator = torch.zeros_like(maximum)
                chunks = []
                for start in range(0, 512, class_chunk):
                    logits = F.conv2d(features.unsqueeze(0), self.head.weight[start:start + class_chunk],
                                      self.head.bias[start:start + class_chunk])[0]
                    exponent = torch.exp(logits - maximum)
                    denominator += exponent.sum(0)
                    chunks.append(exponent)
                cumulative = torch.zeros_like(maximum)
                for start, exponent in zip(range(0, 512, class_chunk), chunks):
                    values = exponent / denominator
                    if output == "cdf":
                        values = values.cumsum(0) + cumulative
                        cumulative = values[-1]
                    array[start:start + values.shape[0], y:y + size, x0:x0 + size] = values.cpu().numpy()
        array.flush()
        return path


def build_model(config: dict, paper: bool = False) -> MetNet2:
    values = dict(config["model"])
    if paper:
        values.update({key: value for key, value in config["paper_model"].items()
                       if key in {"input_channels", "classes", "stacks", "dilations"}})
    dilations = tuple(values.get("dilations", ()))
    if dilations != (1, 2, 4, 8, 16, 32, 64, 128):
        raise ValueError("each dilation stack must use rates 1,2,4,8,16,32,64,128")
    if paper and values["stacks"] != 3:
        raise ValueError("the paper model requires three dilation stacks")
    return MetNet2(**values)


def categorical_nll_chunked(model: MetNet2, x: Tensor, lead: Tensor, target: Tensor,
                            output_size: int = 32, class_chunk: int = 64) -> Tensor:
    """Compute exact categorical NLL while applying the class head in chunks."""
    features = model._features(x, lead, output_size)
    selected, logsumexp = torch.zeros_like(target, dtype=features.dtype), None
    for start in range(0, model.classes, class_chunk):
        end = min(start + class_chunk, model.classes)
        logits = F.conv2d(features, model.head.weight[start:end], model.head.bias[start:end])
        part = torch.logsumexp(logits, dim=1)
        logsumexp = part if logsumexp is None else torch.logaddexp(logsumexp, part)
        mask = (target >= start) & (target < end)
        picked = logits.gather(1, (target - start).clamp(0, end - start - 1).unsqueeze(1)).squeeze(1)
        selected = torch.where(mask, picked, selected)
    return (logsumexp - selected).mean()


def save_checkpoint(path: str | Path, model: nn.Module, model_config: dict) -> None:
    if int(os.environ.get("RANK", "0")) != 0:
        return
    module = model.module if hasattr(model, "module") else model
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = Path(f"{destination}.tmp")
    torch.save({"model": module.state_dict(), "model_config": model_config,
                "format_version": "metnet_2_v1"}, temporary)
    os.replace(temporary, destination)


def load_checkpoint(path: str | Path, model: nn.Module) -> dict:
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    if set(checkpoint) != {"model", "model_config", "format_version"}:
        raise ValueError("checkpoint must contain model, model_config, and format_version")
    model.load_state_dict(checkpoint["model"])
    return checkpoint


def scores(probabilities: np.ndarray, target: np.ndarray,
           thresholds: tuple[float, ...] = (.2, 1., 2., 4., 8.)) -> dict:
    if probabilities.shape[0] != 512 or target.shape != probabilities.shape[1:]:
        raise ValueError("expected probabilities [512,H,W] and target [H,W]")
    cdf = np.cumsum(probabilities.astype(np.float32), axis=0)
    observed_cdf = (np.arange(512)[:, None, None] >= target[None]).astype(np.float32)
    result = {"discrete_crps": float(np.mean(np.sum((cdf - observed_cdf) ** 2, axis=0)))}
    brier, csi = {}, {}
    for threshold in thresholds:
        index = min(511, int(round(threshold / .2)))
        event_probability = 1.0 - cdf[index - 1] if index else np.ones_like(cdf[0])
        observed, forecast = target >= index, event_probability >= .5
        hits = np.logical_and(forecast, observed).sum()
        denominator = hits + np.logical_and(forecast, ~observed).sum() + np.logical_and(~forecast, observed).sum()
        brier[str(threshold)] = float(np.mean((event_probability - observed) ** 2))
        csi[str(threshold)] = float(hits / denominator) if denominator else 1.0
    result.update(brier=brier, csi=csi)
    return result


def write_json(path: str | Path, value: dict) -> None:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")