File size: 6,862 Bytes
fa2b79f | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """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()
|