WRF-ML / scripts /result.py
zhangrenchao's picture
Publish WRF-ML reproduction
d3225f2 verified
Raw
History Blame Contribute Delete
1.23 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.wrf_ml import load_config,write_json,OUTPUT_NAMES
c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);pr=np.sqrt(np.mean((d["profile_prediction"]-d["profile_target"])**2,axis=(0,1)));br=np.sqrt(np.mean((d["boundary_prediction"]-d["boundary_target"])**2,axis=0));write_json(ROOT/c["paths"]["evaluation"],{"sw_flux_rmse":float(pr[0]),"lw_flux_rmse":float(pr[1]),"sw_heating_rate_rmse":float(pr[2]),"lw_heating_rate_rmse":float(pr[3]),"surface_sw_flux_rmse":float(br[0]),"toa_sw_flux_rmse":float(br[1]),"synthetic":True});fig,ax=plt.subplots(1,2,figsize=(9,3.5));ax[0].plot(d["profile_target"][0,:,0],d["vertical_levels"],label="RRTMG proxy");ax[0].plot(d["profile_prediction"][0,:,0],d["vertical_levels"],label="ML");ax[0].legend();ax[0].set(xlabel="SW flux",ylabel="Level");ax[1].bar(np.arange(4),pr);ax[1].set_xticks(np.arange(4),["SWF","LWF","SWH","LWH"]);ax[1].set_title("Profile RMSE");fig.tight_layout();path=ROOT/c["paths"]["figure"];path.parent.mkdir(parents=True,exist_ok=True);fig.savefig(path,dpi=150);print(path)