File size: 841 Bytes
aa529c9 | 1 2 3 4 5 6 7 8 9 10 11 | from pathlib import Path
import sys,numpy as np,torch
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
from model.weatherbench import *
c=load_config(ROOT);d=np.load(ROOT/c["data"]["path"]);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);m=WeatherBenchCNN(**ck["model_config"]);m.load_state_dict(ck["model"]);x=torch.tensor(d["input"][d["split"]==2]);p=[]
with torch.no_grad():
for step in range(20):
x=m(x)
if step in (11,19):p.append(x.numpy())
path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,prediction=np.stack(p,1),target=np.stack((d["target_3d"][d["split"]==2],d["target_5d"][d["split"]==2]),1),initial=d["input"][d["split"]==2],lead_days=np.array([3,5]),latitude=d["latitude"],longitude=d["longitude"]);print(path)
|