File size: 873 Bytes
d3225f2 | 1 2 3 4 5 6 7 8 | from pathlib import Path
import sys,numpy as np,torch
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
from model.wrf_ml import *
c=load_config(ROOT);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);m=RadiationBiLSTM(**ck["model_config"]);m.load_state_dict(ck["model"]);m.eval();x,p,b=synthetic_columns(100,n=6);wrf=x.reshape(2,57,3,10);coupler=SynchronousCoupler(m)
with torch.no_grad():ph_layout,bh_layout=coupler.infer_run(wrf);ph=ph_layout.permute(0,2,1,3).reshape(6,57,4);bh=bh_layout.reshape(6,2)
path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,profile_prediction=ph.numpy(),profile_target=p.numpy(),boundary_prediction=bh.numpy(),boundary_target=b.numpy(),vertical_levels=np.arange(57),coupler_mode=np.array("synchronous_request_response"));print(path)
|