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