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 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 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
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)
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)
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)
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)
