add figure scripts

pcmdi_metrics/precip_variability/scripts_pcmdi/calc_ps_area.freq.mean.regional.py

@@ -0,0 +1,140 @@
+import numpy as np
+import datetime
+import pickle
+import json
+import copy
+import os
+
+def prdday_to_frq3hridx(prdday,frequency):
+ frq3hr=1./(float(prdday)*8.)
+ idx = (np.abs(frequency-frq3hr)).argmin()
+ return int(idx)
+
+def prdday_to_frq1didx(prdday,frequency):
+ frq24hr=1./(float(prdday))
+ idx = (np.abs(frequency-frq24hr)).argmin()
+ return int(idx)
+
+#--------------------
+# User settings here
+#--------------------
+# "3hr" or "day"
+hr=''
+# Input file name (pickle .pkl output from calc_ps_area.mean.py)
+fname = ''
+# Output directory (not including version)
+outdir = ''
+
+#-----------------------
+# The rest of the script
+#-----------------------
+ver = datetime.datetime.now().strftime("v%Y%m%d")
+
+if hr == 'day':
+ frqs_forced=['semi-annual','annual']
+ frqs_unforced=['synoptic','sub-seasonal','seasonal-annual','interannual']
+elif hr == '3hr':
+ frqs_forced=['semi-diurnal','diurnal','semi-annual','annual']
+ frqs_unforced=['sub-daily','synoptic','sub-seasonal','seasonal-annual','interannual'] 
+
+infile = open(fname, 'rb')
+psdm = pickle.load(infile)
+
+psdmfm=copy.deepcopy(psdm)
+for frc in psdm.keys():
+ if (frc=='forced'):
+ frqs=frqs_forced
+ elif (frc=='unforced'): 
+ frqs=frqs_unforced
+ print(frc)
+ for mip in psdm[frc].keys():
+ print(mip)
+ for dat in psdm[frc][mip].keys():
+ frequency=np.array(psdm[frc][mip][dat]['freqs'])
+ del psdm[frc][mip][dat]['freqs']
+ del psdmfm[frc][mip][dat]['freqs']
+
+ for var in psdm[frc][mip][dat].keys():
+ print(dat,var)
+ for idm, dom in enumerate(psdm[frc][mip][dat][var].keys()):
+ am=np.array(psdm[frc][mip][dat][var][dom])
+ del psdmfm[frc][mip][dat][var][dom]
+ psdmfm[frc][mip][dat][var][dom] = {}
+ print(dom)
+ for frq in frqs:
+ print(frq)
+ if (frq=='semi-diurnal'): # pr=0.5day
+ idx=prdday_to_frq1didx(0.5,frequency)
+ amfm = am[idx]
+ elif (frq=='diurnal'): # pr=1day
+ idx=prdday_to_frq1didx(1,frequency)
+ amfm = am[idx]
+ if (frq=='semi-annual'): # 180day=<pr=<183day
+ if hr=='day':
+ idx2=prdday_to_frq1didx(180,frequency)
+ idx1=prdday_to_frq1didx(183,frequency)
+ elif hr=='3hr':
+ idx2=prdday_to_frq3hridx(180,frequency)
+ idx1=prdday_to_frq3hridx(183,frequency)
+ amfm = np.nanmean(am[idx1:idx2+1])
+ elif (frq=='annual'): # 360day=<pr=<366day
+ if hr=='day':
+ idx2=prdday_to_frq1didx(360,frequency)
+ idx1=prdday_to_frq1didx(366,frequency)
+ elif hr=='3hr':
+ idx2=prdday_to_frq3hridx(360,frequency)
+ idx1=prdday_to_frq3hridx(366,frequency)
+ amfm = np.nanmean(am[idx1:idx2+1])
+ elif (frq=='sub-daily'): # pr<1day
+ idx1=prdday_to_frq1didx(1,frequency)
+ amfm = np.nanmean(am[idx1+1:])
+ elif (frq=='synoptic'): # 1day=<pr<20day
+ if hr=='day':
+ idx2=prdday_to_frq1didx(1,frequency)
+ idx1=prdday_to_frq1didx(20,frequency)
+ elif hr=='3hr':
+ idx2=prdday_to_frq3hridx(1,frequency)
+ idx1=prdday_to_frq3hridx(20,frequency)
+ amfm = np.nanmean(am[idx1+1:idx2+1])
+ elif (frq=='sub-seasonal'): # 20day=<pr<90day
+ if hr=='day':
+ idx2=prdday_to_frq1didx(20,frequency)
+ idx1=prdday_to_frq1didx(90,frequency)
+ elif hr=='3hr':
+ idx2=prdday_to_frq3hridx(20,frequency)
+ idx1=prdday_to_frq3hridx(90,frequency)
+ amfm = np.nanmean(am[idx1+1:idx2+1])
+ elif (frq=='seasonal-annual'): # 90day=<pr<365day
+ if hr=='day':
+ idx2=prdday_to_frq1didx(90,frequency)
+ idx1=prdday_to_frq1didx(365,frequency)
+ elif hr=='3hr':
+ idx2=prdday_to_frq3hridx(90,frequency)
+ idx1=prdday_to_frq3hridx(365,frequency)
+ amfm = np.nanmean(am[idx1+1:idx2+1])
+ elif (frq=='interannual'): # 365day=<pr
+ if hr=='day':
+ idx2=prdday_to_frq1didx(365,frequency)
+ elif hr=='3hr':
+ idx2=prdday_to_frq3hridx(365,frequency)
+ amfm = np.nanmean(am[:idx2+1])
+
+ psdmfm[frc][mip][dat][var][dom][frq] = amfm.tolist()
+
+res = os.path.basename(fname).split("_")[2].split(".")[1]
+
+outdir = os.path.join(outdir,ver) # add version to outdir
+if not(os.path.isdir(outdir)):
+ os.makedirs(outdir)
+
+# Save as pickle for figure creation
+outfile=open(outdir+'/PS_pr.{0}_regrid.{1}_area.freq.mean.forced.pkl'.format(hr,res),'wb')
+pickle.dump(psdmfm, outfile)
+outfile.close()
+
+# Write to JSON file
+outfile = open(outdir+'/PS_pr.{0}_regrid.{1}_area.freq.mean.forced.json'.format(hr,res),'w')
+json.dump(psdmfm, outfile, sort_keys=True, indent=4, separators=(',', ': '))
+outfile.close()
+
+

pcmdi_metrics/precip_variability/scripts_pcmdi/calc_ps_area.mean.regional.py

@@ -0,0 +1,98 @@
+import xcdat as xc
+import numpy as np
+import pickle
+import glob
+import os
+import sys
+import datetime
+
+
+#--------------------
+# User settings here
+#--------------------
+# List of domain names
+# e.g. ['CONUS']
+domains=[]
+# Which ensembles are included
+# e.g. ['cmip6','obs']
+mips=[]
+# Input directories. Replace model, realization,
+# and other variables with wildcard as needed
+# Replace domain name references with '%(domain)'
+model_root = ''
+obs_root = ''
+# Output directory (not including version)
+outdir = ''
+
+#-----------------------
+# The rest of the script
+#-----------------------
+ver = datetime.datetime.now().strftime("v%Y%m%d")
+
+frcs=['forced','unforced']
+
+vars=['power','rednoise','sig95']
+
+psdm={}
+
+for ifc, frc in enumerate(frcs):
+ psdm[frc]={}
+ for dom in domains:
+ for im, mip in enumerate(mips):
+ if mip not in psdm[frc]:
+ psdm[frc][mip]={}
+ dir = '/'
+ if mip=="cmip6":
+ root=model_root.replace('%(domain)',dom)
+ elif mip=="obs":
+ root=obs_root.replace('%(domain)',dom)
+ print(root)
+ if(frc=='forced'):
+ file_list = sorted(set(glob.glob(os.path.join(root,'PS*')))-set(glob.glob(os.path.join(root,'PS*_unforced.nc'))) - set(glob.glob(os.path.join(root,'PS*.json'))))
+ elif(frc=='unforced'):
+ file_list = sorted(set(glob.glob(os.path.join(root,'PS*_unforced.nc'))))
+
+ data=[]
+ for file in file_list:
+ print("File (first loop): ",file)
+ if(mip=='obs'):
+ tmp = file.split('/')[-1].split('_')[3]
+ model = tmp.split('.')[0]
+ data.append(model)
+ else:
+ tmp = file.split('/')[-1].split('_')[3]
+ model = tmp.split('.')[0]
+ ens = tmp.split('.')[1]
+ data.append(model+'.'+ens)
+ print(mip, '# of data:', len(data))
+ print("DATA: ",data)
+
+ for id, dat in enumerate(data):
+ #freqs = f[id]['freqs']
+ print("File (second loop): ",file_list[id])
+ frds = xc.open_dataset(file_list[id])
+ if dat not in psdm[frc][mip]:
+ psdm[frc][mip][dat]={}
+ if 'freqs' not in psdm[frc][mip][dat]:
+ psdm[frc][mip][dat]['freqs']=list(frds["freqs"])
+ for iv, var in enumerate(vars):
+ print(frc,mip,dat,var)
+ if var not in psdm[frc][mip][dat]:
+ psdm[frc][mip][dat][var]={}
+ am = frds.spatial.average(var,axis=["X","Y"],weights="generate")[var]
+ psdm[frc][mip][dat][var][dom]=list(am.data)
+ frds.close()
+
+res = os.path.basename(file_list[0]).split("_")[2].split(".")[1]
+hr = os.path.basename(file_list[0]).split("_")[1].split(".")[1]
+
+outdir = os.path.join(outdir,ver) # add version to outdir
+if not(os.path.isdir(outdir)):
+ os.makedirs(outdir)
+
+outfile=open(os.path.join(outdir,'PS_pr.{0}_regrid.{1}_area.mean.pkl'.format(hr,res)),'wb')
+pickle.dump(psdm,outfile)
+outfile.close()
+
+
+