SamudrACE / scripts /result.py
zhangrenchao's picture
Publish SamudrACE reproduction
956cd31 verified
Raw
History Blame Contribute Delete
1.18 kB
from pathlib import Path
import sys,numpy as np
import matplotlib;matplotlib.use("Agg");import matplotlib.pyplot as plt
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
from model.samudrace import load_config,write_json
c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);a=d["atmosphere"];o=d["ocean"];armse=np.sqrt(np.mean((a-a[:,0:1])**2,axis=(0,2,3,4)));ormse=np.sqrt(np.mean((o-o[:,0:1])**2,axis=(0,2,3,4)));heat=o[:,:,0].mean((0,2,3));salt=o[:,:,19].mean((0,2,3));write_json(ROOT/c["paths"]["evaluation"],{"atmosphere_drift_rmse":armse.tolist(),"ocean_drift_rmse":ormse.tolist(),"ocean_heat_proxy":heat.tolist(),"ocean_salt_proxy":salt.tolist(),"is_complete_global":False,"synthetic":True});fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].plot(d["lead_days"],armse,label="atmosphere");ax[0].plot(d["lead_days"],ormse,label="ocean");ax[0].legend();ax[0].set(xlabel="Lead (days)",ylabel="Drift RMSE");im=ax[1].imshow(o[0,-1,76]-o[0,0,76],cmap="coolwarm");ax[1].set_title("SST drift proxy");fig.colorbar(im,ax=ax[1]);fig.tight_layout();path=ROOT/c["paths"]["figure"];path.parent.mkdir(parents=True,exist_ok=True);fig.savefig(path,dpi=150);print(path)