def getTime(self,timeinfo): """ Get the particle time step info """ self.dt = timeinfo[2] self.time_track = othertime.TimeVector(timeinfo[0],timeinfo[1],timeinfo[2],timeformat ='%Y%m%d.%H%M%S') self.time_track_sec = othertime.SecondsSince(self.time_track) self.time_sec = othertime.SecondsSince(self.time) self.time_index = -9999
def __init__(self, xin, yin, zin, tin, xout, yout, zout, tout, mask=None, **kwargs): """ Construct the interpolation components **kwargs are passed straight to interpXYZ """ self.is4D = True if zin == None: self.is4D = False self.nz = 1 else: self.zin = zin self.zout = zout self.nz = zin.shape[0] # Create a 3D mask self.szxy = xin.shape if mask == None: self.mask = np.zeros((self.nz, ) + self.szxy, np.bool) else: self.mask = mask # Horizontal interpolation for each layer self._Fxy = [] for kk in range(self.nz): if self.is4D: mask = self.mask[kk, ...] else: mask = self.mask xyin = np.vstack([xin[~mask].ravel(), yin[~mask].ravel()]).T xyout = np.vstack([xout.ravel(), yout.ravel()]).T self.nxy = xyout.shape[0] self._Fxy.append(interpXYZ(xyin, xyout, **kwargs)) # Just store the other coordinates for now # Convert time to floats self.tin = othertime.SecondsSince(tin) self.tout = othertime.SecondsSince(tout) self.nt = tin.shape[0]
def window_index_time(t, windowsize, overlap): """ Determines the indices for window start and end points of a time vector The window does not need to be evenly spaced Inputs: t - list or array of datetime objects windowsize - length of the window [seconds] overlap - number of overlap points [seconds] Returns: pt1,pt2 the start and end indices of each window """ tsec = othertime.SecondsSince(t) t1 = tsec[0] t2 = t1 + windowsize pt1 = [0] pt2 = [np.searchsorted(tsec, t2)] while t2 < tsec[-1]: t1 = t2 - overlap t2 = t1 + windowsize pt1.append(np.searchsorted(tsec, t1)) pt2.append(np.searchsorted(tsec, t2)) return pt1, pt2
def __init__(self, t, y, **kwargs): self.__dict__.update(kwargs) self.t = t # independent variable (t,x, etc) self.y = y # dependent variable self.shape = self.y.shape self.ndim = len(self.shape) self.tsec = othertime.SecondsSince(self.t, basetime=self.basetime) self.ny = np.size(self.y) # make sure the original data is a masked array if not isinstance(self.y, np.ma.core.MaskedArray): mask = ~np.isfinite(self.y) self.y = np.ma.MaskedArray(self.y, mask=mask) self._checkDT() self.Nt = self.t.shape[0] # Make sure that time is the last dimension if self.y.shape[-1] != self.Nt: self.y = self.y.T
def tide_pred(modfile,lon,lat,time,z=None,conlist=None): """ Performs a tidal prediction at all points in [lon,lat] at times in vector [time] """ # Read and interpolate the constituents u_re, u_im, v_re, v_im, h_re, h_im, omega, conlist = extract_HC(modfile,lon,lat,z=z,conlist=conlist) # Initialise the output arrays sz = lon.shape nx = np.prod(sz) nt = time.shape[0] ncon = omega.shape[0] h_re = h_re.reshape((ncon,nx)) h_im = h_im.reshape((ncon,nx)) u_re = u_re.reshape((ncon,nx)) u_im = u_im.reshape((ncon,nx)) v_re = v_re.reshape((ncon,nx)) v_im = v_im.reshape((ncon,nx)) # Calculate nodal correction to amps and phases #baseyear = time[0].year #t1992 = othertime.SecondsSince(datetime(baseyear,1,1),basetime=datetime(1992,1,1))/86400.0 t1992 = othertime.SecondsSince(time[0],basetime=datetime(1992,1,1))/86400.0 pu,pf,v0u = nodal(t1992+48622.0,conlist) # Calculate the time series tsec = othertime.SecondsSince(time,basetime=datetime(1992,1,1)) # Needs to be referenced to 1992 h=np.zeros((nt,nx)) u=np.zeros((nt,nx)) v=np.zeros((nt,nx)) for nn,om in enumerate(omega): for ii in range(0,nx): h[:,ii] += pf[nn]*h_re[nn,ii] * np.cos(om*tsec + v0u[nn] + pu[nn]) - \ pf[nn]*h_im[nn,ii] * np.sin(om*tsec + v0u[nn] + pu[nn]) u[:,ii] += pf[nn]*u_re[nn,ii] * np.cos(om*tsec + v0u[nn] + pu[nn]) - \ pf[nn]*u_im[nn,ii] * np.sin(om*tsec + v0u[nn] + pu[nn]) v[:,ii] += pf[nn]*v_re[nn,ii] * np.cos(om*tsec + v0u[nn] + pu[nn]) - \ pf[nn]*v_im[nn,ii] * np.sin(om*tsec + v0u[nn] + pu[nn]) szo = (nt,)+sz return h.reshape(szo), u.reshape(szo), v.reshape(szo)
def filter2nc(self, outfile, tstart, tend, substep=12, varlist=None, **kwargs): """ Filters the variables in the list, varlist, and outputs the results to netcdf """ self.__dict__.update(kwargs) if tstart == -1: self.tstep = np.arange(0, self.Nt, 1) else: self.tstep = self.getTstep(tstart, tend) if varlist == None: varlist = ['eta', 'uc', 'vc', 'w'] # Create the output file self.writeNC(outfile) # Create the output variables for vv in varlist: print('Creating variable: %s' % vv) self.create_nc_var(outfile, vv, ugrid[vv]['dimensions'], ugrid[vv]['attributes'],\ dtype=ugrid[vv]['dtype'],zlib=ugrid[vv]['zlib'],complevel=ugrid[vv]['complevel'],fill_value=ugrid[vv]['fill_value']) self.create_nc_var(outfile, 'time', ugrid['time']['dimensions'], ugrid['time']['attributes']) # Loop through and filter each variable (do layer by layer on 3D variables for sake of memory) nc = Dataset(outfile, 'a') # Create the time variable first nctime = othertime.SecondsSince( self.time[self.tstep[0]:self.tstep[-1]:substep]) nc.variables['time'][:] = nctime for vv in varlist: print('Filtering variable: %s' % vv) if len(ugrid[vv]['dimensions']) == 2: dataf = self.__call__(tstart, tend, varname=vv) nc.variables[vv][:] = dataf[::substep, :].copy() elif len(ugrid[vv]['dimensions']) == 3: for kk in range(0, self.Nkmax): print(' layer: %d' % kk) self.klayer = [kk] dataf = self.__call__(tstart, tend, varname=vv) nc.variables[vv][:, kk, :] = dataf[::substep, :].copy() nc.close() print('#####\nComplete - Filtered data written to: \n%s \n#####' % outfile)
def interp(self, timein, method='linear', timeformat='%Y%m%d.%H%M%S', axis=-1): """ Interpolate the data onto an equally spaced vector timein is either: (3x1 tuple) - (tstart,tend,dt) tstart and tend - string with format 'yyyymmdd.HHMMSS' or datetime vector method - method passed to interp1d - use 'nearest' to preserve masking in gap regions """ # Create the time vector try: tstart = timein[0] tend = timein[1] dt = timein[2] tnew = othertime.TimeVector(tstart, tend, dt, timeformat=timeformat) except: tnew = timein dt = (tnew[1] - tnew[0]).total_seconds() if method == 'nearest': # Nearest neighbour doesn't use interp1d to preserve mask self._evenly_dist_data(dt) tnew, output = self.subset(tnew[0], tnew[-1]) else: t = othertime.SecondsSince(tnew, basetime=self.basetime) # Don't include nan points if self.ndim > 1: # Interpolate multidimensional arrays without a mask F = interpolate.interp1d(self.tsec,self.y,kind=method,axis=axis,\ bounds_error=False,fill_value=0) output = F(t) else: #mask = np.isnan(self.y) == False mask = ~self.y.mask F = interpolate.interp1d(self.tsec[mask],self.y[mask],kind=method,axis=axis,\ bounds_error=False,fill_value=0) output = F(t) return tnew, output
def tide_pred_old(modfile,lon,lat,time,z=None,conlist=None): """ ### UNUSED ### Performs a tidal prediction at all points in [lon,lat] at times in vector [time] """ # Read and interpolate the constituents u_re, u_im, v_re, v_im, h_re, h_im, omega, conlist = extract_HC(modfile,lon,lat,z=z,conlist=conlist) # Initialise the output arrays sz = lon.shape nx = np.prod(sz) nt = time.shape[0] ncon = omega.shape[0] h_re = h_re.reshape((ncon,nx)) h_im = h_im.reshape((ncon,nx)) u_re = u_re.reshape((ncon,nx)) u_im = u_im.reshape((ncon,nx)) v_re = v_re.reshape((ncon,nx)) v_im = v_im.reshape((ncon,nx)) # Nodal correction to amps and phases here... baseyear = time[0].year amp, phase = cart2pol(h_re, h_im) amp,phase = nodal_correction(baseyear,conlist, amp, phase) h_re, h_im = pol2cart(amp, phase) amp, phase = cart2pol(u_re, u_im) amp, phase = nodal_correction(baseyear,conlist, amp, phase) u_re, u_im = pol2cart(amp, phase) amp, phase = cart2pol(v_re, v_im) amp, phase = nodal_correction(baseyear,conlist, amp, phase) v_re, v_im = pol2cart(amp, phase) # Calculate the time series tsec = othertime.SecondsSince(time,basetime=datetime(baseyear,1,1)) h=np.zeros((nt,nx)) u=np.zeros((nt,nx)) v=np.zeros((nt,nx)) for nn,om in enumerate(omega): for ii in range(0,nx): h[:,ii] += h_re[nn,ii] * np.cos(om*tsec) + h_im[nn,ii] * np.sin(om*tsec) u[:,ii] += u_re[nn,ii] * np.cos(om*tsec) + u_im[nn,ii] * np.sin(om*tsec) v[:,ii] += v_re[nn,ii] * np.cos(om*tsec) + v_im[nn,ii] * np.sin(om*tsec) szo = (nt,)+sz return h.reshape(szo), u.reshape(szo), v.reshape(szo)
def __init__(self, data, dtime, **kwargs): """ Time series w/ specific IO methods """ self.__dict__.update(kwargs) TimeSeries.__init__(self, data, index=dtime) #super(ObsTimeSeries,self).__init__(data,index=dtime) # Time coordinates self.nt = self.index.shape self.tsec = othertime.SecondsSince(self.index,\ basetime = pd.datetime(self.baseyear,1,1))
def __init__(self, x, y, z, timeinfo, tformat='%Y%m%d.%H%M%S', **kwargs): metfile.__init__(self, mode='create') self.x = x self.y = y self.z = z self.time =\ otime.TimeVector(timeinfo[0],timeinfo[1],timeinfo[2],timeformat=tformat) self.nctime = otime.SecondsSince(self.time) self.varnames = [ 'Uwind', 'Vwind', 'Tair', 'Pair', 'RH', 'rain', 'cloud' ] # Update all of the metdata objects for vv in self.varnames: self.updateMetData(self[vv])
def __call__(self,tstart,tend,varnames=['eta','uc','vc']): """ Actually does the harmonic calculation for the model time steps in tsteps (or at least calls the class that does the calculation) Set tstart = -1 to do all steps """ if tstart == -1: self.tstep=np.arange(0,self.Nt,1) else: self.tstep=self.getTstep(tstart,tend) time = othertime.SecondsSince(self.time[self.tstep]) self.varnames=varnames self._prepDict(varnames) for vv in varnames: if vv in ['ubar','vbar']: ndim=2 if vv=='ubar': self.variable='uc' if vv=='vbar': self.variable='vc' else: ndim = self._returnDim(vv) self.variable=vv if ndim == 2 or self.Nkmax==1: print 'Loading data from %s...'%vv if vv in ['ubar','vbar']: data=self.loadDataBar() else: data=self.loadData() print 'Performing harmonic fit on variable, %s...'%(self.variable) self.Amp[vv], self.Phs[vv], self.Mean[vv] = harmonic_fit(time,data,self.frq,phsbase=self.reftime) elif ndim == 3: for k in range(self.Nkmax): self.klayer=[k] print 'Loading data...' data = self.loadData() print 'Performing harmonic fit on variable, %s, layer = %d of %d...'%(self.variable,self.klayer[0],self.Nkmax) self.Amp[vv][:,k,:], self.Phs[vv][:,k,:], self.Mean[vv][k,:] = harmonic_fit(time,data,self.frq,phsbase=self.reftime)
def calc_agebin(binfile, ncfile, polyfile, ntout): """ Calculate the from a binary file and save to netcdf """ # Load the polygon from a shapefile xypoly, field = readShpPoly(polyfile) # Load the binary file object PTM = PtmBin(binfile) # Count the number of particles from the first time step time, pdata = PTM.read_timestep() N = pdata.shape[0] tsec = othertime.SecondsSince(PTM.time) dt = tsec[1] - tsec[0] # Initialize the age particle object Age = ParticleAge(xypoly[0], N) # Loop through outctr = ntout ncctr = 0 for tt in range(PTM.nt - 1): # Read the current time step time, pdata = PTM.read_timestep(ts=tt) # Update the age variable Age.update_age(pdata['x'][:, 0], pdata['x'][:, 1], pdata['x'][:, 2], dt) # Write to netcdf if outctr == ntout: Age.write_nc(time, ncctr, ncfile=ncfile) ncctr += 1 outctr = 0 outctr += 1 print 'Done.'
def write_nc(self, time, tstep, ncfile=None): """ Writes the particle locations at the output time step, 'tstep' """ if self.verbose: print '\tWriting netcdf output at tstep: %d...' % tstep # Check if the file is opened if not self.__dict__.has_key('_nc'): if ncfile == None: raise Exception, 'must set "ncfile" on call to write_nc().' else: self.create_nc(ncfile) t = othertime.SecondsSince(time, basetime=self.basetime) self._nc.variables['tp'][tstep] = t self._nc.variables['xp'][:, tstep] = self.X self._nc.variables['yp'][:, tstep] = self.Y self._nc.variables['zp'][:, tstep] = self.Z if self.has_age: self._nc.variables['age'][:, tstep] = self.age self._nc.variables['agemax'][:, tstep] = self.agemax
def ProfilePlot(t,y,z,scale=86400,\ axis=0,color=[0.5,0.5,0.5],xlim=None,units='m/s',scalebar=1.0): """ Plot a series of vertical profiles as a time series scale - Sets 1 unit = scale (seconds) See this page on formatting: http://matplotlib.org/examples/pylab_examples/date_index_formatter.html """ from matplotlib import collections from matplotlib.ticker import Formatter class MyFormatter(Formatter): def __init__(self, dates, fmt='%b %d %Y'): self.fmt = fmt self.dates = dates def __call__(self, x, pos=0): 'Return the label for time x s' return datetime.strftime( datetime(1990, 1, 1) + timedelta(seconds=x), self.fmt) tsec = othertime.SecondsSince(t) formatter = MyFormatter(tsec) y = np.swapaxes(y, 0, axis) lines = [] line2 = [] for ii, tt in enumerate(tsec): #xplot = set_scale(y[:,ii],tt) xplot = tt + y[:, ii] * scale lines.append(np.array((xplot, z)).T) line2.append(np.array([[tt, tt], [z[0], z[-1]]]).T) LC1 = collections.LineCollection(lines, colors=color, linewidths=1.5) LC2 = collections.LineCollection(line2, colors='k', linestyles='dashed') # Zero axis ax = plt.gca() ax.add_collection(LC1) ax.add_collection(LC2) ax.set_ylim((z.min(), z.max())) ax.xaxis.set_major_formatter(formatter) if xlim == None: xlim = (tsec[0] - scale / 2, tsec[-1] + scale / 2) else: xlim = othertime.SecondsSince(xlim) ax.set_xlim(xlim) plt.xticks(rotation=17) ### # Add a scale bar ### # Compute the scale bar size in dimensionless units if not scalebar == None: xscale = scalebar * scale / (xlim[-1] - xlim[0]) x0 = 0.1 y0 = 0.8 dy = 0.02 ax.add_line( Line2D([x0, x0 + xscale], [y0, y0], linewidth=0.5, color='k', transform=ax.transAxes)) #Little caps ax.add_line( Line2D([x0, x0], [y0 - dy, y0 + dy], linewidth=0.5, color='k', transform=ax.transAxes)) ax.add_line( Line2D([x0 + xscale, x0 + xscale], [y0 - dy, y0 + dy], linewidth=0.5, color='k', transform=ax.transAxes)) plt.text(x0,y0+0.05,'Scale %3.1f %s'%(scalebar,units),\ transform=ax.transAxes) return ax
def runmpi(ncfile, outfile, tstart, tend, dt, dtout, x, y, z, agepoly=None, method='nearest', is3D=False): # Generate a list of tuples for the time info timevec = othertime.TimeVector(tstart, tend, dtout, timeformat='%Y%m%d.%H%M%S') timevec_sec = othertime.SecondsSince(timevec) timeinfos = [] for ii in range(timevec.shape[0] - 1): if ii == 0: timestart = datetime.strftime(timevec[ii], '%Y%m%d.%H%M%S') else: timestart = datetime.strftime(timevec[ii] + timedelta(seconds=dt), '%Y%m%d.%H%M%S') timeend = datetime.strftime(timevec[ii + 1], '%Y%m%d.%H%M%S') timeinfos.append((timestart, timeend, dt)) # Initialise the particle tracking object print('Initialising the particle tracking object on processor: %d...' % (comm.rank)) sun = SunTrack(ncfile, interp_method='mesh', interp_meshmethod=method, is3D=is3D) # Initialise the age values if not agepoly == None: calcage = True else: calcage = False age = None agemax = None # On rank = 0 only if comm.rank == 0: n = int(x.shape[0]) # Check if the number of processes divides eveny into the array length rem = np.mod(n, size) if rem == 0: #length of each process's portion of the original vector local_n = np.array([n / size]) else: print('Padding array with extra values...') nextra = size - rem xpad = np.zeros((nextra, )) x = np.hstack((x, xpad)) y = np.hstack((y, xpad)) z = np.hstack((z, xpad)) n = n + nextra local_n = np.array([n / size]) print('Size of original vector = %d\nSize of split vector = %d' % (n, local_n)) if calcage: age = np.zeros_like(x) agemax = np.zeros_like(x) # Initialise the output netcdf file sun.initParticleNC(outfile, n, age=calcage) else: x = None y = None z = None local_n = np.array([0]) if calcage: age = None agemax = None comm.Barrier() t_start = MPI.Wtime() # Broadcast the particle tracking object everywhere #sun = comm.bcast(sun, root=0) # !! Doesn't work for this object !! # Scatter the x,y,z locations up amongst all processors #communicate local array size to all processes comm.Bcast(local_n, root=0) #initialize the local particle arrays as numpy arrays x_local = np.zeros(local_n) y_local = np.zeros(local_n) z_local = np.zeros(local_n) if calcage: age_local = np.zeros(local_n) agemax_local = np.zeros(local_n) #divide up vectors comm.Scatter(x, x_local, root=0) comm.Scatter(y, y_local, root=0) comm.Scatter(z, z_local, root=0) if calcage: comm.Scatter(age, age_local, root=0) comm.Scatter(agemax, agemax_local, root=0) else: age_local = None agemax_local = None ### # Testing plot of particle positions ### #if comm.rank ==0: # import matplotlib.pyplot as plt # plt.figure # plt.plot(x_local,y_local,'.') # plt.show() #comm.Barrier() ### # Write out the initial location to netcdf if comm.rank == 0: sun.writeParticleNC(outfile, x, y, z, timevec_sec[0], 0, age=age, agemax=agemax) comm.Barrier() ### # ... Call the particle tracking module on each processor for ii, timeinfo in enumerate(timeinfos): if comm.rank == 0: sun(x_local, y_local, z_local, timeinfo, agepoly=agepoly, age=age_local, agemax=agemax_local, verbose=True) else: sun(x_local, y_local, z_local, timeinfo, agepoly=agepoly, age=age_local, agemax=agemax_local, verbose=False) # Send the particles back to their main array comm.Barrier() comm.Gather(sun.particles['X'], x, root=0) comm.Gather(sun.particles['Y'], y, root=0) comm.Gather(sun.particles['Z'], z, root=0) if calcage: comm.Gather(sun.particles['age'], age, root=0) comm.Gather(sun.particles['agemax'], agemax, root=0) # Write the output to a netcdf file if comm.rank == 0: sun.writeParticleNC(outfile, x, y, z, sun.time_track_sec[-1], ii + 1, age=age, agemax=agemax) comm.Barrier() t_diff = MPI.Wtime() - t_start ### Stop stopwatch ### if comm.rank == 0: print(78 * '=' + '\n' + 78 * '=') print('Completed particle tracking using %d cores in %6.2f seconds.' % (comm.size, t_diff)) print(78 * '=' + '\n' + 78 * '=')
def tide_pred_correc(modfile,lon,lat,time,dbfile,ID,z=None,conlist=None): """ Performs a tidal prediction at all points in [lon,lat] at times in vector [time] Applies an amplitude and phase correction based on a time series """ from timeseries import timeseries, loadDBstation print 'Calculating tidal correction factors from time series...' # Load using the timeseries module t0 = datetime.strftime(time[0],'%Y%m%d.%H%M%S') t1 = datetime.strftime(time[-1],'%Y%m%d.%H%M%S') dt = time[1]-time[0] print t0, t1, dt.total_seconds() timeinfo = (t0,t1,dt.total_seconds()) TS,meta = loadDBstation(dbfile,ID,'waterlevel',timeinfo=timeinfo,filttype='low',cutoff=2*3600,output_meta=True) lonpt=meta['longitude'] latpt=meta['latitude'] print lonpt,latpt # Extract the OTIS tide prediction u_re, u_im, v_re, v_im, h_re, h_im, omega, conlist = extract_HC(modfile,lonpt,latpt) h_amp = np.abs(h_re+1j*h_im)[:,0] h_phs = np.angle(h_re+1j*h_im)[:,0] # Harmonic analysis of observation time series amp, phs, frq, frqnames, htide = TS.tidefit(frqnames=conlist) TS_harm = timeseries(time,htide) residual = TS.y - htide # Calculate the amp and phase corrections dphs = phs - h_phs + np.pi damp = amp/h_amp # Extract the data along the specified points u_re, u_im, v_re, v_im, h_re, h_im, omega, conlist = extract_HC(modfile,lon,lat,z=z,conlist=conlist) h_amp = np.abs(h_re+1j*h_im) h_phs = np.angle(h_re+1j*h_im) u_amp = np.abs(u_re+1j*u_im) u_phs = np.angle(u_re+1j*u_im) v_amp = np.abs(v_re+1j*v_im) v_phs = np.angle(v_re+1j*v_im) # Initialise the output arrays sz = lon.shape nx = np.prod(sz) nt = time.shape[0] h=np.zeros((nt,nx)) u=np.zeros((nt,nx)) v=np.zeros((nt,nx)) # Rebuild the time series #tsec=TS_harm.tsec - TS_harm.tsec[0] tsec = othertime.SecondsSince(time,basetime=time[0]) print tsec[0] for nn,om in enumerate(omega): for ii in range(0,nx): h[:,ii] += damp[nn]*h_amp[nn,ii] * np.cos(om*tsec - (h_phs[nn,ii] + dphs[nn])) u[:,ii] += damp[nn]*u_amp[nn,ii] * np.cos(om*tsec - (u_phs[nn,ii] + dphs[nn])) v[:,ii] += damp[nn]*v_amp[nn,ii] * np.cos(om*tsec - (v_phs[nn,ii] + dphs[nn])) szo = (nt,)+sz return h.reshape(szo), u.reshape(szo), v.reshape(szo), residual
def harmonic_fit(t, X, frq, mask=None, axis=0, phsbase=None): """ Least-squares harmonic fit on an array, X, with frequencies, frq. X - vector [Nt] or array [Nt, (size)] t - vector [Nt] frq - vector [Ncon] mask - array [(size non-time X)] phsbase - phase offset where, dimension with Nt should correspond to axis = axis. """ t = np.asarray(t) # Reshape the array sizes X = X.swapaxes(0, axis) sz = X.shape lenX = np.prod(sz[1:]) if not len(t) == sz[0]: raise 'length of t (%d) must equal dimension of X (%s)' % (len(t), sz[0]) X = np.reshape(X, (sz[0], lenX)) if not mask == None: mask = np.reshape(mask, (lenX, )) else: mask = np.ones((lenX, )) frq = np.array(frq) Nfrq = frq.shape[0] def buildA(t, frq): """ Construct matrix A """ nt = t.shape[0] nf = frq.shape[0] nff = nf * 2 + 1 A = np.ones((nt, nff)) for ff in range(0, nf): A[:, ff * 2 + 1] = np.cos(frq[ff] * t) A[:, ff * 2 + 2] = np.sin(frq[ff] * t) return A def lstsqnumpy(A, y): """ Solve the least square problem Return: the complex amplitude the mean """ N = A.shape[1] b = np.linalg.lstsq(A, y) A = b[0][1::2] B = b[0][2::2] return A + 1j * B, b[0][0::N] def phsamp(C): return np.abs(C), np.angle(C) # Least-squares matrix approach A = buildA(t, frq) C, C0 = lstsqnumpy(A, X) # This works on all columns of X!! Amp, Phs = phsamp(C) # Reference the phase to some time if not phsbase == None: base = othertime.SecondsSince(phsbase) phsoff = phase_offset(frq, t[0], base) phsoff = np.repeat(phsoff.reshape((phsoff.shape[0], 1)), lenX, axis=1) phs = np.mod(Phs + phsoff, 2 * np.pi) # Non-vectorized method (~20x slower) # Amp = np.zeros((Nfrq,lenX)) # Phs = np.zeros((Nfrq,lenX)) # for ii in range(0,lenX): # if mask[ii]==True: # C = lstsqnumpy(A,X[:,ii]) # # Calculate the phase and amplitude # am, ph= phsamp(C) # Amp[:,ii] = am; Phs[:,ii] = ph # reshape the array Amp = np.reshape(Amp, (Nfrq, ) + sz[1:]) Phs = np.reshape(Phs, (Nfrq, ) + sz[1:]) C0 = np.reshape(C0, sz[1:]) # Output back along the original axis return Amp.swapaxes(axis, 0), Phs.swapaxes(axis, 0), C0.swapaxes(axis, 0)