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