def step(self, xstart, ystart, zstart, ufsub, vfsub):
'''
Take some number of steps between a start and end time.
FIGURE OUT HOW TO KEEP TRACK OF TIME FOR EACH SET OF LINES
:param tind: Time index to use for stepping
FILL IN
'''
# Figure out where in time we are
xend, yend, zend, flag,\
ttend, U, V = \
tracmass.step(np.ma.compressed(xstart),
np.ma.compressed(ystart),
np.ma.compressed(zstart),
self.tseas_use, ufsub, vfsub, self.ff,
self.grid['kmt'].astype(int),
self.dzt, self.grid['dxdy'], self.grid['dxv'],
self.grid['dyu'], self.grid['h'], self.nsteps,
self.ah, self.av, self.do3d, self.doturb,
self.doperiodic, self.dostream, self.N,
t0=self.T0,
ut=self.U, vt=self.V)
# return the new positions or the delta lat/lon
return xend, yend, zend, flag, ttend, U, V
def step(self, xstart, ystart, zstart, ufsub, vfsub, T0, U, V):
"""
Take some number of steps between a start and end time.
FIGURE OUT HOW TO KEEP TRACK OF TIME FOR EACH SET OF LINES
FILL IN
Transpose arrays when sent to Fortran.
"""
if T0 is not None:
xend, yend, zend, flag,\
ttend, U, V = tracmass.step(np.ma.compressed(xstart),
np.ma.compressed(ystart),
np.ma.compressed(zstart),
self.tseas_use, ufsub.T, vfsub.T,
self.ff,
self.grid.kmt.astype(int).T,
self.dzt.T, self.grid.dxdy.T,
self.grid.dxv.T,
self.grid.dyu.T, self.grid.h.T,
self.nsteps, self.ah, self.av,
self.do3d, self.doturb,
self.doperiodic, self.dostream,
self.N, t0=np.ma.compressed(T0),
ut=U.T, vt=V.T)
else:
xend, yend, zend, flag,\
ttend, U, V = tracmass.step(np.ma.compressed(xstart),
np.ma.compressed(ystart),
np.ma.compressed(zstart),
self.tseas_use, ufsub.T, vfsub.T,
self.ff,
self.grid.kmt.astype(int).T,
self.dzt.T, self.grid.dxdy.T,
self.grid.dxv.T,
self.grid.dyu.T, self.grid.h.T,
self.nsteps, self.ah, self.av,
self.do3d, self.doturb,
self.doperiodic, self.dostream,
self.N)
# return the new positions or the delta lat/lon
return xend, yend, zend, flag, ttend, U, V
def step(self, xstart, ystart, zstart, ufsub, vfsub, T0, U, V):
"""
Take some number of steps between a start and end time.
FIGURE OUT HOW TO KEEP TRACK OF TIME FOR EACH SET OF LINES
FILL IN
Transpose arrays when sent to Fortran.
"""
if T0 is not None:
xend, yend, zend, flag,\
ttend, U, V = tracmass.step(np.ma.compressed(xstart),
np.ma.compressed(ystart),
np.ma.compressed(zstart),
self.tseas_use, ufsub.T, vfsub.T,
self.ff,
self.grid.kmt.astype(int).T,
self.dzt.T, self.grid.dxdy.T,
self.grid.dxv.T,
self.grid.dyu.T, self.grid.h.T,
self.nsteps, self.ah, self.av,
self.do3d, self.doturb,
self.doperiodic, self.dostream,
self.N, t0=np.ma.compressed(T0),
ut=U.T, vt=V.T)
else:
xend, yend, zend, flag,\
ttend, U, V = tracmass.step(np.ma.compressed(xstart),
np.ma.compressed(ystart),
np.ma.compressed(zstart),
self.tseas_use, ufsub.T, vfsub.T,
self.ff,
self.grid.kmt.astype(int).T,
self.dzt.T, self.grid.dxdy.T,
self.grid.dxv.T,
self.grid.dyu.T, self.grid.h.T,
self.nsteps, self.ah, self.av,
self.do3d, self.doturb,
self.doperiodic, self.dostream,
self.N)
# return the new positions or the delta lat/lon
return xend, yend, zend, flag, ttend, U, V
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
ind = (flag[:] == 0) # indices where the drifters are inside the domain to start # pdb.set_trace() # Linearly interpolate uflux and vflux to step between model outputs, field at ending time for this step ufluxinterp[:,:,:,1] = uflux[:,:,:,0] + (uflux[:,:,:,1]-uflux[:,:,:,0])*((i+1.)/nsteps) vfluxinterp[:,:,:,1] = vflux[:,:,:,0] + (vflux[:,:,:,1]-vflux[:,:,:,0])*((i+1.)/nsteps) # hsinterp[:,:,1] = hs[:,:,0] + (hs[:,:,1]-hs[:,:,0])*((i+1.)/nsteps) dxyzinterp[:,:,:,1] = dxyz[:,:,:,0] + (dxyz[:,:,:,1]-dxyz[:,:,:,0])*((i+1.)/nsteps) # Find drifter locations # pdb.set_trace() # 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: xend[j,ind],yend[j,ind],zend[j,ind],iend[j,ind],jend[j,ind],kend[j,ind],flag[ind] = \ tracmass.step(np.ma.compressed(xstart),np.ma.compressed(ystart), np.ma.compressed(zstart),0.,np.ma.compressed(ia),np.ma.compressed(ja), np.ma.compressed(ka),tseas/float(nsteps),ufluxinterp, vfluxinterp,ff,dxyzinterp)#dz.data,dxdy) # if np.sum(flag)>0: # pdb.set_trace() t0[j+1] = t0[j] + tseas/float(nsteps) # update time in seconds to match drifters # if i == 8: # pdb.set_trace() j = j + 1 nc.close() grid.close() t0 = t0 + t0save # add back in base time in seconds xg=np.concatenate((xstart0,xend)) yg=np.concatenate((ystart0,yend))
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
