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