OneScience-Group/PrecipDD

1"""Run ensemble inference for all held-out synthetic daily maps."""2 3import sys4from pathlib import Path5 6import numpy as np7import torch8 9 10ROOT = Path(__file__).resolve().parents[1]11sys.path.insert(0, str(ROOT))12from model.precipdd import ensemble_predict, load_config, load_ensemble, validate_archive13 14 15def main():16    config = load_config(ROOT / "conf/config.yaml")17    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")18    data = np.load(ROOT / config["paths"]["data"])19    validate_archive(data)20    models, checkpoint = load_ensemble(ROOT / config["paths"]["checkpoint"], device)21    mask = data["split"] == 222    fields = torch.from_numpy(data["precipitation"][mask]).float().to(device)23    prediction = ensemble_predict(models, fields, config["training"]["batch_size"]).cpu().numpy()24    output = ROOT / config["paths"]["predictions"]25    output.parent.mkdir(parents=True, exist_ok=True)26    np.savez_compressed(output, format_version=data["format_version"], prediction=prediction, target=data["agmt"][mask],27                        precipitation=data["precipitation"][mask], year=data["year"][mask], day_of_year=data["day_of_year"][mask],28                        latitude=data["latitude"], longitude=data["longitude"], ensemble_members=np.array(len(models)),29                        checkpoint_world_size=np.array(checkpoint["world_size"]))30    print(f"predictions={output.relative_to(ROOT)} days={len(prediction)} members={len(models)}")31 32 33if __name__ == "__main__":34    main()35