OneScience-Group/NNCAM

1#!/usr/bin/env python32import argparse3from pathlib import Path4 5import numpy as np6import torch7 8from model.nncam import NNCAM, unscale_output9 10 11ROOT = Path(__file__).resolve().parents[1]12 13 14def main():15    parser = argparse.ArgumentParser(description="Run offline NNCAM inference.")16    parser.add_argument("--data", type=Path, default=ROOT / "data/nncam_fake.npz")17    parser.add_argument("--checkpoint", type=Path, default=ROOT / "result/checkpoints/nncam.pt")18    parser.add_argument("--output", type=Path, default=ROOT / "result/output/predictions.npz")19    args = parser.parse_args()20    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True)21    required = {"model", "model_config", "format_version", "normalization"}22    if not required.issubset(checkpoint):23        raise ValueError(f"checkpoint missing {sorted(required - checkpoint.keys())}")24    model = NNCAM(**checkpoint["model_config"])25    model.load_state_dict(checkpoint["model"])26    model.eval()27    with np.load(args.data) as data:28        x, truth, lat, time = (data[name] for name in ("x", "y", "lat", "time"))29    norm = checkpoint["normalization"]30    normalized = (torch.from_numpy(x) - norm["input_mean"]) / norm["input_scale"]31    with torch.no_grad():32        scaled = model(normalized) * norm["target_scale"] + norm["target_mean"]33    prediction = unscale_output(scaled.numpy()).astype(np.float32)34    if prediction.shape != truth.shape or not np.isfinite(prediction).all():35        raise RuntimeError(f"invalid prediction shape or values: {prediction.shape}")36    args.output.parent.mkdir(parents=True, exist_ok=True)37    np.savez_compressed(args.output, input=x, truth=truth, prediction=prediction, lat=lat, time=time)38    print(f"saved {args.output}: prediction={prediction.shape}, finite=true")39 40 41if __name__ == "__main__":42    main()43