| from pathlib import Path |
| import sys,numpy as np,torch |
| R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.climate2weather 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=ScoreUNet(**z['model_config']);m.load_state_dict(z['model']);condition=torch.tensor(d['coarse'][-1:]);up=F.interpolate(condition,size=(128,128),mode='nearest');ens=[] |
| with torch.no_grad(): |
| for j in range(8): |
| x=up+.2*torch.randn_like(up) |
| for k in range(4):x=x+.02*m(x,.2/(k+1));x=x+(up-F.interpolate(observe(x[0])[None],size=(128,128),mode='nearest'))*.2 |
| ens.append(x[0].reshape(3,4,128,128).numpy()) |
| p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,ensemble=ens,target=d['fine'][-1].reshape(3,4,128,128),condition=condition.numpy().reshape(1,3,4,8,8));print(p) |
|
|
