"""Elastic-net logistic calibration of WoFS ensemble storm-track hazards.""" from __future__ import annotations import numpy as np import torch from torch import nn HAZARDS = ("tornado", "hail", "wind") LEAD_GROUPS = ("first_hour", "second_hour") def _pav(values: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """Fit an isotonic map with the pool-adjacent-violators algorithm.""" order = np.argsort(values, kind="stable") x, y = values[order], targets[order].astype(np.float64) starts, ends, sums, counts = [], [], [], [] for index, target in enumerate(y): starts.append(index); ends.append(index); sums.append(float(target)); counts.append(1) while len(sums) > 1 and sums[-2] / counts[-2] > sums[-1] / counts[-1]: ends[-2] = ends[-1] sums[-2] += sums[-1] counts[-2] += counts[-1] starts.pop(); ends.pop(); sums.pop(); counts.pop() xp, yp = [], [] for start, end, total, count in zip(starts, ends, sums, counts): level = total / count xp.extend((float(x[start]), float(x[end]))) yp.extend((level, level)) xp = np.maximum.accumulate(np.asarray(xp, dtype=np.float32)) return xp, np.asarray(yp, dtype=np.float32) class WoFSStormCal(nn.Module): """Two lead-group linear classifiers with portable isotonic calibration.""" input_dim = 113 output_dim = 3 hazards = HAZARDS lead_groups = LEAD_GROUPS ensemble_members = 18 grid_spacing_km = 3 forecast_window_minutes = 30 forecast_interval_minutes = 5 def __init__(self, calibration_points: int = 256): super().__init__() self.calibration_points = int(calibration_points) self.weight = nn.Parameter(torch.empty(2, 3, self.input_dim)) self.bias = nn.Parameter(torch.zeros(2, 3)) nn.init.normal_(self.weight, std=0.01) self.register_buffer("feature_mean", torch.zeros(2, self.input_dim)) self.register_buffer("feature_scale", torch.ones(2, self.input_dim)) grid = torch.linspace(0, 1, self.calibration_points) self.register_buffer("calibration_x", grid.expand(2, 3, -1).clone()) self.register_buffer("calibration_y", grid.expand(2, 3, -1).clone()) self.register_buffer("calibration_length", torch.full((2, 3), self.calibration_points, dtype=torch.long)) @staticmethod def validate_features(features: torch.Tensor) -> None: if features.ndim != 2 or features.shape[1] != 113: raise ValueError(f"features must have shape [N,113], got {tuple(features.shape)}") if not torch.isfinite(features).all(): raise ValueError("features contain NaN or Inf") @staticmethod def validate_lead_group(lead_group: torch.Tensor, samples: int) -> None: if lead_group.ndim != 1 or len(lead_group) != samples: raise ValueError(f"lead_group must have shape [N], got {tuple(lead_group.shape)}") if bool(((lead_group < 0) | (lead_group > 1)).any()): raise ValueError("lead_group values must be 0 (first hour) or 1 (second hour)") def set_normalization(self, mean: torch.Tensor, scale: torch.Tensor) -> None: if mean.shape != (2, 113) or scale.shape != (2, 113): raise ValueError("normalization statistics must both have shape [2,113]") self.feature_mean.copy_(mean) self.feature_scale.copy_(scale.clamp_min(1e-6)) def logits(self, features: torch.Tensor, lead_group: torch.Tensor) -> torch.Tensor: self.validate_features(features) lead_group = lead_group.to(device=features.device, dtype=torch.long) self.validate_lead_group(lead_group, len(features)) normalized = (features - self.feature_mean[lead_group]) / self.feature_scale[lead_group] return torch.einsum("ni,noi->no", normalized, self.weight[lead_group]) + self.bias[lead_group] def _calibrate(self, probabilities: torch.Tensor, lead_group: torch.Tensor) -> torch.Tensor: result = torch.empty_like(probabilities) for group in range(2): mask = lead_group == group if not bool(mask.any()): continue for hazard in range(3): length = int(self.calibration_length[group, hazard]) xp = self.calibration_x[group, hazard, :length] yp = self.calibration_y[group, hazard, :length] value = probabilities[mask, hazard].clamp(xp[0], xp[-1]) upper = torch.searchsorted(xp.contiguous(), value.contiguous()).clamp(1, length - 1) lower = upper - 1 fraction = (value - xp[lower]) / (xp[upper] - xp[lower]).clamp_min(1e-7) result[mask, hazard] = yp[lower] + fraction * (yp[upper] - yp[lower]) return result.clamp(0, 1) def forward(self, features: torch.Tensor, lead_group: torch.Tensor, calibrated: bool = True) -> torch.Tensor: probabilities = torch.sigmoid(self.logits(features, lead_group)) return self._calibrate(probabilities, lead_group.to(probabilities.device)) if calibrated else probabilities @torch.no_grad() def fit_calibration(self, features: torch.Tensor, targets: torch.Tensor, lead_group: torch.Tensor) -> None: if targets.shape != (len(features), 3): raise ValueError(f"targets must have shape [N,3], got {tuple(targets.shape)}") probabilities = torch.sigmoid(self.logits(features, lead_group)).cpu().numpy() target_array, groups = targets.cpu().numpy(), lead_group.cpu().numpy() for group in range(2): for hazard in range(3): mask = groups == group xp, yp = _pav(probabilities[mask, hazard], target_array[mask, hazard]) if len(xp) > self.calibration_points: selected = np.linspace(0, len(xp) - 1, self.calibration_points).round().astype(int) xp, yp = xp[selected], yp[selected] if len(xp) == 1: xp, yp = np.repeat(xp, 2), np.repeat(yp, 2) length = len(xp) self.calibration_x[group, hazard, :length] = torch.from_numpy(xp).to(self.calibration_x) self.calibration_y[group, hazard, :length] = torch.from_numpy(yp).to(self.calibration_y) self.calibration_length[group, hazard] = length def elastic_net_loss(self, logits: torch.Tensor, targets: torch.Tensor, l1: float, l2: float) -> torch.Tensor: if logits.shape != targets.shape or logits.ndim != 2 or logits.shape[1] != 3: raise ValueError("logits and targets must both have shape [N,3]") bce = nn.functional.binary_cross_entropy_with_logits(logits, targets) return bce + float(l1) * self.weight.abs().mean() + 0.5 * float(l2) * self.weight.square().mean()