File size: 2,560 Bytes
b20ca9c | 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 | """Run bounded ensemble reconstruction from a unified checkpoint."""
from pathlib import Path
import argparse
import json
import sys
import numpy as np
import torch
import yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.crai_climateextremes import CRAIClimateExtremes
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
args = parser.parse_args()
config_path = args.config if args.config.is_absolute() else ROOT / args.config
with open(config_path, encoding="utf-8") as handle:
cfg = yaml.safe_load(handle)
data = np.load(ROOT / cfg["data_path"])
inputs = torch.from_numpy(np.concatenate((data["observed"], data["valid_mask"]), axis=1))
checkpoint_path = ROOT / cfg["checkpoint_path"]
if not checkpoint_path.is_file():
raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}; run scripts/train.py first")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
if checkpoint.get("format_version") != "1.0" or not isinstance(checkpoint.get("model"), list):
raise ValueError(f"unsupported checkpoint format: {checkpoint_path}")
model_config = checkpoint.get("model_config", {})
predictions = []
for state in checkpoint["model"]:
model = CRAIClimateExtremes(**model_config).to(device)
model.load_state_dict(state); model.eval()
with torch.no_grad():
predictions.append(model(inputs.to(device)).cpu().numpy())
if not predictions:
raise ValueError(f"checkpoint contains no ensemble members: {checkpoint_path}")
members = np.stack(predictions)
output = ROOT / cfg["output_dir"]
output.mkdir(parents=True, exist_ok=True)
np.savez_compressed(
output / "predictions.npz", prediction=members.mean(0),
ensemble_std=members.std(0), target=data["target"],
observed=data["observed"], valid_mask=data["valid_mask"],
europe_mask=data["europe_mask"], index_ids=data["index_ids"],
index_names=data["index_names"],
)
metadata = {"ensemble_members": len(predictions), "checkpoint_semantics": "member state list in one checkpoint", "output_range": [0, 100]}
(output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
print(f"predicted {inputs.shape[0]} samples with {len(predictions)} members")
if __name__ == "__main__":
main()
|