#!/usr/bin/env python3 """Weekly cost-sensitive EHD classifiers; all outputs are isolated from the regression benchmark.""" import argparse, base64, io, json, os, pickle, socket from pathlib import Path import numpy as np from netCDF4 import Dataset import shapefile import ehd_ai as c BASE=Path(__file__).resolve().parent; OUT=BASE/'outputs/AI_RareEvent_Results'; PARTS=OUT/'parts' REPORT=BASE/'ehd_ai_rare_event_results.html'; SHP='/home/tanpv/GrADS_SCR/VN_SubRegs1/Vietnam7.shp' SETS=('core','full','selected'); NAMES=c.ALL+['sev_mod']+c.STATIC def dirs(): OUT.mkdir(parents=True,exist_ok=True); PARTS.mkdir(exist_ok=True) def clf(seed,w): from xgboost import XGBClassifier return XGBClassifier(n_estimators=450,max_depth=7,learning_rate=.045,subsample=.8,colsample_bytree=.8,min_child_weight=15,reg_lambda=3,n_jobs=20,objective='binary:logistic',eval_metric='logloss',tree_method='hist',scale_pos_weight=w,random_state=seed) def threshold(y,p): best=(-1,None,None) for t in np.unique(np.r_[.5,np.quantile(p,np.linspace(0,1,401))]): f=p>=t; h=np.sum(f&y); m=np.sum(~f&y); a=np.sum(f&~y); score=h/max(h+m+a,1) if score>best[0]: best=(score,float(t),{'csi':float(score),'pod':float(h/max(h+m,1)),'far':float(a/max(h+a,1))}) return best[1],best[2] def train(a): dirs(); X,_,csev=c.sample_data(range(2003,2017),a.points,a.seed,c.ALL+['sev_mod']); V,_,vsev=c.sample_data(range(2017,2019),a.validation,a.seed+1,c.ALL+['sev_mod']) evt=csev>=0; vevt=vsev>=0; li=NAMES.index('lead'); tw=np.minimum((X[:,li]*45).astype(int)//7,5); vw=np.minimum((V[:,li]*45).astype(int)//7,5); pos={n:i for i,n in enumerate(NAMES)} core=c.CORE+c.STATIC; meta={'method':'weekly cost-sensitive direct EHD classification','train_years':'2003-2016','validation_years':'2017-2018','test_years':'2019-2022','models':{}} for w in range(6): tr=tw==w; va=vw==w; weight=float((~evt[tr]).sum()/evt[tr].sum()); ranked=None for fs,names in (('core',core),('full',NAMES)): ix=[pos[n] for n in names]; m=clf(a.seed+w,weight); m.fit(X[tr][:,ix],evt[tr]); p=m.predict_proba(V[va][:,ix])[:,1]; th,score=threshold(vevt[va],p); pickle.dump(m,open(OUT/f'classifier_{fs}_w{w+1}.pkl','wb')); imp=dict(zip(names,map(float,m.feature_importances_))); meta['models'][f'{fs}_w{w+1}']={'features':names,'threshold':th,'validation':score,'weight':weight,'importance':imp} if fs=='full': ranked=sorted(names,key=imp.get,reverse=True) names=list(dict.fromkeys(c.STATIC+['tx','sev_mod']+ranked[:a.selected])); ix=[pos[n] for n in names]; m=clf(a.seed+100+w,weight); m.fit(X[tr][:,ix],evt[tr]); p=m.predict_proba(V[va][:,ix])[:,1]; th,score=threshold(vevt[va],p); pickle.dump(m,open(OUT/f'classifier_selected_w{w+1}.pkl','wb')); meta['models'][f'selected_w{w+1}']={'features':names,'threshold':th,'validation':score,'weight':weight,'importance':dict(zip(names,map(float,m.feature_importances_)))} json.dump(meta,open(OUT/'metadata.json','w'),indent=2); print('rare-event training complete',flush=True) def worker(a): dirs(); rank=int(os.environ.get('SLURM_PROCID',0)); size=int(os.environ.get('SLURM_NTASKS',1)); files=c.paths(range(2019,2023))[rank::size]; lon,lat,mask,yy,xx,*_=c.context(); meta=json.load(open(OUT/'metadata.json')); models={(fs,w):pickle.load(open(OUT/f'classifier_{fs}_w{w+1}.pkl','rb')) for fs in SETS for w in range(6)}; acc={fs:{k:np.zeros((6,)+mask.shape) for k in ('n','hit','miss','fa','cn','brier')} for fs in SETS} for i,path in enumerate(files): fmap,_,sev,lead,py,px=c.read_rows(path,None,c.ALL+['sev_mod']); obs=sev>=0 for w,(lo,hi) in enumerate(c.WEEKS): use=(lead>=lo)&(lead=info['threshold']; A=acc[fs] for k,v in (('n',np.ones(len(y))),('hit',f&o),('miss',~f&o),('fa',f&~o),('cn',~f&~o),('brier',(p-o)**2)): np.add.at(A[k][w],(y,x),v) print(f'rank {rank}: {i+1}/{len(files)}',flush=True) np.savez_compressed(PARTS/f'part_{rank:03d}.npz',**{f'{fs}__{k}':v for fs,A in acc.items() for k,v in A.items()}) def div(x,y): return np.divide(x,y,out=np.full_like(x,np.nan),where=y!=0) def finalize(a): ps=sorted(PARTS.glob('part_*.npz')); expected=int(os.environ.get('EXPECTED_TASKS',a.tasks)); if len(ps)!=expected: raise RuntimeError(f'expected {expected} parts, found {len(ps)}') T={} for p in ps: with np.load(p) as d: for k in d.files:T[k]=T.get(k,0)+d[k] lon,lat,mask,*_=c.context(); R={} for fs in SETS: A={k:T[f'{fs}__{k}'] for k in ('n','hit','miss','fa','cn','brier')}; h,m,f=A['hit'],A['miss'],A['fa']; R[fs]={'pod':div(h,h+m),'far':div(f,h+f),'csi':div(h,h+m+f),'fbi':div(h+f,h+m),'brier':div(A['brier'],A['n'])} total=h+m+f+A['cn']; random_hits=div((h+m)*(h+f),total) R[fs]['pofd']=div(f,f+A['cn']); R[fs]['ets']=div(h-random_hits,h+m+f-random_hits) with Dataset(OUT/'ehd_ai_rare_event_verification.nc','w') as d: d.createDimension('week',6); d.createDimension('lat',len(lat)); d.createDimension('lon',len(lon)); d.createVariable('week','i1',('week',))[:]=range(1,7); d.createVariable('lat','f4',('lat',))[:]=lat; d.createVariable('lon','f4',('lon',))[:]=lon for fs,M in R.items(): for k,z in M.items(): v=d.createVariable(f'{fs}_{k}','f4',('week','lat','lon'),zlib=True,fill_value=np.float32(9.96921e36)); v[:]=np.ma.masked_invalid(np.where(mask,z,np.nan)) html(R,lon,lat,mask) def html(R,lon,lat,mask): import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt meta=json.load(open(OUT/'metadata.json')); rows=[]; images=[] titles={'pod':'POD','pofd':'POFD','far':'FAR','fbi':'Frequency Bias Index','csi':'CSI','ets':'ETS','brier':'Brier score'}; order=('pod','pofd','far','fbi','csi','ets','brier') for fs,M in R.items(): vals={k:np.nanmean(np.where(mask,v,np.nan)) for k,v in M.items()}; rows.append(''+fs+''+''.join(f'{vals[k]:.3f}' for k in order)+'') for k in order: z=np.where(mask,M[k],np.nan); finite=z[np.isfinite(z)] if k=='ets': lim=np.nanpercentile(np.abs(finite),98) if finite.size else 1.; lim=lim if lim>0 else 1.; vmin,vmax,cmap=-lim,lim,'RdBu_r' else: vmax=np.nanpercentile(finite,98) if finite.size else 1.; vmax=vmax if vmax>0 else 1.; vmin,cmap=0.,'viridis' fig,axs=plt.subplots(1,6,figsize=(18,6.35),sharex=True,sharey=True) for w,ax in enumerate(axs): im=ax.pcolormesh(lon,lat,z[w],vmin=vmin,vmax=vmax,cmap=cmap,shading='auto') for s in shapefile.Reader(SHP).shapes(): q=np.asarray(s.points); cuts=list(s.parts)+[len(q)] for x,y in zip(cuts[:-1],cuts[1:]): ax.plot(q[x:y,0],q[x:y,1],'k',lw=.65,zorder=3) ax.set(title=f'W{w+1}',xlabel='Longitude (°E)',xlim=(101.5,110.5),ylim=(7.5,24)); ax.set_aspect('equal',adjustable='box'); ax.tick_params(labelsize=8) if w==0: ax.set_ylabel('Latitude (°N)') label=f'{fs}: {titles[k]}'; fig.suptitle(label,fontsize=15,fontweight='bold',y=.99); fig.colorbar(im,ax=axs,orientation='horizontal',fraction=.07,pad=.16,aspect=55,label=label); fig.subplots_adjust(left=.045,right=.992,bottom=.23,top=.87,wspace=.02) b=io.BytesIO(); fig.savefig(b,format='png',dpi=135,bbox_inches='tight'); plt.close(fig); images.append((label,base64.b64encode(b.getvalue()).decode())) screening=[] for fs in SETS: screening.append(f'

{fs}

') for w in range(6): info=meta['models'][f'{fs}_w{w+1}']; top=sorted(info['importance'].items(),key=lambda z:z[1],reverse=True)[:10] screening.append(f"

W{w+1} (threshold {info['threshold']:.3f}): "+', '.join(f'{n} ({v:.3f})' for n,v in top)+'

') metric_rows=''.join(f'
{label}: W1 to W6
' for label,uri in images) text=f'''Stage 2 Rare-event EHD AI

Stage 2: AI prediction of rare extreme-heat events

Design. Weekly cost-sensitive direct EHD classifiers are compared using physical core, all-variable, and importance-selected feature sets. Training: {meta['train_years']}; threshold selection: {meta['validation_years']}; untouched test: {meta['test_years']}. Initializations—not individual rows—define the temporal split. Decision thresholds are selected separately by week and feature set to maximize validation CSI. Scores below are spatial means across W1–W6; POD/FAR/CSI/FBI concern EHD occurrence and Brier is the mean squared probability error.

Test summary

{''.join(f'' for k in order)}{''.join(rows)}
feature_set{titles[k]}

Feature screening and thresholds

{''.join(screening)}

Importance is normalized tree gain and can share credit unpredictably among correlated variables; it was computed without test years. Selected models always retain lead, annual-cycle terms, coordinates, model Tx and model severity.

Weekly verification maps

{metric_rows}

Interpretation cautions

This rare-event benchmark uses separate weekly cost-sensitive classifiers. Thresholds were tuned only on the validation period; the 2019–2022 test period remained untouched until final verification. Each row uses one shared color scale across W1–W6. AUC is not reported because probability histograms were not retained by the completed distributed verification run.

''' (OUT/'ehd_ai_rare_event_results.html').write_text(text); REPORT.write_text(text) def preflight(a): dirs(); assert len(c.paths(range(2003,2023)))==2080; print(json.dumps({'host':socket.gethostname(),'output':str(OUT),'old_output_untouched':str(c.RUN)},indent=2)) def main(): p=argparse.ArgumentParser(); p.add_argument('command',choices=('preflight','train','worker','finalize')); p.add_argument('--points',type=int,default=32); p.add_argument('--validation',type=int,default=64); p.add_argument('--selected',type=int,default=18); p.add_argument('--seed',type=int,default=20260827); p.add_argument('--tasks',type=int,default=120); a=p.parse_args(); globals()[a.command](a) if __name__=='__main__':main()