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