def get_fits(shots, dt=1):
if plt.is_string_like(shots) and ',' in shots:
shots = [[utc_ns(nls), utc_ns(nls) + int(dt * 1000000000)]
for nls in shots.split(',')]
i_data = [
dev.acq.getdata(sh, ichans, exceptions=exceptions) for sh in shots
]
vsweep_data = [dev.acq.getdata(sh, vsweep) for sh in shots]
ichan_names = np.array([ch.name for ch in i_data[0].channels])
gain_used = np.array(
[[i_dat.params[ch.name]['gain_used'] for ch in i_dat.channels]
for i_dat in i_data])
if gain_used.std(axis=0).sum() != 0:
print('inconsistent gain_used')
gain_used = gain_used[0]
npts = len(vsweep_data[0].timebase)
cm_part = np.array([np.sum(dat.signal, axis=1) for dat in i_data]).T / npts
v_sum = np.array([float(np.sum(vdat.signal))
for vdat in vsweep_data]) / npts
fits = [np.polyfit(v_sum, cm, 1, full=True) for cm in cm_part]
if len(fits[0][1]) == 0:
res = [None for fit in fits]
else:
res = np.array([
np.sqrt(fit[1][0] / len(shots)) / np.max(np.abs(cm_part))
for fit in fits
])
return (dict(fits=[fit[0] for fit in fits],
res=res,
gain_used=gain_used,
names=ichan_names))
def prepare_for_model_step(self, tind, nc, flag, xend, yend, zend, j, nsubstep, T0):
'''
Already in a step, get ready to actually do step
'''
xstart = xend[:,j*self.N]
ystart = yend[:,j*self.N]
zstart = zend[:,j*self.N]
# 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)
if T0 is not None:
T0 = np.ma.masked_where(flag[:]==1,T0)
# Move previous new time step to old time step info
self.uf[:,:,:,0] = self.uf[:,:,:,1].copy()
self.vf[:,:,:,0] = self.vf[:,:,:,1].copy()
self.dzt[:,:,:,0] = self.dzt[:,:,:,1].copy()
self.zrt[:,:,:,0] = self.zrt[:,:,:,1].copy()
self.zwt[:,:,:,0] = self.zwt[:,:,:,1].copy()
# Read stuff in for next time loop
if is_string_like(self.z0): # isoslice case
self.uf[:,:,:,1],self.vf[:,:,:,1],self.dzt[:,:,:,1],self.zrt[:,:,:,1],self.zwt[:,:,:,1] = tracpy.inout.readfields(tind, self.grid, nc, self.z0, self.zpar, zparuv=self.zparuv)
else: # 3d case
self.uf[:,:,:,1],self.vf[:,:,:,1],self.dzt[:,:,:,1],self.zrt[:,:,:,1],self.zwt[:,:,:,1] = tracpy.inout.readfields(tind, self.grid, nc)
# Find the fluxes of the immediately bounding range for the desired time step, which can be less than 1 model output
# SHOULD THIS BE PART OF SELF TOO? Leave uf and vf as is, though, because they may be used for interpolating the
# input fluxes for substeps.
ufsub = np.ones(self.uf.shape)*np.nan
vfsub = np.ones(self.vf.shape)*np.nan
# for earlier bounding flux info
rp = nsubstep/self.nsubsteps # weighting for later time step
rm = 1 - rp # timing for earlier time step
ufsub[:,:,:,0] = rm*self.uf[:,:,:,0] + rp*self.uf[:,:,:,1]
vfsub[:,:,:,0] = rm*self.vf[:,:,:,0] + rp*self.vf[:,:,:,1]
# for later bounding flux info
rp = (nsubstep+1)/self.nsubsteps # weighting for later time step
rm = 1 - rp # timing for earlier time step
ufsub[:,:,:,1] = rm*self.uf[:,:,:,0] + rp*self.uf[:,:,:,1]
vfsub[:,:,:,1] = rm*self.vf[:,:,:,0] + rp*self.vf[:,:,:,1]
# Change the horizontal indices from python to fortran indexing
# (vertical are zero-based in tracmass)
xstart, ystart = tracpy.tools.convert_indices('py2f',xstart,ystart)
return xstart, ystart, zstart, ufsub, vfsub, T0
def prepare_for_model_run(self, date, lon0, lat0):
'''
Get everything ready so that we can get to the simulation.
'''
# Convert date to number
date = netCDF.date2num(date, self.time_units)
# Figure out what files will be used for this tracking
nc, tinds = tracpy.inout.setupROMSfiles(self.currents_filename, date, self.ff, self.tout, tstride=self.tstride)
# Read in grid parameters into dictionary, grid, if haven't already
if self.grid is None:
self._readgrid()
# If dostream==1, do transport calculations and initialize to an empty array
if self.U is None and self.dostream:
self.U = np.ma.zeros(grid['xu'].shape, order='F')
self.V = np.ma.zeros(grid['xv'].shape, order='F')
# Interpolate to get starting positions in grid space
if self.usespherical: # convert from assumed input lon/lat coord locations to grid space
xstart0, ystart0, _ = tracpy.tools.interpolate2d(lon0, lat0, self.grid, 'd_ll2ij')
else: # assume input seed locations are in projected/idealized space and change to index space
xstart0, ystart0, _ = tracpy.tools.interpolate2d(lon0, lat0, self.grid, 'd_xy2ij')
# Do z a little lower down
# Initialize seed locations
ia = np.ceil(xstart0)
ja = np.ceil(ystart0)
# 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)*self.N+1))*np.nan
yend = np.ones((ia.size,(len(tinds)-1)*self.N+1))*np.nan
zend = np.ones((ia.size,(len(tinds)-1)*self.N+1))*np.nan
zp = np.ones((ia.size,(len(tinds)-1)*self.N+1))*np.nan
ttend = np.zeros((ia.size,(len(tinds)-1)*self.N+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
lx = self.grid['xr'].shape[0]
ly = self.grid['xr'].shape[1]
lk = self.grid['sc_r'].size
if is_string_like(self.z0): # isoslice case
# Now that we have the grid, initialize the info for the two bounding model
# steps using the grid size
self.uf = np.asfortranarray(np.ones((lx-1, ly, lk-1, 2)))*np.nan
self.vf = np.asfortranarray(np.ones((lx, ly-1, lk-1, 2)))*np.nan
self.dzt = np.asfortranarray(np.ones((lx, ly, lk-1, 2)))*np.nan
self.zrt = np.asfortranarray(np.ones((lx, ly, lk-1, 2)))*np.nan
self.zwt = np.asfortranarray(np.ones((lx, ly, lk, 2)))*np.nan
self.uf[:,:,:,1], self.vf[:,:,:,1], \
self.dzt[:,:,:,1], self.zrt[:,:,:,1], \
self.zwt[:,:,:,1] = tracpy.inout.readfields(tinds[0], self.grid, nc, self.z0, self.zpar, zparuv=self.zparuv)
else: # 3d case
# Now that we have the grid, initialize the info for the two bounding model
# steps using the grid size
self.uf = np.asfortranarray(np.ones((lx-1, ly, lk-1, 2)))*np.nan
self.vf = np.asfortranarray(np.ones((lx, ly-1, lk-1, 2)))*np.nan
self.dzt = np.asfortranarray(np.ones((lx, ly, lk-1, 2)))*np.nan
self.zrt = np.asfortranarray(np.ones((lx, ly, lk-1, 2)))*np.nan
self.zwt = np.asfortranarray(np.ones((lx, ly, lk, 2)))*np.nan
self.uf[:,:,:,1], self.vf[:,:,:,1], \
self.dzt[:,:,:,1], self.zrt[:,:,:,1], \
self.zwt[:,:,:,1] = tracpy.inout.readfields(tinds[0], self.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(self.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 self.zpar == 'fromMSL':
print 'zpar==''fromMSL'' not implemented yet...'
# for i in xrange(ia.size):
# # pdb.set_trace()
# ind = (self.grid['zwt0'][ia[i],ja[i],:]<=self.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 (self.z0[i] != self.grid['zwt0'][ia[i],ja[i],ka[i]]) and (ka[i] != self.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(self.z0[i]-self.grid['zwt0'][ia[i],ja[i],ka[i]]) \
# /abs(self.grid['zwt0'][ia[i],ja[i],ka[i]-1]-self.grid['zwt0'][ia[i],ja[i],ka[i]])
elif self.zpar == 'fromZeta':
# In this case, the starting z values of the drifters are found in grid space as z0 below
# the z surface for each drifter
for i in xrange(ia.size):
ind = (self.zwt[ia[i],ja[i],:,1]<=self.z0[i])
ka[i] = find(ind)[-1] # find value that is just shallower than starting vertical position
if (self.z0[i] != self.zwt[ia[i],ja[i],ka[i],1]) and (ka[i] != self.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(self.z0[i]-self.zwt[ia[i],ja[i],ka[i],1]) \
/abs(self.zwt[ia[i],ja[i],ka[i]-1,1]-self.zwt[ia[i],ja[i],ka[i],1])
# Find initial cell depths to concatenate to beginning of drifter tracks later
zsave = tracpy.tools.interpolate3d(xstart0, ystart0, zstart0, self.zwt[:,:,:,1])
# Initialize x,y,z with initial seeded positions
xend[:,0] = xstart0
yend[:,0] = ystart0
zend[:,0] = zstart0
return tinds, nc, t0save, xend, yend, zend, zp, ttend, flag
def extract(self, dictionary=False, varnames=None, inds=None, limit=None, strict=0, masked=1, debug=0):
""" extract the listed variables into the dictionary (local by default)
selecting those at indices <inds> (all be default
variables must be strings, either an array, or separated by commas
if the dictionary is False, return them in a tuple instead
Note: returning a list requires you to make the order consistent
if varnames is None - extract all.
e.g. if da is a dictionary or arrays
da = DA('mydata.npz')
da.extract('shot,beta')
plot(shot,beta)
(shot,beta,n_e) = da.extract(['shot','beta','n_e'], \
inds=np.where(da['beta']>3)[0])
# makes a tuple of 3 arrays of data for high beta.
Note syntax of where()! It is evaluted in your variable space.
to extract one var, need trailing "," (tuple notation) e.g.
(allbeta,) = D54.extract('beta',locals())
which can be abbreviated to
allbeta, = D54.extract('beta',locals())
"""
start_mem = report_mem(msg="extract")
if debug == 0:
debug = self.debug
if varnames is None:
varnames = self.da.keys() # all variables
if plt.is_string_like(varnames):
varlist = varnames.split(",")
else:
varlist = varnames
val_tuple = ()
if inds is None:
inds = np.arange(self.len)
if len(np.shape(inds)) == 2:
inds = inds[0] # trick to catch when you forget [0] on where
if limit != None and len(inds) > abs(limit):
if limit < 0:
print("repeatably"),
np.random.seed(0) # if positive, should be random
# negative will repeat the sequence
else:
print("randomly"),
print("decimating from sample of {n} and".format(n=len(inds))),
ir = np.where(np.random.random(len(inds)) < float(abs(limit)) / len(inds))[0]
inds = inds[ir]
if len(inds) < 500:
print("*** {n} is a very small number to extract????".format(n=len(inds)))
if self.verbose > 0:
print("extracting a sample of {n} ".format(n=len(inds)))
for k in varlist:
if k in self.da: # be careful that self keys is up to date!
# this is normally OK if you use self.update
debug_(debug, key="extract")
# used to refer to da[k] twice - two reads if npz
# if masking is enabled and it is a legitimate key for masking, use the masked version
if masked and hasattr(self, "masked") and k in self.masked.keys():
print("extracting masked values for {k} - use masked=0 to get raw values".format(k=k))
dak = self.masked[k]
else:
dak = self.da[k] # We know we want it - let's
# hope space is not wasted
if hasattr(dak, "keys"): # used to be self.da[k]
allvals = dak
else:
allvals = np.array(dak)
if len(np.shape(allvals)) == 0:
sel_vals = allvals
else:
if (len(np.shape(allvals)) > 0) and (len(allvals) < np.max(inds)):
print("{k} does not match the size of the other arrays - extracting all".format(k=k))
sel_vals = allvals
else:
sel_vals = allvals[inds]
if dictionary is False:
val_tuple += (sel_vals,)
else:
dictionary.update({k: sel_vals})
else:
print("variable {k} not found in {ks}".format(k=k, ks=np.sort(self.da.keys())))
report_mem(start_mem)
if dictionary is False:
return val_tuple
