| from pathlib import Path |
| import sys,numpy as np,torch |
| ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT)) |
| from model.sfno_bvmc import * |
| c=load_config(ROOT);torch.manual_seed(c["seed"]);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);base=CompactSFNO(**ck["model_config"]);base.load_state_dict(ck["model"]);models=[] |
| for j in range(c["ensemble"]["checkpoint_members"]): |
| m=CompactSFNO(**ck["model_config"]);m.load_state_dict(ck["model"]) |
| with torch.no_grad(): |
| for p in m.parameters():p.add_(torch.randn_like(p)*(j+1)*1e-4) |
| m.eval();models.append(m) |
| x=state(c["data"]["tile_origins"][0],c["data"]["tile_size"],30,0)[None];members=[] |
| with torch.no_grad(): |
| for m in models: |
| plus,minus=centered_bred_vectors(m,x,c["ensemble"]["bred_cycles"],c["ensemble"]["perturbation_norm"]) |
| for perturb in (plus,minus): |
| current=x+perturb;seq=[] |
| for step in range(c["ensemble"]["forecast_steps"]): |
| out=m(current);seq.append(out[0].numpy());current=torch.cat((out,current[:,-3:]),dim=1) |
| members.append(seq) |
| targets=np.stack([state(c["data"]["tile_origins"][0],32,30,s+1,74).numpy() for s in range(c["ensemble"]["forecast_steps"])]);path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,ensemble=np.asarray(members),target=targets,lead_hours=np.arange(1,c["ensemble"]["forecast_steps"]+1)*6,checkpoint_indices=np.repeat(np.arange(len(models)),2),bred_sign=np.tile(np.array([1,-1]),len(models)),logical_shape=np.array([74,721,1440]),is_complete_global=np.bool_(False));print(path) |
|
|
