WeatherBench / scripts /result.py
zhangrenchao's picture
Publish WeatherBench reproduction
aa529c9 verified
Raw
History Blame Contribute Delete
1.28 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.weatherbench import load_config,weighted_rmse,weighted_acc,write_json
c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);p=d["prediction"];t=d["target"];clim=t.mean((0,1),keepdims=True);metrics={}
for k,name in enumerate(("Z500","T850")):metrics[name]={"rmse":[weighted_rmse(p[:,j,k:k+1],t[:,j,k:k+1],d["latitude"]) for j in range(2)],"acc":[weighted_acc(p[:,j,k:k+1],t[:,j,k:k+1],clim[:,:,k:k+1],d["latitude"]) for j in range(2)],"persistence_rmse":[weighted_rmse(d["initial"][:,k:k+1],t[:,j,k:k+1],d["latitude"]) for j in range(2)]}
metrics["synthetic"]=True;write_json(ROOT/c["paths"]["evaluation"],metrics);fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].plot(d["lead_days"],metrics["Z500"]["rmse"],"o-",label="CNN");ax[0].plot(d["lead_days"],metrics["Z500"]["persistence_rmse"],"s-",label="persistence");ax[0].legend();im=ax[1].imshow(p[0,1,0]-t[0,1,0],cmap="coolwarm");ax[1].set_title("5-day Z500 error");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)