def plot_domains(wrfouts, labels, latlons, outpath, colour=0):
"""
wrfouts : list of wrfout file paths
latlons : dictionary of Nlim,Elim,Slim,Wlim
for plot
"""
fig, ax = plt.subplots(1)
# Create basemap first of all
#basemap_res = getattr(self.C,'basemap_res',self.D.basemap_res)
basemap_res = 'h'
m = Basemap(projection='merc',
llcrnrlon=latlons['Wlim'],
llcrnrlat=latlons['Slim'],
urcrnrlon=latlons['Elim'],
urcrnrlat=latlons['Nlim'],
lat_0=latlons['lat0'],
lon_0=latlons['lon0'],
resolution=basemap_res,
area_thresh=500,
ax=ax)
m.drawcoastlines()
m.drawstates()
m.drawcountries()
if not isinstance(colour, collections.Sequence):
colours = [
'k',
] * len(wrfouts)
else:
colours = colour
# Get corners of each domain
for gridlabel, fpath, colour in zip(labels, wrfouts, colours):
W = WRFOut(fpath)
print("Plotting domain {0} for {1}".format(gridlabel, fpath))
#Nlim, Elim, Slim, Wlim = W.get_limits()
x, y = m(W.lons, W.lats)
xl = len(x[0, :])
midpt = len(y[0, :]) / 2
ax.annotate(gridlabel,
color=colour,
fontsize=10,
xy=(x[0, -(0.12 * xl)], y[0, midpt]),
bbox=dict(fc='white'),
alpha=1,
va='center',
ha='left')
m.plot(x[0, :], y[0, :], colour, lw=2)
ax.plot(x[:, 0], y[:, 0], colour, lw=2)
ax.plot(x[len(y) - 1, :], y[len(y) - 1, :], colour, lw=2)
ax.plot(x[:, len(x) - 1], y[:, len(x) - 1], colour, lw=2)
# fpath = os.path.join(self.C.output_root,'domains.png')
fname = 'domains.png'
fpath = os.path.join(outpath, fname)
fig.savefig(fpath)
print("Saved to " + fpath)
def get_wrfout(self, wrf_sd=0, wrf_nc=0, dom=0, path_only=0):
"""Returns the WRFOut instance, given arguments:
Optional inputs:
wrf_sd : subdirectory for wrf file
wrf_nc : filename for wrf file
dom : domain for wrf file
path_only : only return absolute path
"""
# Check configuration to see if wrfout files should be
# sought inside subdirectories.
descend = getattr(self.C, 'wrf_folders_descend', 1)
if wrf_sd and wrf_nc:
wrfpath = os.path.join(self.C.wrfout_root, wrf_sd, wrf_nc)
elif wrf_sd:
wrfdir = os.path.join(self.C.wrfout_root, wrf_sd)
# print wrfdir
wrfpath = utils.wrfout_files_in(wrfdir,
dom=dom,
unambiguous=1,
descend=descend)
else:
wrfdir = os.path.join(self.C.wrfout_root)
wrfpath = utils.wrfout_files_in(wrfdir,
dom=dom,
unambiguous=1,
descend=descend)
if path_only:
return wrfpath
else:
return WRFOut(wrfpath)
def spaghetti(self, t, lv, va, contour, wrfouts, outpath, da=0, dom=0):
"""
wrfouts : list of wrfout files
Only change dom if there are multiple domains.
"""
m, x, y = self.basemap_setup()
time_idx = self.W.get_time_idx(t)
colours = utils.generate_colours(M, len(wrfouts))
# import pdb; pdb.set_trace()
if lv == 2000:
lv_idx = None
else:
print("Only support surface right now")
raise Exception
lat_sl, lon_sl = self.get_limited_domain(da)
slices = {'t': time_idx, 'lv': lv_idx, 'la': lat_sl, 'lo': lon_sl}
# self.ax.set_color_cycle(colours)
ctlist = []
for n, wrfout in enumerate(wrfouts):
self.W = WRFOut(wrfout)
data = self.W.get(va, slices)[0, ...]
# m.contour(x,y,data,levels=[contour,])
ct = m.contour(x,
y,
data,
colors=[
colours[n],
],
levels=[
contour,
],
label=wrfout.split('/')[-2])
print("Plotting contour level {0} for {1} from file \n {2}".format(
contour, va, wrfout))
# ctlist.append(ct)
# self.ax.legend()
# labels = [w.split('/')[-2] for w in wrfouts]
# print labels
# self.fig.legend(handles=ctlist)
# plt.legend(handles=ctlist,labels=labels)
#labels,ncol=3, loc=3,
# bbox_to_anchor=[0.5,1.5])
datestr = utils.string_from_time('output', t, tupleformat=0)
lv_na = utils.get_level_naming(va, lv)
naming = ['spaghetti', va, lv_na, datestr]
if dom:
naming.append(dom)
fname = self.create_fname(*naming)
self.save(outpath, fname)
def colocate_WRF_map(self,wrfdir):
searchfor = os.path.join(wrfdir,'wrfout*')
# pdb.set_trace()
files = glob.glob(searchfor)
W = WRFOut(files[0])
limits = {}
limits['Nlim'] = W.lats.max()
limits['Slim'] = W.lats.min()
limits['Wlim'] = W.lons.min()
limits['Elim'] = W.lons.max()
return limits
def std(ncfiles,va,tidx,lvidx):
"""
Find standard deviation in along axis of ensemble
members. Returns matrix x-y for plotting
"""
for n, nc in enumerate(ncfiles):
W = WRFOut(nc)
slices = {'lv':lvidx, 't':tidx}
va_array = W.get(va,slices)
dims = va_array.shape
if n==0:
all_members = N.zeros([len(ncfiles),1,1,dims[-2],dims[-1]])
all_members[n,0,0,:,:] = va_array[...]
std = N.std(all_members,axis=0).reshape([dims[-2],dims[-1]])
# pdb.set_trace()
return std
def plot_skewT(self,
plot_time,
plot_latlon,
dom=1,
save_output=0,
composite=0):
wrfouts = self.wrfout_files_in(self.C.wrfout_root)
for wrfout in wrfouts:
if not composite:
W = WRFOut(wrfout)
ST = SkewT(self.C, W)
ST.plot_skewT(plot_time, plot_latlon, dom, save_output)
nice_time = utils.string_from_time('title', plot_time)
print("Plotted Skew-T for time {0} at {1}".format(
nice_time, plot_latlon))
else:
#ST = SkewT(self.C)
pass
def plot_variable2D(self,va,pt,en,lv,p2p,na=0,da=0):
"""Plot a longitude--latitude cross-section (bird's-eye-view).
Use Basemap to create geographical data
========
REQUIRED
========
va = variable(s)
pt = plot time(s)
nc = ensemble member(s)
lv = level(s)
p2p = path to plots
========
OPTIONAL
========
da = smaller domain area(s), needs dictionary || DEFAULT = 0
na = naming scheme for plot files || DEFAULT = get what you're given
"""
va = self.get_sequence(va)
pt = self.get_sequence(pt,SoS=1)
en = self.get_sequence(en)
lv = self.get_sequence(lv)
da = self.get_sequence(da)
perms = self.make_iterator(va,pt,en,lv,da)
# Find some way of looping over wrfout files first, avoiding need
# to create new W instances
# print("Beginning plotting of {0} figures.".format(len(list(perms))))
#pdb.set_trace()
for x in perms:
va,pt,en,lv,da = x
W = WRFOut(en) # wrfout file class using path
F = BirdsEye(self.C,W,p2p) # 2D figure class
F.plot2D(va,pt,en,lv,da,na) # Plot/save figure
pt_s = utils.string_from_time('title',pt)
print("Plotting from file {0}: \n variable = {1}"
" time = {2}, level = {3}, area = {4}.".format(en,va,pt_s,lv,da))
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)
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
raise Exception
# Look up the method to use depending on type of plot
PLOTS = {'sum_z':self.DE_z, 'sum_xyz':self.DE_xyz}
# Creates sequence of times
ts = self.get_sequence(times)
# Dictionary of data
DATA = {}
# 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)
def plot_diff_energy(self,
ptype,
energy,
time,
folder,
fname,
p2p,
plotname,
V,
no_title=0,
ax=0):
"""
folder : directory holding computed data
fname : naming scheme of required files
p2p : root directory for plots
V : constant values to contour at
"""
sw = 0
DATA = self.load_data(folder, fname, format='pickle')
if isinstance(time, collections.Sequence):
time = calendar.timegm(time)
#for n,t in enumerate(times):
for pn, perm in enumerate(DATA):
f1 = DATA[perm]['file1']
f2 = DATA[perm]['file2']
if sw == 0:
# Get times and info about nc files
# First time to save power
W1 = WRFOut(f1)
permtimes = DATA[perm]['times']
sw = 1
# Find array for required time
x = N.where(N.array(permtimes) == time)[0][0]
data = DATA[perm]['values'][x][0]
if not pn:
stack = data
else:
stack = N.dstack((data, stack))
stack_average = N.average(stack, axis=2)
if ax:
kwargs1 = {'ax': ax}
kwargs2 = {'save': 0}
#birdseye plot with basemap of DKE/DTE
F = BirdsEye(self.C, W1, **kwargs1) # 2D figure class
#F.plot2D(va,t,en,lv,da,na) # Plot/save figure
tstr = utils.string_from_time('output', time)
fname_t = ''.join((plotname, '_{0}'.format(tstr)))
# fpath = os.path.join(p2p,fname_t)
fig_obj = F.plot_data(stack_average,
'contourf',
p2p,
fname_t,
time,
V,
no_title=no_title,
**kwargs2)
if ax:
return fig_obj
