| from pathlib import Path | |
| import sys,numpy as np,torch | |
| R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.fuxi_weather import * | |
| c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);m=FuXiWeather(**z['model_config']);m.load_state_dict(z['model']);x=m.analysis(*map(torch.tensor,(d['background'][:2],d['observation'][:2],d['mask'][:2])));seq=[] | |
| with torch.no_grad(): | |
| for s in range(c['inference']['steps']):x=m.forecast(x,s>=c['inference']['switch_step']);seq.append(x.numpy()) | |
| target=np.stack([[sample(i,s+1)[0].numpy() for s in range(12)] for i in range(2)]);p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,prediction=np.stack(seq,1),target=target,lead_hours=np.arange(1,13)*6,logical_shape=np.array([20,721,1440]),is_complete_global=False);print(p) | |