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