| from pathlib import Path | |
| import sys,numpy as np,torch | |
| ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT)) | |
| from model.samudrace import * | |
| c=load_config(ROOT);d=np.load(ROOT/c["data"]["path"]);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);m=SamudrACE(**ck["model_config"]);m.load_state_dict(ck["model"]);a=torch.tensor(d["atmosphere"]);o=torch.tensor(d["ocean"]);ap=[];op=[] | |
| with torch.no_grad(): | |
| for _ in range(4):a,o=m.coupled_step(a,o,20);ap.append(a.numpy());op.append(o.numpy()) | |
| path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,atmosphere=np.stack(ap,1),ocean=np.stack(op,1),origins=d["origins"],lead_days=np.arange(1,5)*5,logical_shape=d["logical_shape"],is_complete_global=False);print(path) | |