File size: 4,323 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 | """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()
|