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 _load_data(nc_fpath, vrbl, npyfpaths):
W = WRFOut(nc_fpath)
data = W.get(vrbl=vrbl)[:, 0, :, :]
assert data.shape[0] == len(npyfpaths)
for tidx, npyfpath in enumerate(npyfpaths):
utils.trycreate(npyfpath, debug=False)
_save(arr=data[tidx, :, :], file=npyfpath)
return
def std_ttest(ncfiles1,
ncfiles2,
vrbl,
utc=False,
level=False,
other=False,
th=0):
"""
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 < th] = th
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 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, 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 __init__(self, inherit=None):
if isinstance(inherit, WRFOut):
self.__dict__ = inherit.__dict__.copy()
elif isinstance(inherit, WRF_native_grid):
self.__dict__ = inherit.__dict__.copy()
elif isinstance(inherit, VerifGrid):
self.__dict__ = inherit.__dict__.copy()
elif isinstance(inherit, str):
inherit = WRFOut(inherit)
self.__dict__ = inherit.__dict__.copy()
# self.lat_0 = inherit.nc.CEN_LAT
# self.lon_0 = inherit.nc.CEN_LON
else:
raise Exception("Inherited object is {}".format(type(inherit)))
return
def load_wrfout(self,fpath):
self.W = WRFOut(fpath)
return
cmap = S.cm
lvs = S.clvs
elif vrbl == 'WSPD10MAX':
cmap = 'gist_earth_r'
lvs = N.arange(2.5, 22.5, 2.5)
elif vrbl == 'W':
lvidx = 20
else:
raise Exception
fpath = os.path.join(folder, mem, fname)
# ax.set_title("Dom {} for the {} run".format(dom,ss),fontsize=6)
ax.set_title("{} for netCDF{}".format(mem, ncno))
if ("onedom" in ss) and (dom == 'd02'):
continue
W = WRFOut(fpath)
data = W.get(vrbl=vrbl, utc=utc)[0, lvidx, :, :]
ax.contourf(data, cmap=cmap, levels=lvs)
ax.grid('off')
ax.axis('off')
fname = "NCOruns_{:03d}min_{}_{}.png".format(fcstmin, vrbl, dom)
fpath = os.path.join(outdir, fname)
utils.trycreate(fpath)
fig.tight_layout()
fig.savefig(fpath)
plt.close(fig)
print("Saved to", fpath)
class BirdsEye(Figure):
""" Top-down 2D plots.
Todo:
* multiple plots on same figure - like a 'hold' - maybe by
using set and get decorators. Or stopping save/close.
* Massive rewrite with cartopy, now basemap is depreciated.
* Move most options to mplkwargs (e.g. cmap, levels).
Args:
ax,fig (optional): If None (default), then a new figure/axis object
will be created in the instance (via parent class
:class:`~evac.plot.figure.Figure`.
ideal (bool): If True, options for an idealised plot will be used.
For example, no lat/lon.
proj (str,bool): If `str`, such as `'merc'`, this will set the
instance's plotting projection.
mplargs, mplkwargs: Arguments to pass to matplotlib.
Returns:
A figure is generated (if no fig/ax is specified), or plotting (onto an
existing figure).
"""
def __init__(self,ax=None,fig=None,ideal=False,proj='merc',
mplargs=[],mplkwargs={},):#W=None):
""" Setting ideal to True removes the geography (e.g. for idealised
plots)
"""
self.proj = proj
self.ideal = ideal
super(BirdsEye,self).__init__(ax=ax,fig=fig,mplargs=mplargs,
mplkwargs=mplkwargs)
def plot2D(self,data,fname=False,outdir=False,plottype='contourf',
save=True,smooth=1,lats=False,lons=False,
clvs=None,cmap=None,title=False,cb=True,
locations=False,m=False,x=None,y=None,
Nlim=False,Elim=False,Slim=False,Wlim=False,
color='k',inline=False,cblabel=False,ideal=False,
drawcounties=False,mplargs=[],mplkwargs={},
extend=False,nx=None,ny=None,cen_lat=None,cen_lon=None,
lat_ts=None,W=None,hold=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
"""
if isinstance(cmap,str):
cmap = M.cm.get_cmap(cmap)
save = getattr(self,'save_opt',save)
hold = getattr(self,'hold_opt',hold)
if plottype not in ("pcolor","pcolormesh"):
mplkwargs['levels'] = clvs
mplkwargs['cmap'] = cmap
# WRFOut instance for help plotting.
self.W = W
if save:
assert (fname is not False) and (outdir is not False)
# INITIALISE
self.data = data
if self.ideal:
ideal = True
if ideal:
self.bmap = self.ax
# pdb.set_trace()
self.y = N.arange(len(data[:,0]))
self.x = N.arange(len(data[0,:]))
# Can all this be moved into Figure() or whatever?
# def basemap_from_newgrid(self,Nlim=None,Elim=None,Slim=None,Wlim=None,proj='merc',
# lat_ts=None,resolution='i',nx=None,ny=None,
# tlat1=30.0,tlat2=60.0,cen_lat=None,cen_lon=None,
# lllon=None,lllat=None,urlat=None,urlon=None,
# drawcounties=False):
# self.basemap_from_newgrid(Nlim=Nlim,Elim=Elim,Slim=Slim,Wlim=Wlim,xx=x,yy=y,
# drawcounties=drawcounties,nx=nx,ny=ny,lons=lons,lats=lats,
# lat_ts=lat_ts)
elif (not x) and (not y):
if (not m):
if not Nlim:
self.bmap,self.x,self.y = self.basemap_setup(smooth=smooth,lats=lats,
lons=lons,drawcounties=drawcounties)#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,drawcounties=drawcounties)
else:
self.bmap = m
self.x, self.y = self.bmap(lons,lats)
else:
self.bmap = m
self.x = x
self.y = y
mplargs = [self.x,self.y,self.data,] + mplargs
# TODO: move this logic to other methods
if plottype == 'contour':
mplkwargs['colors'] = color
mplkwargs['inline'] = inline
# This won't work in Python 3
# f1 = self.bmap.contour(*mplargs,**mplkwargs)
# Fixed:
f1 = self.bmap.contour(*mplargs,**mplkwargs)
# This is repeated below
if inline:
plt.clabel(f1,inline=True,fmt='%d',color='black',fontsize=9)
elif plottype == 'contourf':
if isinstance(extend,str):
mplkwargs['extend'] = extend
# f1 = self.bmap.contourf(*mplargs,**mplkwargs)
f1 = self.ax.contourf(*mplargs,**mplkwargs)
elif plottype == 'pcolor':
# f1 = self.bmap.pcolor(*mplargs,**mplkwargs)
f1 = self.ax.pcolor(*mplargs,**mplkwargs)
elif plottype == 'pcolormesh':
# f1 = self.bmap.pcolormesh(*mplargs,**mplkwargs)
f1 = self.ax.pcolormesh(*mplargs,**mplkwargs)
elif plottype == 'scatter':
# f1 = self.bmap.scatter(*mplargs,**mplkwargs)
f1 = self.ax.scatter(*mplargs,**mplkwargs)
# elif plottype == 'quiver':
# f1 = self.bmap.quiver(*mplargs,**mplkwargs)
else:
print("Specify correct plot type.")
raise Exception
if ideal:
self.ax.set_aspect('equal')
if isinstance(locations,dict):
self.plot_locations(locations)
if isinstance(title,str):
self.ax.set_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)
def plot_locations(self,locations):
for k,v in locations.items():
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
return
def do_contour_plot(self):
pass
def do_contourf_plot(self):
pass
def do_pcolormesh_plot(self):
pass
def plot_streamlines(self,U,V,outdir,fname,lats=False,lons=False,smooth=1,
title=False,lw_speed=False,density=1.8,ideal=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
"""
if ideal:
m = plt
x = N.arange(len(U[:,0]))
y = N.arange(len(U[0,:]))
else:
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,str):
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
class BirdsEye(Figure):
""" Top-down 2D plots.
Todo:
* Massive rewrite with cartopy, now basemap is depreciated.
* Move most options to mplkwargs (e.g. cmap, levels).
Args:
ax,fig (optional): If None (default), then a new figure/axis object
will be created in the instance (via parent class
:class:`~evac.plot.figure.Figure`.
ideal (bool): If True, options for an idealised plot will be used.
For example, no lat/lon, geography data.
grid: instance of Grid to assist cartopy setup.
mplargs, mplkwargs: Arguments to pass to matplotlib.
Returns:
A figure is generated (if no fig/ax is specified), or plotting (onto an
existing figure).
"""
def __init__(self,
fpath,
ax=None,
fig=None,
ideal=False,
grid=None,
proj='merc',
use_basemap=True,
figsize=None):
# ccrs_proj = grid.get_cartopy_proj(proj)
super().__init__(fpath=fpath,
ax=ax,
fig=fig,
grid=grid,
use_basemap=use_basemap,
figsize=figsize)
# proj=ccrs_proj,
self.ideal = ideal
self.use_basemap = use_basemap
self.proj = proj
def plot2D(
self,
data,
plottype='contourf',
locations=False,
save=True,
smooth=1,
title=False,
cb=True,
# lats=False,lons=False,m=False,x=None,y=None,
# Nlim=False,Elim=False,Slim=False,Wlim=False,
# nx=None,ny=None,cen_lat=None,cen_lon=None,
# lat_ts=None,W=None,
my_basemap=None,
inline=False,
cblabel=False,
ideal=False,
draw_geography=True,
mplkwargs=None,
extend=False,
hold=False,
lats=None,
lons=None,
proj='merc',
# color='k',
*args,
**kwargs):
"""
Generic method that plots any matrix of data on a map
Notes:
mplkwargs settings include:
* levels sets contour levels
* cmap sets colormap
* color sets plot line colours
Todo:
* just_one_colorbar integration
Args:
data: 2D matrix of data
plottype (str): matplotlib function for plotting
smooth (int): Gaussian smooth by this many grid spaces
title (str): title on plot
save (bool): whether to save to file
locations: plot locations on map. Format is {'label':(lat,lon)}.
cb (bool,str): if True, plot colorbar. If string, this
should be an absolute path to plot one cb separately
cblabel (str): label for colorbar
ideal (bool): if True, plot without looking for geographical data
drawcounties (bool): if True, plot US counties on map
mplkwargs (dict): dictionary of keyword arguments to pass
to plotting method
kwargs: Alias for mplkwargs.
"""
if (proj is None) and (not ideal):
proj = getattr(self, 'proj', 'merc')
if mplkwargs is None:
mplkwargs = {}
mplkwargs.update(**kwargs)
save = getattr(self, 'save_opt', save)
hold = getattr(self, 'hold_opt', hold)
self.data = data
# if idealised, don't use geographical map
if ideal:
self.bmap = self.ax
# pdb.set_trace()
self.y = N.arange(len(data[:, 0]))
self.x = N.arange(len(data[0, :]))
else:
if lats is None:
assert self.grid is not None
# self.lats2d, self.lons2d = N.meshgrid(self.grid.lats,self.grid.lons)
self.lats2d = self.grid.lats
self.lons2d = self.grid.lons
elif lats.ndim == 1:
# self.lats2d, self.lons2d = N.meshgrid(lats,lons)
self.lons2d, self.lats2d = N.meshgrid(lons, lats)
elif lats.ndim == 2:
self.lats2d = lats
self.lons2d = lons
else:
raise Exception("lats and lons not valid.")
assert self.lats2d.ndim == 2
if my_basemap is not None:
self.bmap = my_basemap
self.x, self.y = self.bmap(self.lons2d, self.lats2d)
elif self.use_basemap:
if not self.basemap_done:
self.basemap_setup(proj=proj,
draw_geography=draw_geography)
self.bmap = self.m
self.x, self.y = self.m(self.lons2d, self.lats2d)
# self.y = self.grid.yy
else:
raise Exception("Need to make Cartopy work")
self.x = self.lats2d
self.y = self.lons2d
# mplkwargs['X'] = self.x
# mplkwargs['Y'] = self.y
# mplkwargs['Z'] = self.data
if plottype == 'scatter':
mplargs = [*self.data]
else:
mplargs = [self.x, self.y, self.data]
# pdb.set_trace()
# TODO: move this logic to other methods
if plottype == 'contour':
cb = False
f1 = self.bmap.contour(*mplargs, **mplkwargs)
if inline:
plt.clabel(f1,
inline=True,
fmt='%d',
color='black',
fontsize=9)
elif plottype == 'contourf':
f1 = self.bmap.contourf(*mplargs, **mplkwargs)
elif plottype == 'pcolor':
f1 = self.bmap.pcolor(*mplargs, **mplkwargs)
elif plottype == 'pcolormesh':
f1 = self.bmap.pcolormesh(*mplargs, **mplkwargs)
elif plottype == 'scatter':
cb = False
f1 = self.bmap.scatter(*mplargs, **mplkwargs)
elif plottype == 'quiver':
f1 = self.bmap.quiver(*mplargs, **mplkwargs)
elif plottype == 'plot':
cb = False
f1 = self.bmap.plot(*mplargs, **mplkwargs)
else:
print("Specify correct plot type.")
raise Exception
if isinstance(locations, dict):
self.plot_locations(locations)
if isinstance(title, str):
self.ax.set_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()
plt.close(self.fig)
return
def plot_locations(self, locations):
for k, v in locations.items():
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
return
def plot_streamlines(self,
U,
V,
outdir,
fname,
lats=False,
lons=False,
smooth=1,
title=False,
lw_speed=False,
density=1.8,
ideal=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
"""
if ideal:
m = plt
x = N.arange(len(U[:, 0]))
y = N.arange(len(U[0, :]))
else:
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, str):
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 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, str):
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
def get_data(caseutc,
fmt,
vrbl=None,
validutc=None,
initutc=None,
mem=None,
latlon_only=False,
load_latlons=False):
""" Whatever
"""
casestr = utils.string_from_time('dir', caseutc, strlen='day')
if mem:
assert initutc
assert (vrbl in FCST_VRBLS) or (vrbl in other_list)
initstr = utils.string_from_time('dir', initutc, strlen='hour')
memdir = os.path.join(ensroot, initstr, mem)
latf, lonf = get_llf_fpath(casestr, fmt)
if latlon_only:
if load_latlons:
return N.load(latf), N.load(lonf)
else:
return latf, lonf
else:
data_fpath = get_tempdata_fpath(caseutc, fmt, vrbl, validutc, initutc,
mem)
if fmt == "d01_raw":
# Load lat/lon grid
latf, lonf = get_llf_fpath(casestr, "d01_raw")
# Load 3-km wrfout domain as is, on 2D
fname = get_wrfout_fname(initutc, dom=1)
fpath = os.path.join(memdir, fname)
W = WRFOut(fpath)
# if vrbl in ("Wmax","UH02","UH25","RAINNC"):
data = W.get(utc=validutc, vrbl=vrbl)[0, 0, :, :]
elif fmt == "d02_raw":
# Load lat/lon grid
latf, lonf = get_llf_fpath(casestr, "d02_raw")
# Load 1-km wrfout domain as is, on 2D
fname = get_wrfout_fname(initutc, dom=2)
fpath = os.path.join(memdir, fname)
W = WRFOut(fpath)
# if vrbl in ("Wmax","UH02","UH25","RAINNC"):
data = W.get(utc=validutc, vrbl=vrbl)[0, 0, :, :]
elif fmt in ("mrms_aws_raw", "mrms_dz_raw"):
latf, lonf = get_llf_fpath(casestr, fmt)
# Load MRMS data, 2D
FILES, utcs, key, fdir = return_all_mrms(caseutc, vrbl)
t = closest(utcs, validutc, return_val=True)
fpath = os.path.join(fdir, FILES[t])
mrms_nc = Dataset(fpath)
# d01_nc = get_random_d01(caseutc)
data = _unsparsify(
vrbl=key,
mrms_nc=mrms_nc,
) #d01_nc=d01_nc)
with N.errstate(invalid='ignore'):
data[data < -100] = N.nan
data[data > 100] = N.nan
# pdb.set_trace()
elif fmt == "stageiv_raw":
latf, lonf = get_llf_fpath(casestr, "stageiv_raw")
# Load stage iv data
data = ST4.get(utc=validutc)[0, 0, :, :]
# pdb.set_trace()
elif fmt == 'nexrad_raw':
latf, lonf = get_llf_fpath(casestr, "nexrad_raw")
# Load stage iv data
data = RADARS.get(utc=validutc)
else:
print("fmt not valid.")
assert True is False
# This is now done in _save function
# if N.ma.is_masked(data):
# data = data.data
# pdb.set_trace()
_save(arr=data, file=data_fpath)
if load_latlons:
return data_fpath, N.load(latf), N.load(lonf)
else:
return data_fpath, latf, lonf