OneScience-Group/CRAI-ClimateExtremes

1"""Run bounded ensemble reconstruction from a unified checkpoint."""2 3from pathlib import Path4import argparse5import json6import sys7import numpy as np8import torch9import yaml10 11ROOT = Path(__file__).resolve().parents[1]12sys.path.insert(0, str(ROOT))13 14from model.crai_climateextremes import CRAIClimateExtremes15 16 17def main():18    parser = argparse.ArgumentParser()19    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")20    args = parser.parse_args()21    config_path = args.config if args.config.is_absolute() else ROOT / args.config22    with open(config_path, encoding="utf-8") as handle:23        cfg = yaml.safe_load(handle)24    data = np.load(ROOT / cfg["data_path"])25    inputs = torch.from_numpy(np.concatenate((data["observed"], data["valid_mask"]), axis=1))26    checkpoint_path = ROOT / cfg["checkpoint_path"]27    if not checkpoint_path.is_file():28        raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}; run scripts/train.py first")29    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")30    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)31    if checkpoint.get("format_version") != "1.0" or not isinstance(checkpoint.get("model"), list):32        raise ValueError(f"unsupported checkpoint format: {checkpoint_path}")33    model_config = checkpoint.get("model_config", {})34    predictions = []35    for state in checkpoint["model"]:36        model = CRAIClimateExtremes(**model_config).to(device)37        model.load_state_dict(state); model.eval()38        with torch.no_grad():39            predictions.append(model(inputs.to(device)).cpu().numpy())40    if not predictions:41        raise ValueError(f"checkpoint contains no ensemble members: {checkpoint_path}")42    members = np.stack(predictions)43    output = ROOT / cfg["output_dir"]44    output.mkdir(parents=True, exist_ok=True)45    np.savez_compressed(46        output / "predictions.npz", prediction=members.mean(0),47        ensemble_std=members.std(0), target=data["target"],48        observed=data["observed"], valid_mask=data["valid_mask"],49        europe_mask=data["europe_mask"], index_ids=data["index_ids"],50        index_names=data["index_names"],51    )52    metadata = {"ensemble_members": len(predictions), "checkpoint_semantics": "member state list in one checkpoint", "output_range": [0, 100]}53    (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")54    print(f"predicted {inputs.shape[0]} samples with {len(predictions)} members")55 56 57if __name__ == "__main__":58    main()59