def background(grid=None):
"""
Plot basic TXLA shelf background: coastline, bathymetry, meridians, etc
Can optionally input grid (so it doesn't have to be loaded again)
"""
matplotlib.rcParams.update({'font.size': 20})#,'font.weight': 'bold'})
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
# Do plot
grid['basemap'].drawcoastlines()
grid['basemap'].fillcontinents('0.8')
grid['basemap'].drawparallels(np.arange(18, 35), dashes=(1, 0), linewidth=0.15, labels=[1, 0, 0, 0])
grid['basemap'].drawmeridians(np.arange(-100, -80), dashes=(1, 0), linewidth=0.15, labels=[0, 0, 0, 1])
hold('on')
contour(grid['xr'], grid['yr'], grid['h'], np.hstack(([10,20],np.arange(50,500,50))), colors='lightgrey', linewidths=0.25)
# Outline numerical domain
plot(grid['xr'][0,:], grid['yr'][0,:], 'k:')
plot(grid['xr'][-1,:], grid['yr'][-1,:], 'k:')
plot(grid['xr'][:,0], grid['yr'][:,0], 'k:')
plot(grid['xr'][:,-1], grid['yr'][:,-1], 'k:')
def background(grid=None,
ax=None,
pars=np.arange(18, 35),
mers=np.arange(-100, -80),
hlevs=np.hstack(([10, 20], np.arange(50, 500, 50))),
col='lightgrey',
fig=None,
outline=True,
merslabels=[0, 0, 0, 1],
parslabels=[1, 0, 0, 0]):
"""
Plot basic TXLA shelf background: coastline, bathymetry, meridians, etc
Can optionally input grid (so it doesn't have to be loaded again)
pars parallels to plot
mers meridians to plot
hlevs which depth contours to plot
"""
# matplotlib.rcParams.update({'font.size': 18})#,'font.weight': 'bold'})
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
if fig is None:
fig = gcf()
if ax is None:
ax = gca()
# Do plot
grid['basemap'].drawcoastlines(ax=ax)
grid['basemap'].fillcontinents('0.8', ax=ax)
grid['basemap'].drawparallels(pars,
dashes=(1, 1),
linewidth=0.15,
labels=parslabels,
ax=ax)
grid['basemap'].drawmeridians(mers,
dashes=(1, 1),
linewidth=0.15,
labels=merslabels,
ax=ax)
# hold('on')
ax.contour(grid['xr'],
grid['yr'],
grid['h'],
hlevs,
colors=col,
linewidths=0.5)
if outline:
# Outline numerical domain
ax.plot(grid['xr'][0, :], grid['yr'][0, :], 'k:')
ax.plot(grid['xr'][-1, :], grid['yr'][-1, :], 'k:')
ax.plot(grid['xr'][:, 0], grid['yr'][:, 0], 'k:')
ax.plot(grid['xr'][:, -1], grid['yr'][:, -1], 'k:')
def background(
grid=None,
ax=None,
pars=np.arange(18, 35),
mers=np.arange(-100, -80),
hlevs=np.hstack(([10, 20], np.arange(50, 500, 50))),
col="lightgrey",
halpha=1,
fig=None,
outline=[1, 1, 0, 1],
merslabels=[0, 0, 0, 1],
parslabels=[1, 0, 0, 0],
):
"""
Plot basic TXLA shelf background: coastline, bathymetry, meridians, etc
Can optionally input grid (so it doesn't have to be loaded again)
Args:
pars: parallels to plot
mers: meridians to plot
hlevs: which depth contours to plot
outline: west, east, north, south lines (left, right, top, bottom)
"""
# matplotlib.rcParams.update({'font.size': 18})#,'font.weight': 'bold'})
if grid is None:
loc = "http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc"
grid = inout.readgrid(loc)
if fig is None:
fig = plt.gcf()
if ax is None:
ax = plt.gca()
# Do plot
grid.proj.drawcoastlines(ax=ax)
grid.proj.fillcontinents("0.8", ax=ax)
grid.proj.drawparallels(pars, dashes=(1, 1), linewidth=0.15, labels=parslabels, ax=ax)
grid.proj.drawmeridians(mers, dashes=(1, 1), linewidth=0.15, labels=merslabels, ax=ax)
ax.contour(grid.x_rho, grid.y_rho, grid.h, hlevs, colors=col, linewidths=0.5, alpha=halpha)
# Outline numerical domain
# if outline: # backward compatibility
# outline = [1,1,1,1]
if outline[0]: # left
ax.plot(grid.x_rho[:, 0], grid.y_rho[:, 0], "k:")
if outline[1]: # right
ax.plot(grid.x_rho[:, -1], grid.y_rho[:, -1], "k:")
if outline[2]: # top
ax.plot(grid.x_rho[-1, :], grid.y_rho[-1, :], "k:")
if outline[3]: # bottom
ax.plot(grid.x_rho[0, :], grid.y_rho[0, :], "k:")
def background(grid=None, ax=None, pars=np.arange(18, 35),
mers=np.arange(-100, -80),
hlevs=np.hstack(([10, 20], np.arange(50, 500, 50))),
col='lightgrey', halpha=1, fig=None, outline=[1, 1, 0, 1],
merslabels=[0, 0, 0, 1], parslabels=[1, 0, 0, 0]):
"""
Plot basic TXLA shelf background: coastline, bathymetry, meridians, etc
Can optionally input grid (so it doesn't have to be loaded again)
Args:
pars: parallels to plot
mers: meridians to plot
hlevs: which depth contours to plot
outline: west, east, north, south lines (left, right, top, bottom)
"""
# matplotlib.rcParams.update({'font.size': 18})#,'font.weight': 'bold'})
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
if fig is None:
fig = plt.gcf()
if ax is None:
ax = plt.gca()
# Do plot
grid.proj.drawcoastlines(ax=ax)
grid.proj.fillcontinents('0.8', ax=ax)
grid.proj.drawparallels(pars, dashes=(1, 1),
linewidth=0.15, labels=parslabels, ax=ax)
grid.proj.drawmeridians(mers, dashes=(1, 1),
linewidth=0.15, labels=merslabels, ax=ax)
ax.contour(grid.x_rho, grid.y_rho, grid.h, hlevs, colors=col,
linewidths=0.5, alpha=halpha)
# Outline numerical domain
# if outline: # backward compatibility
# outline = [1,1,1,1]
if outline[0]: # left
ax.plot(grid.x_rho[:, 0], grid.y_rho[:, 0], 'k:')
if outline[1]: # right
ax.plot(grid.x_rho[:, -1], grid.y_rho[:, -1], 'k:')
if outline[2]: # top
ax.plot(grid.x_rho[-1, :], grid.y_rho[-1, :], 'k:')
if outline[3]: # bottom
ax.plot(grid.x_rho[0, :], grid.y_rho[0, :], 'k:')
def background(grid=None, ax=None, pars=np.arange(18, 35), mers=np.arange(-100, -80),
hlevs=np.hstack(([10,20],np.arange(50,500,50))),
col='lightgrey', fig=None, outline=[1, 1, 0, 1], merslabels=[0, 0, 0, 1],
parslabels=[1, 0, 0, 0]):
"""
Plot basic TXLA shelf background: coastline, bathymetry, meridians, etc
Can optionally input grid (so it doesn't have to be loaded again)
pars parallels to plot
mers meridians to plot
hlevs which depth contours to plot
outline west, east, north, south lines (left, right, top, bottom)
"""
# matplotlib.rcParams.update({'font.size': 18})#,'font.weight': 'bold'})
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
if fig is None:
fig = gcf()
if ax is None:
ax = gca()
# Do plot
grid['basemap'].drawcoastlines(ax=ax)
grid['basemap'].fillcontinents('0.8',ax=ax)
grid['basemap'].drawparallels(pars, dashes=(1, 1),
linewidth=0.15, labels=parslabels, ax=ax)
grid['basemap'].drawmeridians(mers, dashes=(1, 1),
linewidth=0.15, labels=merslabels, ax=ax)
# hold('on')
ax.contour(grid['xr'], grid['yr'], grid['h'], hlevs,
colors=col, linewidths=0.5)
# Outline numerical domain
# if outline: # backward compatibility
# outline = [1,1,1,1]
if outline[0]:
ax.plot(grid['xr'][0,:], grid['yr'][0,:], 'k:')
if outline[1]:
ax.plot(grid['xr'][-1,:], grid['yr'][-1,:], 'k:')
if outline[2]:
ax.plot(grid['xr'][:,-1], grid['yr'][:,-1], 'k:')
if outline[3]:
ax.plot(grid['xr'][:,0], grid['yr'][:,0], 'k:')
def tracks(lonp,latp,fname,grid=None):
"""
Plot tracks as lines with starting points in green and ending points in red.
Inputs:
lonp,latp Drifter track positions [time x ndrifters]
fname Plot name to save
"""
# pdb.set_trace()
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
# Change positions from lon/lat to x/y
xp,yp = grid['basemap'](lonp,latp)
# Need to retain nan's since basemap changes them to values
ind = np.isnan(lonp)
xp[ind] = np.nan
yp[ind] = np.nan
figure(figsize=(12,10))
background(grid) # Plot coastline and such
# pdb.set_trace()
# Starting marker
plot(xp[:,0],yp[:,0],'o',color='g',markersize=3,label='_nolegend_',alpha=0.4)
# Plot tracks
plot(xp.T,yp.T,'-',color='grey',linewidth=.2)
# Find final positions of drifters
xpc,ypc = tools.find_final(xp,yp)
plot(xpc,ypc,'o',color='r',label='_nolegend_')
# pdb.set_trace()
# Legend, of sorts
ax = gca()
xtext = 0.45; ytext = 0.18;
text(xtext, ytext, 'starting location', fontsize=16, color='green',
alpha=.8, transform = ax.transAxes)
text(xtext, ytext-.03, 'track', fontsize=16, color='grey', transform = ax.transAxes)
text(xtext, ytext-.03*2, 'ending location', fontsize=16, color='red',
transform = ax.transAxes)
# xtext, ytext = grid['basemap'](-94,24) # text location
# text(xtext,ytext,'starting location',fontsize=16,color='green',alpha=.8)
# text(xtext,ytext-30000,'track',fontsize=16,color='grey')#,alpha=.8)
# text(xtext,ytext-60000,'ending location',fontsize=16,color='red')#,alpha=.8)
# # get psi mask from rho mask
# # maskp = grid['mask'][1:,1:]*grid['mask'][:-1,1:]* \
# # grid['mask'][1:,:-1]*grid['mask'][:-1,:-1]
# # ind = maskp
# # ind[ind==0] = np.nan
# # plot(grid['xpsi']*ind,grid['ypsi']*ind,'k', \
# # (grid['xpsi']*ind).T,(grid['ypsi']*ind).T,'k')
# plot(grid['xpsi'],grid['ypsi'],'k', \
# (grid['xpsi']).T,(grid['ypsi']).T,'k')
# 16 is (lower) one that is near islands, 41 is higher one
# show()
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'tracks.png',bbox_inches='tight')
def hist(lonp, latp, fname, tind='final', which='contour', vmax=None, fig=None, ax=None, \
bins=(40,40), N=10, grid=None, xlims=None, ylims=None, C=None, Title=None,
weights=None, Label='Final drifter location (%)', isll=True, binscale=None):
"""
Plot histogram of given track data at time index tind.
Inputs:
lonp,latp Drifter track positions in lon/lat [time x ndrifters]
fname Plot name to save
tind (optional) Default is 'final', in which case the final
position of each drifter in the array is found
and plotted. Alternatively, a time index
can be input and drifters at that time will be plotted.
Note that once drifters hit the outer numerical boundary,
they are nan'ed out so this may miss some drifters.
which (optional) 'contour', 'pcolor', 'hexbin', 'hist2d'
for type of plot used. Default 'hexbin'.
bins (optional) Number of bins used in histogram. Default (15,25).
N (optional) Number of contours to make. Default 10.
grid (optional) grid as read in by inout.readgrid()
xlims (optional) value limits on the x axis
ylims (optional) value limits on the y axis
isll Default True. Inputs are in lon/lat. If False, assume they
are in projected coords.
Note: Currently assuming we are plotting the final location
of each drifter regardless of tind.
"""
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
if isll: # if inputs are in lon/lat, change to projected x/y
# Change positions from lon/lat to x/y
xp, yp = grid['basemap'](lonp, latp)
# Need to retain nan's since basemap changes them to values
ind = np.isnan(lonp)
xp[ind] = np.nan
yp[ind] = np.nan
else:
xp = lonp
yp = latp
if fig is None:
fig = figure(figsize=(11, 10))
else:
fig = fig
background(grid) # Plot coastline and such
# pdb.set_trace()
if tind == 'final':
# Find final positions of drifters
xpc, ypc = tools.find_final(xp, yp)
elif is_numlike(tind):
xpc = xp[:, tind]
ypc = yp[:, tind]
else: # just plot what is input if some other string
xpc = xp.flatten()
ypc = yp.flatten()
if which == 'contour':
# Info for 2d histogram
H, xedges, yedges = np.histogram2d(xpc, ypc,
range=[[grid['xr'].min(), \
grid['xr'].max()], \
[grid['yr'].min(), \
grid['yr'].max()]],
bins=bins)
# Contour Plot
XE, YE = np.meshgrid(op.resize(xedges, 0), op.resize(yedges, 0))
d = (H / H.sum()) * 100
# # from http://matplotlib.1069221.n5.nabble.com/question-about-contours-and-clim-td21111.html
# locator = ticker.MaxNLocator(50) # if you want no more than 10 contours
# locator.create_dummy_axis()
# locator.set_bounds(0,1)#d.min(),d.max())
# levs = locator()
con = contourf(XE, YE, d.T,
N) #,levels=levs)#(0,15,30,45,60,75,90,105,120))
con.set_cmap('YlOrRd')
if Title is not None:
set_title(Title)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3725, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(con, cax=cax, orientation='horizontal')
cb.set_label('Final drifter location (percent)')
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'histcon.png', bbox_inches='tight')
# savefig('figures/' + fname + 'histcon.pdf',bbox_inches='tight')
elif which == 'pcolor':
# Info for 2d histogram
H, xedges, yedges = np.histogram2d(xpc, ypc,
range=[[grid['xr'].min(), \
grid['xr'].max()], \
[grid['yr'].min(), \
grid['yr'].max()]],
bins=bins, weights=weights)
# print H.T.max()
# pdb.set_trace()
# Pcolor plot
# C is the z value plotted, and is normalized by the total number of drifters
if C is None:
C = (H.T / H.sum()) * 100
else:
# or, provide some other weighting
C = (H.T / C) * 100
p = pcolor(xedges, yedges, C, cmap='YlOrRd')
if Title is not None:
set_title(Title)
# Set x and y limits
# pdb.set_trace()
if xlims is not None:
xlim(xlims)
if ylims is not None:
ylim(ylims)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(p, cax=cax, orientation='horizontal')
cb.set_label('Final drifter location (percent)')
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'histpcolor.png', bbox_inches='tight')
# savefig('figures/' + fname + 'histpcolor.pdf',bbox_inches='tight')
elif which == 'hexbin':
if ax is None:
ax = gca()
else:
ax = ax
if C is None:
# C with the reduce_C_function as sum is what makes it a percent
C = np.ones(len(xpc)) * (1. / len(xpc)) * 100
else:
C = C * np.ones(len(xpc)) * 100
hb = hexbin(xpc,
ypc,
C=C,
cmap='YlOrRd',
gridsize=bins[0],
extent=(grid['xpsi'].min(), grid['xpsi'].max(),
grid['ypsi'].min(), grid['ypsi'].max()),
reduce_C_function=sum,
vmax=vmax,
axes=ax,
bins=binscale)
# Set x and y limits
# pdb.set_trace()
if xlims is not None:
xlim(xlims)
if ylims is not None:
ylim(ylims)
if Title is not None:
ax.set_title(Title)
# Want colorbar at the given location relative to axis so this works regardless of # of subplots,
# so convert from axis to figure coordinates
# To do this, first convert from axis to display coords
# transformations: http://matplotlib.org/users/transforms_tutorial.html
ax_coords = [0.35, 0.25, 0.6,
0.02] # axis: [x_left, y_bottom, width, height]
disp_coords = ax.transAxes.transform([
(ax_coords[0], ax_coords[1]),
(ax_coords[0] + ax_coords[2], ax_coords[1] + ax_coords[3])
]) # display: [x_left,y_bottom,x_right,y_top]
inv = fig.transFigure.inverted(
) # inverter object to go from display coords to figure coords
fig_coords = inv.transform(
disp_coords) # figure: [x_left,y_bottom,x_right,y_top]
# actual desired figure coords. figure: [x_left, y_bottom, width, height]
fig_coords = [
fig_coords[0, 0], fig_coords[0, 1],
fig_coords[1, 0] - fig_coords[0, 0],
fig_coords[1, 1] - fig_coords[0, 1]
]
# Inlaid colorbar
cax = fig.add_axes(fig_coords)
# # Horizontal colorbar below plot
# cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(cax=cax, orientation='horizontal')
cb.set_label(Label)
# pdb.set_trace()
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'histhexbin.png', bbox_inches='tight')
# savefig('figures/' + fname + 'histhexbin.pdf',bbox_inches='tight')
elif which == 'hist2d':
# pdb.set_trace()
hist2d(xpc,
ypc,
bins=40,
range=[[grid['xr'].min(), grid['xr'].max()],
[grid['yr'].min(), grid['yr'].max()]],
normed=True)
set_cmap('YlOrRd')
# Set x and y limits
# pdb.set_trace()
if xlims is not None:
xlim(xlims)
if ylims is not None:
ylim(ylims)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(cax=cax, orientation='horizontal')
cb.set_label('Final drifter location (percent)')
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'hist2d.png', bbox_inches='tight')
def transport(name,
fmod=None,
Title=None,
dmax=None,
N=7,
extraname=None,
llcrnrlon=-98.5,
llcrnrlat=22.5,
urcrnrlat=31.0,
urcrnrlon=-87.5,
colormap='Blues',
fig=None,
ax=None):
'''
Make plot of zoomed-in area near DWH spill of transport of drifters over
time.
FILL IN
Inputs:
name
U
V
lon0
lat0
T0
'''
# (name=None, U, V, lon0, lat0, T0, dmax, extraname, Title, N,
# llcrnrlon, llcrnrlat, urcrnrlat, urcrnrlon, colormap):
# Load in transport information
U, V, lon0, lat0, T0 = inout.loadtransport(name, fmod=fmod)
# Smaller basemap parameters.
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc,
llcrnrlon=llcrnrlon,
llcrnrlat=llcrnrlat,
urcrnrlat=urcrnrlat,
urcrnrlon=urcrnrlon)
S = np.sqrt(op.resize(U, 1)**2 + op.resize(V, 0)**2)
Splot = (S / T0) * 100
if dmax is None:
dmax = Splot.max()
else:
dmax = dmax
# from http://matplotlib.1069221.n5.nabble.com/question-about-contours-and-clim-td21111.html
locator = ticker.MaxNLocator(N) # if you want no more than 10 contours
locator.create_dummy_axis()
locator.set_bounds(0, dmax) #d.min(),d.max())
levs = locator()
if fig is None:
fig = figure(figsize=(11, 10))
else:
fig = fig
background(grid=grid)
c = contourf(grid['xpsi'],
grid['ypsi'],
Splot,
cmap=colormap,
extend='max',
levels=levs)
title(Title)
# # Add initial drifter location (all drifters start at the same location)
# lon0 = lon0.mean()
# lat0 = lat0.mean()
# x0, y0 = grid['basemap'](lon0, lat0)
# plot(x0, y0, 'go', markersize=10)
if ax is None:
ax = gca()
else:
ax = ax
# Want colorbar at the given location relative to axis so this works regardless of # of subplots,
# so convert from axis to figure coordinates
# To do this, first convert from axis to display coords
# transformations: http://matplotlib.org/users/transforms_tutorial.html
ax_coords = [0.35, 0.25, 0.6,
0.02] # axis: [x_left, y_bottom, width, height]
disp_coords = ax.transAxes.transform([
(ax_coords[0], ax_coords[1]),
(ax_coords[0] + ax_coords[2], ax_coords[1] + ax_coords[3])
]) # display: [x_left,y_bottom,x_right,y_top]
inv = fig.transFigure.inverted(
) # inverter object to go from display coords to figure coords
fig_coords = inv.transform(
disp_coords) # figure: [x_left,y_bottom,x_right,y_top]
# actual desired figure coords. figure: [x_left, y_bottom, width, height]
fig_coords = [
fig_coords[0, 0], fig_coords[0,
1], fig_coords[1, 0] - fig_coords[0, 0],
fig_coords[1, 1] - fig_coords[0, 1]
]
# Inlaid colorbar
cax = fig.add_axes(fig_coords)
# cax = fig.add_axes([0.39, 0.25, 0.49, 0.02])
# cax = fig.add_axes([0.49, 0.25, 0.39, 0.02])
cb = colorbar(cax=cax, orientation='horizontal')
cb.set_label('Normalized drifter transport (%)')
if extraname is None:
savefig('figures/' + name + '/transport', bbox_inches='tight')
else:
savefig('figures/' + name + '/' + extraname + 'transport',
bbox_inches='tight')
def tracks(lonp,
latp,
fname,
grid=None,
fig=None,
ax=None,
Title=None,
mers=None,
pars=None):
"""
Plot tracks as lines with starting points in green and ending points in red.
Inputs:
lonp,latp Drifter track positions [time x ndrifters]
fname Plot name to save
"""
# pdb.set_trace()
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
if fig is None:
figure(figsize=(12, 10))
else:
fig = fig
if ax is None:
ax = gca()
else:
ax = ax
# Change positions from lon/lat to x/y
xp, yp = grid['basemap'](lonp, latp)
# Need to retain nan's since basemap changes them to values
ind = np.isnan(lonp)
xp[ind] = np.nan
yp[ind] = np.nan
if mers is not None:
background(grid, ax=ax, mers=mers,
pars=pars) # Plot coastline and such
else:
background(grid, ax=ax) # Plot coastline and such
# pdb.set_trace()
# Starting marker
ax.plot(xp[:, 0],
yp[:, 0],
'o',
color='g',
markersize=3,
label='_nolegend_',
alpha=0.4)
# Plot tracks
ax.plot(xp.T, yp.T, '-', color='grey', linewidth=.2)
# Find final positions of drifters
xpc, ypc = tools.find_final(xp, yp)
ax.plot(xpc, ypc, 'o', color='r', label='_nolegend_')
# pdb.set_trace()
if Title is not None:
ax.set_title(Title)
# Legend, of sorts
# ax = gca()
xtext = 0.45
ytext = 0.18
text(xtext,
ytext,
'starting location',
fontsize=16,
color='green',
alpha=.8,
transform=ax.transAxes)
text(xtext,
ytext - .03,
'track',
fontsize=16,
color='grey',
transform=ax.transAxes)
text(xtext,
ytext - .03 * 2,
'ending location',
fontsize=16,
color='red',
transform=ax.transAxes)
# xtext, ytext = grid['basemap'](-94,24) # text location
# text(xtext,ytext,'starting location',fontsize=16,color='green',alpha=.8)
# text(xtext,ytext-30000,'track',fontsize=16,color='grey')#,alpha=.8)
# text(xtext,ytext-60000,'ending location',fontsize=16,color='red')#,alpha=.8)
# # get psi mask from rho mask
# # maskp = grid['mask'][1:,1:]*grid['mask'][:-1,1:]* \
# # grid['mask'][1:,:-1]*grid['mask'][:-1,:-1]
# # ind = maskp
# # ind[ind==0] = np.nan
# # plot(grid['xpsi']*ind,grid['ypsi']*ind,'k', \
# # (grid['xpsi']*ind).T,(grid['ypsi']*ind).T,'k')
# plot(grid['xpsi'],grid['ypsi'],'k', \
# (grid['xpsi']).T,(grid['ypsi']).T,'k')
# 16 is (lower) one that is near islands, 41 is higher one
# show()
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'tracks.png', bbox_inches='tight')
def test1(loc=None, nsteps=None, ff=None, ah=None, grid=None, nlon=None, nlat=None, doturb=None, name=None): ''' A drifter test using TXLA model output. The comparison case for this simulation is 2D (do3d=0) with no turbulence/diffusion (doturb=0). Drifters are started at the surface and run forward for ten days (ndays=10) from 11/25/09 (in date). Compare results with figure in examples/test1.png. Optional inputs for making tests easy to run: loc 'thredds' or 'local', default = 'thredds' nsteps Number of particle steps to record between model outputs Default = 5 ff Backward (-1) or forward (1) in time. Default is forward (1). ah Horizontal viscosity, default = 5 grid If input, will not redo this step. Default is to load in grid. nlon, nlat Number of drifters to use in the lon/lat direction in seed array Default = 110, 98 (10 km spacing) doturb What, if any, subgrid parameterization to use. Default is 'none' name Specific name for track and figure files. Default is 'temp' ''' # Location of TXLA model output # file and then grid. # 0150 file goes from (2009, 11, 19, 12, 0) to (2009, 12, 6, 0, 0) if loc is None or loc == 'thredds': loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc' elif loc is 'local': # Location of TXLA model output if 'rainier' in os.uname(): loc = '/Users/kthyng/Documents/research/postdoc/' # for model outputs elif 'hafen.tamu.edu' in os.uname(): loc = '/home/kthyng/shelf/' # for model outputs # Initialize parameters if nsteps is None: nsteps = 5 else: nsteps = nsteps ndays = .5 #1 #16 if ff is None: ff = 1 else: ff = ff # Start date date = datetime(2009,11, 25, 0) # date = datetime(2009,11, 20, 0) # Time between outputs # Dt = 14400. # in seconds (4 hours), nc.variables['dt'][:] tseas = 4*3600 # 4 hours between outputs, in seconds, time between model outputs if ah is None: ah = 5. #100. else: ah = ah av = 1.e-5 # m^2/s, or try 5e-6 # grid = netCDF.Dataset(loc+'grid.nc') # lonr = grid.variables['lon_rho'][:] # latr = grid.variables['lat_rho'][:] if grid is None: grid = inout.readgrid(loc) else: grid = grid ## Input starting locations as real space lon,lat locations # lon0,lat0 = np.meshgrid(-95.498218005315309,23.142258627126882) # [0,0] (SE) corner # lon0,lat0 = np.meshgrid(-97.748582291691989,23.000027311710628) # [-1,0] (SW) corner # lon0,lat0 = np.meshgrid(-87.757124031927574,29.235771320764623) # [0,-1] (NE) corner # lon0,lat0 = np.meshgrid(-88.3634073986196,30.388542615201313) # [-1,-1] (NW) corner # lon0,lat0 = np.meshgrid(np.linspace(-94,-93,10),np.linspace(28,29,10)) # grid outside Galveston Bay # lon0,lat0 = np.meshgrid(np.linspace(-95,-91,100),np.linspace(28,29,50)) # rectangle outside Galveston # lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,1100),np.linspace(22.5,31,980)) # whole domain, 1 km # lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,220),np.linspace(22.5,31,196)) # whole domain, 5 km # # FOR TEST1: # lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,110),np.linspace(22.5,31,98)) # whole domain, 10 km # lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,21),np.linspace(22.5,31,20)) # whole domain, 50 km if nlon is None: nlon = 110 else: nlon = nlon if nlat is None: nlat = 98 else: nlat = nlat lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,nlon),np.linspace(22.5,31,nlat)) # whole domain, 10 km # Eliminate points that are outside domain or in masked areas lon0,lat0 = tools.check_points(lon0,lat0,grid) ## Choose method for vertical placement of drifters # Also update makefile accordingly. Choose the twodim flag for isoslice. # See above for more notes, but do the following two lines for an isoslice z0 = 's' #'salt' #'s' #'z' #'salt' #'s' zpar = 29 #30 #29 #-10 #grid['km']-1 # 30 #grid['km']-1 # Do the following two for a 3d simulation # z0 = np.ones(xstart0.shape)*-40 # below the surface # zpar = 'fromMSL' # pdb.set_trace() # for 3d flag, do3d=0 makes the run 2d and do3d=1 makes the run 3d do3d = 0 # turbulence/diffusion flag. doturb=0 means no turb/diffusion, # doturb=1 means adding parameterized turbulence # doturb=2 means adding diffusion on a circle # doturb=3 means adding diffusion on an ellipse (anisodiffusion) if doturb is None: doturb = 0 else: doturb = doturb # simulation name, used for saving results into netcdf file if name is None: name = 'temp' #'5_5_D5_F' else: name = name return loc,nsteps,ndays,ff,date,tseas,ah,av,lon0,lat0,z0,zpar,do3d,doturb,name,grid
def test1(loc=None,
nsteps=None,
ff=None,
ah=None,
grid=None,
nlon=None,
nlat=None,
doturb=None,
name=None):
'''
A drifter test using TXLA model output.
The comparison case for this simulation is 2D (do3d=0)
with no turbulence/diffusion (doturb=0).
Drifters are started at the surface and run forward
for ten days (ndays=10) from 11/25/09 (in date). Compare results with figure in examples/test1.png.
Optional inputs for making tests easy to run:
loc 'thredds' or 'local', default = 'thredds'
nsteps Number of particle steps to record between model outputs
Default = 5
ff Backward (-1) or forward (1) in time. Default is forward (1).
ah Horizontal viscosity, default = 5
grid If input, will not redo this step. Default is to load in grid.
nlon, nlat Number of drifters to use in the lon/lat direction in seed array
Default = 110, 98 (10 km spacing)
doturb What, if any, subgrid parameterization to use. Default is 'none'
name Specific name for track and figure files. Default is 'temp'
'''
# Location of TXLA model output
# file and then grid.
# 0150 file goes from (2009, 11, 19, 12, 0) to (2009, 12, 6, 0, 0)
if loc is None or loc == 'thredds':
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
elif loc is 'local':
# Location of TXLA model output
if 'rainier' in os.uname():
loc = '/Users/kthyng/Documents/research/postdoc/' # for model outputs
elif 'hafen.tamu.edu' in os.uname():
loc = '/home/kthyng/shelf/' # for model outputs
# Initialize parameters
if nsteps is None:
nsteps = 5
else:
nsteps = nsteps
ndays = .5 #1 #16
if ff is None:
ff = 1
else:
ff = ff
# Start date
date = datetime(2009, 11, 25, 0)
# date = datetime(2009,11, 20, 0)
# Time between outputs
# Dt = 14400. # in seconds (4 hours), nc.variables['dt'][:]
tseas = 4 * 3600 # 4 hours between outputs, in seconds, time between model outputs
if ah is None:
ah = 5. #100.
else:
ah = ah
av = 1.e-5 # m^2/s, or try 5e-6
# grid = netCDF.Dataset(loc+'grid.nc')
# lonr = grid.variables['lon_rho'][:]
# latr = grid.variables['lat_rho'][:]
if grid is None:
grid = inout.readgrid(loc)
else:
grid = grid
## Input starting locations as real space lon,lat locations
# lon0,lat0 = np.meshgrid(-95.498218005315309,23.142258627126882) # [0,0] (SE) corner
# lon0,lat0 = np.meshgrid(-97.748582291691989,23.000027311710628) # [-1,0] (SW) corner
# lon0,lat0 = np.meshgrid(-87.757124031927574,29.235771320764623) # [0,-1] (NE) corner
# lon0,lat0 = np.meshgrid(-88.3634073986196,30.388542615201313) # [-1,-1] (NW) corner
# lon0,lat0 = np.meshgrid(np.linspace(-94,-93,10),np.linspace(28,29,10)) # grid outside Galveston Bay
# lon0,lat0 = np.meshgrid(np.linspace(-95,-91,100),np.linspace(28,29,50)) # rectangle outside Galveston
# lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,1100),np.linspace(22.5,31,980)) # whole domain, 1 km
# lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,220),np.linspace(22.5,31,196)) # whole domain, 5 km
# # FOR TEST1:
# lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,110),np.linspace(22.5,31,98)) # whole domain, 10 km
# lon0,lat0 = np.meshgrid(np.linspace(-98.5,-87.5,21),np.linspace(22.5,31,20)) # whole domain, 50 km
if nlon is None:
nlon = 110
else:
nlon = nlon
if nlat is None:
nlat = 98
else:
nlat = nlat
lon0, lat0 = np.meshgrid(np.linspace(-98.5, -87.5, nlon),
np.linspace(22.5, 31,
nlat)) # whole domain, 10 km
# Eliminate points that are outside domain or in masked areas
lon0, lat0 = tools.check_points(lon0, lat0, grid)
## Choose method for vertical placement of drifters
# Also update makefile accordingly. Choose the twodim flag for isoslice.
# See above for more notes, but do the following two lines for an isoslice
z0 = 's' #'salt' #'s' #'z' #'salt' #'s'
zpar = 29 #30 #29 #-10 #grid['km']-1 # 30 #grid['km']-1
# Do the following two for a 3d simulation
# z0 = np.ones(xstart0.shape)*-40 # below the surface
# zpar = 'fromMSL'
# pdb.set_trace()
# for 3d flag, do3d=0 makes the run 2d and do3d=1 makes the run 3d
do3d = 0
# turbulence/diffusion flag. doturb=0 means no turb/diffusion,
# doturb=1 means adding parameterized turbulence
# doturb=2 means adding diffusion on a circle
# doturb=3 means adding diffusion on an ellipse (anisodiffusion)
if doturb is None:
doturb = 0
else:
doturb = doturb
# simulation name, used for saving results into netcdf file
if name is None:
name = 'temp' #'5_5_D5_F'
else:
name = name
return loc, nsteps, ndays, ff, date, tseas, ah, av, lon0, lat0, z0, zpar, do3d, doturb, name, grid
def run(loc, nsteps, ndays, ff, date, tseas, ah, av, lon0, lat0, z0, \
zpar, do3d, doturb, name, grid=None, dostream=0, \
T0=None, U=None, V=None):
'''
To re-compile tracmass fortran code, type "make clean" and "make f2py", which will give
a file tracmass.so, which is the module we import above. Then in ipython, "run run.py"
xend,yend,zend are particle locations at next step
some variables are not specifically because f2py is hiding them from me:
imt, jmt, km, ntractot
Look at tracmass.step to see what it is doing and making optional at the end.
Do this by importing tracmass and then tracmass.step?
I am assuming here that the velocity field at two times are being input into tracmass
such that the output is the position for the drifters at the time corresponding to the
second velocity time. Each drifter may take some number of steps in between, but those
are not saved.
loc Path to directory of grid and output files
nsteps Number of steps to do between model outputs (iter in tracmass)
ndays number of days to track the particles from start date
ff ff=1 to go forward in time and ff=-1 for backward in time
date Start date in datetime object
tseas Time between outputs in seconds
ah Horizontal diffusion in m^2/s.
See project values of 350, 100, 0, 2000. For -turb,-diffusion
av Vertical diffusion in m^2/s.
do3d for 3d flag, do3d=0 makes the run 2d and do3d=1 makes the run 3d
doturb turbulence/diffusion flag.
doturb=0 means no turb/diffusion,
doturb=1 means adding parameterized turbulence
doturb=2 means adding diffusion on a circle
doturb=3 means adding diffusion on an ellipse (anisodiffusion)
lon0 Drifter starting locations in x/zonal direction.
lat0 Drifter starting locations in y/meridional direction.
z0/zpar For 3D drifter movement, turn off twodim flag in makefile.
Then z0 should be an array of initial drifter depths.
The array should be the same size as lon0 and be negative
for under water. Currently drifter depths need to be above
the seabed for every x,y particle location for the script to run.
To do 3D but start at surface, use z0=zeros(ia.shape) and have
either zpar='fromMSL'
choose fromMSL to have z0 starting depths be for that depth below the base
time-independent sea level (or mean sea level).
choose 'fromZeta' to have z0 starting depths be for that depth below the
time-dependent sea surface. Haven't quite finished the 'fromZeta' case.
For 2D drifter movement, turn on twodim flag in makefile.
Then:
set z0 to 's' for 2D along a terrain-following slice
and zpar to be the index of s level you want to use (0 to km-1)
set z0 to 'rho' for 2D along a density surface
and zpar to be the density value you want to use
Can do the same thing with salinity ('salt') or temperature ('temp')
The model output doesn't currently have density though.
set z0 to 'z' for 2D along a depth slice
and zpar to be the constant (negative) depth value you want to use
To simulate drifters at the surface, set z0 to 's'
and zpar = grid['km']-1 to put them in the upper s level
z0='s' is currently not working correctly!!!
In the meantime, do surface using the 3d set up option but with 2d flag set
xp x-locations in x,y coordinates for drifters
yp y-locations in x,y coordinates for drifters
zp z-locations (depths from mean sea level) for drifters
t time for drifter tracks
name Name of simulation to be used for netcdf file containing final tracks
grid (optional) Grid information, as read in by tracpy.inout.readgrid().
The following inputs are for calculating Lagrangian stream functions
dostream Calculate streamfunctions (1) or not (0). Default is 0.
U0, V0 (optional) Initial volume transports of drifters (m^3/s)
U, V (optional) Array aggregating volume transports as drifters move [imt-1,jmt], [imt,jmt-1]
'''
tic_start = time.time()
tic_initial = time.time()
# Units for time conversion with netCDF.num2date and .date2num
units = 'seconds since 1970-01-01'
# Number of model outputs to use
# Adding one index so that all necessary indices are captured by this number.
# Then the run loop uses only the indices determined by tout instead of needing
# an extra one beyond
tout = np.int((ndays*(24*3600))/tseas + 1)
# Convert date to number
date = netCDF.date2num(date, units)
# Figure out what files will be used for this tracking
nc, tinds = inout.setupROMSfiles(loc, date, ff, tout)
# Read in grid parameters into dictionary, grid
if grid is None:
grid = inout.readgrid(loc, nc)
else: # don't need to reread grid
grid = grid
# Interpolate to get starting positions in grid space
xstart0, ystart0, _ = tools.interpolate2d(lon0, lat0, grid, 'd_ll2ij')
# Do z a little lower down
# Initialize seed locations
ia = np.ceil(xstart0) #[253]#,525]
ja = np.ceil(ystart0) #[57]#,40]
# don't use nan's
ind2 = ~np.isnan(ia) * ~np.isnan(ja)
ia = ia[ind2]
ja = ja[ind2]
xstart0 = xstart0[ind2]
ystart0 = ystart0[ind2]
dates = nc.variables['ocean_time'][:]
t0save = dates[tinds[0]] # time at start of drifter test from file in seconds since 1970-01-01, add this on at the end since it is big
# Initialize drifter grid positions and indices
xend = np.ones((ia.size,len(tinds)*nsteps))*np.nan
yend = np.ones((ia.size,len(tinds)*nsteps))*np.nan
zend = np.ones((ia.size,len(tinds)*nsteps))*np.nan
zp = np.ones((ia.size,len(tinds)*nsteps))*np.nan
iend = np.ones((ia.size,len(tinds)*nsteps))*np.nan
jend = np.ones((ia.size,len(tinds)*nsteps))*np.nan
kend = np.ones((ia.size,len(tinds)*nsteps))*np.nan
ttend = np.ones((ia.size,len(tinds)*nsteps))*np.nan
t = np.zeros((len(tinds)*nsteps+1))
flag = np.zeros((ia.size),dtype=np.int) # initialize all exit flags for in the domain
# Initialize vertical stuff and fluxes
# Read initial field in - to 'new' variable since will be moved
# at the beginning of the time loop ahead
if is_string_like(z0): # isoslice case
ufnew,vfnew,dztnew,zrtnew,zwtnew = inout.readfields(tinds[0],grid,nc,z0,zpar)
else: # 3d case
ufnew,vfnew,dztnew,zrtnew,zwtnew = inout.readfields(tinds[0],grid,nc)
## Find zstart0 and ka
# The k indices and z grid ratios should be on a wflux vertical grid,
# which goes from 0 to km since the vertical velocities are defined
# at the vertical cell edges. A drifter's grid cell is vertically bounded
# above by the kth level and below by the (k-1)th level
if is_string_like(z0): # then doing a 2d isoslice
# there is only one vertical grid cell, but with two vertically-
# bounding edges, 0 and 1, so the initial ka value is 1 for all
# isoslice drifters.
ka = np.ones(ia.size)
# for s level isoslice, place drifters vertically at the center
# of the grid cell since that is where the u/v flux info is from.
# For a rho/temp/density isoslice, we treat it the same way, such
# that the u/v flux info taken at a specific rho/temp/density value
# is treated as being at the center of the grid cells vertically.
zstart0 = np.ones(ia.size)*0.5
else: # 3d case
# Convert initial real space vertical locations to grid space
# first find indices of grid cells vertically
ka = np.ones(ia.size)*np.nan
zstart0 = np.ones(ia.size)*np.nan
if zpar == 'fromMSL':
for i in xrange(ia.size):
# pdb.set_trace()
ind = (grid['zwt0'][ia[i],ja[i],:]<=z0[i])
# check to make sure there is at least one true value, so the z0 is shallower than the seabed
if np.sum(ind):
ka[i] = find(ind)[-1] # find value that is just shallower than starting vertical position
# if the drifter starting vertical location is too deep for the x,y location, complain about it
else: # Maybe make this nan or something later
print 'drifter vertical starting location is too deep for its x,y location. Try again.'
if (z0[i] != grid['zwt0'][ia[i],ja[i],ka[i]]) and (ka[i] != grid['km']): # check this
ka[i] = ka[i]+1
# Then find the vertical relative position in the grid cell by adding on the bit of grid cell
zstart0[i] = ka[i] - abs(z0[i]-grid['zwt0'][ia[i],ja[i],ka[i]]) \
/abs(grid['zwt0'][ia[i],ja[i],ka[i]-1]-grid['zwt0'][ia[i],ja[i],ka[i]])
# elif zpar == 'fromZeta':
# for i in xrange(ia.size):
# pdb.set_trace()
# ind = (zwtnew[ia[i],ja[i],:]<=z0[i])
# ka[i] = find(ind)[-1] # find value that is just shallower than starting vertical position
# if (z0[i] != zwtnew[ia[i],ja[i],ka[i]]) and (ka[i] != grid['km']): # check this
# ka[i] = ka[i]+1
# # Then find the vertical relative position in the grid cell by adding on the bit of grid cell
# zstart0[i] = ka[i] - abs(z0[i]-zwtnew[ia[i],ja[i],ka[i]]) \
# /abs(zwtnew[ia[i],ja[i],ka[i]-1]-zwtnew[ia[i],ja[i],ka[i]])
# Find initial cell depths to concatenate to beginning of drifter tracks later
zsave = tools.interpolate3d(xstart0, ystart0, zstart0, zwtnew)
toc_initial = time.time()
# j = 0 # index for number of saved steps for drifters
tic_read = np.zeros(len(tinds))
toc_read = np.zeros(len(tinds))
tic_zinterp = np.zeros(len(tinds))
toc_zinterp = np.zeros(len(tinds))
tic_tracmass = np.zeros(len(tinds))
toc_tracmass = np.zeros(len(tinds))
# pdb.set_trace()
xr3 = grid['xr'].reshape((grid['xr'].shape[0],grid['xr'].shape[1],1)).repeat(zwtnew.shape[2],axis=2)
yr3 = grid['yr'].reshape((grid['yr'].shape[0],grid['yr'].shape[1],1)).repeat(zwtnew.shape[2],axis=2)
# Loop through model outputs. tinds is in proper order for moving forward
# or backward in time, I think.
for j,tind in enumerate(tinds[:-1]):
# pdb.set_trace()
# Move previous new time step to old time step info
ufold = ufnew
vfold = vfnew
dztold = dztnew
zrtold = zrtnew
zwtold = zwtnew
tic_read[j] = time.time()
# Read stuff in for next time loop
if is_string_like(z0): # isoslice case
ufnew,vfnew,dztnew,zrtnew,zwtnew = inout.readfields(tinds[j+1],grid,nc,z0,zpar)
else: # 3d case
ufnew,vfnew,dztnew,zrtnew,zwtnew = inout.readfields(tinds[j+1],grid,nc)
toc_read[j] = time.time()
# print "readfields run time:",toc_read-tic_read
print j
# flux fields at starting time for this step
if j != 0:
xstart = xend[:,j*nsteps-1]
ystart = yend[:,j*nsteps-1]
zstart = zend[:,j*nsteps-1]
# mask out drifters that have exited the domain
xstart = np.ma.masked_where(flag[:]==1,xstart)
ystart = np.ma.masked_where(flag[:]==1,ystart)
zstart = np.ma.masked_where(flag[:]==1,zstart)
ind = (flag[:] == 0) # indices where the drifters are still inside the domain
else: # first loop, j==0
xstart = xstart0
ystart = ystart0
zstart = zstart0
# TODO: Do a check to make sure all drifter starting locations are within domain
ind = (flag[:] == 0) # indices where the drifters are inside the domain to start
# Find drifter locations
# only send unmasked values to step
if not np.ma.compressed(xstart).any(): # exit if all of the drifters have exited the domain
break
else:
# Combine times for arrays for input to tracmass
# from [ixjxk] to [ixjxkxt]
# Change ordering for these three arrays here instead of in readfields since
# concatenate does not seem to preserve ordering
uflux = np.asfortranarray(np.concatenate((ufold.reshape(np.append(ufold.shape,1)), \
ufnew.reshape(np.append(ufnew.shape,1))), \
axis=ufold.ndim))
vflux = np.asfortranarray(np.concatenate((vfold.reshape(np.append(vfold.shape,1)), \
vfnew.reshape(np.append(vfnew.shape,1))), \
axis=vfold.ndim))
dzt = np.asfortranarray(np.concatenate((dztold.reshape(np.append(dztold.shape,1)), \
dztnew.reshape(np.append(dztnew.shape,1))), \
axis=dztold.ndim))
# Change the horizontal indices from python to fortran indexing
# (vertical are zero-based in tracmass)
xstart, ystart = tools.convert_indices('py2f',xstart,ystart)
# km that is sent to tracmass is determined from uflux (see tracmass?)
# so it will be the correct value for whether we are doing the 3D
# or isoslice case.
# vec = np.arange(j*nsteps,j*nsteps+nsteps) # indices for storing new track locations
tic_tracmass[j] = time.time()
# pdb.set_trace()
if dostream: # calculate Lagrangian stream functions
xend[ind,j*nsteps:j*nsteps+nsteps],\
yend[ind,j*nsteps:j*nsteps+nsteps],\
zend[ind,j*nsteps:j*nsteps+nsteps], \
iend[ind,j*nsteps:j*nsteps+nsteps],\
jend[ind,j*nsteps:j*nsteps+nsteps],\
kend[ind,j*nsteps:j*nsteps+nsteps],\
flag[ind],\
ttend[ind,j*nsteps:j*nsteps+nsteps], U, V = \
tracmass.step(np.ma.compressed(xstart),\
np.ma.compressed(ystart),
np.ma.compressed(zstart),
tseas, uflux, vflux, ff, \
grid['kmt'].astype(int), \
dzt, grid['dxdy'], grid['dxv'], \
grid['dyu'], grid['h'], nsteps, \
ah, av, do3d, doturb, dostream, \
t0=T0[ind],
ut=U, vt=V)
else: # don't calculate Lagrangian stream functions
xend[ind,j*nsteps:j*nsteps+nsteps],\
yend[ind,j*nsteps:j*nsteps+nsteps],\
zend[ind,j*nsteps:j*nsteps+nsteps], \
iend[ind,j*nsteps:j*nsteps+nsteps],\
jend[ind,j*nsteps:j*nsteps+nsteps],\
kend[ind,j*nsteps:j*nsteps+nsteps],\
flag[ind],\
ttend[ind,j*nsteps:j*nsteps+nsteps], _, _ = \
tracmass.step(np.ma.compressed(xstart),\
np.ma.compressed(ystart),
np.ma.compressed(zstart),
tseas, uflux, vflux, ff, \
grid['kmt'].astype(int), \
dzt, grid['dxdy'], grid['dxv'], \
grid['dyu'], grid['h'], nsteps, \
ah, av, do3d, doturb, dostream)
toc_tracmass[j] = time.time()
# pdb.set_trace()
# Change the horizontal indices from python to fortran indexing
xend[ind,j*nsteps:j*nsteps+nsteps], \
yend[ind,j*nsteps:j*nsteps+nsteps] \
= tools.convert_indices('f2py', \
xend[ind,j*nsteps:j*nsteps+nsteps], \
yend[ind,j*nsteps:j*nsteps+nsteps])
# Calculate times for the output frequency
if ff == 1:
t[j*nsteps+1:j*nsteps+nsteps+1] = t[j*nsteps] + np.linspace(tseas/nsteps,tseas,nsteps) # update time in seconds to match drifters
else:
t[j*nsteps+1:j*nsteps+nsteps+1] = t[j*nsteps] - np.linspace(tseas/nsteps,tseas,nsteps) # update time in seconds to match drifters
# Skip calculating real z position if we are doing surface-only drifters anyway
if z0 != 's' and zpar != grid['km']-1:
tic_zinterp[j] = time.time()
# Calculate real z position
r = np.linspace(1./nsteps,1,nsteps) # linear time interpolation constant that is used in tracmass
for n in xrange(nsteps): # loop through time steps
# interpolate to a specific output time
# pdb.set_trace()
zwt = (1.-r[n])*zwtold + r[n]*zwtnew
zp[ind,j*nsteps:j*nsteps+nsteps], dt = tools.interpolate3d(xend[ind,j*nsteps:j*nsteps+nsteps], \
yend[ind,j*nsteps:j*nsteps+nsteps], \
zend[ind,j*nsteps:j*nsteps+nsteps], \
zwt)
toc_zinterp[j] = time.time()
nc.close()
t = t + t0save # add back in base time in seconds
# pdb.set_trace()
# Add on to front location for first time step
xg=np.concatenate((xstart0.reshape(xstart0.size,1),xend),axis=1)
yg=np.concatenate((ystart0.reshape(ystart0.size,1),yend),axis=1)
# Concatenate zp with initial real space positions
zp=np.concatenate((zsave[0].reshape(zstart0.size,1),zp),axis=1)
# Delaunay interpolation
# xp, yp, dt = tools.interpolate(xg,yg,grid,'d_ij2xy')
# lonp, latp, dt = tools.interpolation(xg,yg,grid,'d_ij2ll')
## map coordinates interpolation
# xp2, yp2, dt = tools.interpolate(xg,yg,grid,'m_ij2xy')
# tic = time.time()
lonp, latp, dt = tools.interpolate2d(xg,yg,grid,'m_ij2ll',mode='constant',cval=np.nan)
# print '2d interp time=', time.time()-tic
# pdb.set_trace()
runtime = time.time()-tic_start
print "run time:\t\t\t", runtime
print "---------------------------------------------"
print "Time spent on:"
initialtime = toc_initial-tic_initial
print "\tInitial stuff: \t\t%4.2f (%4.2f%%)" % (initialtime, (initialtime/runtime)*100)
readtime = np.sum(toc_read-tic_read)
print "\tReading in fields: \t%4.2f (%4.2f%%)" % (readtime, (readtime/runtime)*100)
zinterptime = np.sum(toc_zinterp-tic_zinterp)
print "\tZ interpolation: \t%4.2f (%4.2f%%)" % (zinterptime, (zinterptime/runtime)*100)
tractime = np.sum(toc_tracmass-tic_tracmass)
print "\tTracmass: \t\t%4.2f (%4.2f%%)" % (tractime, (tractime/runtime)*100)
# Save results to netcdf file
if dostream:
inout.savetracks(lonp, latp, zp, t, name, nsteps, ff, tseas, ah, av, \
do3d, doturb, loc, T0, U, V)
return lonp, latp, zp, t, grid, T0, U, V
else:
inout.savetracks(lonp, latp, zp, t, name, nsteps, ff, tseas, ah, av, \
do3d, doturb, loc)
return lonp, latp, zp, t, grid
def hist(
lonp,
latp,
fname,
tind="final",
which="contour",
vmax=None,
fig=None,
ax=None,
bins=(40, 40),
N=10,
grid=None,
xlims=None,
ylims=None,
C=None,
Title=None,
weights=None,
Label="Final drifter location (%)",
isll=True,
binscale=None,
):
"""
Plot histogram of given track data at time index tind.
Args:
lonp,latp: Drifter track positions in lon/lat [time x ndrifters]
fname: Plot name to save
tind (Optional): Default is 'final', in which case the final
position of each drifter in the array is found and plotted.
Alternatively, a time index can be input and drifters at that time
will be plotted. Note that once drifters hit the outer numerical
boundary, they are nan'ed out so this may miss some drifters.
which (Optional[str]): 'contour', 'pcolor', 'hexbin', 'hist2d' for
type of plot used. Default 'hexbin'.
bins (Optional): Number of bins used in histogram. Default (15,25).
N (Optional[int]): Number of contours to make. Default 10.
grid (Optional): grid as read in by inout.readgrid()
xlims (Optional): value limits on the x axis
ylims (Optional): value limits on the y axis
isll: Default True. Inputs are in lon/lat. If False, assume they
are in projected coords.
Note:
Currently assuming we are plotting the final location of each drifter
regardless of tind.
"""
if grid is None:
loc = "http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc"
grid = inout.readgrid(loc)
if isll: # if inputs are in lon/lat, change to projected x/y
# Change positions from lon/lat to x/y
xp, yp = grid.proj(lonp, latp)
# Need to retain nan's since basemap changes them to values
ind = np.isnan(lonp)
xp[ind] = np.nan
yp[ind] = np.nan
else:
xp = lonp
yp = latp
if fig is None:
fig = plt.figure(figsize=(11, 10))
else:
fig = fig
background(grid) # Plot coastline and such
if tind == "final":
# Find final positions of drifters
xpc, ypc = tools.find_final(xp, yp)
elif isinstance(tind, int):
xpc = xp[:, tind]
ypc = yp[:, tind]
else: # just plot what is input if some other string
xpc = xp.flatten()
ypc = yp.flatten()
if which == "contour":
# Info for 2d histogram
H, xedges, yedges = np.histogram2d(
xpc, ypc, range=[[grid.x_rho.min(), grid.x_rho.max()], [grid.y_rho.min(), grid.y_rho.max()]], bins=bins
)
# Contour Plot
XE, YE = np.meshgrid(op.resize(xedges, 0), op.resize(yedges, 0))
d = (H / H.sum()) * 100
# # from http://matplotlib.1069221.n5.nabble.com/question-about-contours-and-clim-td21111.html
# locator = ticker.MaxNLocator(50) # if you want no more than 10 contours
# locator.create_dummy_axis()
# locator.set_bounds(0,1)#d.min(),d.max())
# levs = locator()
con = fig.contourf(XE, YE, d.T, N) # ,levels=levs)#(0,15,30,45,60,75,90,105,120))
con.set_cmap("YlOrRd")
if Title is not None:
plt.set_title(Title)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3725, 0.25, 0.48, 0.02]) # colorbar axes
cb = fig.colorbar(con, cax=cax, orientation="horizontal")
cb.set_label("Final drifter location (percent)")
# Save figure into a local directory called figures. Make directory
# if it doesn't exist.
if not os.path.exists("figures"):
os.makedirs("figures")
fig.savefig("figures/" + fname + "histcon.png", bbox_inches="tight")
elif which == "pcolor":
# Info for 2d histogram
H, xedges, yedges = np.histogram2d(
xpc,
ypc,
range=[[grid.x_rho.min(), grid.x_rho.max()], [grid.y_rho.min(), grid.y_rho.max()]],
bins=bins,
weights=weights,
)
# Pcolor plot
# C is the z value plotted, and is normalized by the total number of
# drifters
if C is None:
C = (H.T / H.sum()) * 100
else:
# or, provide some other weighting
C = (H.T / C) * 100
p = plt.pcolor(xedges, yedges, C, cmap="YlOrRd")
if Title is not None:
plt.set_title(Title)
# Set x and y limits
if xlims is not None:
plt.xlim(xlims)
if ylims is not None:
plt.ylim(ylims)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) # colorbar axes
cb = fig.colorbar(p, cax=cax, orientation="horizontal")
cb.set_label("Final drifter location (percent)")
# Save figure into a local directory called figures. Make directory
# if it doesn't exist.
if not os.path.exists("figures"):
os.makedirs("figures")
fig.savefig("figures/" + fname + "histpcolor.png", bbox_inches="tight")
# savefig('figures/' + fname + 'histpcolor.pdf',bbox_inches='tight')
elif which == "hexbin":
if ax is None:
ax = plt.gca()
else:
ax = ax
if C is None:
# C with the reduce_C_function as sum is what makes it a percent
C = np.ones(len(xpc)) * (1.0 / len(xpc)) * 100
else:
C = C * np.ones(len(xpc)) * 100
hb = plt.hexbin(
xpc,
ypc,
C=C,
cmap="YlOrRd",
gridsize=bins[0],
extent=(grid.x_psi.min(), grid.x_psi.max(), grid.y_psi.min(), grid.y_psi.max()),
reduce_C_function=sum,
vmax=vmax,
axes=ax,
bins=binscale,
)
# Set x and y limits
if xlims is not None:
plt.xlim(xlims)
if ylims is not None:
plt.ylim(ylims)
if Title is not None:
ax.set_title(Title)
# Want colorbar at the given location relative to axis so this works
# regardless of # of subplots, so convert from axis to figure
# coordinates. To do this, first convert from axis to display coords
# transformations:
# http://matplotlib.org/users/transforms_tutorial.html
# axis: [x_left, y_bottom, width, height]
ax_coords = [0.35, 0.25, 0.6, 0.02]
# display: [x_left,y_bottom,x_right,y_top]
disp_coords = ax.transAxes.transform(
[(ax_coords[0], ax_coords[1]), (ax_coords[0] + ax_coords[2], ax_coords[1] + ax_coords[3])]
)
# inverter object to go from display coords to figure coords
inv = fig.transFigure.inverted()
# figure: [x_left,y_bottom,x_right,y_top]
fig_coords = inv.transform(disp_coords)
# actual desired figure coords. figure:
# [x_left, y_bottom, width, height]
fig_coords = [
fig_coords[0, 0],
fig_coords[0, 1],
fig_coords[1, 0] - fig_coords[0, 0],
fig_coords[1, 1] - fig_coords[0, 1],
]
# Inlaid colorbar
cax = fig.add_axes(fig_coords)
# # Horizontal colorbar below plot
# cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) # colorbar axes
cb = fig.colorbar(hb, cax=cax, orientation="horizontal")
cb.set_label(Label)
# Save figure into a local directory called figures. Make directory
# if it doesn't exist.
if not os.path.exists("figures"):
os.makedirs("figures")
fig.savefig("figures/" + fname + "histhexbin.png", bbox_inches="tight")
# savefig('figures/' + fname + 'histhexbin.pdf',bbox_inches='tight')
elif which == "hist2d":
plt.hist2d(
xpc,
ypc,
bins=40,
range=[[grid.x_rho.min(), grid.x_rho.max()], [grid.y_rho.min(), grid.y_rho.max()]],
normed=True,
)
plt.set_cmap("YlOrRd")
# Set x and y limits
if xlims is not None:
xlim(xlims)
if ylims is not None:
ylim(ylims)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) # colorbar axes
cb = fig.colorbar(cax=cax, orientation="horizontal")
cb.set_label("Final drifter location (percent)")
# Save figure into a local directory called figures. Make directory
# if it doesn't exist.
if not os.path.exists("figures"):
os.makedirs("figures")
fig.savefig("figures/" + fname + "hist2d.png", bbox_inches="tight")
def transport(
name,
fmod=None,
Title=None,
dmax=None,
N=7,
extraname=None,
llcrnrlon=-98.5,
llcrnrlat=22.5,
urcrnrlat=31.0,
urcrnrlon=-87.5,
colormap="Blues",
fig=None,
ax=None,
):
"""
Make plot of zoomed-in area near DWH spill of transport of drifters over
time.
FILL IN
Args:
name
U
V
lon0
lat0
T0
"""
# Load in transport information
U, V, lon0, lat0, T0 = inout.loadtransport(name, fmod=fmod)
# Smaller basemap parameters.
loc = "http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc"
grid = inout.readgrid(loc, llcrnrlon=llcrnrlon, llcrnrlat=llcrnrlat, urcrnrlat=urcrnrlat, urcrnrlon=urcrnrlon)
S = np.sqrt(op.resize(U, 1) ** 2 + op.resize(V, 0) ** 2)
Splot = (S / T0) * 100
if dmax is None:
dmax = Splot.max()
else:
dmax = dmax
# from http://matplotlib.1069221.n5.nabble.com/question-about-contours-and-clim-td21111.html
locator = ticker.MaxNLocator(N) # if you want no more than 10 contours
locator.create_dummy_axis()
locator.set_bounds(0, dmax) # d.min(),d.max())
levs = locator()
if fig is None:
fig = plt.figure(figsize=(11, 10))
else:
fig = fig
background(grid=grid)
c = fig.contourf(grid.xpsi, grid.ypsi, Splot, cmap=colormap, extend="max", levels=levs)
plt.title(Title)
# # Add initial drifter location (all drifters start at the same location)
# lon0 = lon0.mean()
# lat0 = lat0.mean()
# x0, y0 = grid['basemap'](lon0, lat0)
# plot(x0, y0, 'go', markersize=10)
if ax is None:
ax = plt.gca()
else:
ax = ax
# Want colorbar at the given location relative to axis so this works
# regardless of # of subplots,
# so convert from axis to figure coordinates
# To do this, first convert from axis to display coords
# transformations: http://matplotlib.org/users/transforms_tutorial.html
ax_coords = [0.35, 0.25, 0.6, 0.02] # axis: [x_left, y_bottom, width, height]
# display: [x_left,y_bottom,x_right,y_top]
disp_coords = ax.transAxes.transform(
[(ax_coords[0], ax_coords[1]), (ax_coords[0] + ax_coords[2], ax_coords[1] + ax_coords[3])]
)
# inverter object to go from display coords to figure coords
inv = fig.transFigure.inverted()
# figure: [x_left,y_bottom,x_right,y_top]
fig_coords = inv.transform(disp_coords)
# actual desired figure coords. figure: [x_left, y_bottom, width, height]
fig_coords = [
fig_coords[0, 0],
fig_coords[0, 1],
fig_coords[1, 0] - fig_coords[0, 0],
fig_coords[1, 1] - fig_coords[0, 1],
]
# Inlaid colorbar
cax = fig.add_axes(fig_coords)
# cax = fig.add_axes([0.39, 0.25, 0.49, 0.02])
# cax = fig.add_axes([0.49, 0.25, 0.39, 0.02])
cb = fig.colorbar(cax=cax, orientation="horizontal")
cb.set_label("Normalized drifter transport (%)")
if extraname is None:
fig.savefig("figures/" + name + "/transport", bbox_inches="tight")
else:
fig.savefig("figures/" + name + "/" + extraname + "transport", bbox_inches="tight")
def tracks(lonp, latp, fname, grid=None, fig=None, ax=None, Title=None, mers=None, pars=None):
"""
Plot tracks as lines with starting points in green and ending points in
red.
Args:
lonp,latp: Drifter track positions [time x ndrifters]
fname: Plot name to save
"""
if grid is None:
loc = "http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc"
grid = inout.readgrid(loc)
if fig is None:
fig = plt.figure(figsize=(12, 10))
else:
fig = fig
if ax is None:
ax = plt.gca()
else:
ax = ax
# Change positions from lon/lat to x/y
xp, yp = grid.proj(lonp, latp)
# Need to retain nan's since basemap changes them to values
ind = np.isnan(lonp)
xp[ind] = np.nan
yp[ind] = np.nan
if mers is not None:
# Plot coastline and such
background(grid, ax=ax, mers=mers, pars=pars)
else:
background(grid, ax=ax) # Plot coastline and such
# Starting marker
ax.plot(xp[:, 0], yp[:, 0], "o", color="g", markersize=3, label="_nolegend_", alpha=0.4)
# Plot tracks
ax.plot(xp.T, yp.T, "-", color="grey", linewidth=0.2)
# Find final positions of drifters
xpc, ypc = tools.find_final(xp, yp)
ax.plot(xpc, ypc, "o", color="r", label="_nolegend_")
if Title is not None:
ax.set_title(Title)
# Legend, of sorts
# ax = gca()
xtext = 0.45
ytext = 0.18
ax.text(xtext, ytext, "starting location", fontsize=16, color="green", alpha=0.8, transform=ax.transAxes)
ax.text(xtext, ytext - 0.03, "track", fontsize=16, color="grey", transform=ax.transAxes)
ax.text(xtext, ytext - 0.03 * 2, "ending location", fontsize=16, color="red", transform=ax.transAxes)
# xtext, ytext = grid['basemap'](-94,24) # text location
# text(xtext,ytext,'starting location',fontsize=16,color='green',alpha=.8)
# text(xtext,ytext-30000,'track',fontsize=16,color='grey')#,alpha=.8)
# text(xtext,ytext-60000,'ending location',fontsize=16,color='red')#,alpha=.8)
# # get psi mask from rho mask
# # maskp = grid['mask'][1:,1:]*grid['mask'][:-1,1:]* \
# # grid['mask'][1:,:-1]*grid['mask'][:-1,:-1]
# # ind = maskp
# # ind[ind==0] = np.nan
# # plot(grid['xpsi']*ind,grid['ypsi']*ind,'k', \
# # (grid['xpsi']*ind).T,(grid['ypsi']*ind).T,'k')
# plot(grid['xpsi'],grid['ypsi'],'k', \
# (grid['xpsi']).T,(grid['ypsi']).T,'k')
# 16 is (lower) one that is near islands, 41 is higher one
# show()
# Save figure into a local directory called figures. Make directory if it
# doesn't exist.
if not os.path.exists("figures"):
os.makedirs("figures")
fig.savefig("figures/" + fname + "tracks.png", bbox_inches="tight")
def transport(name, fmod=None, Title=None, dmax=None, N=7, extraname=None,
llcrnrlon=-98.5, llcrnrlat=22.5, urcrnrlat=31.0, urcrnrlon=-87.5,
colormap='Blues'):
'''
Make plot of zoomed-in area near DWH spill of transport of drifters over
time.
FILL IN
Inputs:
name
U
V
lon0
lat0
T0
'''
# (name=None, U, V, lon0, lat0, T0, dmax, extraname, Title, N,
# llcrnrlon, llcrnrlat, urcrnrlat, urcrnrlon, colormap):
# Load in transport information
U, V, lon0, lat0, T0 = inout.loadtransport(name,fmod=fmod)
# Smaller basemap parameters.
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc, llcrnrlon=llcrnrlon, llcrnrlat=llcrnrlat,
urcrnrlat=urcrnrlat, urcrnrlon=urcrnrlon)
S = np.sqrt(op.resize(U,1)**2+op.resize(V,0)**2)
Splot = (S/T0)*100
if dmax is None:
dmax = Splot.max()
else:
dmax = dmax
# from http://matplotlib.1069221.n5.nabble.com/question-about-contours-and-clim-td21111.html
locator = ticker.MaxNLocator(N) # if you want no more than 10 contours
locator.create_dummy_axis()
locator.set_bounds(0,dmax)#d.min(),d.max())
levs = locator()
fig = figure(figsize=(12,10))
background(grid=grid)
c = contourf(grid['xpsi'], grid['ypsi'], Splot,
cmap=colormap, extend='max', levels=levs)
title(Title)
# Add initial drifter location (all drifters start at the same location)
lon0 = lon0.mean()
lat0 = lat0.mean()
x0, y0 = grid['basemap'](lon0, lat0)
plot(x0, y0, 'go', markersize=10)
# Inlaid colorbar
cax = fig.add_axes([0.49, 0.25, 0.39, 0.02])
# cax = fig.add_axes([0.5, 0.2, 0.35, 0.02])
cb = colorbar(cax=cax,orientation='horizontal')
cb.set_label('Normalized drifter transport (%)')
if extraname is None:
savefig('figures/' + name + '/transport', bbox_inches='tight')
else:
savefig('figures/' + name + '/transport' + extraname, bbox_inches='tight')
def hist(lonp, latp, fname, tind='final', which='contour', \
bins=(40,40), N=10, grid=None, xlims=None, ylims=None):
"""
Plot histogram of given track data at time index tind.
Inputs:
lonp,latp Drifter track positions in lon/lat [time x ndrifters]
fname Plot name to save
tind (optional) Default is 'final', in which case the final
position of each drifter in the array is found
and plotted. Alternatively, a time index
can be input and drifters at that time will be plotted.
Note that once drifters hit the outer numerical boundary,
they are nan'ed out so this may miss some drifters.
which (optional) 'contour', 'pcolor', 'hexbin', 'hist2d'
for type of plot used. Default 'hexbin'.
bins (optional) Number of bins used in histogram. Default (15,25).
N (optional) Number of contours to make. Default 10.
grid (optional) grid as read in by inout.readgrid()
xlims (optional) value limits on the x axis
ylims (optional) value limits on the y axis
Note: Currently assuming we are plotting the final location
of each drifter regardless of tind.
"""
if grid is None:
loc = 'http://barataria.tamu.edu:8080/thredds/dodsC/NcML/txla_nesting6.nc'
grid = inout.readgrid(loc)
# Change positions from lon/lat to x/y
xp, yp = grid['basemap'](lonp, latp)
# Need to retain nan's since basemap changes them to values
ind = np.isnan(lonp)
xp[ind] = np.nan
yp[ind] = np.nan
fig = figure(figsize=(12,10))
background(grid) # Plot coastline and such
# pdb.set_trace()
if tind == 'final':
# Find final positions of drifters
xpc, ypc = tools.find_final(xp, yp)
elif is_numlike(tind):
xpc = xp[:,tind]
ypc = yp[:,tind]
else: # just plot what is input if some other string
xpc = xp.flatten()
ypc = yp.flatten()
if which == 'contour':
# Info for 2d histogram
H, xedges, yedges = np.histogram2d(xpc, ypc,
range=[[grid['xr'].min(), \
grid['xr'].max()], \
[grid['yr'].min(), \
grid['yr'].max()]],
bins=bins)
# Contour Plot
XE, YE = np.meshgrid(op.resize(xedges,0), op.resize(yedges,0))
d = (H/H.sum())*100
# # from http://matplotlib.1069221.n5.nabble.com/question-about-contours-and-clim-td21111.html
# locator = ticker.MaxNLocator(50) # if you want no more than 10 contours
# locator.create_dummy_axis()
# locator.set_bounds(0,1)#d.min(),d.max())
# levs = locator()
con = contourf(XE, YE, d.T, N)#,levels=levs)#(0,15,30,45,60,75,90,105,120))
con.set_cmap('YlOrRd')
# Horizontal colorbar below plot
cax = fig.add_axes([0.3725, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(con, cax=cax, orientation='horizontal')
cb.set_label('Final drifter location (percent)')
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'histcon.png',bbox_inches='tight')
# savefig('figures/' + fname + 'histcon.pdf',bbox_inches='tight')
elif which == 'pcolor':
# Info for 2d histogram
H, xedges, yedges = np.histogram2d(xpc, ypc,
range=[[grid['xr'].min(), \
grid['xr'].max()], \
[grid['yr'].min(), \
grid['yr'].max()]],
bins=bins)
# Pcolor plot
p = pcolor(xedges, yedges, (H.T/H.sum())*100, cmap='YlOrRd')
# Set x and y limits
# pdb.set_trace()
if xlims is not None:
xlim(xlims)
if ylims is not None:
ylim(ylims)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(p, cax=cax, orientation='horizontal')
cb.set_label('Final drifter location (percent)')
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'histpcolor.png', bbox_inches='tight')
# savefig('figures/' + fname + 'histpcolor.pdf',bbox_inches='tight')
elif which == 'hexbin':
# C with the reduce_C_function as sum is what makes it a percent
C = np.ones(len(xpc))*(1./len(xpc))*100
hb = hexbin(xpc, ypc, C=C, cmap='YlOrRd', gridsize=bins[0],
extent=(grid['xr'].min(), grid['xr'].max(),
grid['yr'].min(), grid['yr'].max()),
reduce_C_function=sum)
# Set x and y limits
# pdb.set_trace()
if xlims is not None:
xlim(xlims)
if ylims is not None:
ylim(ylims)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(cax=cax, orientation='horizontal')
cb.set_label('Final drifter location (percent)')
# pdb.set_trace()
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'histhexbin.png', bbox_inches='tight')
# savefig('figures/' + fname + 'histhexbin.pdf',bbox_inches='tight')
elif which == 'hist2d':
# pdb.set_trace()
hist2d(xpc, ypc, bins=40,
range=[[grid['xr'].min(), grid['xr'].max()],
[grid['yr'].min(), grid['yr'].max()]], normed=True)
set_cmap('YlOrRd')
# Set x and y limits
# pdb.set_trace()
if xlims is not None:
xlim(xlims)
if ylims is not None:
ylim(ylims)
# Horizontal colorbar below plot
cax = fig.add_axes([0.3775, 0.25, 0.48, 0.02]) #colorbar axes
cb = colorbar(cax=cax,orientation='horizontal')
cb.set_label('Final drifter location (percent)')
# Save figure into a local directory called figures. Make directory if it doesn't exist.
if not os.path.exists('figures'):
os.makedirs('figures')
savefig('figures/' + fname + 'hist2d.png',bbox_inches='tight')
def run(loc, nsteps, ndays, ff, date, tseas, ah, av, lon0, lat0, z0, \
zpar, do3d, doturb, name, grid=None, dostream=0, \
T0=None, U=None, V=None, zparuv=None, tseas_use=None):
'''
To re-compile tracmass fortran code, type "make clean" and "make f2py", which will give
a file tracmass.so, which is the module we import above. Then in ipython, "run run.py"
xend,yend,zend are particle locations at next step
some variables are not specifically because f2py is hiding them from me:
imt, jmt, km, ntractot
Look at tracmass.step to see what it is doing and making optional at the end.
Do this by importing tracmass and then tracmass.step?
I am assuming here that the velocity field at two times are being input into tracmass
such that the output is the position for the drifters at the time corresponding to the
second velocity time. Each drifter may take some number of steps in between, but those
are not saved.
loc Path to directory of grid and output files
nsteps Number of steps to do between model outputs (iter in tracmass)
ndays number of days to track the particles from start date
ff ff=1 to go forward in time and ff=-1 for backward in time
date Start date in datetime object
tseas Time between outputs in seconds
ah Horizontal diffusion in m^2/s.
See project values of 350, 100, 0, 2000. For -turb,-diffusion
av Vertical diffusion in m^2/s.
do3d for 3d flag, do3d=0 makes the run 2d and do3d=1 makes the run 3d
doturb turbulence/diffusion flag.
doturb=0 means no turb/diffusion,
doturb=1 means adding parameterized turbulence
doturb=2 means adding diffusion on a circle
doturb=3 means adding diffusion on an ellipse (anisodiffusion)
lon0 Drifter starting locations in x/zonal direction.
lat0 Drifter starting locations in y/meridional direction.
z0/zpar For 3D drifter movement, turn off twodim flag in makefile.
Then z0 should be an array of initial drifter depths.
The array should be the same size as lon0 and be negative
for under water. Currently drifter depths need to be above
the seabed for every x,y particle location for the script to run.
To do 3D but start at surface, use z0=zeros(ia.shape) and have
either zpar='fromMSL'
choose fromMSL to have z0 starting depths be for that depth below the base
time-independent sea level (or mean sea level).
choose 'fromZeta' to have z0 starting depths be for that depth below the
time-dependent sea surface. Haven't quite finished the 'fromZeta' case.
For 2D drifter movement, turn on twodim flag in makefile.
Then:
set z0 to 's' for 2D along a terrain-following slice
and zpar to be the index of s level you want to use (0 to km-1)
set z0 to 'rho' for 2D along a density surface
and zpar to be the density value you want to use
Can do the same thing with salinity ('salt') or temperature ('temp')
The model output doesn't currently have density though.
set z0 to 'z' for 2D along a depth slice
and zpar to be the constant (negative) depth value you want to use
To simulate drifters at the surface, set z0 to 's'
and zpar = grid['km']-1 to put them in the upper s level
z0='s' is currently not working correctly!!!
In the meantime, do surface using the 3d set up option but with 2d flag set
zparuv (optional) Use this if the k index for the model output fields (e.g, u, v) is different
from the k index in the grid. This might happen if, for example, only the surface current
were saved, but the model run originally did have many layers. This parameter
represents the k index for the u and v output, not for the grid.
tseas_use (optional) Desired time between outputs in seconds, as opposed to the actual time between outputs
(tseas). Should be >= tseas since this is just an ability to use model output at less
frequency than is available, probably just for testing purposes or matching other models.
Should to be a multiple of tseas (or will be rounded later).
xp x-locations in x,y coordinates for drifters
yp y-locations in x,y coordinates for drifters
zp z-locations (depths from mean sea level) for drifters
t time for drifter tracks
name Name of simulation to be used for netcdf file containing final tracks
grid (optional) Grid information, as read in by tracpy.inout.readgrid().
The following inputs are for calculating Lagrangian stream functions
dostream Calculate streamfunctions (1) or not (0). Default is 0.
U0, V0 (optional) Initial volume transports of drifters (m^3/s)
U, V (optional) Array aggregating volume transports as drifters move [imt-1,jmt], [imt,jmt-1]
'''
tic_start = time.time()
tic_initial = time.time()
# Units for time conversion with netCDF.num2date and .date2num
units = 'seconds since 1970-01-01'
# If tseas_use isn't set, use all available model output
if tseas_use is None:
tseas_use = tseas
# Number of model outputs to use (based on tseas, actual amount of model output)
# This should not be updated with tstride since it represents the full amount of
# indices in the original model output. tstride will be used separately to account
# for the difference.
# Adding one index so that all necessary indices are captured by this number.
# Then the run loop uses only the indices determined by tout instead of needing
# an extra one beyond
tout = np.int((ndays * (24 * 3600)) / tseas + 1)
# Calculate time outputs stride. Will be 1 if want to use all model output.
tstride = int(tseas_use / tseas) # will round down
# pdb.set_trace()
# Convert date to number
date = netCDF.date2num(date, units)
# Figure out what files will be used for this tracking
nc, tinds = inout.setupROMSfiles(loc, date, ff, tout, tstride=tstride)
# Read in grid parameters into dictionary, grid
if grid is None:
grid = inout.readgrid(loc, nc)
else: # don't need to reread grid
grid = grid
# Interpolate to get starting positions in grid space
xstart0, ystart0, _ = tools.interpolate2d(lon0, lat0, grid, 'd_ll2ij')
# Do z a little lower down
# Initialize seed locations
ia = np.ceil(xstart0) #[253]#,525]
ja = np.ceil(ystart0) #[57]#,40]
# don't use nan's
# pdb.set_trace()
ind2 = ~np.isnan(ia) * ~np.isnan(ja)
ia = ia[ind2]
ja = ja[ind2]
xstart0 = xstart0[ind2]
ystart0 = ystart0[ind2]
dates = nc.variables['ocean_time'][:]
t0save = dates[tinds[
0]] # time at start of drifter test from file in seconds since 1970-01-01, add this on at the end since it is big
# Initialize drifter grid positions and indices
xend = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
yend = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
zend = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
zp = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
iend = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
jend = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
kend = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
ttend = np.ones((ia.size, (len(tinds) - 1) * nsteps)) * np.nan
t = np.zeros(((len(tinds) - 1) * nsteps + 1))
flag = np.zeros(
(ia.size), dtype=np.int) # initialize all exit flags for in the domain
# Initialize vertical stuff and fluxes
# Read initial field in - to 'new' variable since will be moved
# at the beginning of the time loop ahead
if is_string_like(z0): # isoslice case
ufnew, vfnew, dztnew, zrtnew, zwtnew = inout.readfields(tinds[0],
grid,
nc,
z0,
zpar,
zparuv=zparuv)
else: # 3d case
ufnew, vfnew, dztnew, zrtnew, zwtnew = inout.readfields(
tinds[0], grid, nc)
## Find zstart0 and ka
# The k indices and z grid ratios should be on a wflux vertical grid,
# which goes from 0 to km since the vertical velocities are defined
# at the vertical cell edges. A drifter's grid cell is vertically bounded
# above by the kth level and below by the (k-1)th level
if is_string_like(z0): # then doing a 2d isoslice
# there is only one vertical grid cell, but with two vertically-
# bounding edges, 0 and 1, so the initial ka value is 1 for all
# isoslice drifters.
ka = np.ones(ia.size)
# for s level isoslice, place drifters vertically at the center
# of the grid cell since that is where the u/v flux info is from.
# For a rho/temp/density isoslice, we treat it the same way, such
# that the u/v flux info taken at a specific rho/temp/density value
# is treated as being at the center of the grid cells vertically.
zstart0 = np.ones(ia.size) * 0.5
else: # 3d case
# Convert initial real space vertical locations to grid space
# first find indices of grid cells vertically
ka = np.ones(ia.size) * np.nan
zstart0 = np.ones(ia.size) * np.nan
if zpar == 'fromMSL':
for i in xrange(ia.size):
# pdb.set_trace()
ind = (grid['zwt0'][ia[i], ja[i], :] <= z0[i])
# check to make sure there is at least one true value, so the z0 is shallower than the seabed
if np.sum(ind):
ka[i] = find(
ind
)[-1] # find value that is just shallower than starting vertical position
# if the drifter starting vertical location is too deep for the x,y location, complain about it
else: # Maybe make this nan or something later
print 'drifter vertical starting location is too deep for its x,y location. Try again.'
if (z0[i] != grid['zwt0'][ia[i], ja[i], ka[i]]) and (
ka[i] != grid['km']): # check this
ka[i] = ka[i] + 1
# Then find the vertical relative position in the grid cell by adding on the bit of grid cell
zstart0[i] = ka[i] - abs(z0[i]-grid['zwt0'][ia[i],ja[i],ka[i]]) \
/abs(grid['zwt0'][ia[i],ja[i],ka[i]-1]-grid['zwt0'][ia[i],ja[i],ka[i]])
# elif zpar == 'fromZeta':
# for i in xrange(ia.size):
# pdb.set_trace()
# ind = (zwtnew[ia[i],ja[i],:]<=z0[i])
# ka[i] = find(ind)[-1] # find value that is just shallower than starting vertical position
# if (z0[i] != zwtnew[ia[i],ja[i],ka[i]]) and (ka[i] != grid['km']): # check this
# ka[i] = ka[i]+1
# # Then find the vertical relative position in the grid cell by adding on the bit of grid cell
# zstart0[i] = ka[i] - abs(z0[i]-zwtnew[ia[i],ja[i],ka[i]]) \
# /abs(zwtnew[ia[i],ja[i],ka[i]-1]-zwtnew[ia[i],ja[i],ka[i]])
# Find initial cell depths to concatenate to beginning of drifter tracks later
zsave = tools.interpolate3d(xstart0, ystart0, zstart0, zwtnew)
toc_initial = time.time()
# j = 0 # index for number of saved steps for drifters
tic_read = np.zeros(len(tinds))
toc_read = np.zeros(len(tinds))
tic_zinterp = np.zeros(len(tinds))
toc_zinterp = np.zeros(len(tinds))
tic_tracmass = np.zeros(len(tinds))
toc_tracmass = np.zeros(len(tinds))
# pdb.set_trace()
xr3 = grid['xr'].reshape(
(grid['xr'].shape[0], grid['xr'].shape[1], 1)).repeat(zwtnew.shape[2],
axis=2)
yr3 = grid['yr'].reshape(
(grid['yr'].shape[0], grid['yr'].shape[1], 1)).repeat(zwtnew.shape[2],
axis=2)
# Loop through model outputs. tinds is in proper order for moving forward
# or backward in time, I think.
for j, tind in enumerate(tinds[:-1]):
# pdb.set_trace()
# Move previous new time step to old time step info
ufold = ufnew
vfold = vfnew
dztold = dztnew
zrtold = zrtnew
zwtold = zwtnew
tic_read[j] = time.time()
# Read stuff in for next time loop
if is_string_like(z0): # isoslice case
ufnew, vfnew, dztnew, zrtnew, zwtnew = inout.readfields(
tinds[j + 1], grid, nc, z0, zpar, zparuv=zparuv)
else: # 3d case
ufnew, vfnew, dztnew, zrtnew, zwtnew = inout.readfields(
tinds[j + 1], grid, nc)
toc_read[j] = time.time()
# print "readfields run time:",toc_read-tic_read
print j
# pdb.set_trace()
# flux fields at starting time for this step
if j != 0:
xstart = xend[:, j * nsteps - 1]
ystart = yend[:, j * nsteps - 1]
zstart = zend[:, j * nsteps - 1]
# mask out drifters that have exited the domain
xstart = np.ma.masked_where(flag[:] == 1, xstart)
ystart = np.ma.masked_where(flag[:] == 1, ystart)
zstart = np.ma.masked_where(flag[:] == 1, zstart)
ind = (flag[:] == 0
) # indices where the drifters are still inside the domain
else: # first loop, j==0
xstart = xstart0
ystart = ystart0
zstart = zstart0
# TODO: Do a check to make sure all drifter starting locations are within domain
ind = (
flag[:] == 0
) # indices where the drifters are inside the domain to start
# Find drifter locations
# only send unmasked values to step
if not np.ma.compressed(xstart).any(
): # exit if all of the drifters have exited the domain
break
else:
# Combine times for arrays for input to tracmass
# from [ixjxk] to [ixjxkxt]
# Change ordering for these three arrays here instead of in readfields since
# concatenate does not seem to preserve ordering
uflux = np.asfortranarray(np.concatenate((ufold.reshape(np.append(ufold.shape,1)), \
ufnew.reshape(np.append(ufnew.shape,1))), \
axis=ufold.ndim))
vflux = np.asfortranarray(np.concatenate((vfold.reshape(np.append(vfold.shape,1)), \
vfnew.reshape(np.append(vfnew.shape,1))), \
axis=vfold.ndim))
dzt = np.asfortranarray(np.concatenate((dztold.reshape(np.append(dztold.shape,1)), \
dztnew.reshape(np.append(dztnew.shape,1))), \
axis=dztold.ndim))
# Change the horizontal indices from python to fortran indexing
# (vertical are zero-based in tracmass)
xstart, ystart = tools.convert_indices('py2f', xstart, ystart)
# km that is sent to tracmass is determined from uflux (see tracmass?)
# so it will be the correct value for whether we are doing the 3D
# or isoslice case.
# vec = np.arange(j*nsteps,j*nsteps+nsteps) # indices for storing new track locations
tic_tracmass[j] = time.time()
# pdb.set_trace()
if dostream: # calculate Lagrangian stream functions
xend[ind,j*nsteps:j*nsteps+nsteps],\
yend[ind,j*nsteps:j*nsteps+nsteps],\
zend[ind,j*nsteps:j*nsteps+nsteps], \
iend[ind,j*nsteps:j*nsteps+nsteps],\
jend[ind,j*nsteps:j*nsteps+nsteps],\
kend[ind,j*nsteps:j*nsteps+nsteps],\
flag[ind],\
ttend[ind,j*nsteps:j*nsteps+nsteps], U, V = \
tracmass.step(np.ma.compressed(xstart),\
np.ma.compressed(ystart),
np.ma.compressed(zstart),
tseas_use, uflux, vflux, ff, \
grid['kmt'].astype(int), \
dzt, grid['dxdy'], grid['dxv'], \
grid['dyu'], grid['h'], nsteps, \
ah, av, do3d, doturb, dostream, \
t0=T0[ind],
ut=U, vt=V)
else: # don't calculate Lagrangian stream functions
xend[ind,j*nsteps:j*nsteps+nsteps],\
yend[ind,j*nsteps:j*nsteps+nsteps],\
zend[ind,j*nsteps:j*nsteps+nsteps], \
iend[ind,j*nsteps:j*nsteps+nsteps],\
jend[ind,j*nsteps:j*nsteps+nsteps],\
kend[ind,j*nsteps:j*nsteps+nsteps],\
flag[ind],\
ttend[ind,j*nsteps:j*nsteps+nsteps], _, _ = \
tracmass.step(np.ma.compressed(xstart),\
np.ma.compressed(ystart),
np.ma.compressed(zstart),
tseas_use, uflux, vflux, ff, \
grid['kmt'].astype(int), \
dzt, grid['dxdy'], grid['dxv'], \
grid['dyu'], grid['h'], nsteps, \
ah, av, do3d, doturb, dostream)
toc_tracmass[j] = time.time()
# pdb.set_trace()
# Change the horizontal indices from python to fortran indexing
xend[ind,j*nsteps:j*nsteps+nsteps], \
yend[ind,j*nsteps:j*nsteps+nsteps] \
= tools.convert_indices('f2py', \
xend[ind,j*nsteps:j*nsteps+nsteps], \
yend[ind,j*nsteps:j*nsteps+nsteps])
# Calculate times for the output frequency
if ff == 1:
t[j * nsteps + 1:j * nsteps +
nsteps + 1] = t[j * nsteps] + np.linspace(
tseas_use / nsteps, tseas_use,
nsteps) # update time in seconds to match drifters
else:
t[j * nsteps + 1:j * nsteps +
nsteps + 1] = t[j * nsteps] - np.linspace(
tseas_use / nsteps, tseas_use,
nsteps) # update time in seconds to match drifters
# Skip calculating real z position if we are doing surface-only drifters anyway
if z0 != 's' and zpar != grid['km'] - 1:
tic_zinterp[j] = time.time()
# Calculate real z position
r = np.linspace(
1. / nsteps, 1, nsteps
) # linear time interpolation constant that is used in tracmass
for n in xrange(nsteps): # loop through time steps
# interpolate to a specific output time
# pdb.set_trace()
zwt = (1. - r[n]) * zwtold + r[n] * zwtnew
zp[ind,j*nsteps:j*nsteps+nsteps], dt = tools.interpolate3d(xend[ind,j*nsteps:j*nsteps+nsteps], \
yend[ind,j*nsteps:j*nsteps+nsteps], \
zend[ind,j*nsteps:j*nsteps+nsteps], \
zwt)
toc_zinterp[j] = time.time()
nc.close()
t = t + t0save # add back in base time in seconds
# pdb.set_trace()
# Add on to front location for first time step
xg = np.concatenate((xstart0.reshape(xstart0.size, 1), xend), axis=1)
yg = np.concatenate((ystart0.reshape(ystart0.size, 1), yend), axis=1)
# Concatenate zp with initial real space positions
zp = np.concatenate((zsave[0].reshape(zstart0.size, 1), zp), axis=1)
# Delaunay interpolation
# xp, yp, dt = tools.interpolate(xg,yg,grid,'d_ij2xy')
# lonp, latp, dt = tools.interpolation(xg,yg,grid,'d_ij2ll')
## map coordinates interpolation
# xp2, yp2, dt = tools.interpolate(xg,yg,grid,'m_ij2xy')
# tic = time.time()
lonp, latp, dt = tools.interpolate2d(xg,
yg,
grid,
'm_ij2ll',
mode='constant',
cval=np.nan)
# print '2d interp time=', time.time()-tic
# pdb.set_trace()
runtime = time.time() - tic_start
print "============================================="
print ""
print "Simulation name: ", name
print ""
print "============================================="
print "run time:\t\t\t", runtime
print "---------------------------------------------"
print "Time spent on:"
initialtime = toc_initial - tic_initial
print "\tInitial stuff: \t\t%4.2f (%4.2f%%)" % (initialtime,
(initialtime / runtime) *
100)
readtime = np.sum(toc_read - tic_read)
print "\tReading in fields: \t%4.2f (%4.2f%%)" % (readtime,
(readtime / runtime) *
100)
zinterptime = np.sum(toc_zinterp - tic_zinterp)
print "\tZ interpolation: \t%4.2f (%4.2f%%)" % (zinterptime,
(zinterptime / runtime) *
100)
tractime = np.sum(toc_tracmass - tic_tracmass)
print "\tTracmass: \t\t%4.2f (%4.2f%%)" % (tractime,
(tractime / runtime) * 100)
print "============================================="
# Save results to netcdf file
if dostream:
inout.savetracks(lonp, latp, zp, t, name, nsteps, ff, tseas_use, ah, av, \
do3d, doturb, loc, T0, U, V)
return lonp, latp, zp, t, grid, T0, U, V
else:
inout.savetracks(lonp, latp, zp, t, name, nsteps, ff, tseas_use, ah, av, \
do3d, doturb, loc)
return lonp, latp, zp, t, grid
