def __init__(self,Wctrl_fpath,Wmod_fpath,vrbl,utc,lv=False,
accum_hr=False,radar_datadir=False,thresh=False,
footprint=500):
self.utc = utc
self.C = {}
self.M = {}
self.thresh = thresh
self.footprint = footprint
self.M['WRFOut'] = WRFOut(Wmod_fpath)
self.dx = self.M['WRFOut'].dx
self.compute_d(self.M['WRFOut'])
if Wctrl_fpath is False and (vrbl=='REFL_comp' or vrbl=='cref'):
use_radar_obs = True
self.C['data'] = self.get_radar_verif(utc,radar_datadir)
else:
use_radar_obs = False
Wctrl = WRFOut(Wctrl_fpath)
# Get 2D grids for ctrl and model
if vrbl == 'accum_precip':
if not accum_hr:
raise Exception("Need to set accumulation hours.")
self.C['data'] = Wctrl.compute_accum_rain(utc,accum_hr)[0,0,:,:]
self.M['data'] = self.M['WRFOut'].compute_accum_rain(utc,accum_hr)[0,0,:,:]
else:
self.M['data'] = self.M['WRFOut'].get(vrbl,level=lv,utc=utc)[0,0,:,:]
if not use_radar_obs:
self.C['data'] = Wctrl.get(vrbl,level=lv,utc=utc)[0,0,:,:]
# Set negative values to 0
# if vrbl == 'REFL_comp':
self.C['data'][self.C['data']<0] = 0
self.M['data'][self.M['data']<0] = 0
self.vrbl = vrbl
"""
fig, ax = plt.subplots(1)
ax.pcolor(self.C['data'])
fig.savefig('/home/jrlawson/public_html/bowecho/SALtests/obs_pcolor.png')
plt.close(fig)
fig, ax = plt.subplots(1)
ax.pcolor(self.M['data'])
fig.savefig('/home/jrlawson/public_html/bowecho/SALtests/mod_pcolor.png')
plt.close(fig)
import pdb; pdb.set_trace()
"""
self.identify_objects()
self.compute_amplitude()
self.compute_location()
self.compute_structure()
print("S = {0} A = {1} L = {2}".format(self.S,self.A,self.L))
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 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 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 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 std(ncfiles,vrbl,utc=False,level=False,other=False,axis=0):
"""
Find standard deviation in along axis of ensemble
members. Returns matrix x-y for plotting
"""
for n, nc in enumerate(ncfiles):
# print("Ensemble member {0} loaded.".format(n))
W = WRFOut(nc)
vrbl_array = W.get(vrbl,utc=utc,level=level,other=other)
if vrbl=='cref':
vrbl_array[vrbl_array<0] = 0
if n==0:
dims = [len(ncfiles),] + list(vrbl_array.shape)
all_members = N.zeros(dims)
all_members[n,...] = vrbl_array[...]
# import pdb; pdb.set_trace()
std = N.std(all_members,axis=axis)
return std
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_ttest(ncfiles1,ncfiles2,vrbl,utc=False,level=False,other=False):
"""
Find standard deviation in along axis of ensemble
members. Returns matrix x-y for plotting. Returns sig test
"""
std = []
std_ave = []
for ncfiles in (ncfiles1, ncfiles2):
for n, nc in enumerate(ncfiles):
# print("Ensemble member {0} loaded.".format(n))
W = WRFOut(nc)
vrbl_array = W.get(vrbl,utc=utc,level=level,other=other)
if vrbl=='cref':
vrbl_array[vrbl_array<0] = 0
if n==0:
dims = [len(ncfiles),] + list(vrbl_array.shape)
all_members = N.zeros(dims)
all_members[n,...] = vrbl_array[...]
# sample[n,...] = vrbl_array[
# import pdb; pdb.set_trace()
sample=all_members
std.append(N.std(sample,axis=0)[0,0,:,:])
std_ave.append(N.std(all_members, axis=None))
# N.random.shuffle(std[0])
# N.random.shuffle(std[1])
choice0 = N.random.choice(std[0].flatten(),size=1000)
choice1 = N.random.choice(std[1].flatten(),size=1000)
# import pdb; pdb.set_trace()
# tstat, pvalue = scipy.stats.ttest_rel(std[0],std[1],axis=None)
tstat, pvalue = scipy.stats.ttest_rel(choice0,choice1,axis=None)
return std_ave[0], std_ave[1], tstat, pvalue
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 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))
class BirdsEye(Figure):
def __init__(self,config,wrfout,ax=0,fig=0):
super(BirdsEye,self).__init__(config,wrfout,ax=ax,fig=fig)
def get_contouring(self,vrbl='user',lv='user',**kwargs):
"""
Returns colourmap and contouring levels
Options keyword arguments:
V : manually override contour levels
"""
data = self.data.reshape((self.la_n,self.lo_n))
# List of args and dictionary of kwargs
plotargs = [self.x,self.y,data]
plotkwargs = kwargs
# cmap = getattr(kwargs,'cmap',plt.cm.jet)
if vrbl=='user':
pass
else:
S = Scales(vrbl,lv)
if S.cm:
plotkwargs['cmap'] = S.cm
if isinstance(S.clvs,N.ndarray):
plotkwargs['levels'] = S.clvs
# if self.mplcommand == 'contour':
# multiplier = S.get_multiplier(vrbl,lv)
if 'clvs' in kwargs:
if isinstance(kwargs['clvs'],N.ndarray):
plotkwargs['levels'] = kwargs['clvs']
kwargs.pop('clvs')
# pdb.set_trace()
return plotargs, plotkwargs
def plot_data(self,data,mplcommand,p2p,fname,pt,no_title=1,save=1,**kwargs):
"""
Generic method that plots any matrix of data on a map
Inputs:
data : lat/lon matrix of data
vrbl : variable type for contouring convention
m : basemap instance
mplcommand : contour or contourf etc
p2p : path to plots
fname : filename for plot
V : scale for contours
no_title : switch to turn off title
save : whether to save to file
"""
# INITIALISE
# self.fig = plt.figure()
# self.fig = self.figsize(8,8,self.fig) # Create a default figure size if not set by user
# self.fig.set_size_inches(5,5)
self.bmap,self.x,self.y = self.basemap_setup()#ax=self.ax)
self.mplcommand = mplcommand
self.data = data
self.la_n = self.data.shape[-2]
self.lo_n = self.data.shape[-1]
# if plottype == 'contourf':
# f1 = self.bmap.contourf(*plotargs,**plotkwargs)
# elif plottype == 'contour':
# plotkwargs['colors'] = 'k'
# f1 = self.bmap.contour(*plotargs,**plotkwargs)
# scaling_func = M.ticker.FuncFormatter(lambda x, pos:'{0:d}'.format(int(x*multiplier)))
# plt.clabel(f1, inline=1, fmt=scaling_func, fontsize=9, colors='k')
plotargs, plotkwargs = self.get_contouring(**kwargs)
# pdb.set_trace()
if self.mplcommand == 'contour':
f1 = self.bmap.contour(*plotargs,**plotkwargs)
elif self.mplcommand == 'contourf':
f1 = self.bmap.contourf(*plotargs,**plotkwargs)
elif self.mplcommand == 'pcolor':
f1 = self.bmap.pcolor(*plotargs,**plotkwargs)
elif self.mplcommand == 'pcolormesh':
f1 = self.bmap.pcolormesh(*plotargs,**plotkwargs)
elif self.mplcommand == 'scatter':
f1 = self.bmap.scatter(*plotargs,**plotkwargs)
else:
print("Specify plot type.")
raise Exception
# LABELS, TITLES etc
"""
Change these to hasattr!
"""
#if self.C.plot_titles:
if not no_title:
title = utils.string_from_time('title',pt,tupleformat=0)
plt.title(title)
plot_colorbar = 1
if plot_colorbar:
self.fig.colorbar(f1,orientation='horizontal')
# plt.show(self.fig)
# div0 = make_axes_locatable(self.ax)
# cax0 = div0.append_axes("bottom", size="20%", pad=0.05)
# cb0 = self.fig.colorbar(f1, cax=cax0)
# SAVE FIGURE
datestr = utils.string_from_time('output',pt,tupleformat=0)
# self.fname = self.create_fname(fpath) # No da variable here
if save:
self.save(p2p,fname)
plt.close(self.fig)
return f1
# print("Plot saved to {0}.".format(os.path.join(p2p,fname)))
#def plot2D(self,va,**kwargs):
def plot2D(self,vrbl,t,lv,dom,outpath,bounding=0,smooth=1,
plottype='contourf',save=1,return_data=0):
"""
Inputs:
vrbl : variable string
t : date/time in (YYYY,MM,DD,HH,MM,SS) or datenum format
If tuple of two dates, it's start time and
end time, e.g. for finding max/average.
lv : level
dom : domain
outpath : absolute path to output
bounding : list of four floats (Nlim, Elim, Slim, Wlim):
Nlim : northern limit
Elim : eastern limit
Slim : southern limit
Wlim : western limit
smooth : smoothing. 1 is off. integer greater than one is
the degree of smoothing, to be specified.
save : whether to save to file
"""
# INITIALISE
self.fig.set_size_inches(8,8)
self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth)
self.mplcommand = plottype
# Make sure smooth=0 is corrected to 1
# They are both essentially 'off'.
if smooth==0:
smooth = 1
# Get indices for time, level, lats, lons
if isinstance(t,collections.Sequence) and len(t)!=6:
# List of two dates, start and end
# pdb.set_trace()
it_idx = self.W.get_time_idx(t[0])
ft_idx = self.W.get_time_idx(t[1])
assert ft_idx > it_idx
tidx = slice(it_idx,ft_idx,None)
title = "range"
datestr = "range"
else:
tidx = self.W.get_time_idx(t)
title = utils.string_from_time('title',t)
datestr = utils.string_from_time('output',t)
# Until pressure coordinates are fixed TODO
lvidx = 0
latidx, lonidx = self.get_limited_domain(bounding,smooth=smooth)
# if vc == 'surface':
# lv_idx = 0
# elif lv == 'all':
# lv_idx = 'all'
# else:
# print("Need to sort other levels")
# raise Exception
# FETCH DATA
ncidx = {'t': tidx, 'lv': lvidx, 'la': latidx, 'lo': lonidx}
self.data = self.W.get(vrbl,ncidx)#,**vardict)
self.la_n = self.data.shape[-2]
self.lo_n = self.data.shape[-1]
# COLORBAR, CONTOURING
plotargs, plotkwargs = self.get_contouring(vrbl,lv)
# S = Scales(vrbl,lv)
# multiplier = S.get_multiplier(vrbl,lv)
# if S.cm:
# plotargs = (self.x,self.y,data.reshape((la_n,lo_n)),S.clvs)
# cmap = S.cm
# elif isinstance(S.clvs,N.ndarray):
# if plottype == 'contourf':
# plotargs = (self.x,self.y,data.reshape((la_n,lo_n)),S.clvs)
# cmap = plt.cm.jet
# else:
# plotargs = (self.x,self.y,data.reshape((la_n,lo_n)),S.clvs)
# else:
# plotargs = (self.x,self.y,data.reshape((la_n,lo_n)))
# cmap = plt.cm.jet
# pdb.set_trace()
if self.mplcommand == 'contourf':
# f1 = self.bmap.contourf(*plotargs,cmap=cmap)
f1 = self.bmap.contourf(*plotargs,**plotkwargs)
elif self.mplcommand == 'contour':
plotkwargs['colors'] = 'k'
f1 = self.bmap.contour(*plotargs,**kwargs)
scaling_func = M.ticker.FuncFormatter(lambda x, pos:'{0:d}'.format(int(x*multiplier)))
plt.clabel(f1, inline=1, fmt=scaling_func, fontsize=9, colors='k')
# LABELS, TITLES etc
if self.C.plot_titles:
plt.title(title)
if self.mplcommand == 'contourf' and self.C.colorbar:
plt.colorbar(f1,orientation='horizontal')
# SAVE FIGURE
# pdb.set_trace()
lv_na = utils.get_level_naming(vrbl,lv)
naming = [vrbl,lv_na,datestr]
if dom:
naming.append(dom)
self.fname = self.create_fname(*naming)
if save:
self.save(outpath,self.fname)
plt.close()
if isinstance(self.data,N.ndarray):
return self.data.reshape((self.la_n,self.lo_n))
def plot_streamlines(self,lv,pt,outpath,da=0):
m,x,y = self.basemap_setup()
time_idx = self.W.get_time_idx(pt)
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}
if lv == 2000:
u = self.W.get('U10',slices)[0,:,:]
v = self.W.get('V10',slices)[0,:,:]
else:
u = self.W.get('U',slices)[0,0,:,:]
v = self.W.get('V',slices)[0,0,:,:]
# pdb.set_trace()
#div = N.sum(N.dstack((N.gradient(u)[0],N.gradient(v)[1])),axis=2)*10**4
#vort = (N.gradient(v)[0] - N.gradient(u)[1])*10**4
#pdb.set_trace()
lv_na = utils.get_level_naming('wind',lv=2000)
m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],u,v,
density=2.5,linewidth=0.75,color='k')
#div_Cs = N.arange(-30,31,1)
#divp = m.contourf(x,y,vort,alpha=0.6)
#divp = m.contour(x,y,vort)
#plt.colorbar(divp,orientation='horizontal')
if self.C.plot_titles:
title = utils.string_from_time('title',pt)
m.title(title)
datestr = utils.string_from_time('output',pt)
na = ('streamlines',lv_na,datestr)
fname = self.create_fname(*na)
self.save(outpath,fname)
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)
class BirdsEye(Figure):
def __init__(self, wrfout, ax=0, fig=0):
super(BirdsEye, self).__init__(wrfout, ax=ax, fig=fig)
def get_plot_arguments(self, cmap=False, clvs=False):
"""
Returns colourmap and contouring levels
Options keyword arguments:
clvs : manually override contour levels
"""
data = self.data.reshape((self.la_n, self.lo_n))
# List of args and dictionary of kwargs
plotargs = [self.x, self.y, data]
plotkwargs = {}
# if self.mplcommand == 'contour':
# multiplier = S.get_multiplier(vrbl,lv)
if clvs is not False:
plotkwargs['levels'] = clvs
if cmap is not False:
# cmap = eval('M.cm.{0}'.format(cmap))
plotkwargs['cmap'] = cmap
# import pdb; pdb.set_trace()
return plotargs, plotkwargs
# Old plot_data
def plot2D(self,
data,
fname,
outdir,
plottype='contourf',
save=1,
smooth=1,
lats=False,
lons=False,
clvs=False,
cmap=False,
title=False,
colorbar=True):
"""
Generic method that plots any matrix of data on a map
Inputs:
data : 2D matrix of data
outdir : path to plots
outf : filename for output (with or without .png)
Optional:
plottype : matplotlib function for plotting
smooth : Gaussian smooth by this many grid spaces
clvs : scale for contours
title : title on plot
save : whether to save to file
"""
# INITIALISE
self.data = data
self.bmap, self.x, self.y = self.basemap_setup(
smooth=smooth,
lats=lats,
lons=lons,
) #ax=self.ax)
self.la_n = self.data.shape[-2]
self.lo_n = self.data.shape[-1]
plotargs, plotkwargs = self.get_plot_arguments(clvs=clvs, cmap=cmap)
# import pdb; pdb.set_trace()
if plottype == 'contour':
f1 = self.bmap.contour(*plotargs, **plotkwargs)
elif plottype == 'contourf':
f1 = self.bmap.contourf(*plotargs, **plotkwargs)
elif plottype == 'pcolor':
f1 = self.bmap.pcolor(*plotargs, **plotkwargs)
elif plottype == 'pcolormesh':
f1 = self.bmap.pcolormesh(*plotargs, **plotkwargs)
elif plottype == 'scatter':
f1 = self.bmap.scatter(*plotargs, **plotkwargs)
else:
print("Specify correct plot type.")
raise Exception
if isinstance(title, basestring):
plt.title(title)
if colorbar:
self.fig.colorbar(f1, orientation='vertical')
if save:
self.save(outdir, fname)
plt.close(self.fig)
def plot_streamlines(self,
U,
V,
outdir,
fname,
lats=False,
lons=False,
smooth=1,
title=False,
lw_speed=False):
"""
Plot streamlines.
U : U-component of wind (nx x ny)
V : V-component of wind (same dimensions)
lw_speed : linewidth is proportional to wind speed
"""
m, x, y = self.basemap_setup()
if lw_speed:
wind = N.sqrt(U**2 + V**2)
lw = 5 * wind / wind.max()
else:
lw = 1
if smooth > 1:
U = stats.gauss_smooth(U, smooth)
V = stats.gauss_smooth(V, smooth)
m.streamplot(x[self.W.x_dim / 2, :],
y[:, self.W.y_dim / 2],
U,
V,
density=1.8,
linewidth=lw,
color='k',
arrowsize=3)
if isinstance(title, basestring):
self.ax.set_title(title)
self.save(outdir, fname)
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 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
class BirdsEye(Figure):
def __init__(self,nc=False,ax=0,fig=0):
super(BirdsEye,self).__init__(nc=nc,ax=ax,fig=fig)
def get_plot_arguments(self,cmap=False,clvs=False,color=False):
"""
Returns colourmap and contouring levels
Options keyword arguments:
clvs : manually override contour levels
"""
# import pdb; pdb.set_trace()
# data = self.data.reshape((self.la_n,self.lo_n))
data = self.data
# List of args and dictionary of kwargs
plotargs = [self.x,self.y,data]
plotkwargs = {}
# if self.mplcommand == 'contour':
# multiplier = S.get_multiplier(vrbl,lv)
if clvs is not False:
plotkwargs['levels'] = clvs
if cmap is not False:
# cmap = eval('M.cm.{0}'.format(cmap))
plotkwargs['cmap'] = cmap
return plotargs, plotkwargs
def axes_of_dilatation(self,xdata,ydata,fname,outdir,
lats=False,lons=False,smooth=False, locations=False,
x=False,y=False,m=False,
Nlim=False,Elim=False,Slim=False,Wlim=False):
if x is False and y is False and m is False:
if not Nlim:
self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
lons=lons,)#ax=self.ax)
else:
self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
lons=lons,Nlim=Nlim,Elim=Elim,
Slim=Slim,Wlim=Wlim)
else:
self.bmap = m
self.x = x
self.y = y
# self.la_n = self.data.shape[-2]
# self.lo_n = self.data.shape[-1]
# import pdb; pdb.set_trace()
# b = 10
# self.bmap.quiver(self.x[::b],self.y[::b],xdata[::b,::b],ydata[::b,::b],headwidth=0, units='xy',scale=10)
self.bmap.streamplot(self.x,self.y,xdata,ydata)
# m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],U,V,
# density=1.8,linewidth=lw,color='k',arrowsize=3)
self.save(outdir,fname)
plt.close(self.fig)
# Old plot_data
def plot2D(self,data,fname,outdir,plottype='contourf',
save=True,smooth=1,lats=False,lons=False,
clvs=False,cmap=False,title=False,cb=True,
locations=False,m=False,x=False,y=False,
Nlim=False,Elim=False,Slim=False,Wlim=False,
color='k',inline=False,lw=False,extend=False,
cblabel=False):
"""
Generic method that plots any matrix of data on a map
Inputs:
data : 2D matrix of data
outdir : path to plots
outf : filename for output (with or without .png)
Optional:
plottype : matplotlib function for plotting
smooth : Gaussian smooth by this many grid spaces
clvs : scale for contours
title : title on plot
save : whether to save to file
:param locations: Locations to plot on the basemap.
Format: locations = {'label':(lat,lon),etc}
:type locations: dict
"""
# INITIALISE
self.data = data
if x is False and y is False and m is False:
if not Nlim:
self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
lons=lons,)#ax=self.ax)
else:
self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
lons=lons,Nlim=Nlim,Elim=Elim,
Slim=Slim,Wlim=Wlim)
else:
self.bmap = m
self.x = x
self.y = y
# self.la_n = self.data.shape[-2]
# self.lo_n = self.data.shape[-1]
plotargs, plotkwargs = self.get_plot_arguments(clvs=clvs,cmap=cmap,color=color)
# import pdb; pdb.set_trace()
if plottype == 'contour':
plotkwargs['colors'] = color
plotkwargs['inline'] = inline
if lw:
plotkwargs['lw'] = lw
f1 = self.bmap.contour(*plotargs,**plotkwargs)
if inline:
plt.clabel(f1,inline=True,fmt='%d',color='black',fontsize=9)
elif plottype == 'contourf':
if isinstance(extend,str):
plotkwargs['extend'] = extend
f1 = self.bmap.contourf(*plotargs,**plotkwargs)
elif plottype == 'pcolor':
f1 = self.bmap.pcolor(*plotargs,**plotkwargs)
elif plottype == 'pcolormesh':
f1 = self.bmap.pcolormesh(*plotargs,**plotkwargs)
elif plottype == 'scatter':
f1 = self.bmap.scatter(*plotargs,**plotkwargs)
# elif plottype == 'quiver':
# f1 = self.bmap.quiver(*plotargs,**plotkwargs)
else:
print("Specify correct plot type.")
raise Exception
if isinstance(locations,dict):
for k,v in locations.iteritems():
if isinstance(v,tuple) and len(v) == 2:
xpt, ypt = self.bmap(v[1],v[0])
# bbox_style = {'boxstyle':'square','fc':'white','alpha':0.5}
self.bmap.plot(xpt,ypt,'ko',markersize=3,zorder=100)
self.ax.text(xpt,ypt,k,ha='left',fontsize=7)
# self.ax.text(xpt-15000,ypt,k,bbox=bbox_style,ha='left',fontsize=7)
else:
print("Not a valid location argument.")
raise Exception
if isinstance(title,basestring):
plt.title(title)
if cb != False:
if cb==True:
cb1 = plt.colorbar(f1,orientation='vertical',ax=self.ax)
elif cb=='horizontal':
cb1 = plt.colorbar(f1,orientation='horizontal',ax=self.ax)
elif cb == 'only':
save = False
self.fig,self.ax = plt.subplots(figsize=(4,0.8))
cb1 = plt.colorbar(f1,cax=self.ax,orientation='horizontal')
if isinstance(cblabel,str):
cb1.set_label(cblabel)
self.save(outdir,fname+'_cb')
else:
cb1 = plt.colorbar(f1,orientation='vertical',cax=cb)
if cb and isinstance(cblabel,str):
cb1.set_label(cblabel)
if save:
self.save(outdir,fname)
plt.close(self.fig)
else:
return f1
def plot_streamlines(self,U,V,outdir,fname,lats=False,lons=False,smooth=1,
title=False,lw_speed=False,density=1.8):
"""
Plot streamlines.
U : U-component of wind (nx x ny)
V : V-component of wind (same dimensions)
lw_speed : linewidth is proportional to wind speed
"""
m,x,y = self.basemap_setup(lats=lats,lons=lons)
if lw_speed:
wind = N.sqrt(U**2 + V**2)
lw = 5*wind/wind.max()
else:
lw = 1
if smooth>1:
U = stats.gauss_smooth(U,smooth)
V = stats.gauss_smooth(V,smooth)
# m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],U,V,
# density=density,linewidth=lw,color='k',arrowsize=3)
m.streamplot(x,y,U,V,
density=density,linewidth=lw,color='k',arrowsize=3)
if isinstance(title,basestring):
self.ax.set_title(title)
self.save(outdir,fname)
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 make_subplot_label(ax,label):
ax.text(0.05,0.85,label,transform=ax.transAxes,
bbox={'facecolor':'white'},fontsize=15,zorder=1000)
return
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
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
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)
class BirdsEye(Figure):
def __init__(self,wrfout,ax=0,fig=0):
super(BirdsEye,self).__init__(wrfout,ax=ax,fig=fig)
def get_plot_arguments(self,cmap=False,clvs=False):
"""
Returns colourmap and contouring levels
Options keyword arguments:
clvs : manually override contour levels
"""
data = self.data.reshape((self.la_n,self.lo_n))
# List of args and dictionary of kwargs
plotargs = [self.x,self.y,data]
plotkwargs = {}
# if self.mplcommand == 'contour':
# multiplier = S.get_multiplier(vrbl,lv)
if clvs is not False:
plotkwargs['levels'] = clvs
if cmap is not False:
# cmap = eval('M.cm.{0}'.format(cmap))
plotkwargs['cmap'] = cmap
# import pdb; pdb.set_trace()
return plotargs, plotkwargs
# Old plot_data
def plot2D(self,data,fname,outdir,plottype='contourf',
save=1,smooth=1,lats=False,lons=False,
clvs=False,cmap=False,title=False,colorbar=True):
"""
Generic method that plots any matrix of data on a map
Inputs:
data : 2D matrix of data
outdir : path to plots
outf : filename for output (with or without .png)
Optional:
plottype : matplotlib function for plotting
smooth : Gaussian smooth by this many grid spaces
clvs : scale for contours
title : title on plot
save : whether to save to file
"""
# INITIALISE
self.data = data
self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
lons=lons,)#ax=self.ax)
self.la_n = self.data.shape[-2]
self.lo_n = self.data.shape[-1]
plotargs, plotkwargs = self.get_plot_arguments(clvs=clvs,cmap=cmap)
# import pdb; pdb.set_trace()
if plottype == 'contour':
f1 = self.bmap.contour(*plotargs,**plotkwargs)
elif plottype == 'contourf':
f1 = self.bmap.contourf(*plotargs,**plotkwargs)
elif plottype == 'pcolor':
f1 = self.bmap.pcolor(*plotargs,**plotkwargs)
elif plottype == 'pcolormesh':
f1 = self.bmap.pcolormesh(*plotargs,**plotkwargs)
elif plottype == 'scatter':
f1 = self.bmap.scatter(*plotargs,**plotkwargs)
else:
print("Specify correct plot type.")
raise Exception
if isinstance(title,basestring):
plt.title(title)
if colorbar:
self.fig.colorbar(f1,orientation='vertical')
if save:
self.save(outdir,fname)
plt.close(self.fig)
def plot_streamlines(self,U,V,outdir,fname,lats=False,lons=False,smooth=1,
title=False,lw_speed=False):
"""
Plot streamlines.
U : U-component of wind (nx x ny)
V : V-component of wind (same dimensions)
lw_speed : linewidth is proportional to wind speed
"""
m,x,y = self.basemap_setup()
if lw_speed:
wind = N.sqrt(U**2 + V**2)
lw = 5*wind/wind.max()
else:
lw = 1
if smooth>1:
U = stats.gauss_smooth(U,smooth)
V = stats.gauss_smooth(V,smooth)
m.streamplot(x[self.W.x_dim/2,:],y[:,self.W.y_dim/2],U,V,
density=1.8,linewidth=lw,color='k',arrowsize=3)
if isinstance(title,basestring):
self.ax.set_title(title)
self.save(outdir,fname)
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)
