#!/export/hpc-lw-hmtdev1/strojniak/miniconda3/bin/python ####### # Sarah Python /export/hpc-lw-hmtdev1/strojniak/miniconda3/bin/python #Jimmy python /export/wpc-lw-hmtdev4/jcorreia/miniconda3/bin/python ################### ###################IMPORTANT#################################################### # # this relies on gen_npz_ens.py running with the data download for each model # gen_npz generates precip npz which this code needs to do the 1,3,6,24 hr means as the grib files are read / stored. # # this code must also point to hmt_functions_ens.py which is where all methods are defined # # Code Last Update: Dec 29 2021 -- SMT # #Updated : Sarah Trojniak Apr 14 2022 # new directory paths and names # updated webpath ###################################################################################### import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import numpy as np import os from datetime import datetime, timedelta import xarray as xr import cfgrib import pygrib from pyproj import Proj, transform import cartopy.crs as ccrs import cartopy from matplotlib.colors import LinearSegmentedColormap import matplotlib.colors as mcolors import argparse, copy import subprocess, pickle import json from matplotlib.path import Path import scipy.ndimage as ndimage from skimage.morphology import dilation,disk from scipy.spatial import cKDTree import sys #func_loc = '/export/wpc-lw-hmtdev1/strojniak/ffair/code/' #sys.path.insert(0, func_loc) from hmt_functions_ensv2 import * #np.set_printoptions(threshold=sys.maxsize) dta = datetime.utcnow() web = 'hpc@vm-lnx-rzdm06:/home/people/hpc/www/htdocs/hmt/hmt_webpages/images/' web_aws = 's3://s3-east-www.wpc.woc.noaa.gov/public/hmt/hmt_webpages/images/' transfer = True transfer_aws = True mask = True debug = False ovr_snow = False add_wwe = False # plot_ensv1.py -m 'HREF' -c '00' -t '0' -s 'True' -v 's' # init argparse parser = argparse.ArgumentParser() parser.add_argument('-m', '--model', help='Model', type=str) parser.add_argument('-c', '--cycle', help='Cycle', type=str) parser.add_argument('-t', '--time', help='time', type=str) parser.add_argument('-s', '--skip', help='skip', type=str) parser.add_argument('-n', '--snow', help='snow', type=str,required = False) parser.add_argument('-v', '--system', help='system: j-jimmy s-sarah h-hmtvm w-wwe', type=str) #v for virtual machine parser.add_argument('-p', '--ptype', help='ptype', type=str,required = False) #read in command line arguments args = parser.parse_args() model = args.model cycle = args.cycle #date/cycle if args.time == '0': dta = datetime.utcnow() elif args.time == '1': dta = datetime.utcnow()-timedelta(hours=24) elif len(args.time) > 1: dta = datetime.strptime(args.time,'%Y%m%d') else: print('not recognized run time') raise ValueError if args.skip == 'True': #plot mostly all the things add_precip = True else: add_precip = False if args.snow == '1': add_snow = True add_cate = True else: add_snow = False if args.ptype == '1': add_ptype = True else: add_ptype = False #what system are we working on? if args.system == 'j': sys = 'jimmy' elif args.system == 'h': sys = 'hmtvm' elif args.system == 's': sys = 'sarah' elif args.system == 'w': sys = 'wwe' else: sys = 'sarah' ###### System Options ###### work, ds, mask_loc = get_system(sys) ###work is where the data is -- ds is where to save things print(dta,model,args.time,cycle,work, ds) # unique to precip/snow pqq = [3,6,12,24] #start and accum freq pqt = [1,1,6,6] if ovr_snow or add_snow: ccc = [.1] #inches for contouring on top of qpf # not invoked, in hmt_functions its specified, # just enter extent=extent in plot calls for testing # llb = -120. lla = 20. lld = -70. llc = 50. extent = (llb,lld,lla,llc) #(-120.,-70.,20.,50.) is the default ############################## du = datetime.strftime(dta,'%Y%m%d') dd = datetime.strftime(dta,'%Y%m%d')+cycle dj = datetime.strftime(dta,'%y%j')+cycle+'00' ###### Model Options ###### dct,ix,iy,its,ite,wdir = model_dict(model, work, du) print(dct,ix,iy,its,ite,wdir) goods = [] fils = [] #look for the npz file with data not just lat long mll = np.load(mask_loc+model.lower()+'.npz') xf = mll['xf'] if xf[0,0] >= 0.: xf = xf - 360. print('negative longitudes') yf = mll['yf'] print('SHAPE: ',np.shape(xf),np.shape(yf)) central = mll['central'] if model == 'HREF' or model == 'RRFSCE': parms = ['ffri','prob','eas','mean'] for i in range(its,ite,1): #file construction if model == 'HREF': foo = ['href.t'+str(cycle)+'z.conus.'+parm+'.f'+fhrs(i)+'.grib2' for f, parm in enumerate(parms)] fils.extend(foo) elif model == 'SSEF': #parms = ['ens', 'sam06', 'sam24'] parms = ['ens'] for f, parm in enumerate(parms): foo = ['fv3lam-'+parm+'_'+str(dd)+'f'+fHHH(i)+'.grib2' for j,i in enumerate(range(its,ite,1))] fils.extend(foo) elif model == 'HRRRE' or model == 'HRRREMP': fils = [dj+str(fhrs(i))+'00' for j,i in enumerate(range(its,ite,1))] else: fils = [model.lower()+'_ens_'+str(dd)+fhrs(i)+'.grib2' for j,i in enumerate(range(its,ite,1))] ix,iy = 1799,1059 means = ['mean', 'pmm', 'lpmm'] snmeans = ['snmean', 'snpmm', 'snlpmm'] ####get *npz files for means for m, mean in enumerate(means): if model == 'HREF': # if mean == 'mean': # mll = np.load(work+model.lower()+'/'+model.lower()+'_'+mean+'_'+dd+'.npz') mll = np.load(dct+model.lower()+'_'+mean+'_'+dd+'.npz') # else: # mll = np.load(work+'/'+model.lower()+'/'+model.lower()+'_'+mean+'_'+dd+'.npz') # elif model == 'SSEF': # mll = np.load(work+'/'+model.lower()+'/'+model.lower()+'_'+mean+'_'+dd+'_test.npz') else: mll = np.load(dct+model.lower()+'_'+mean+'_'+dd+'.npz') # mll = np.load(work+model.lower()+'/'+model.lower()+'_'+mean+'_'+dd+'.npz') print("keys", mll.files) xf = mll['xf'] central = mll['central'] if xf[0,0] >= 0.: xf = xf - 360. yf = mll['yf'] print('SHAPE: ',np.shape(xf),np.shape(yf)) central = mll['central'] if mean == 'mean': if add_snow: snow = mll['snow'] # print("Snow mll: ",snow.shape) precip = mll['fcsts'] # print("Mean mll: ",precip.shape) if mean == 'pmm': if add_snow: if model == 'SSEF': psnow = mll['snow'] # print("PMM Snow mll: ",psnow.shape) pprecip = mll['fcsts'] # print("PMM mll: ",pprecip.shape) if mean == 'lpmm': if add_snow: if model == 'SSEF': lsnow = mll['snow'] # print("LPMM Snow mll: ",lsnow.shape) lprecip = mll['fcsts'] # print("LMM mll: ",lprecip.shape) if add_snow: if model == 'HREF' or model == 'RRFSCE': # mll = np.load(work+model.lower()+'/'+model.lower()+'_categorical_'+dd+'.npz') mll = np.load(dct+'/'+model.lower()+'_categorical_'+dd+'.npz') crain = mll['crain'] cfrzr = mll['cfrzr'] cicep = mll['cicep'] csnow = mll['csnow'] else: add_cate = False if mask: mgrid = masking(mask_loc,xf,yf) #define some stuff for loops#### ####get *npz files for means if add_precip: plist = [precip, pprecip, lprecip] #array names for means if add_snow: if model == 'SSEF': slist = [snow,psnow,lsnow] #array names for snow means else: slist = [snow] count = 0 #count for means #####set the parameterNames you will be looking for in for loop below #add_precip nam where 194 is ffg in ffri but also the code for ice pellets and storm surface is fro ARI in ffri nam = ['Total precipitation', 'Flash flood runoff (Encoded as an accumulation over a floating subinterval of time)', 'Storm surface runoff', '194'] #add_snow snam snam =[] for k,fil in enumerate(fils): good = True ik = k check = k+1 try: grbs = pygrib.open(dct+fil) print(grbs) except: print('File not found:',dct+fil) good = False goods.append(good) #extract the initial time and confirm its a match for expected time if good: try: validTime = datetime.strftime(grbs[1].analDate,'%Y%m%d%H') if validTime == dd: #print(validTime,dd) great = 1 else: print('Not the correct data file: ',validTime,dd) raise ValueError except: print('empty grib file') continue if model == 'HREF' or model == 'RRFSCE': if check % 4 == 0: print(k % 4) ik = count if add_precip: for p, precip in enumerate(plist): vari = means[p] print("K is: ",k,"\n ik is ",ik,"\n count is ",count,) nvar= 'precip' intervals,mmap,norm = get_colors(nvar) plot_hr(precip[ik]-precip[ik-1],nvar,vari,k,count,1,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) for o,p in enumerate(pqq): if (ik+1) >= p and (ik+1) % pqt[o] == 0: plot_hr(precip[ik]-precip[ik-p],nvar,vari,k,count,p,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) if add_snow: for s, sn in enumerate(slist): if snmeans[s] == 'snmean': print(np.unique(snow[ik+1])) vari = snmeans[s] nvar = 'snow' intervals,mmap,norm = get_colors(nvar) #snow[snow==9999] = np.nan snow[snow==9999] = 0.0 plot_hr(snow[ik]-snow[ik-1],nvar,vari,k,count,1,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) for o,p in enumerate(pqq): if (ik+1) >= p and (ik+1) % pqt[o] == 0: plot_hr(snow[ik]-snow[ik-p],nvar,vari,k,count,p,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) #### if running gen_npy_enstest.py comment out above and comment in below plot_hr call. This is because in the enstest snow[0] is set to zeros not f01 values #### # plot_hr(snow[ik+1]-snow[ik],nvar,vari,k,count,1,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) else: print("This model doesn't plot: ",sn) if add_cate: nvar = 'categorical' crain[crain==9999] = np.nan cfrzr[cfrzr==9999] = np.nan cicep[cicep==9999] = np.nan csnow[csnow==9999] = np.nan plot_categorical(crain[ik],cfrzr[ik],cicep[ik],csnow[ik],nvar,k,count,model,dd,its,tcentral=central,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) else: print("add_cate is: ",add_cate) count += 1 else: print("not plotting a mean") else: # ik = k if add_precip: for p, precip in enumerate(plist): vari = means[p] # print("K is: ",k,"\n ik is ",ik) nvar= 'precip' intervals,mmap,norm = get_colors(nvar) plot_hr(precip[ik]-precip[ik-1],nvar,vari,k,count,1,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) for o,p in enumerate(pqq): if (ik) >= p and (ik) % pqt[o] == 0: plot_hr(precip[ik]-precip[ik-p],nvar,vari,k,count,p,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) if add_snow: for s, sn in enumerate(slist): # if sn == 'smean': print(np.unique(snow[ik+1])) vari = snmeans[s] nvar = 'snow' intervals,mmap,norm = get_colors(nvar) test = sn[ik]-sn[ik-1] print(np.unique(test)) plot_hr(test,nvar,vari,k,count,1,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) for o,p in enumerate(pqq): if (ik) >= p and (ik) % pqt[o] == 0: plot_hr(sn[ik]-sn[ik-p],nvar,vari,k,count,p,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) if add_cate: # if add_cate: nvar = 'categorical' # intervals,mmap,norm = get_colors(nvar) crain[crain==9999] = np.nan cfrzr[cfrzr==9999] = np.nan cicep[cicep==9999] = np.nan csnow[csnow==9999] = np.nan # print(np.unique(crain[ik])) plot_categorical(crain[ik],cfrzr[ik],cicep[ik],csnow[ik],nvar,k,count,model,dd,its,tcentral=central,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) else: print("add_cate: ",add_cate) for grb in grbs: vtime = grb.endStep print("this is valid time of grib: ",vtime) start = grb.startStep end = grb.endStep # print("LOOOK HERE: ",grb.parameterName) if add_precip: nvar = 'precip' print(fil) if grb.parameterName == nam[0]: #total preciptation print(grb.parameterName) try: # print("this is Derived : ",grb.derivedForecast,type(grb.derivedForecast),grb) df = grb.derivedForecast except: print("no dervied value") df = 'nodf' print(fil) print(df) if df == 'nodf': nvar = 'probability' try: if model == 'HREF' or model == 'RRFSCE' or model == 'SSEF': sf = grb.scaleFactorOfUpperLimit sc = grb.scaledValueOfUpperLimit/10.**sf if 'ffri' in fil: vari = 'ari' # print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari now: ",vari) elif 'eas' in fil: vari = 'eas' # print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari now: ",vari) else: vari = 'prob' # print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari now: ",vari) else: sf = grb.scaleFactorOfLowerLimit sc = grb.scaledValueOfLowerLimit/10.**sf vari = 'prob' # print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari now: ",vari) except: print("Not a prob of exceeding") vari = 'notprob' else: vari = 'somemean' #else: print("Check next nam") print(fil) elif grb.parameterName == nam[1]: #Flash flood runoff (Encoded as an accumulation over a floating subinterval of time) nvar = 'probability' print(grb.parameterName) try: #scl.append(grb.scaledValueOfLowerLimit) sf = grb.scaleFactorOfLowerLimit sc = grb.scaledValueOfLowerLimit/10.**sf vari = 'ffg' print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari now: ",vari) except: vari ='foo' sc = '' sf = '' #else: print("Check next nam") print(fil) elif grb.parameterName == nam[2]: #Storm Runoff nvar = 'probability' print(grb.parameterName) try: #scl.append(grb.scaledValueOfLowerLimit) sf = grb.scaleFactorOfLowerLimit sc = grb.scaledValueOfLowerLimit/10.**sf vari = 'ari' print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari now: ",vari) except: vari ='foo' sc = '' sf = '' #else: print("Check next nam") print(fil) elif grb.parameterName == nam[3] and 'ffri' in fil: #194 nvar = 'probability' print(grb.parameterName) sf = grb.scaleFactorOfUpperLimit sc = grb.scaledValueOfUpperLimit/10.**sf if sc != 0: vari = 'ffg' print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari: ",vari) else: print(grb.parameterName) print("Not a precip prob") vari = 'foo' if vari == 'prob' or vari == 'ffg' or vari == 'ari' or vari == 'eas': fcst = grb.values # print("First Unique: ",np.unique(fcst)) intervals,mmap,norm = get_colors(nvar) if vari == 'prob' or vari == 'eas': try: if sc == 12.7: p = str(0.5) if sc == 25.4: p = str(1) if sc == 50.8: p = str(2) if sc == 76.2: p = str(3) if sc == 127: p = str(5) if sc == 6.35: p = str(0.25) if sc == 0.254: p = str(0.01) if sc == 203.2: p = str(8) except: print("dont't plot this prob: ",sc) p = 'foo' vari = 'foo' continue elif vari == 'ari': if sc == 2: p = str(2)+'yr' if sc == 5: p = str(5)+'yr' if sc == 10: p = str(10)+'yr' if sc == 100: p = str(100)+'yr' else: p = 'foo' if p!= 'foo': plot_probs(fcst,nvar,vari,p,k,vtime,start,end,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) if add_snow: if model == 'HREF' or model == 'RRFSCE': snam = ["Water equivalent of accumulated snow depth"] #,"192","193","194","195"] elif model == 'SSEF': snam = ["Total snowfall"] else: snam = ["Water equivalent of accumulated snow depth"] print(grb.parameterName,fil) if grb.parameterName == snam[0]: nvar = 'sprobability' print(grb.parameterName) if model == 'HREF' or model == 'RRFSCE': if 'prob' in fil: vari = 'sprob' elif 'eas' in fil: vari = 'sprobeas' else: continue ###in mm sf = grb.scaleFactorOfUpperLimit sc = grb.scaledValueOfUpperLimit/10.**sf # print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari: ",vari) elif model == 'SSEF': start = grb.startStep end = grb.endStep dif = end - start print(dif) grok = grb.typeOfGeneratingProcess if grok == 194: #in m? vari = 'sprob' sf = grb.scaleFactorOfUpperLimit sc = grb.upperLimit # if grok == 4 and dif == 1: # nvar = 'snow' # vari = snmeans[0] # p = 'foo' # sn = grb.values*100/2.54 # intervals,mmap,norm = get_colors(nvar) # sn[sn==9999] = 0.0 # sn[np.isnan(sn)] = 0 # print(sn) # plot_hr(sn,nvar,vari,k,count,p,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) else: vari = 'foo' sf = '' sc = '' else: sf = grb.scaleFactorOfUpperLimit sc = grb.scaledValueOfUpperLimit/10.**sf print("this is sf: ",str(sf)," This is sc: ",str(sc)," This is vari: ",vari) if vari == 'sprob' or vari == 'sprobeas': fcst = grb.values # print(np.unique(fcst)) intervals,mmap,norm = get_colors(nvar) try: if sc == 2.54: p = str(0.1) if sc == 7.62: p = str(0.3) if sc == 15.24: p = str(0.6) if sc == 25.4 or sc == 0.0254: p = str(1) if sc == 30.48 or sc == 0.03048: p = str(1.4) if sc == 50.8 or sc == 0.0508: p =str(2) if sc == 76.2 or sc == 0.0762: p = str(3) print("this is p ",p) except: print("dont't plot this prob: ",sc) p = 'foo' if p != 'foo': plot_probs(fcst,nvar,vari,p,k,vtime,start,end,model,dd,its,tcentral=central,cmap=mmap,intervals=intervals,norm=norm,work=work,ds=ds,mask=mask,mgrid=mgrid,debug=debug,web=web,wdir=wdir,transfer=transfer,transfer_aws=transfer_aws,web_aws=web_aws,xf=xf,yf=yf) else: print("Not ",snam[0]) vari = 'foo'