File size: 8,079 Bytes
7a7efc9 | 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 134 135 136 137 138 139 | """Station-space IMPROVER-style post-processing with explicit local grid patches."""
from __future__ import annotations
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
VARIABLES = ("temperature", "dewpoint", "wind_speed")
class ImproverAIFS(nn.Module):
"""Calibrate three sources at 569 station-centred target-grid neighborhoods."""
def __init__(self, config: dict):
super().__init__()
self.variables = tuple(config.get("variables", VARIABLES))
self.lead_hours = int(config.get("lead_hours", 240))
self.station_count = int(config.get("station_count", 569))
self.fuzzy_widths = torch.tensor(config["fuzzy_widths"], dtype=torch.float32)
self.recursive_coefficient = float(config.get("recursive_coefficient", 0.18))
self.recursive_iterations = int(config.get("recursive_iterations", 1))
self.register_buffer("knot_hours", torch.linspace(0, self.lead_hours, 11))
self.blend_logits = nn.Parameter(torch.zeros(len(self.variables), 11, 3))
self.register_buffer("bias", torch.zeros(3, self.lead_hours + 1, len(self.variables), self.station_count))
max_thresholds = max(map(len, config["thresholds"]))
knots = torch.zeros(3, len(self.variables), max_thresholds, 7, 2)
knots[..., 0] = torch.linspace(0, 1, 7)
knots[..., 1] = torch.linspace(0, 1, 7)
self.register_buffer("calibration_knots", knots)
self.register_buffer("threshold_counts", torch.tensor([len(v) for v in config["thresholds"]]))
@torch.no_grad()
def fit_bias_chunk(self, forecasts: torch.Tensor, analyses: torch.Tensor, elevation_delta: torch.Tensor, station_slice: slice) -> None:
"""Fit each source from all 30 preceding daily 1200 UTC histories."""
adjusted = forecasts.clone() # history, lead, variable, source, station
adjusted[:, :, 0] += -0.0098 * elevation_delta.view(1, 1, 1, -1)
self.bias[..., station_slice] = (adjusted - analyses.unsqueeze(3)).mean(0).permute(2, 0, 1, 3)
def correct_expected(self, forecasts: torch.Tensor, elevation_delta: torch.Tensor, source: int, station_slice: slice) -> torch.Tensor:
"""Correct station-centre forecasts shaped [date/history, lead, variable, station]."""
corrected = forecasts.clone()
corrected[:, :, 0] += -0.0098 * elevation_delta.view(1, 1, -1)
return corrected - self.bias[source, :, :, station_slice].unsqueeze(0)
def fuzzy_threshold(self, expected: torch.Tensor, thresholds: list[torch.Tensor]) -> list[torch.Tensor]:
outputs = []
widths = self.fuzzy_widths.to(expected.device)
for variable, values in enumerate(thresholds):
forecast = expected[:, :, variable].unsqueeze(2)
threshold = values.to(expected.device).view(1, 1, -1, 1)
outputs.append(((forecast - threshold + widths[variable]) / (2 * widths[variable])).clamp(0, 1))
return outputs
@staticmethod
def neighborhood(patch_probabilities: torch.Tensor) -> torch.Tensor:
"""Reduce pre-extracted real-semantic 3x3 target-grid patches to stations."""
if patch_probabilities.shape[-2:] != (3, 3):
raise ValueError("station neighborhood must be an authoritative 3x3 target-grid patch")
return patch_probabilities.mean(dim=(-2, -1))
def recursive_filter(self, patch_probabilities: torch.Tensor) -> torch.Tensor:
"""Apply separable filtering only within each station's extracted 3x3 patch."""
output = patch_probabilities.clone()
coefficient = self.recursive_coefficient
for _ in range(self.recursive_iterations):
for axis in (-2, -1):
for index in range(1, 3):
current, previous = [slice(None)] * output.ndim, [slice(None)] * output.ndim
current[axis], previous[axis] = index, index - 1
output[tuple(current)] = (1 - coefficient) * output[tuple(current)] + coefficient * output[tuple(previous)]
for index in range(1, -1, -1):
current, following = [slice(None)] * output.ndim, [slice(None)] * output.ndim
current[axis], following[axis] = index, index + 1
output[tuple(current)] = (1 - coefficient) * output[tuple(current)] + coefficient * output[tuple(following)]
return output
@torch.no_grad()
def fit_reliability(self, probabilities: list[torch.Tensor], analyses: torch.Tensor, thresholds: list[torch.Tensor], source: int) -> None:
"""Fit seven-bin mappings on configured station samples with full time axes."""
defaults = torch.linspace(0, 1, 7, device=analyses.device)
for variable, probability in enumerate(probabilities):
truth = analyses[:, :, variable].unsqueeze(2) > thresholds[variable].view(1, 1, -1, 1)
count = probability.shape[2]
p = probability.permute(2, 0, 1, 3).reshape(count, -1)
y = truth.permute(2, 0, 1, 3).reshape(count, -1).float()
ids = torch.bucketize(p, torch.linspace(1 / 7, 6 / 7, 6, device=p.device))
xs, ys = [], []
for bin_index in range(7):
mask = ids == bin_index
samples = mask.sum(1)
denominator = samples.clamp_min(1)
xs.append(torch.where(samples > 0, (p * mask).sum(1) / denominator, defaults[bin_index]))
ys.append(torch.where(samples > 0, (y * mask).sum(1) / denominator, defaults[bin_index]))
xs = torch.stack(xs, dim=1).cummax(1).values
ys = torch.stack(ys, dim=1).cummax(1).values.clamp(0, 1)
self.calibration_knots[source, variable, :count, :, 0] = xs
self.calibration_knots[source, variable, :count, :, 1] = ys
def calibrate(self, probabilities: list[torch.Tensor], source: int) -> list[torch.Tensor]:
outputs = []
for variable, probability in enumerate(probabilities):
count = probability.shape[2]
knots = self.calibration_knots[source, variable, :count]
x, y = knots[..., 0].contiguous(), knots[..., 1].contiguous()
values = probability.permute(2, 0, 1, 3).reshape(count, -1).contiguous()
ids = torch.searchsorted(x, values).clamp(1, 6)
x0, x1 = x.gather(1, ids - 1), x.gather(1, ids)
y0, y1 = y.gather(1, ids - 1), y.gather(1, ids)
calibrated = y0 + (values - x0) * (y1 - y0) / (x1 - x0).clamp_min(1e-6)
calibrated = calibrated.reshape(count, probability.shape[0], probability.shape[1], probability.shape[3]).permute(1, 2, 0, 3)
outputs.append(calibrated.clamp(0, 1).cummin(dim=2).values)
return outputs
def blend_weights(self) -> torch.Tensor:
"""Interpolate the paper's 11 valid-time knots to 241 hourly leads."""
return F.interpolate(self.blend_logits.softmax(-1).permute(0, 2, 1), size=241, mode="linear", align_corners=True).permute(0, 2, 1)
def blend_expected(self, model_expected: torch.Tensor) -> torch.Tensor:
weights = self.blend_weights().permute(1, 0, 2).view(1, 241, 3, 3, 1)
return (model_expected * weights).sum(3)
def blend_probabilities(self, source_probabilities: list[list[torch.Tensor]]) -> list[torch.Tensor]:
weights = self.blend_weights()
return [sum(source_probabilities[s][v] * weights[v, :, s].view(1, 241, 1, 1) for s in range(3)).clamp(0, 1)
for v in range(3)]
def crps_from_thresholds(probabilities: np.ndarray, thresholds: np.ndarray, truth: np.ndarray) -> float:
total = np.zeros_like(truth, dtype=np.float32)
previous = (1.0 - probabilities[:, :, 0] - (thresholds[0] >= truth)) ** 2
for index in range(1, len(thresholds)):
current = (1.0 - probabilities[:, :, index] - (thresholds[index] >= truth)) ** 2
total += 0.5 * (previous + current) * (thresholds[index] - thresholds[index - 1])
previous = current
return float(total.mean())
|