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