from pathlib import Path
import sys,numpy as np,torch;R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.weatherbench2 import *
c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],weights_only=True);m=WB2Baseline(**z['model_config']);m.load_state_dict(z['model']);base=m(torch.tensor(d['input'])).detach().numpy();ens=np.stack([base+np.random.default_rng(i).normal(0,.02,base.shape) for i in range(8)]);p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez(p,ensemble=ens,target=d['target'],latitude=d['latitude']);print(p)