def composite_profile(self,va,plot_time,plot_latlon,wrfouts,outpath,
dom=1,mean=1,std=1,xlim=0,ylim=0,fig=0,ax=0,
locname=0,ml=-2):
"""
Loop over wrfout files.
Get profile of variable
Plot all members
Optional standard deviation
Optional mean
If ax, save image to that axis
ml : member level. negative number that corresponds to the
folder in absolute string for naming purposes.
"""
# Set up figure
if isinstance(fig,M.figure.Figure):
self.fig = fig
self.ax = ax
else:
self.fig, self.ax = plt.subplots()
# Plot settings
if xlim:
xmin, xmax, xint = xlim
if ylim:
if ylim[1] > ylim[0]:
# top to bottom
P_top, P_bot, dp = [y*100.0 for y in ylim]
else:
P_bot, P_top, dp = [y*100.0 for y in ylim]
else:
P_bot = 100000.0
P_top = 20000.0
dp = 10000.0
# plevs = N.arange(P_bot,P_top,dp)
# Get wrfout prototype for information
W = WRFOut(wrfouts[0])
lat, lon = plot_latlon
datestr = utils.string_from_time('output',plot_time)
t_idx = W.get_time_idx(plot_time,)
y, x, exact_lat, exact_lon = utils.getXY(W.lats1D,W.lons1D,lat,lon)
slices = {'t': t_idx, 'la': y, 'lo': x}
#var_slices = {'t': t_idx, 'lv':0, 'la':y, 'lo':x}
# Initialise array
# Number of levels and ensembles:
nPlevs = W.z_dim
data = W.get(va,slices)
nvarlevs = data.shape[1]
nens = len(wrfouts)
# 2D: (profile,member)
profile_arr = N.zeros((nvarlevs,nens))
composite_P = N.zeros((nPlevs,nens))
# Set up legend
labels = []
colourlist = utils.generate_colours(M,nens)
# M.rcParams['axes.color_cycle'] = colourlist
# Collect profiles
for n,wrfout in enumerate(wrfouts):
W = WRFOut(wrfout)
# Get pressure levels
composite_P[:,n] = W.get('pressure',slices)[0,:,0,0]
#elev = self.W.get('HGT',H_slices)
#pdb.set_trace()
# Grab variable
profile_arr[:,n] = W.get(va,slices)[0,:,0,0]
# Plot variable on graph
self.ax.plot(profile_arr[:,n],composite_P[:,n],color=colourlist[n])
member = wrfout.split('/')[ml]
labels.append(member)
# if locname=='KOAX': pdb.set_trace()
# Compute mean, std etc
if mean:
profile_mean = N.mean(profile_arr,axis=1)
profile_mean_P = N.mean(composite_P,axis=1)
self.ax.plot(profile_mean,profile_mean_P,color='black')
if std:
# Assume mean P across ensemble is correct level
# to plot standard deviation
profile_std = N.std(profile_arr,axis=1)
std_upper = profile_mean + profile_std
std_lower = profile_mean - profile_std
self.ax.plot(std_upper,profile_mean_P,'k--')
self.ax.plot(std_lower,profile_mean_P,'k--')
if not locname:
fname = '_'.join(('profile_comp',va,datestr,'{0:03d}'.format(x),'{0:03d}'.format(y))) + '.png'
else:
fname = '_'.join(('profile_comp',va,datestr,locname)) + '.png'
# Set semi-log graph
self.ax.set_yscale('log')
# Plot limits, ticks
yticks = N.arange(P_bot,P_top+dp*100.0,-100*100.0)
# yticks = N.arange(P_bot,P_top+dp,-dp*100)
# self.ax.set_yticks(yticks)
ylabels = ["%4u" %(p/100.0) for p in yticks]
self.ax.set_yticks(yticks)
self.ax.set_yticklabels(ylabels)
# import pdb; pdb.set_trace()
# self.ax.yaxis.tick_right()
#plt.axis([-20,50,105000.0,20000.0])
#plt.xlabel(r'Temperature ($^{\circ}$C) at 1000 hPa')
#plt.xticks(xticks,['' if tick%10!=0 else str(tick) for tick in xticks])
self.ax.set_ylabel('Pressure (hPa)')
#yticks = N.arange(self.P_bot,P_t-1,-10**4)
#plt.yticks(yticks,yticks/100)
if xlim:
self.ax.set_xlim([xmin,xmax])
xticks = N.arange(xmin,xmax+xint,xint)
self.ax.set_xticks(xticks)
# Flip y axis
# Limits are already set
self.ax.set_ylim([P_bot,P_top])
# plt.tight_layout(self.fig)
#plt.autoscale(enable=1,axis='x')
#ax = plt.gca()
#ax.relim()
#ax.autoscale_view()
#plt.draw()
self.ax.legend(labels,loc=2,fontsize=6)
self.save(outpath,fname)
plt.close(self.fig)
# Get all permutations of files
for perm in itertools.combinations(files,2):
DATA[perm] = {}
f1, f2 = perm
W1 = WRFOut(f1)
W2 = WRFOut(f2)
# Make sure times are the same in both files
if not W1.wrf_times == W2.wrf_times:
print("Times are not identical between input files.")
raise Exception
# Find indices of each time
t_idx = []
for t in ts:
t_idx.append(W1.get_time_idx(t))
DATA[perm]['times'] = ts
DATA[perm]['values'] = PLOTS[ptype](nc0,nc1,t_idx,
energy,lower,upper)
if d_return and not d_save:
return DATA
elif d_save and not d_return:
self.save(DATA,d_save,d_fname)
elif d_return and d_save:
self.save(DATA,d_save,d_fname)
return DATA
def DE_xyz(self,nc0,nc1,t_idx,energy,*args):
"""
def compute_diff_energy(ptype,
energy,
files,
times,
upper=None,
lower=None,
d_save=True,
d_return=True,
d_fname='diff_energy_data'):
"""
This method computes difference kinetic energy (DKE)
or different total energy (DTE, including temp)
between WRFout files for a given depth of the
atmosphere, at given time intervals
:param ptype: 'sum_z' or 'sum_xyz'.
'sum_z' integrates vertically between lower and
upper hPa and creates a time series.
'sum_xyz' integrates over the 3D space (again between
the upper and lower bounds) and creates 2D arrays.
:param energy: 'kinetic' or 'total'
:param upper: upper limit of vertical integration
:param lower: lower limit of vertical integration
:param files: abs paths to all wrfout files
:param times: times for computations - tuple format
:param d_save: save dictionary to folder (path to folder)
:param d_return: return dictionary (True or False)
:param d_fname: custom filename
:returns: N.ndarray -- time series or list of 2D arrays
"""
if d_save and not isinstance(d_save, basestring):
d_save = os.environ['HOME']
# First, save or output? Can't be neither!
if not d_save and not d_return:
print("Pick save or output, otherwise it's a waste of computer"
"power")
raise Exception
print("Saving pickle file to {0}".format(d_save))
# Look up the method to use depending on type of plot
PLOTS = {'sum_z': self.DE_z, 'sum_xyz': self.DE_xyz}
print('Get sequence of time')
# Creates sequence of times
ts = utils.get_sequence(times)
# Dictionary of data
DATA = {}
print('Get permutations')
# Get all permutations of files
nperm = len(list(itertools.combinations(files, 2)))
print('Start loop')
# pdb.set_trace()
for n, perm in enumerate(itertools.combinations(files, 2)):
print("No. {0} from {1} permutations".format(n, nperm))
perm_start = time.time()
DATA[str(n)] = {}
f1, f2 = perm
W1 = WRFOut(f1)
W2 = WRFOut(f2)
print('WRFOuts loaded.')
#pdb.set_trace()
# Make sure times are the same in both files
if not N.all(N.array(W1.wrf_times) == N.array(W2.wrf_times)):
print("Times are not identical between input files.")
raise Exception
else:
print(
"Passed check for identical timestamps between "
"NetCDF files")
# Find indices of each time
print('Finding time indices')
t_idx = []
for t in ts:
t_idx.append(W1.get_time_idx(t))
print("Calculating values now...")
DATA[str(n)]['times'] = ts
DATA[str(n)]['values'] = []
for t in t_idx:
DATA[str(n)]['values'].append(PLOTS[ptype](W1.nc, W2.nc, t, energy,
lower, upper))
DATA[str(n)]['file1'] = f1
DATA[str(n)]['file2'] = f2
print "Calculation #{0} took {1:2.2f} seconds.".format(
n,
time.time() - perm_start)
if d_return and not d_save:
return DATA
elif d_save and not d_return:
#self.save_data(DATA,d_save,d_fname)
self.save_data(DATA, d_save, d_fname)
#self.json_data(DATA,d_save,d_fname)
return
elif d_return and d_save:
#self.save_data(DATA,d_save,d_fname)
self.save_data(DATA, d_save, d_fname)
#self.json_data(DATA,d_save,d_fname)
return DATA
def compute_diff_energy(ptype,energy,files,times,upper=None,lower=None,
d_save=True,d_return=True,d_fname='diff_energy_data'):
"""
This method computes difference kinetic energy (DKE)
or different total energy (DTE, including temp)
between WRFout files for a given depth of the
atmosphere, at given time intervals
:param ptype: '2D' or '3D'.
'2D' integrates vertically between lower and
upper hPa and creates a time series.
'3D' integrates over the 3D space (again between
the upper and lower bounds) and creates 2D arrays.
:param energy: 'DKE' or 'DTE'
:param upper: upper limit of vertical integration
:param lower: lower limit of vertical integration
:param files: abs paths to all wrfout files
:param times: times for computations - tuple format
:param d_save: save dictionary to folder (path to folder)
:param d_return: return dictionary (True or False)
:param d_fname: custom filename
:returns: N.ndarray -- time series or list of 2D arrays
"""
if d_save and not isinstance(d_save,basestring):
d_save = os.environ['HOME']
# First, save or output? Can't be neither!
if not d_save and not d_return:
print("Pick save or output, otherwise it's a waste of computer"
"power")
raise Exception
print("Saving pickle file to {0}".format(d_save))
# Look up the method to use depending on type of plot
PLOTS = {'1D':DE_z, '3D':DE_xyz}
print('Get sequence of time')
# Creates sequence of times
ts = utils.get_sequence(times)
# Dictionary of data
DATA = {}
print('Get permutations')
# Get all permutations of files
nperm = len(list(itertools.combinations(files,2)))
print('Start loop')
# pdb.set_trace()
for n, perm in enumerate(itertools.combinations(files,2)):
print("No. {0} from {1} permutations".format(n,nperm))
perm_start = time.time()
DATA[str(n)] = {}
f1, f2 = perm
W1 = WRFOut(f1)
W2 = WRFOut(f2)
print('WRFOuts loaded.')
#pdb.set_trace()
# Make sure times are the same in both files
if not N.all(N.array(W1.wrf_times) == N.array(W2.wrf_times)):
print("Times are not identical between input files.")
raise Exception
else:
print("Passed check for identical timestamps between "
"NetCDF files")
# Find indices of each time
print('Finding time indices')
t_idx = []
for t in ts:
t_idx.append(W1.get_time_idx(t))
print("Calculating values now...")
DATA[str(n)]['times'] = ts
DATA[str(n)]['values'] = []
for t in t_idx:
DATA[str(n)]['values'].append(PLOTS[ptype](W1.nc,W2.nc,t,
energy,lower,upper))
DATA[str(n)]['file1'] = f1
DATA[str(n)]['file2'] = f2
# import pdb; pdb.set_trace()
print "Calculation #{0} took {1:2.1f} seconds.".format(n,time.time()-perm_start)
if d_return and not d_save:
return DATA
elif d_save and not d_return:
#self.save_data(DATA,d_save,d_fname)
utils.save_data(DATA,d_save,d_fname)
#self.json_data(DATA,d_save,d_fname)
return
elif d_return and d_save:
#self.save_data(DATA,d_save,d_fname)
utils.save_data(DATA,d_save,d_fname)
#self.json_data(DATA,d_save,d_fname)
return DATA
