Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,148 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Weather Postprocessing
+- Probabilistic Forecasting
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">IMPROVER-AIFS</span></strong>
+</p>
+
+# Model Introduction
+
+IMPROVER-AIFS corrects systematic biases and improves the reliability of artificial-intelligence weather forecasts. It converts deterministic AIFS forecasts into spatially processed and statistically calibrated deterministic and probabilistic products and can blend them with conventional numerical weather predictions for near-surface weather postprocessing and operational product generation.
+
+Paper: Statistical Postprocessing Yields Accurate Probabilistic Forecasts from Artificial Intelligence Weather Models  
+https://doi.org/10.1175/AIES-D-25-0037.1
+
+# Model Description
+
+The method reproduced by IMPROVER-AIFS was proposed by a research team at the Australian Bureau of Meteorology and uses the IMPROVER postprocessing system led by the Met Office. The paper performs calibration and evaluation with ECMWF AIFS, HRES, and ENS forecasts, MSAS gridded analyses, and Bureau Jive automatic weather-station observations. The model is suitable for deterministic correction, probabilistic calibration, and multi-model blended forecasting of surface temperature, surface dewpoint temperature, and 10 m wind speed.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Deterministic forecast postprocessing | Apply elevation and historical bias correction to AIFS-like weather forecasts. |
+| Probabilistic forecasting and calibration | Produce threshold probabilities for three variables with spatial smoothing and reliability calibration. |
+| Multi-model blending | Smoothly blend deterministic and probabilistic outputs from AIFS-, HRES-, and ENS-like inputs. |
+| ModelScope/OneCode execution | Validate training, inference, evaluation, visualization, and checkpoint workflows in ModelScope or OneCode. |
+| Multi-GPU training | Launch distributed data-parallel training with `torchrun`. |
+
+# Usage Instructions
+
+## 1.OneCode
+
+Experience intelligent, one-click AI4S programming through the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Download and Installation
+
+```bash
+hf download OneScience-Group/IMPROVER-AIFS --local-dir ./IMPROVER-AIFS
+cd IMPROVER-AIFS
+```
+
+### Environment Dependencies
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for connectivity validation with the default small-sample configuration.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+**DCU Environment**
+
+```bash
+ # Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+ # Activate Conda first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+This repository uses a small number of structured synthetic samples to validate the engineering workflow, with AIFS-, HRES-, and ENS-like forecasts and corresponding MSAS-analysis and Bureau Jive station-observation semantics. Each sample retains 30 history days, 241 hourly leads, three variables, `61/47/49` thresholds, all 569 stations, and a `3x3` neighborhood for each station; only the number of valid dates and station samples used for parameter fitting is reduced. The paper does not publish the complete target Albers grid dimensions, so this implementation does not infer them; synthetic data validate postprocessing, training, inference, and evaluation only and do not represent the paper data distribution or performance.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+For single-device training, use:
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training outputs are saved to:
+
+```text
+result/checkpoints/improver_aifs.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not include weights under `weight/`. The paper does not provide a confirmed official checkpoint, and the current engineering checkpoint is not claimed to be compatible with external weights.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads a checkpoint containing both valid-time folds and generates chunked deterministic and probabilistic blends for two valid dates, 241 leads, and all 569 stations. Numerical results and station metadata are saved to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation results include errors for the raw, postprocessed, and blended forecasts and probabilistic forecast quality, together with target, prediction, and error comparisons at representative lead times. Synthetic-data results validate the engineering workflow only and do not represent paper performance.
+
+```text
+result/evaluation/metrics.json
+result/evaluation/multi_lead_temperature.png
+```
+
+# Official OneScience Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+This repository is an independent engineering reproduction of the public IMPROVER-AIFS specifications.
+
+Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.

conf/config.yaml

@@ -0,0 +1,70 @@
+seed: 42
+paper_model:
+  name: ECMWF deterministic AIFS
+  role: deterministic forecast input
+  native_grid: N320_reduced_Gaussian
+  native_grid_degrees_approx: 0.25
+  preprocessed_grid: regular_latlon_0.25_degree
+  native_output_interval_hours: 6
+protocol:
+  forecast_date_start: 2024-03-01
+  forecast_date_end: 2024-07-23
+  initialization_utc: "1200"
+  history_days: 30
+  lead_hours: 240
+  output_interval_hours: 1
+  valid_time_folds: 2
+  blend_knots: 11
+  probability_bins: 7
+  sources:
+    AIFS: {native_grid: N320_reduced_Gaussian, native_interval_hours: 6, preprocessed_grid: regular_latlon_0.25_degree}
+    HRES: {accepted_grid_degrees: 0.1}
+    ENS: {accepted_grid_degrees: 0.2}
+  analysis: {name: MSAS, resolution_arcminutes: 2.5}
+  verification: {name: Bureau_Jive, station_count: 569, interval_hours: 1}
+data:
+  root: data
+  format_version: aiweather_station_patch_v2
+  valid_dates: [2024-03-01, 2024-07-23]
+  history_days: 30
+  lead_hours: 240
+  output_interval_hours: 1
+  station_count: 569
+  station_chunk_size: 128
+  variables: [temperature, dewpoint, wind_speed]
+  units: [degC, degC, m_s-1]
+  sources: [AIFS, HRES, ENS]
+  target_representation: station_center_with_preextracted_target_grid_neighborhood
+  neighborhood_shape: [3, 3]
+  target_grid_adapter:
+    status: required_for_real_data
+    projection: Albers
+    nx: null
+    ny: null
+    coordinates: null
+    note: Full target Albers dimensions are not public and must be supplied by an authoritative data adapter.
+model:
+  variables: [temperature, dewpoint, wind_speed]
+  lead_hours: 240
+  station_count: 569
+  fuzzy_widths: [2.0, 2.0, 1.5]
+  thresholds:
+    - [-20, -19, -18, -17, -16, -15, -14, -13, -12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 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]
+    - [-20, -19, -18, -17, -16, -15, -14, -13, -12, -11, -10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 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]
+    - [0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7, 7.5, 8, 8.5, 9, 9.5, 10, 10.5, 11, 11.5, 12, 12.5, 13, 13.5, 14, 14.5, 15, 15.5, 16, 16.5, 17, 17.5, 18, 18.5, 19, 19.5, 20, 20.5, 21, 21.5, 22, 22.5, 23, 23.5, 24]
+  neighborhood_cells: 3
+  recursive_coefficient: 0.18
+  recursive_iterations: 1
+train:
+  epochs: 2
+  learning_rate: 0.08
+  station_chunk_size: 128
+  calibration_stations: 8
+  blend_stations: 8
+runtime:
+  device: auto
+paths:
+  checkpoint: result/checkpoints/improver_aifs.pt
+  training_metrics: result/training/metrics.json
+  inference: result/output/predictions.npz
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,33 @@
+{
+  "format_version": "aiweather_station_patch_v2",
+  "model": "IMPROVER-AIFS",
+  "model_type": "statistical-weather-postprocessing",
+  "architectures": ["ImproverAIFS"],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "deterministic-to-probabilistic-weather-postprocessing",
+  "implementation": {
+    "entry_point": "model/improver_aifs.py",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "paper": {"arxiv": "2504.12672", "doi": "10.1175/AIES-D-25-0037.1", "method": "Bureau of Meteorology IMPROVER applied to ECMWF deterministic AIFS"},
+  "scientific_dimensions": {
+    "forecast_dates": ["2024-03-01", "2024-07-23"],
+    "conceptual_date_range": ["2024-03-01", "2024-07-23"],
+    "initialization_utc": "1200",
+    "history_days": 30,
+    "lead_hours": [0, 240, 1],
+    "stations": 569,
+    "sources": ["AIFS", "HRES", "ENS"],
+    "variables": ["temperature", "dewpoint", "10 m wind speed"],
+    "threshold_counts": [61, 47, 49],
+    "probability_bins": 7,
+    "valid_time_folds": 2,
+    "blend_knots": 11,
+    "station_neighborhood": [3, 3]
+  },
+  "target_grid_adapter": {"projection": "Albers", "nx": null, "ny": null, "status": "not_public_do_not_infer"}
+}

configuration.json

@@ -0,0 +1,15 @@
+{
+  "format_version": "aiweather_station_patch_v2",
+  "framework": "PyTorch",
+  "task": "weather_forecast_statistical_postprocessing",
+  "model": "IMPROVER-AIFS",
+  "input_format": "lazy [valid_date, history_day=30, lead=241, variable=3, source=3, station=569, patch_y=3, patch_x=3]",
+  "output_format": "[valid_date, lead=241, variable=3, station=569] plus threshold probabilities",
+  "protocol": "improver_aifs_trotta_2025_station_validation",
+  "fidelity_scope": "paper-confirmed 569-station validation object and pre-extracted 3x3 target-grid neighborhoods; not the unpublished full operational Albers grid",
+  "target_grid_adapter": {"projection": "Albers", "nx": null, "ny": null, "required": true},
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py"
+}

model/improver_aifs.py

@@ -0,0 +1,138 @@
+"""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())

scripts/fake_data.py

@@ -0,0 +1,70 @@
+"""Write a lazy 569-station protocol manifest; fields are generated per chunk."""
+
+from pathlib import Path
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def station_metadata(count: int):
+    index = np.arange(count, dtype=np.float32)
+    latitude = -43.5 + 33.0 * ((index * 0.61803398875) % 1)
+    longitude = 113.0 + 40.0 * ((index * 0.41421356237) % 1)
+    elevation = 20 + 1450 * ((index * 0.2718281828) % 1)
+    return latitude, longitude, elevation.astype(np.float32)
+
+
+def make_split(path: Path, config: dict):
+    data = config["data"]
+    dates = np.asarray(data["valid_dates"], dtype="U10")
+    lat, lon, elevation = station_metadata(data["station_count"])
+    history_dates = np.empty((len(dates), data["history_days"]), dtype="U10")
+    for i, date in enumerate(dates.astype("datetime64[D]")):
+        history_dates[i] = (date - np.arange(data["history_days"], 0, -1)).astype("U10")
+    np.savez(path, format_version=np.asarray(data["format_version"]), valid_dates=dates,
+             history_dates=history_dates, initialization_utc=np.asarray("1200"), lead_hours=np.arange(241),
+             station_id=np.asarray([f"JIVE-{i:04d}" for i in range(data["station_count"])]),
+             station_latitude=lat, station_longitude=lon, station_elevation_m=elevation,
+             variables=np.asarray(data["variables"]), units=np.asarray(data["units"]), sources=np.asarray(data["sources"]),
+             neighborhood_shape=np.asarray([3, 3]), target_grid_projection=np.asarray("Albers"),
+             target_grid_nx=np.asarray(None), target_grid_ny=np.asarray(None),
+             representation=np.asarray(data["target_representation"]))
+
+
+def generate_chunk(meta, date_index: int, station_selector, history: bool, seed: int, include_patch: bool = True):
+    """Emulate lazy reads after authoritative station patches have been extracted."""
+    station = np.arange(len(meta["station_id"]))[station_selector]
+    days = 30 if history else 1
+    lead = np.arange(241, dtype=np.float32)[None, :, None, None]
+    day = np.arange(days, dtype=np.float32)[:, None, None, None]
+    lat = meta["station_latitude"][station][None, None, None, :]
+    lon = meta["station_longitude"][station][None, None, None, :]
+    phase = date_index * 0.7 + day * 0.11
+    temp = 20 - 0.35 * (lat + 25) + 4 * np.sin(2 * np.pi * (lead + 12) / 24 + phase)
+    dew = temp - 5 - 1.5 * np.cos(np.deg2rad(lon) + lead / 48)
+    wind = 5 + 1.2 * np.abs(np.sin(np.deg2rad(lon) + lead / 18 + phase))
+    truth = np.concatenate((temp, dew, wind), axis=2).astype(np.float32)
+    source = np.arange(3, dtype=np.float32)[None, None, None, :, None]
+    forecast = truth[:, :, :, None] + (source - 0.6) * np.asarray([1.1, 0.8, 0.5], np.float32)[None, None, :, None, None]
+    forecast += (lead[..., None] / 240) * np.asarray([0.8, -0.5, 0.6], np.float32)[None, None, :, None, None]
+    rng = np.random.default_rng(seed + date_index * 1000 + int(station[0]))
+    forecast += rng.normal(0, 0.08, forecast.shape).astype(np.float32)
+    offsets = np.asarray([[-0.18, -0.10, -0.04], [-0.08, 0.0, 0.09], [0.03, 0.12, 0.20]], np.float32)
+    elevation_delta = meta["station_elevation_m"][station] - np.mean(meta["station_elevation_m"])
+    if not include_patch:
+        return forecast.astype(np.float32), truth, elevation_delta.astype(np.float32)
+    patches = forecast[..., None, None] + offsets
+    return patches.astype(np.float32), truth, elevation_delta.astype(np.float32)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    output = ROOT / config["data"]["root"]
+    output.mkdir(parents=True, exist_ok=True)
+    make_split(output / "protocol.npz", config)
+    print("generated=data/protocol.npz dates=2 history_days=30 leads=241 stations=569 patch=3x3 nx=unknown ny=unknown")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,70 @@
+"""Run fold-held-out inference for every valid date, lead, variable, and station."""
+
+import sys
+from pathlib import Path
+import numpy as np
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.improver_aifs import ImproverAIFS
+from fake_data import generate_chunk
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    meta = np.load(ROOT / config["data"]["root"] / "protocol.npz", allow_pickle=True)
+    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location="cpu", weights_only=False)
+    fold_states = checkpoint.get("model", {}).get("fold_states", [])
+    if checkpoint["format_version"] != config["data"]["format_version"] or len(fold_states) != 2:
+        raise ValueError("checkpoint must contain two valid-time folds")
+    dates, leads, stations = 2, 241, config["data"]["station_count"]
+    raw = np.empty((dates, leads, 3, stations), np.float32)
+    truth = np.empty_like(raw); expected = np.empty_like(raw)
+    probabilities = [np.empty((dates, leads, len(v), stations), np.float32) for v in config["model"]["thresholds"]]
+    thresholds = [torch.tensor(v, dtype=torch.float32) for v in config["model"]["thresholds"]]
+    chunk = config["data"]["station_chunk_size"]
+    for fold in range(2):
+        model = ImproverAIFS(checkpoint["model_config"]); model.load_state_dict(fold_states[fold]); model.eval()
+        for start in range(0, stations, chunk):
+            section = slice(start, min(start + chunk, stations))
+            patches, analyses, elevation = generate_chunk(meta, fold, section, False, int(config["seed"]) + 100)
+            source_expected, source_probabilities = [], []
+            with torch.no_grad():
+                for source in range(3):
+                    patch = torch.from_numpy(patches[:, :, :, source])
+                    patch[:, :, 0] += (-0.0098 * torch.from_numpy(elevation)).view(1, 1, -1, 1, 1)
+                    patch -= model.bias[source, :, :, section].unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                    centre = patch[..., 1, 1]
+                    source_expected.append(centre)
+                    raw_probabilities = []
+                    for variable, values in enumerate(thresholds):
+                        width = model.fuzzy_widths[variable]
+                        p = ((patch[:, :, variable].unsqueeze(2) - values.view(1, 1, -1, 1, 1, 1) + width) / (2 * width)).clamp(0, 1)
+                        raw_probabilities.append(model.neighborhood(model.recursive_filter(p)))
+                    source_probabilities.append(model.calibrate(raw_probabilities, source))
+                blend = model.blend_expected(torch.stack(source_expected, dim=3))
+                blend_probability = model.blend_probabilities(source_probabilities)
+            raw[fold, ..., section] = patches[0, :, :, 0, :, 1, 1]
+            truth[fold, ..., section] = analyses[0]
+            expected[fold, ..., section] = blend[0].numpy()
+            for variable in range(3): probabilities[variable][fold, ..., section] = blend_probability[variable][0].numpy()
+    arrays = [raw, truth, expected, *probabilities]
+    if not all(np.isfinite(array).all() for array in arrays):
+        raise ValueError("inference produced non-finite values")
+    if not all(((array >= 0) & (array <= 1)).all() for array in probabilities):
+        raise ValueError("probabilities must be within [0, 1]")
+    output = ROOT / config["paths"]["inference"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    payload = {"raw_aifs": raw, "analyses": truth, "blend_expected": expected, "valid_dates": meta["valid_dates"], "lead_hours": meta["lead_hours"],
+               "station_id": meta["station_id"], "station_latitude": meta["station_latitude"], "station_longitude": meta["station_longitude"],
+               "variables": meta["variables"], "units": meta["units"], "fold_id": np.arange(2)}
+    for variable, name in enumerate(meta["variables"]):
+        payload[f"thresholds_{name}"] = thresholds[variable].numpy(); payload[f"probability_blend_{name}"] = probabilities[variable]
+    np.savez_compressed(output, **payload)
+    print(f"saved={output.relative_to(ROOT)} expected_shape={expected.shape} folds=2 stations=569 thresholds=61/47/49")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,43 @@
+"""Evaluate all 569 Jive stations hourly and plot station summaries."""
+
+import json
+import sys
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.improver_aifs import crps_from_thresholds
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = np.load(ROOT / config["paths"]["inference"])
+    raw, blend, truth = data["raw_aifs"], data["blend_expected"], data["analyses"]
+    if raw.shape != (2, 241, 3, 569):
+        raise ValueError(f"evaluation requires [2,241,3,569], received {raw.shape}")
+    metrics = {"note": "engineering validation, not paper performance", "verification_stations": 569, "hourly_leads": 241, "valid_time_folds": 2, "by_variable": {}}
+    for variable, name in enumerate(data["variables"]):
+        probability, thresholds = data[f"probability_blend_{name}"], data[f"thresholds_{name}"]
+        metrics["by_variable"][str(name)] = {
+            "raw_mse": float(np.mean((raw[:, :, variable] - truth[:, :, variable]) ** 2)),
+            "blend_mse": float(np.mean((blend[:, :, variable] - truth[:, :, variable]) ** 2)),
+            "blend_bias": float(np.mean(blend[:, :, variable] - truth[:, :, variable])),
+            "blend_crps": crps_from_thresholds(probability, thresholds, truth[:, :, variable]),
+        }
+    output = ROOT / config["paths"]["evaluation_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    fig, axes = plt.subplots(2, 2, figsize=(11, 7), constrained_layout=True)
+    for axis, lead in zip(axes.flat, [0, 24, 120, 240]):
+        axis.scatter(data["station_longitude"], data["station_latitude"], c=np.abs(blend[0, lead, 0] - truth[0, lead, 0]), s=8, cmap="viridis")
+        axis.set_title(f"569-station temperature absolute error | {lead} h")
+        axis.set_xlabel("longitude"); axis.set_ylabel("latitude")
+    fig.savefig(output / "multi_lead_temperature.png", dpi=150); plt.close(fig)
+    print("saved=result/evaluation stations=569 hourly_leads=241 folds=2 mse_crps_bias=ok")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,119 @@
+"""Fit two valid-time folds without caching full-station threshold probabilities."""
+
+import json
+import os
+import sys
+from pathlib import Path
+import numpy as np
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.improver_aifs import ImproverAIFS
+from fake_data import generate_chunk
+
+
+def fit_fold(config, meta, fold, rank):
+    model = ImproverAIFS(config["model"])
+    stations = int(config["data"]["station_count"])
+    chunk_size = int(config["train"]["station_chunk_size"])
+
+    # Bias uses every station but only station-centre values, so no full-grid cache exists.
+    for start in range(0, stations, chunk_size):
+        section = slice(start, min(start + chunk_size, stations))
+        centres, analyses, elevation = generate_chunk(
+            meta, fold, section, True, int(config["seed"]), include_patch=False
+        )
+        model.fit_bias_chunk(
+            torch.from_numpy(centres), torch.from_numpy(analyses),
+            torch.from_numpy(elevation), section,
+        )
+
+    sample_count = max(int(config["train"]["calibration_stations"]), int(config["train"]["blend_stations"]))
+    sample_stations = np.linspace(0, stations - 1, sample_count, dtype=np.int64)
+    patches, analyses, elevation = generate_chunk(
+        meta, fold, sample_stations, True, int(config["seed"]), include_patch=True
+    )
+    truth = torch.from_numpy(analyses)
+    thresholds = [torch.tensor(values, dtype=torch.float32) for values in config["model"]["thresholds"]]
+    source_expected = []
+    for source in range(3):
+        patch = torch.from_numpy(patches[:, :, :, source])
+        patch[:, :, 0] += (-0.0098 * torch.from_numpy(elevation)).view(1, 1, -1, 1, 1)
+        patch -= model.bias[source, :, :, sample_stations].unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+        source_expected.append(patch[..., 1, 1])
+        probabilities = []
+        for variable, values in enumerate(thresholds):
+            width = model.fuzzy_widths[variable]
+            probability_patch = ((patch[:, :, variable].unsqueeze(2) - values.view(1, 1, -1, 1, 1, 1) + width) / (2 * width)).clamp(0, 1)
+            probabilities.append(model.neighborhood(model.recursive_filter(probability_patch)))
+        model.fit_reliability(probabilities, truth, thresholds, source)
+        del probabilities
+
+    blend_count = int(config["train"]["blend_stations"])
+    blend_input = torch.stack(source_expected, dim=3)[..., :blend_count]
+    blend_truth = truth[..., :blend_count]
+    optimizer = torch.optim.Adam([model.blend_logits], lr=float(config["train"]["learning_rate"]))
+    history = []
+    for epoch in range(int(config["train"]["epochs"])):
+        loss = torch.mean((model.blend_expected(blend_input) - blend_truth) ** 2)
+        if not torch.isfinite(loss):
+            raise ValueError("non-finite blend loss")
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+        history.append({"epoch": epoch + 1, "blend_mse_loss": float(loss.detach())})
+    record = {
+        "fold": fold, "held_out_valid_date": str(meta["valid_dates"][fold]),
+        "history_days": 30, "bias_stations": stations,
+        "calibration_stations": int(config["train"]["calibration_stations"]),
+        "blend_stations": blend_count, "rank": rank, "history": history,
+    }
+    print(f"rank={rank} fold={fold} history_days=30 bias_stations=569 calibration_stations={sample_count} loss={history[-1]['blend_mse_loss']:.6f}")
+    return model.state_dict(), record
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    meta = np.load(ROOT / config["data"]["root"] / "protocol.npz", allow_pickle=True)
+    if str(meta["format_version"]) != config["data"]["format_version"] or meta["history_dates"].shape != (2, 30):
+        raise ValueError("protocol requires two valid dates and complete 30-day histories")
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    if distributed:
+        torch.distributed.init_process_group("gloo")
+    rank = torch.distributed.get_rank() if distributed else 0
+    world_size = torch.distributed.get_world_size() if distributed else 1
+    local = []
+    for fold in range(2):
+        if fold % world_size == rank:
+            state, record = fit_fold(config, meta, fold, rank)
+            local.append((fold, state, record))
+    if distributed:
+        gathered = [None] * world_size
+        torch.distributed.all_gather_object(gathered, local)
+        combined = [item for rank_items in gathered for item in rank_items]
+    else:
+        combined = local
+    if rank == 0:
+        combined.sort(key=lambda item: item[0])
+        if [item[0] for item in combined] != [0, 1]:
+            raise RuntimeError("DDP ranks did not produce both valid-time folds")
+        fold_states = [item[1] for item in combined]
+        records = [item[2] for item in combined]
+        checkpoint = ROOT / config["paths"]["checkpoint"]
+        metrics = ROOT / config["paths"]["training_metrics"]
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        metrics.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({
+            "model": {"fold_states": fold_states, "fold_semantics": "valid-time 2-fold"},
+            "model_config": config["model"],
+            "format_version": config["data"]["format_version"],
+        }, checkpoint)
+        metrics.write_text(json.dumps({"folds": records, "distributed_world_size": world_size}, indent=2) + "\n")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()