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