def fit_peaks(data, p_init, model_type='Gaussian'): '''Fit a composite spectrum to the data, using the Levenberg-Marquardt least-squares fitter. data -> 1D cross-section p_init -> list of initial estimates for parameters of each peak ''' data = RobustData(data, index=True) data.replace_nans() mean, median, stdev = data.stats() fit_spectrum = data.fit_to_model(composite_model(p_init, model_type)) return (data.x, filter(lambda x: x.amplitude > median, fit_spectrum), filter(lambda x: x.amplitude < median, fit_spectrum))
def fit_trace(data, p_init, tracedir=-1, width=10, model_type='Gaussian', full=False): '''Now we're going to fit the whole trace iteratively, working out from the middle to get the gaussian centroids.''' nx, ny = data.shape tracedir = abs(tracedir) n_wav = data.shape[tracedir] cen = n_wav/2 + (n_wav % 2) all_models = [None]*n_wav trace_centers = np.zeros((len(p_init), n_wav)) on_axis = np.arange(-width,width) peak = ['x_0','mean'][model_type == 'Gaussian'] window = lambda c: ((np.tile(np.clip(on_axis+c, 0, nx-1)[...,np.newaxis], (1,ny)), np.tile(np.arange(ny)[np.newaxis,...],(on_axis.size,1))), (np.tile(np.arange(nx)[...,np.newaxis],(1,on_axis.size)), np.tile(np.clip(on_axis+c, 0, ny-1)[np.newaxis,...], (nx,1))))[tracedir] def fit_section(center, p0): xr, yr = window(center) section = Robust2D(data[xr, yr]).combine(axis=(tracedir+1)%2) return section.fit_to_model(p0, x=section.index) #first we grab the center section fit_profile = fit_section(cen, composite_model(p_init, model_type)) up, down = fit_profile.copy(), fit_profile.copy() all_models[cen] = fit_profile trace_centers[:,cen] = np.array([getattr(component, peak).value for component in fit_profile]) for c in xrange(1, cen+1): #now work outward from the middle if cen + c < [nx, ny][tracedir]: up = fit_section(cen+c, up) all_models[cen+c] = up.copy() trace_centers[:,cen+c] = np.array([getattr(component, peak).value for component in up]) if cen - c >= 0: down = fit_section(cen-c, down) all_models[cen-c] = down.copy() trace_centers[:,cen-c] = np.array([getattr(component, peak).value for component in down]) if full: return all_models pos = np.array([x.amplitude > np.median(data) for x in all_models[cen]]) return trace_centers[pos,:], trace_centers[~pos,:]