File size: 1,549 Bytes
950fc23 | 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 | """Run ensemble inference for all held-out synthetic daily maps."""
import sys
from pathlib import Path
import numpy as np
import torch
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.precipdd import ensemble_predict, load_config, load_ensemble, validate_archive
def main():
config = load_config(ROOT / "conf/config.yaml")
device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
data = np.load(ROOT / config["paths"]["data"])
validate_archive(data)
models, checkpoint = load_ensemble(ROOT / config["paths"]["checkpoint"], device)
mask = data["split"] == 2
fields = torch.from_numpy(data["precipitation"][mask]).float().to(device)
prediction = ensemble_predict(models, fields, config["training"]["batch_size"]).cpu().numpy()
output = ROOT / config["paths"]["predictions"]
output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output, format_version=data["format_version"], prediction=prediction, target=data["agmt"][mask],
precipitation=data["precipitation"][mask], year=data["year"][mask], day_of_year=data["day_of_year"][mask],
latitude=data["latitude"], longitude=data["longitude"], ensemble_members=np.array(len(models)),
checkpoint_world_size=np.array(checkpoint["world_size"]))
print(f"predictions={output.relative_to(ROOT)} days={len(prediction)} members={len(models)}")
if __name__ == "__main__":
main()
|