#!/usr/bin/env python3
import argparse,csv,html,json,logging
from pathlib import Path
import numpy as np
import matplotlib;matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.io import netcdf_file
ROOT=Path('/home/tanpv/Progs/CLI/Dr'); LEADS=(1,3,5); MONTHS=range(1,13); SPLITS=('training','validation'); RCMS=('RegCM','clWRF')
METHODS=('Ridge1','Ridge2','ANN1','ANN2'); SPATIAL=('national','regional'); SETS={'pr':('SP1','SP2','SP3'),'t2':('ST1','ST2','ST3')}
def names(var):return [f'{r}_{m}_{s}_{p}' for r in RCMS for m in METHODS for s in SETS[var] for p in SPATIAL]+['Stack_Ensemble']
def mask_lon_lat():
with netcdf_file(Path('/home/tanpv/Progs/Dr_VN_Ope/DATA/Mask_VN_0.25_Corrected.nc'),'r',mmap=False) as d:
v=d.variables['reg'];dims=list(v.dimensions);reg=np.array(v[:],copy=True)
if dims==['lat','lon']:reg=reg.T
with netcdf_file(Path('/home/tanpv/Data/562_VNU/0.25_VN/For_BC/lead1_RegCM_01_BL_Test_VN.nc'),'r',mmap=False) as d:lon=np.array(d.variables['lon'][:],copy=True);lat=np.array(d.variables['lat'][:],copy=True)
return reg,lon,lat
def fields(obs,pred,sel,var):
o=np.where(sel[None,None,:],obs,np.nan);p=np.where(sel[None,None,:],pred,np.nan);ok=np.isfinite(o)&np.isfinite(p);n=ok.sum(2)
oo=np.where(ok,o,np.nan);pp=np.where(ok,p,np.nan);mae=np.nanmean(np.where(ok,np.abs(p-o),np.nan),2);om=np.nanmean(oo,2);pm=np.nanmean(pp,2)
cov=np.nansum((oo-om[...,None])*(pp-pm[...,None]),2);den=np.sqrt(np.nansum((oo-om[...,None])**2,2)*np.nansum((pp-pm[...,None])**2,2));corr=np.divide(cov,den,out=np.full_like(cov,np.nan),where=(den>0)&(n>=3))
if var=='pr':
base=np.nanmean(oo,2);mae=np.divide(mae,base,out=np.full_like(mae,np.nan),where=base>1e-6)*100
return mae,corr
def load_month(run,var,lead,month):
zs=[np.load(ROOT/'outputs/production_output/shards'/run/f'{r}_{var}_LT{lead}_M{month:02d}.npz') for r in RCMS];stack=np.concatenate([z['pred'] for z in zs],0)
return zs[0]['obs'],stack,np.nanmean(stack,0),zs[0]['train'],zs[0]['valid']
def calculate(run,out,log):
reg,lon,lat=mask_lon_lat();valid=(reg>=1)&(reg<=7);saved={}
for var in ('pr','t2'):
nm=names(var);err=np.full((49,3,12,2,len(lon),len(lat)),np.nan,np.float32);cor=np.full_like(err,np.nan)
for li,lead in enumerate(LEADS):
for mi,month in enumerate(MONTHS):
obs,stack,ens,tr,va=load_month(run,var,lead,month)
for si,sel in enumerate((tr,va)):
for k in range(49):
e,c=fields(obs,stack[k] if k<48 else ens,sel,var);err[k,li,mi,si]=np.where(valid,e,np.nan);cor[k,li,mi,si]=np.where(valid,c,np.nan)
log.info('metrics var=%s LT=%s month=%02d',var,lead,month)
np.savez_compressed(out/f'{var}_verification_fields.npz',error=err,corr=cor,names=np.array(nm),leads=np.array(LEADS),splits=np.array(SPLITS),lon=lon,lat=lat);saved[var]=(err,cor,nm,lon,lat)
return saved
def plot_maps(saved,out,log):
figdir=out/'figures/verification_maps';figdir.mkdir(parents=True,exist_ok=True);rows=[]
for var,(err,cor,nm,lon,lat) in saved.items():
for metric,data in (('RMAE' if var=='pr' else 'MAE',err),('CORR',cor)):
for k,name in enumerate(nm):
vals=data[k];finite=vals[np.isfinite(vals)];vmin,vmax=((-1,1) if metric=='CORR' else (0,float(np.nanpercentile(finite,98))))
fig,axs=plt.subplots(6,12,figsize=(22,10),sharex=True,sharey=True,constrained_layout=True);mesh=None
for si,split in enumerate(SPLITS):
for li,lead in enumerate(LEADS):
for mi,month in enumerate(MONTHS):
ax=axs[si*3+li,mi];mesh=ax.pcolormesh(lon,lat,vals[li,mi,si].T,shading='nearest',cmap='RdBu_r' if metric=='CORR' else 'YlOrRd',vmin=vmin,vmax=vmax)
if li==0:ax.set_title(f'M{month:02d}',fontsize=6)
if mi==0:ax.set_ylabel(f'{split}\nLT{lead}',fontsize=6)
ax.tick_params(labelsize=4,length=2)
unit='%' if metric=='RMAE' else ('°C' if metric=='MAE' else 'coefficient');fig.colorbar(mesh,ax=axs,orientation='horizontal',fraction=.025,pad=.02,label=f'{metric} ({unit})');fig.suptitle(f'{name} – {var} {metric}: training and validation',fontsize=11)
fn=f'{var}_{metric}_{name}.png';fig.savefig(figdir/fn,dpi=150);plt.close(fig);rows.append((var,metric,name,fn));log.info('map %s',fn)
return rows
def statuses(run,out):
rows=[]
for p in sorted((ROOT/'outputs/production_output/shards'/run).glob('*.json')):
for r in json.loads(p.read_text()):rows.append({'unit':p.stem,**r})
keys=sorted(set().union(*(x.keys() for x in rows)))
with open(out/'model_status_and_metrics.csv','w',newline='') as f:w=csv.DictWriter(f,fieldnames=keys);w.writeheader();w.writerows(rows)
return rows
def make_html(run,out,maps,status):
summary={}
for r in status:
key=(r['member'].split('_')[0],r.get('status','unknown'));summary[key]=summary.get(key,0)+1
trs=''.join(f'
{html.escape(m)}
{html.escape(s)}
{n}
' for (m,s),n in sorted(summary.items()))
indexed=[]
for i,(var,metric,name,fn) in enumerate(maps,1):
anchor=f'map-{i:03d}'
label=f'{var} — {metric} — {name}'
indexed.append((anchor,label,var,metric,name,fn))
toc=''.join(f'