| """PPNN NN-aux-emb model and proper scoring rules.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| FORMAT_VERSION = "ppnn_nn_aux_emb_v1" |
| MEMBER_COUNT = 50 |
| VARIABLE_COUNT = 18 |
| STATION_COUNT = 537 |
| LEAD_HOURS = 48 |
| AUX_FEATURE_COUNT = 4 |
| CONTINUOUS_FEATURE_COUNT = VARIABLE_COUNT * 2 + AUX_FEATURE_COUNT |
| EMBEDDING_DIM = 2 |
| INPUT_FEATURE_COUNT = CONTINUOUS_FEATURE_COUNT + EMBEDDING_DIM |
|
|
|
|
| class PPNN(nn.Module): |
| """NN-aux-emb: ensemble moments, auxiliary predictors, and station embedding.""" |
|
|
| def __init__(self, hidden_size: int = 32, station_count: int = STATION_COUNT, eps: float = 1e-4): |
| super().__init__() |
| if station_count != STATION_COUNT: |
| raise ValueError(f"station_count must be the paper value {STATION_COUNT}") |
| self.station_count = station_count |
| self.eps = float(eps) |
| self.station_embedding = nn.Embedding(station_count, EMBEDDING_DIM) |
| self.hidden = nn.Linear(INPUT_FEATURE_COUNT, hidden_size) |
| self.output = nn.Linear(hidden_size, 2) |
|
|
| def forward(self, continuous: torch.Tensor, station_index: torch.Tensor): |
| if continuous.ndim != 2 or continuous.shape[-1] != CONTINUOUS_FEATURE_COUNT: |
| raise ValueError(f"continuous features must have shape [batch,{CONTINUOUS_FEATURE_COUNT}]") |
| if station_index.ndim != 1 or station_index.shape[0] != continuous.shape[0]: |
| raise ValueError("station_index must have shape [batch]") |
| if station_index.numel() and (station_index.min() < 0 or station_index.max() >= self.station_count): |
| raise ValueError("station index outside [0, 536]") |
| features = torch.cat((continuous, self.station_embedding(station_index)), dim=-1) |
| mu, raw_sigma = self.output(torch.relu(self.hidden(features))).unbind(dim=-1) |
| return mu, raw_sigma.abs() + self.eps |
|
|
|
|
| def ensemble_features(ensemble: torch.Tensor, auxiliary: torch.Tensor) -> torch.Tensor: |
| """Reduce [batch, 50, 18] to 36 moments and append four auxiliary features.""" |
| if ensemble.ndim != 3 or ensemble.shape[1:] != (MEMBER_COUNT, VARIABLE_COUNT): |
| raise ValueError(f"ensemble must have shape [batch,{MEMBER_COUNT},{VARIABLE_COUNT}]") |
| if auxiliary.ndim != 2 or auxiliary.shape != (ensemble.shape[0], AUX_FEATURE_COUNT): |
| raise ValueError(f"auxiliary must have shape [batch,{AUX_FEATURE_COUNT}]") |
| moments = torch.cat((ensemble.mean(dim=1), ensemble.std(dim=1, correction=1)), dim=-1) |
| return torch.cat((moments, auxiliary), dim=-1) |
|
|
|
|
| def gaussian_crps(mu: torch.Tensor, sigma: torch.Tensor, target: torch.Tensor) -> torch.Tensor: |
| """Closed-form CRPS for a Gaussian predictive distribution.""" |
| if torch.any(sigma <= 0): |
| raise ValueError("sigma must be strictly positive") |
| z = (target - mu) / sigma |
| pdf = torch.exp(-0.5 * z.square()) / math.sqrt(2.0 * math.pi) |
| cdf = 0.5 * (1.0 + torch.erf(z / math.sqrt(2.0))) |
| return sigma * (z * (2.0 * cdf - 1.0) + 2.0 * pdf - 1.0 / math.sqrt(math.pi)) |
|
|
|
|
| def ensemble_crps(ensemble: torch.Tensor, target: torch.Tensor) -> torch.Tensor: |
| """Empirical ensemble CRPS without constructing a member-pair matrix.""" |
| if ensemble.shape[-1] != MEMBER_COUNT: |
| raise ValueError(f"last ensemble dimension must be {MEMBER_COUNT}") |
| sorted_members = ensemble.sort(dim=-1).values |
| weights = (2 * torch.arange(1, MEMBER_COUNT + 1, device=ensemble.device) - MEMBER_COUNT - 1).to(ensemble.dtype) |
| pair_term = (sorted_members * weights).sum(dim=-1) / (MEMBER_COUNT * MEMBER_COUNT) |
| return (ensemble - target.unsqueeze(-1)).abs().mean(dim=-1) - pair_term |
|
|
