def spline_1D(xarr,yarr,invar,breakpoints,order=4): """ Returns the spline fit to xarr, yarr with inverse variance invar for the specified breakpoint array (and optionally spline order) Assumes independent data points. """ B_spline_matrix = basic_spline(xarr,breakpoints,order=order) coeffs, chi = sf.chi_fit(yarr,B_spline_matrix.T,np.diag(invar)) spline_interp = np.dot(B_spline_matrix.T,coeffs)[:,0] return spline_interp
def spline_2D(img_matrix,invar_matrix,h_breakpoints,v_breakpoints,order=4,return_coeffs=False): """ Returns a 2D spline interpolation of an image. Option to also return scaled spline_coeffs Assumes a pixel-based image (ie, integer horizontal and vertical steps between points) Assumes pixel errors are independent Note, computational time scales up very quickly, suggest modest image sizes (~100 x 100 px or less) """ hpix = np.shape(img_matrix)[1] #horizontal pixels vpix = np.shape(img_matrix)[0] #vertical pixels harr = np.arange(hpix) varr = np.arange(vpix) ### Given h, v arrays and breakpoints, find splines along both directions h_splines = basic_spline(harr,h_breakpoints,order=order) v_splines = basic_spline(varr,v_breakpoints,order=order) h_sp_len = np.shape(h_splines)[0] v_sp_len = np.shape(v_splines)[0] dim1 = hpix*vpix dim2 = h_sp_len * v_sp_len ### Use the h and v splines to construct a 2D profile matrix for linear fitting profile_matrix = np.zeros((vpix,hpix,dim2)) for i in range(h_sp_len): for j in range(v_sp_len): k = i*v_sp_len + j h_sp_tmp = np.reshape(h_splines[i],(1,hpix)) v_sp_tmp = np.reshape(v_splines[j],(vpix,1)) profile_matrix[:,:,k] = np.dot(v_sp_tmp,h_sp_tmp) #Reshape the profile matrix and input image and invar matrices for chi^2 fitting profile_matrix = np.reshape(profile_matrix,(dim1,dim2)) data = np.reshape(img_matrix,(dim1,)) noise = np.diag(np.reshape(invar_matrix,(dim1,))) #Chi^2 fit a_coeffs, chi = sf.chi_fit(data,profile_matrix,noise) #Evaluate spline fit and reshape back to 2D array spline_fit = np.dot(profile_matrix,a_coeffs) spline_fit = np.reshape(spline_fit,(vpix,hpix)) spline_integral = np.sum(spline_fit) a_coeffs_scale = a_coeffs/spline_integral if return_coeffs: return spline_fit, a_coeffs_scale, spline_integral else: return spline_fit
def fit_spline_psf(raw_img,hcenters,vcenters,sigmas,powers,readnoise, gain,plot_results=False,verbose=False): """ function to fit parameters for radial bspline psf. """ ### 1. Estimate spline amplitudes, centers, w/ circular model actypix = raw_img.shape[1] #r_breakpoints = [0, 1.2, 2.5, 3.7, 5, 8, 10] ## 2.3, 3 #r_breakpoints = np.hstack(([0, 1.5, 2.4, 3],np.arange(3.5,10,0.5))) #For cpad=8 ########################################################################## ############## All this hardcoded stuff should be flexible (TODO) ######## ########################################################################## # r_breakpoints = np.hstack(([0, 1.5, 2.4, 3],np.arange(3.5,6.6,1))) #For cpad=5 r_breakpoints = np.hstack(([0, 1.5, 2.4, 3],np.arange(3.5,8.6,1))) #For cpad=6 #r_breakpoints = np.hstack(([0, 1.2, 2.3, 3],np.arange(3.5,10,0.5))) #For cpad=8 theta_orders = [0] cpad = 6 bp_space = 2 #beakpoint spacing in pixels invar = 1/(raw_img+readnoise**2) ### Initial spline coeff guess spl_coeffs, s_scale, fit_params, new_hcenters, new_vcenters = spline_coeff_fit(raw_img,hcenters,vcenters,invar,r_breakpoints,sigmas,powers,theta_orders=theta_orders,cpad=cpad,bp_space=bp_space,return_new_centers=True) #''' ### 2. Set up and initialize while loop (other steps embedded within loop) num_bases = spl_coeffs.shape[1] new_hscale = (new_hcenters-actypix/2)/actypix peak_mask = np.ones((len(new_hscale)),dtype=bool) #Can be used to mask "bad" peaks params1 = lmfit.Parameters() ### Loop to add horizontal/vertical centers for j in range(len(new_hscale)): harr = np.arange(-cpad,cpad+1)+int(np.floor(new_hcenters[j])) varr = np.arange(-cpad,cpad+1)+int(np.floor(new_vcenters[j])) ### Shouldn't need +1... params1.add('vc{}'.format(j), value = new_vcenters[j]-varr[0]) params1.add('hc{}'.format(j), value = new_hcenters[j]-harr[0]) ### and add initial ellitical parameter guesses (for quadratic variation) params1.add('q0', value=0.9, min=0, max=1) params1.add('PA0', value=0, min=-np.pi, max=np.pi) params1.add('q1', value=0, min=-1, max=1) params1.add('PA1', value=0, min=-np.pi, max=np.pi) params1.add('q2', value=0, min=-1, max=1) params1.add('PA2', value=0, min=-np.pi, max=np.pi) params = lmfit.Parameters() params.add('hc', value = params1['hc0'].value) params.add('vc', value = params1['vc0'].value) params.add('q', value = 1, min=0, max=1) params.add('PA', value=0, min=-np.pi, max=np.pi) ### Start while loop - iterate until convergence chi_new = np.ones((sum(peak_mask))) #Can build this from first fit if desired chi_old = np.zeros((sum(peak_mask))) chi_min = 100 coeff_matrix_min = np.zeros((3,np.shape(spl_coeffs)[1])).T params_min = lmfit.Parameters() dlt_chi = 1e-3 #difference between successive chi_squared values to cut off mx_loops = 50 #eventually must cutoff loop_cnt = 0 fit_bg = False ## True fits a constant background at each subimage while abs(np.sum(chi_new)-np.sum(chi_old)) > dlt_chi and loop_cnt < mx_loops: if verbose: print("starting loop {}".format(loop_cnt)) print(" chi_old mean = {}".format(np.mean(chi_old))) print(" chi_new mean = {}".format(np.mean(chi_new))) print(" delta_chi = {}".format((np.sum(chi_new)-np.sum(chi_old)))) chi_old = np.copy(chi_new) ### 3. Build profile, data, and noise matrices at each pixel point and sum dim_s = (2*cpad+1)**2 dim_h = sum(peak_mask)*dim_s profile_matrix = np.zeros((dim_h,3*num_bases+fit_bg*len(new_hscale))) #hardcoded for quadratic # last_profile = np.zeros((dim_s,3*num_bases+fit_bg)) data_array = np.zeros((dim_h)) noise_array = np.zeros((dim_h)) data_for_fitting = np.zeros((2*cpad+1,2*cpad+1,len(new_hscale))) invar_for_fitting = np.zeros((2*cpad+1,2*cpad+1,len(new_hscale))) d_scale = np.zeros(len(new_hscale)) # Will build from data # bg_data = np.zeros(len(new_hscale)) for k in range(len(new_hscale)): ### Slice subset of image data around each peak harr = np.arange(-cpad,cpad+1)+int(np.floor(new_hcenters[k])) varr = np.arange(-cpad,cpad+1)+int(np.floor(new_vcenters[k])) harr = harr[harr>=0] harr = harr[harr<raw_img.shape[1]] varr = varr[varr>=0] varr = varr[varr<raw_img.shape[0]] data_for_fitting[:,:,k] = raw_img[varr[0]:varr[-1]+1,harr[0]:harr[-1]+1]#/s_scale[k] # invar_for_fitting[:,:,k] = invar[varr[0]:varr[-1]+1,harr[0]:harr[-1]+1]#/s_scale[k] d_scale[k] = np.sum(data_for_fitting[:,:,k]) invar_for_fitting[:,:,k] = s_scale[k]/(abs(data_for_fitting[:,:,k])+readnoise**2/s_scale[k]) # rarr = sf.make_rarr(np.arange(2*cpad+1),np.arange(2*cpad+1),cpad,cpad) # bg_mask = rarr > 3 # bg_data[k] = poisson_bg(data_for_fitting[:,:,k],mask=bg_mask) ### bound s_scale to (hopefully) prevent runaway growth # for k in range(len(new_hscale)): # sig_factor = 1 #Constrain s_scale to be within this man stddevs # d_min = d_scale[k]-np.sqrt(d_scale[k])*sig_factor # d_max = d_scale[k]+np.sqrt(d_scale[k])*sig_factor # if s_scale[k] < d_min: # s_scale[k] = d_min # elif s_scale[k] > d_max: # s_scale[k] = d_max # s_scale *= np.sum(d_scale)/np.sum(s_scale) for k in range(len(new_hscale)): ### Pull in best center estimates params['hc'].value = params1['hc{}'.format(k)].value params['vc'].value = params1['vc{}'.format(k)].value ### Pull in best elliptical parameter estimates if loop_cnt == 0: params['q'].value = 1 else: params['q'].value = params1['q0'].value + params1['q1'].value*new_hscale[k] + params1['q2'].value*new_hscale[k]**2 params['PA'].value = params1['PA0'].value + params1['PA1'].value*new_hscale[k] + params1['PA2'].value*new_hscale[k]**2 ### Scale data # data_for_fitting[:,:,k] -= bg_data[k] ### remove bg first data_for_fitting[:,:,k] /= s_scale[k] # invar_for_fitting[:,:,k] *= s_scale[k] ### Setup arrays for spline analysis r_arr, theta_arr, dim1, r_inds = spline.build_rarr_thetaarr(data_for_fitting[:,:,k],params) ### Build data, noise, and profile array data_array[k*dim_s:(k+1)*dim_s] = np.ravel(data_for_fitting[:,:,k])[r_inds] #scaled, sorted data array noise_array[k*dim_s:(k+1)*dim_s] = np.ravel(invar_for_fitting[:,:,k])[r_inds] profile_base = spline.build_radial_profile(r_arr,theta_arr,r_breakpoints,theta_orders,(2*cpad+1)**2,order=4) profile_matrix[k*dim_s:(k+1)*dim_s,0:num_bases] = profile_base profile_matrix[k*dim_s:(k+1)*dim_s,num_bases:2*num_bases] = profile_base*new_hscale[k] profile_matrix[k*dim_s:(k+1)*dim_s,2*num_bases:3*num_bases] = profile_base*(new_hscale[k]**2) if fit_bg: profile_matrix[k*dim_s:(k+1)*dim_s,3*num_bases+k*fit_bg] = 1 # plt.imshow(profile_matrix,interpolation='none') # plt.show() ### 4. Using matrices from step 3. perform chi^2 fitting for coefficients next_coeffs, next_chi = sf.chi_fit(data_array,profile_matrix,np.diag(noise_array)) if fit_bg: bg_array = next_coeffs[3*num_bases:] # print bg_array*s_scale trunc_coeffs = next_coeffs[0:3*num_bases] else: trunc_coeffs = np.copy(next_coeffs) dd2 = int(np.size(trunc_coeffs)/3) coeff_matrix = trunc_coeffs.reshape(3,dd2).T # if fit_bg: ### Don't save background fit term # bg_array = coeff_matrix[:,-1] # print bg_array*s_scale # coeff_matrix = coeff_matrix[:,:-1] # last_coeffs = np.dot(coeff_matrix,(np.vstack((ones(len(new_hscale)),new_hscale,new_hscale**2)))) ### Check each of the profiles with next_coeffs + adjust scale factor profile_matrix = np.zeros((dim_s,3*num_bases+fit_bg*len(new_hscale))) #hardcoded for quadratic data_array = np.zeros((dim_h)) noise_array = np.zeros((dim_h)) chi2_first = np.zeros(len(new_hscale)) # fit_sums = 0 # print("Temp fit sums:") for k in range(len(new_hscale)): ### Pull in best center estimates params['hc'].value = params1['hc{}'.format(k)].value params['vc'].value = params1['vc{}'.format(k)].value ### Pull in best elliptical parameter estimates if loop_cnt == 0: params['q'].value = 1 else: params['q'].value = params1['q0'].value + params1['q1'].value*new_hscale[k] + params1['q2'].value*new_hscale[k]**2 params['PA'].value = params1['PA0'].value + params1['PA1'].value*new_hscale[k] + params1['PA2'].value*new_hscale[k]**2 ### Setup arrays for spline analysis r_arr, theta_arr, dim1, r_inds = spline.build_rarr_thetaarr(data_for_fitting[:,:,k],params) ### Build data, noise, and profile array data_array[k*dim_s:(k+1)*dim_s] = np.ravel(data_for_fitting[:,:,k])[r_inds] #scaled, sorted data array noise_array[k*dim_s:(k+1)*dim_s] = np.ravel(invar_for_fitting[:,:,k])[r_inds] profile_base = spline.build_radial_profile(r_arr,theta_arr,r_breakpoints,theta_orders,(2*cpad+1)**2,order=4) profile_matrix[:,0:num_bases] = profile_base profile_matrix[:,num_bases:2*num_bases] = profile_base*new_hscale[k] profile_matrix[:,2*num_bases:3*num_bases] = profile_base*(new_hscale[k]**2) if fit_bg: profile_matrix[:,3*num_bases:] = 0 profile_matrix[:,3*num_bases+k*fit_bg] = 1 tmp_fit = np.dot(profile_matrix,next_coeffs) # print np.sum(tmp_fit) # fit_sums += np.sum(tmp_fit) resort_inds = np.argsort(r_inds) tmp_fit = np.reshape(tmp_fit[resort_inds],data_for_fitting[:,:,k].shape) # plt.figure("Arc, iteration {}".format(k)) ## plt.imshow(np.hstack((tmp_fit,small_img/s_scale[k])),interpolation='none') chi2_first[k] = np.sum(((tmp_fit-data_for_fitting[:,:,k])**2)*invar_for_fitting[:,:,k])#*s_scale[k]**2 # plt.imshow((tmp_fit-small_img/s_scale[k])*small_inv,interpolation='none') # plt.show() # plt.close() # print "chi2 first:", chi2_first # next_coeffs *= fit_sums/(k+1) # s_scale /= fit_sums/(k+1) ### Optional place to check coefficients variation over order #for i in range(8): # plt.plot(new_hscale,last_coeffs[i]) # #plt.show() #plt.close() #first_fit = np.dot(last_profile,next_coeffs) #print next_coeffs #print params['vc'].value #print params['hc'].value #print r_arr[0:10] #print profile_base[0] #print profile_matrix[0,:]/(k+1) #print last_profile[0] #print first_fit[0] #resort_inds = np.argsort(r_inds) #scale1 = np.max(small_img)/np.max(first_fit) ##print scale1, scale, scale1/scale #first_fit = np.reshape(first_fit[resort_inds],small_img.shape) #print np.sum(first_fit), k, scale1, s_scale[k] #first_fit /= np.sum(first_fit) ##plt.imshow(first_fit,interpolation='none') #plt.imshow(np.hstack((small_img/s_scale[k],first_fit,(small_img/s_scale[k]-first_fit)*small_inv)),interpolation='none') #plt.show() #plt.imshow((small_img/s_scale[k]-first_fit)*small_inv,interpolation='none') #plt.show() #test_xs = (np.arange(xpix)-xpix/2)/xpix #for i in range(num_bases): # test_ys = next_coeffs[i]+next_coeffs[num_bases+i]*test_xs+next_coeffs[2*num_bases+i]*test_xs**2 # plt.plot(test_xs,test_ys) #plt.show() ### 5. Now do a nonlinear fit for hc, vc, q, and PA #data_for_lmfit = np.zeros((np.size(small_img),len(new_hscale))) #invar_for_lmfit = np.zeros((np.size(small_img),len(new_hscale))) # for k in range(len(new_hscale)): # harr = np.arange(-cpad,cpad+1)+int(np.floor(new_hcenters[k])) # varr = np.arange(-cpad,cpad+1)+int(np.floor(new_vcenters[k])) # data_for_lmfit[:,:,k] = raw_img[varr[0]:varr[-1]+1,harr[0]:harr[-1]+1]/s_scale[k] # invar_for_lmfit[:,:,k] = invar[varr[0]:varr[-1]+1,harr[0]:harr[-1]+1]*(s_scale[k]) # r_arr, theta_arr, dim1, r_inds = spline.build_rarr_thetaarr(small_img,params) # data_for_lmfit[:,k] = np.ravel(small_img)[r_inds]/s_scale[k] # invar_for_lmfit[:,k] = np.ravel(small_inv)[r_inds]/np.sqrt(s_scale[k]) # resort_inds = np.argsort(r_inds) # plt.imshow(np.resize(data_for_lmfit[:,k][resort_inds],np.shape(small_img))) # plt.show() # plt.close() ### Make proper inputs for minimizer function #centers = np.vstack((new_hcenters,new_vcenters)).T args = (data_for_fitting,invar_for_fitting,r_breakpoints,new_hscale,next_coeffs) kws = dict() kws['theta_orders'] = theta_orders kws['fit_bg'] = fit_bg minimizer_results = lmfit.minimize(spline.spline_poly_residuals,params1,args=args,kws=kws) ### Re-initialize params1, put in elliptical values. Will add hc/vc at end ### (using mask, so #of values for centers will differ) params1['q0'].value = minimizer_results.params['q0'].value params1['q1'].value = minimizer_results.params['q1'].value params1['q2'].value = minimizer_results.params['q2'].value params1['PA0'].value = minimizer_results.params['PA0'].value params1['PA1'].value = minimizer_results.params['PA1'].value params1['PA2'].value = minimizer_results.params['PA2'].value #hc_ck = minimizer_results.params['hc0'].value + minimizer_results.params['hc1'].value*new_hscale + minimizer_results.params['hc2'].value*new_hscale**2 #vc_ck = minimizer_results.params['vc0'].value + minimizer_results.params['vc1'].value*new_hscale + minimizer_results.params['vc2'].value*new_hscale**2 q_ck = minimizer_results.params['q0'].value + minimizer_results.params['q1'].value*new_hscale + minimizer_results.params['q2'].value*new_hscale**2 PA_ck = minimizer_results.params['PA0'].value + minimizer_results.params['PA1'].value*new_hscale + minimizer_results.params['PA2'].value*new_hscale**2 # print q_ck # print PA_ck ### Convert so q is less than 1 if np.max(q_ck) > 1: q_ck_tmp = 1/q_ck #change axis definition if np.max(q_ck_tmp) > 1: print "q array always over 1!" else: q_ck = q_ck_tmp PA_ck = PA_ck + np.pi/2 #change axis definition q_coeffs = np.polyfit(new_hscale,q_ck,2) PA_coeffs = np.polyfit(new_hscale,PA_ck,2) params1['q0'].value = q_coeffs[2] params1['q1'].value = q_coeffs[1] params1['q2'].value = q_coeffs[0] params1['PA0'].value = PA_coeffs[2] params1['PA1'].value = PA_coeffs[1] params1['PA2'].value = PA_coeffs[0] # print q_ck # print PA_ck #plt.plot(np.arange(5),np.arange(5)) #plt.show() #plt.plot(hc_ck,vc_ck,new_hcenters,new_vcenters) #plt.show() #ecc = minimizer_results.params['q'].value #pos_ang = minimizer_results.params['PA'].value ### Check to see if elliptical values worked out well chi_new = np.zeros(len(new_hscale)) for i in range(len(new_hscale)): params['vc'].value = minimizer_results.params['vc{}'.format(i)].value params['hc'].value = minimizer_results.params['hc{}'.format(i)].value # harr = np.arange(-cpad,cpad+1)+int(np.floor(new_hcenters[i])) # varr = np.arange(-cpad,cpad+1)+int(np.floor(new_vcenters[i])) # params['vc'].value = new_vcenters[i]-varr[0]+1 # params['hc'].value = new_hcenters[i]-harr[0] x_coord = new_hscale[i] img_matrix = data_for_fitting[:,:,i] invar_matrix = invar_for_fitting[:,:,i] q = params1['q0'].value + params1['q1'].value*x_coord + params1['q2'].value*x_coord**2 PA = params1['PA0'].value + params1['PA1'].value*x_coord + params1['PA2'].value*x_coord**2 params['q'].value = q params['PA'].value = PA sp_coeffs = np.dot(coeff_matrix,np.array(([1,new_hscale[i],new_hscale[i]**2]))) if fit_bg: sp_coeffs = np.hstack((sp_coeffs,bg_array[i])) # r_arr, theta_arr, dim1, r_inds = spline.build_rarr_thetaarr(small_img,params) # profile_base = spline.build_radial_profile(r_arr,theta_arr,r_breakpoints,theta_orders,(2*cpad+1)**2,order=4) fitted_image = spline.spline_2D_radial(img_matrix,invar_matrix,r_breakpoints,params,theta_orders,order=4,return_coeffs=False,spline_coeffs=sp_coeffs,sscale=None,fit_bg=fit_bg) ### Update s_scale chi_new[i] = np.sum(((img_matrix-fitted_image)**2)*invar_matrix)*s_scale[i]/(np.size(img_matrix)-len(sp_coeffs)-2)#*s_scale[i]**2 # print chi_new[i] # print s_scale[i] # print np.max(invar_matrix)*3.63**2/s_scale[i] # print chi_new[i]*s_scale[i] # print chi_new[i]*s_scale[i]**2 ### Set new scale - drive sum of image toward unity s_scale[i] = s_scale[i]*np.sum(fitted_image) # plt.imshow(np.hstack((img_matrix,fitted_image)),interpolation='none')#,(img_matrix-fitted_image)*invar_matrix)),interpolation='none') # plt.imshow(invar_matrix,interpolation='none') # plt.plot(img_matrix[:,5]) # plt.plot(fitted_image[:,5]) # plt.show() # plt.close() #print chi2_first #print chi2_second #print s_scale #print s_scale2 ### Mask/eliminate points with high chi2 peak_mask = sf.sigma_clip(chi_new,sigma=3,max_iters=1) if sum(peak_mask) < 4: print("Too few peaks for fitting") exit(0) # break ### Update new_hscale, s_scale, new_h/vcenters s_scale = s_scale[peak_mask] cnts = len(new_hscale) new_hscale = np.zeros((sum(peak_mask))) lp_idx = 0 for j in range(cnts): if not peak_mask[j]: if verbose: print "skipping point {}".format(j) continue else: harr = np.arange(-cpad,cpad+1)+int(np.floor(new_hcenters[j])) params1.add('hc{}'.format(lp_idx), value = minimizer_results.params['hc{}'.format(j)].value) params1.add('vc{}'.format(lp_idx), value = minimizer_results.params['vc{}'.format(j)].value) new_hscale[lp_idx] = (params1['hc{}'.format(lp_idx)].value+harr[0]-1-actypix/2)/actypix lp_idx += 1 new_hcenters = new_hcenters[peak_mask] new_vcenters = new_vcenters[peak_mask] ### Record minimum values (some subsequent iterations give higher chi2) if loop_cnt == 0: coeff_matrix_min = np.copy(coeff_matrix) params_min = lmfit.Parameters(params1) if np.sum(chi_new) < chi_min: if verbose: print "Better fit on loop ", loop_cnt chi_min = np.sum(chi_new) coeff_matrix_min = np.copy(coeff_matrix) params_min = lmfit.Parameters(params1) loop_cnt += 1 ### End of loop if verbose: print("End of Loop") ### Check that q, PA, aren't driving toward unphysical answers # test_hscale = np.arange(-1,1,0.01) #q = params_min['q0'].value + params_min['q1'].value*test_hscale + params_min['q2'].value*test_hscale**2 #PA = params_min['PA0'].value + params_min['PA1'].value*test_hscale + params_min['PA2'].value*test_hscale**2 #bg = coeff_matrix_min[0,-1] + coeff_matrix_min[1,-1]*test_hscale + coeff_matrix_min[2,-1]*test_hscale**2 #plt.plot(test_hscale,q) #plt.show() #plt.plot(test_hscale,PA) #plt.show() #plt.plot(test_hscale,bg) #plt.show() #plt.close() if plot_results: ### Plot final answers for evaluation for i in range(len(new_hscale)): params['vc'].value = minimizer_results.params['vc{}'.format(i)].value params['hc'].value = minimizer_results.params['hc{}'.format(i)].value # harr = np.arange(-cpad,cpad+1)+int(np.floor(new_hcenters[i])) # varr = np.arange(-cpad,cpad+1)+int(np.floor(new_vcenters[i])) # params['vc'].value = new_vcenters[i]-varr[0]+1 # params['hc'].value = new_hcenters[i]-harr[0] x_coord = new_hscale[i] img_matrix = data_for_fitting[:,:,i] invar_matrix = invar_for_fitting[:,:,i] q = params_min['q0'].value + params_min['q1'].value*x_coord + params_min['q2'].value*x_coord**2 PA = params_min['PA0'].value + params_min['PA1'].value*x_coord + params_min['PA2'].value*x_coord**2 params['q'].value = q params['PA'].value = PA sp_coeffs = np.dot(coeff_matrix_min,np.array(([1,new_hscale[i],new_hscale[i]**2]))) if fit_bg: sp_coeffs = np.hstack((sp_coeffs,bg_array[i])) # r_arr, theta_arr, dim1, r_inds = spline.build_rarr_thetaarr(small_img,params) # profile_base = spline.build_radial_profile(r_arr,theta_arr,r_breakpoints,theta_orders,(2*cpad+1)**2,order=4) fitted_image = spline.spline_2D_radial(img_matrix,invar_matrix,r_breakpoints,params,theta_orders,order=4,return_coeffs=False,spline_coeffs=sp_coeffs,sscale=None,fit_bg=fit_bg) ### Update s_scale # print chi_new[i] # print chi_new[i]*s_scale[i] # print chi_new[i]*s_scale[i]**2 chi_sq_red = np.sum(((img_matrix-fitted_image))**2*invar_matrix)/(np.size(img_matrix)-len(sp_coeffs)-2)*(s_scale[i]) print "Reduced Chi^2 on iteration ", i, " is: ", chi_sq_red # plt.plot(fitted_image[:,cpad]/np.max(fitted_image[:,cpad])) # plt.plot(np.sum(fitted_image,axis=1)/np.max(np.sum(fitted_image,axis=1))) # plt.show() # plt.imshow(np.hstack((img_matrix,fitted_image,(img_matrix-fitted_image))),interpolation='none') plt.imshow((img_matrix-fitted_image)*invar_matrix,interpolation='none') # plt.imshow((img_matrix-fitted_image)*invar_matrix,interpolation='none') plt.show() plt.close() centers, ellipse = params_to_array(params_min) results = np.hstack((np.ravel(coeff_matrix_min),np.ravel(ellipse))) return results
def fit_trace(x,y,ccd,form='gaussian'): """quadratic fit (in x) to trace around x,y in ccd x,y are integer pixel values input "form" can be set to quadratic or gaussian """ x = int(x) y = int(y) if form=='quadratic': xpad = 2 xvals = np.arange(-xpad,xpad+1) def make_chi_profile(x,y,ccd): xpad = 2 xvals = np.arange(-xpad,xpad+1) zvals = ccd[x+xvals,y] profile = np.ones((2*xpad+1,3)) #Quadratic fit profile[:,1] = xvals profile[:,2] = xvals**2 noise = np.diag((1/zvals)) return zvals, profile, noise zvals, profile, noise = make_chi_profile(x,y,ccd) coeffs, chi = sf.chi_fit(zvals,profile,noise) # print x # print xvals # print x+xvals # print zvals # plt.errorbar(x+xvals,zvals,yerr=sqrt(zvals)) # plt.plot(x+xvals,coeffs[2]*xvals**2+coeffs[1]*xvals+coeffs[0]) # plt.show() chi_max = 100 if chi>chi_max: #print("bad fit, chi^2 = {}".format(chi)) #try adacent x xl = x-1 xr = x+1 zl, pl, nl = make_chi_profile(xl,y,ccd) zr, pr, nr = make_chi_profile(xr,y,ccd) cl, chil = sf.chi_fit(zl,pl,nl) cr, chir = sf.chi_fit(zr,pr,nr) if chil<chi and chil<chir: # plt.errorbar(xvals-1,zl,yerr=sqrt(zl)) # plt.plot(xvals-1,cl[2]*(xvals-1)**2+cl[1]*(xvals-1)+cl[0]) # plt.show() xnl = -cl[1]/(2*cl[2]) znl = cl[2]*xnl**2+cl[1]*xnl+cl[0] return xl+xnl, znl, chil elif chir<chi and chir<chil: xnr = -cr[1]/(2*cr[2]) znr = cr[2]*xnr**2+cr[1]*xnr+cr[0] # plt.errorbar(xvals+1,zr,yerr=sqrt(zr)) # plt.plot(xvals+1,cr[2]*(xvals+1)**2+cr[1]*(xvals+1)+cr[0]) # plt.show() return xr+xnr, znr, chir else: ca = coeffs[2] cb = coeffs[1] xc = -cb/(2*ca) zc = ca*xc**2+cb*xc+coeffs[0] return x+xc, zc, chi else: ca = coeffs[2] cb = coeffs[1] xc = -cb/(2*ca) zc = ca*xc**2+cb*xc+coeffs[0] return x+xc, zc, chi elif form=='gaussian': xpad = 7 xvals = np.arange(-xpad,xpad+1) xinds = x+xvals xvals = xvals[(xinds>=0)*(xinds<np.shape(ccd)[0])] zvals = ccd[x+xvals,y] params, errarr = sf.gauss_fit(xvals,zvals) xc = x+params[1] #offset plus center zc = params[2] #height (intensity) # pxn = np.linspace(xvals[0],xvals[-1],1000) fit = sf.gaussian(xvals,abs(params[0]),params[1],params[2],params[3],params[4]) chi = sum((fit-zvals)**2/zvals) return xc, zc, chi
else: xtrace[j,i], Itrace[j,i], chi_vals[j,i] = nan, nan, nan #Finally fit x vs. y on traces. Start with quadratic for simple + close enough trace_coeffs = np.zeros((3,num_fibers)) trace_intense_coeffs = np.zeros((3,num_fibers)) for i in range(num_fibers): #Given orientation makes more sense to swap x/y mask = ~np.isnan(xtrace[i,:]) if len(mask[mask])<3: continue profile = np.ones((len(ytrace[i,:][mask]),3)) #Quadratic fit profile[:,1] = (ytrace[i,:][mask]-ypix/2)/ypix #scale data to get better fit profile[:,2] = ((ytrace[i,:][mask]-ypix/2)/ypix)**2 noise = np.diag(chi_vals[i,:][mask]) tmp_coeffs, junk = sf.chi_fit(xtrace[i,:][mask],profile,noise) trace_coeffs[0,i] = tmp_coeffs[0] trace_coeffs[1,i] = tmp_coeffs[1] trace_coeffs[2,i] = tmp_coeffs[2] tmp_coeffs2, junk = sf.chi_fit(Itrace[i,:][mask],profile,noise) trace_intense_coeffs[0,i] = tmp_coeffs2[0] trace_intense_coeffs[1,i] = tmp_coeffs2[1] trace_intense_coeffs[2,i] = tmp_coeffs2[2] #Plot to check traces # fig,ax = plt.subplots() # ax.pcolorfast(ccd) # for i in range(num_fibers): # ys = (np.arange(ypix)-ypix/2)/ypix # xs = trace_coeffs[2,i]*ys**2+trace_coeffs[1,i]*ys+trace_coeffs[0,i] # yp = np.arange(ypix)
def refine_trace_centers(ccd, t_coeffs, i_coeffs, s_coeffs, p_coeffs, fact=10, readnoise=3.63, verbose=False): """ Uses estimated centers from fibers flats as starting point, then fits from there to find traces based on science ccd frame. INPUTS: ccd - image on which to fit traces t/i/s/p_coeffs - modified gaussian coefficients from fiberflat fact - do 1/fact of the available points """ num_fibers = t_coeffs.shape[0] hpix = ccd.shape[1] vpix = ccd.shape[0] ### First fit vc parameters for traces rough_pts = int(np.ceil(hpix/fact)) vc_ccd = np.zeros((num_fibers,rough_pts)) hc_ccd = np.zeros((num_fibers,rough_pts)) inv_chi = np.zeros((num_fibers,rough_pts)) yspec = np.arange(hpix) if verbose: print("Refining trace centers") for i in range(num_fibers): if verbose: print("Running on index {}".format(i)) # slit_num = np.floor((i)/args.telescopes) for j in range(0,hpix,fact): jadj = int(np.floor(j/fact)) yj = (yspec[j]-hpix/2)/hpix hc_ccd[i,jadj] = yspec[j] vc = t_coeffs[2,i]*yj**2+t_coeffs[1,i]*yj+t_coeffs[0,i] # Ij = i_coeffs[2,i]*yj**2+i_coeffs[1,i]*yj+i_coeffs[0,i] sigj = s_coeffs[2,i]*yj**2+s_coeffs[1,i]*yj+s_coeffs[0,i] powj = s_coeffs[2,i]*yj**2+s_coeffs[1,i]*yj+s_coeffs[0,i] if np.isnan(vc): vc_ccd[i,jadj] = np.nan inv_chi[i,jadj] = 0 else: xpad = 7 xvals = np.arange(-xpad,xpad+1) xj = int(vc) xwindow = xj+xvals xvals = xvals[(xwindow>=0)*(xwindow<vpix)] zorig = ccd[xj+xvals,yspec[j]] if len(zorig)<1: vc_ccd[i,jadj] = np.nan inv_chi[i,jadj] = 0 continue invorig = 1/(abs(zorig)+readnoise**2) if np.max(zorig)<20: vc_ccd[i,jadj] = np.nan inv_chi[i,jadj] = 0 else: mn_new, hght, bg = fit_mn_hght_bg(xvals, zorig, invorig, sigj, vc-xj-1, sigj, powj=powj) fitorig = sf.gaussian(xvals,sigj,mn_new,hght,power=powj) if j == 715: print mn_new plt.plot(xvals,zorig,xvals,fitorig) plt.show() plt.close() inv_chi[i,jadj] = 1/sum((zorig-fitorig)**2*invorig) vc_ccd[i,jadj] = mn_new+xj+1 tmp_poly_ord = 10 trace_coeffs_ccd = np.zeros((tmp_poly_ord+1,num_fibers)) for i in range(num_fibers): mask = ~np.isnan(vc_ccd[i,:]) profile = np.ones((len(hc_ccd[i,:][mask]),tmp_poly_ord+1)) #Quadratic fit for order in range(tmp_poly_ord): profile[:,order+1] = ((hc_ccd[i,:][mask]-hpix/2)/hpix)**(order+1) noise = np.diag(inv_chi[i,:][mask]) if len(vc_ccd[i,:][mask])>3: tmp_coeffs, junk = sf.chi_fit(vc_ccd[i,:][mask],profile,noise) else: tmp_coeffs = np.nan*np.ones((tmp_poly_ord+1)) trace_coeffs_ccd[:,i] = tmp_coeffs return trace_coeffs_ccd
def extract_1D(ccd, t_coeffs, i_coeffs=None, s_coeffs=None, p_coeffs=None, readnoise=1, gain=1, return_model=False, verbose=False): """ Function to extract using optimal extraction method. This could benefit from a lot of cleaning up INPUTS: ccd - ccd image to extract t_coeffs - estimate of trace coefficients (from 'find_t_coeffs') i/s/p_coeffs - optional intensity, sigma, power coefficients readnoise, gain - of the ccd return_model - set True to return model of image based on extraction OUTPUTS: spec - extracted spectrum (n x hpix) where n is number of traces spec_invar - inverse variance at each point in extracted spectrum spec_mask - mask for invalid/suspect points in spectrum image_model - only if return_model = True. """ # def extract(ccd,t_coeffs,i_coeffs=None,s_coeffs=None,p_coeffs=None,readnoise=1,gain=1,return_model=False,fact,verbose=False): # """ Extraction. # """ ### t_coeffs are from fiber flat - need to shift based on actual exposure #################################################### ### Prep Needed variables/empty arrays ######### #################################################### ### CCD dimensions and number of fibers hpix = np.shape(ccd)[1] vpix = np.shape(ccd)[0] num_fibers = np.shape(t_coeffs)[1] #################################################### ##### First refine horizontal centers (fit ##### ##### traces from data ccd using fiber flat ##### ##### as initial estimate) ##### #################################################### ta = time.time() ### Start time of trace refinement fact = 20 #do 1/fact * available points ### Empty arrays rough_pts = int(np.ceil(hpix/fact)) xc_ccd = np.zeros((num_fibers,rough_pts)) yc_ccd = np.zeros((num_fibers,rough_pts)) inv_chi = np.zeros((num_fibers,rough_pts)) if verbose: print("Refining trace centers") for i in range(num_fibers): for j in range(0,hpix,fact): ### set coordinates, gaussian parameters from coeffs jadj = int(np.floor(j/fact)) yj = (j-hpix/2)/hpix yc_ccd[i,jadj] = j xc = t_coeffs[2,i]*yj**2+t_coeffs[1,i]*yj+t_coeffs[0,i] # Ij = i_coeffs[2,i]*yj**2+i_coeffs[1,i]*yj+i_coeffs[0,i] #May use later for normalization sigj = s_coeffs[2,i]*yj**2+s_coeffs[1,i]*yj+s_coeffs[0,i] powj = p_coeffs[2,i]*yj**2+p_coeffs[1,i]*yj+p_coeffs[0,i] ### Don't try to fit any bad trace sections if np.isnan(xc): xc_ccd[i,jadj] = np.nan inv_chi[i,jadj] = 0 else: ### Take subset of ccd of interest, xpad pixels to each side of peak xpad = 7 xvals = np.arange(-xpad,xpad+1) xj = int(xc) xwindow = xj+xvals xvals = xvals[(xwindow>=0)*(xwindow<vpix)] zorig = gain*ccd[xj+xvals,j] ### If empty slice, don't try to fit if len(zorig)<1: xc_ccd[i,jadj] = np.nan inv_chi[i,jadj] = 0 continue invorig = 1/(abs(zorig)+readnoise**2) ### inverse variance ### Don't try to fit profile for very low SNR peaks if np.max(zorig)<20: xc_ccd[i,jadj] = np.nan inv_chi[i,jadj] = 0 else: ### Fit for center (mn_new), amongst other values # mn_new, hght, bg = fit_mn_hght_bg(xvals,zorig,invorig,sigj,xc-xj-1,sigj,powj=powj) mn_new, hght, bg = linear_mn_hght_bg(xvals,zorig,invorig,sigj,xc-xj-1,power=powj) fitorig = sf.gaussian(xvals,sigj,mn_new,hght,power=powj) inv_chi[i,jadj] = 1/sum((zorig-fitorig)**2*invorig) ### Shift from relative to absolute center xc_ccd[i,jadj] = mn_new+xj+1 ##################################################### #### Now with new centers, refit trace coefficients # ##################################################### tmp_poly_ord = 6 ### Use a higher order for a closer fit over entire trace t_coeffs_ccd = np.zeros((tmp_poly_ord+1,num_fibers)) for i in range(num_fibers): #Given orientation makes more sense to swap x/y mask = ~np.isnan(xc_ccd[i,:]) ### Mask bad points ### build profile matrix over good points profile = np.ones((len(yc_ccd[i,:][mask]),tmp_poly_ord+1)) for order in range(tmp_poly_ord): profile[:,order+1] = ((yc_ccd[i,:][mask]-hpix/2)/hpix)**(order+1) noise = np.diag(inv_chi[i,:][mask]) if len(xc_ccd[i,:][mask])>(tmp_poly_ord+1): ### Chi^2 fit tmp_coeffs, junk = sf.chi_fit(xc_ccd[i,:][mask],profile,noise) else: ### if not enough points to fit, call entire trace bad tmp_coeffs = np.nan*np.ones((tmp_poly_ord+1)) t_coeffs_ccd[:,i] = tmp_coeffs tb = time.time() ### Start time of extraction/end of trace refinement if verbose: print("Trace refinement time = {}s".format(tb-ta)) ### Uncomment below to see plot of traces # for i in range(num_fibers): # ys = (np.arange(hpix)-hpix/2)/hpix # xs = t_coeffs_ccd[2,i]*ys**2+t_coeffs_ccd[1,i]*ys+t_coeffs_ccd[0,i] # yp = np.arange(hpix) # plt.plot(yp,xs) # plt.show() # plt.close() ########################################################### ##### Finally, full extraction with refined traces ######## ########################################################### ### Make empty arrays for return values spec = np.zeros((num_fibers,hpix)) spec_invar = np.zeros((num_fibers,hpix)) spec_mask = np.ones((num_fibers,hpix),dtype=bool) chi2red_array = np.zeros((num_fibers,hpix)) if return_model: image_model = np.zeros((np.shape(ccd))) ### Used for evaluation ### Run once for each fiber for i in range(num_fibers): #slit_num = np.floor((i)/4)#args.telescopes) # Use with slit flats if verbose: print("extracting trace {}".format(i+1)) ### in each fiber loop run through each trace for j in range(hpix): yj = (j-hpix/2)/hpix xc = np.poly1d(t_coeffs_ccd[::-1,i])(yj) # Ij = i_coeffs[2,i]*yj**2+i_coeffs[1,i]*yj+i_coeffs[0,i] sigj = s_coeffs[2,i]*yj**2+s_coeffs[1,i]*yj+s_coeffs[0,i] powj = p_coeffs[2,i]*yj**2+p_coeffs[1,i]*yj+p_coeffs[0,i] ### If trace center is undefined mask the point if np.isnan(xc): spec_mask[i,j] = False else: ### Set values to use in extraction xpad = 5 ### can't be too big or traces start to overlap xvals = np.arange(-xpad,xpad+1) xj = int(xc) xwindow = xj+xvals xvals = xvals[(xwindow>=0)*(xwindow<vpix)] zorig = gain*ccd[xj+xvals,j] fitorig = sf.gaussian(xvals,sigj,xc-xj-1,hght,power=powj) ### If too short, don't fit, mask point if len(zorig)<1: spec[i,j] = 0 spec_mask[i,j] = False continue invorig = 1/(abs(zorig)+readnoise**2) ### don't try to extract for very low signal if np.max(zorig)<20: continue else: ### Do nonlinear fit for center, height, and background mn_new, hght, bg = fit_mn_hght_bg(xvals,zorig,invorig,sigj,xc-xj-1,sigj/8,powj=powj) # mn_new, hght, bg = linear_mn_hght_bg(xvals,zorig,invorig,sigj,xc-xj-1,power=powj) ### Use fitted values to make best fit arrays fitorig = sf.gaussian(xvals,sigj,mn_new,hght,power=powj) xprecise = np.linspace(xvals[0],xvals[-1],100) fitprecise = sf.gaussian(xprecise,sigj,mn_new,hght,power=powj) ftmp = sum(fitprecise)*np.mean(np.ediff1d(xprecise)) #Following if/else handles failure to fit if ftmp==0: fitnorm = np.zeros(len(zorig)) else: fitnorm = fitorig/ftmp ### Get extracted flux and error fstd = sum(fitnorm*zorig*invorig)/sum(fitnorm**2*invorig) invorig = 1/(readnoise**2 + abs(fstd*fitnorm)) chi2red = np.sum((fstd*fitnorm+bg-zorig)**2*invorig)/(len(zorig)-3) ### Now set up to do cosmic ray rejection rej_min = 0 loop_count=0 while rej_min==0: pixel_reject = cosmic_ray_reject(zorig,fstd,fitnorm,invorig,S=bg,threshhold=0.3*np.mean(zorig),verbose=True) rej_min = np.min(pixel_reject) ### Once no pixels are rejected, re-find extracted flux if rej_min==0: ### re-index arrays to remove rejected points zorig = zorig[pixel_reject==1] invorig = invorig[pixel_reject==1] xvals = xvals[pixel_reject==1] ### re-do fit (can later cast this into a separate function) mn_new, hght, bg = fit_mn_hght_bg(xvals,zorig,invorig,sigj,xc-xj-1,sigj/8,powj=powj) # mn_new, hght, bg = linear_mn_hght_bg(xvals,zorig,invorig,sigj,xc-xj-1,power=powj) fitorig = sf.gaussian(xvals,sigj,mn_new,hght,power=powj) xprecise = np.linspace(xvals[0],xvals[-1],100) fitprecise = sf.gaussian(xprecise,sigj,mn_new,hght,power=powj) ftmp = sum(fitprecise)*np.mean(np.ediff1d(xprecise)) fitnorm = fitorig/ftmp fstd = sum(fitnorm*zorig*invorig)/sum(fitnorm**2*invorig) invorig = 1/(readnoise**2 + abs(fstd*fitnorm)) chi2red = np.sum((fstd*fitnorm+bg-zorig)**2*invorig)/(len(zorig)-3) ### if more than 3 points are rejected, mask the extracted flux if loop_count>3: spec_mask[i,j] = False break loop_count+=1 ### Set extracted spectrum value, inverse variance spec[i,j] = fstd spec_invar[i,j] = sum(fitnorm**2*invorig) chi2red_array[i,j] = chi2red if return_model and not np.isnan(fstd): ### Build model, if desired image_model[xj+xvals,j] += (fstd*fitnorm+bg)/gain ### If a nan came out of the above routine, zero it and mask if np.isnan(spec[i,j]): spec[i,j] = 0 spec_mask[i,j] = False if verbose: print("Average reduced chi^2 = {}".format(np.mean(chi2red))) if return_model: return spec, spec_invar, spec_mask, image_model else: return spec, spec_invar, spec_mask
def find_trace_coeffs(image,pord,fiber_space,num_points=None,num_fibers=None,vertical=False,return_all_coeffs=True,skip_peaks=0): """ Polynomial fitting for trace coefficients. Packs into interval [-1,1] INPUTS: image - 2D ccd image on which you'd like to find traces pord - polynomial order to fit trace positions fiber_space - estimate of fiber spacing in pixels num_points - number of cross sections to average for trace (if None, will set to 1/20 length or 2*pord (whichever is greater)) num_fibers - number of fibers (if None, will auto-detect) vertical - True if traces run vertical. False if horizontal. return_all_coeffs - if False, only returns trace_poly_coeffs OUTPUTS: t_coeffs - nx(pord+1) array where n is number of detected traces gives fitted coeffs for each trace i_coeffs - intensity along each trace s_coeffs - sigma (for gaussian fit) along each trace p_coeffs - power (for pseudo gaussian fit) along each trace """ def find_peaks(array,bg_cutoff=None,mx_peaks=None,skip_peaks=0): """ Finds peaks of a 1D array. Assumes decent signal to noise ratio, no anomalies Assumes separation of at least 5 units between peaks """ ###find initial peaks (center is best in general, but edge is okay here) xpix = len(array) px = 2 pcnt = 0 skcnt = 0 peaks = np.zeros(len(array)) if mx_peaks is None: mx_peaks = len(array) if bg_cutoff is None: bg_cutoff = 0.5*np.mean(array) while px<xpix: # if trct>=num_fibers: # break # y = yvals[0] if array[px-1]>bg_cutoff and array[px]<array[px-1] and array[px-1]>array[px-2]: #not good for noisy if skcnt < skip_peaks: skcnt += 1 #Increment skip counter continue else: peaks[pcnt] = px-1 px += 5 #jump past peak pcnt+=1 #increment peak counts skcnt += 1 #increment skip counter if pcnt >= mx_peaks: break else: px+=1 peaks = peaks[0:pcnt] return peaks def find_fibers(image): """ Estimates number of fibers. Assumes roughly square image. THIS DOESN'T WORK RIGHT YET """ shrt_ax = int(min(np.shape(image))/2) crsx_pts = min(10,shrt_ax) tr_ul_lr = np.zeros(crsx_pts) #traces cutting upper left to lower right tr_ur_ll = np.zeros(crsx_pts) #traces cutting upper right to lower left for i in range(crsx_pts): ul_lr = np.ravel(image)[2*i:np.size(image):(len(image[0])+1)] ur_ll = np.ravel(image)[(len(image[0])-1-2*i):np.size(image):(len(image[0])-1)] tr_ul_lr[i] = len(find_peaks(ul_lr)) tr_ur_ll[i] = len(find_peaks(ur_ll)) tr_ul_lr = int(np.median(tr_ul_lr)) tr_ur_ll = int(np.median(tr_ur_ll)) if tr_ul_lr > tr_ur_ll: return tr_ul_lr, True else: return tr_ur_ll, False ### Force horizontal if vertical: image = image.T ### Find number of fibers and direction (upper left to lower right or opposite) tmp_num_fibers, fiber_dir = find_fibers(image) if num_fibers is None: num_fibers=tmp_num_fibers ### Select number of points to use in tracing if num_points is None: num_points = max(2*pord,int(np.shape(image)[1]/20)) ### Make all empty arrays for trace finding xpix = np.shape(image)[0] ypix = np.shape(image)[1] yspace = int(np.floor(ypix/(num_points+1))) yvals = yspace*(1+np.arange(num_points)) xtrace = np.nan*np.ones((num_fibers,num_points)) #xpositions of traces ytrace = np.zeros((num_fibers,num_points)) #ypositions of traces sigtrace = np.zeros((num_fibers,num_points)) #standard deviation along trace powtrace = np.zeros((num_fibers,num_points)) #pseudo-gaussian power along trace Itrace = np.zeros((num_fibers,num_points)) #relative intensity of flat at trace chi_vals = np.zeros((num_fibers,num_points)) #returned from fit_trace bg_cutoff = 1.05*np.median(image) #won't fit values below this intensity ### Put in initial peak guesses peaks = find_peaks(image[:,yvals[0]],bg_cutoff=bg_cutoff,mx_peaks=num_fibers,skip_peaks=skip_peaks) xtrace[:len(peaks)-1,0] = peaks[:-1] ### have to cut off last point - trace wanders off ccd ytrace[:,0] = yvals[0]*np.ones(len(ytrace[:,0])) ###From initial peak guesses fit for more precise location for i in range(num_fibers): y = yvals[0] if not np.isnan(xtrace[i,0]): xtrace[i,0], Itrace[i,0], sigtrace[i,0], powtrace[i,0], chi_vals[i,0] = fit_trace(xtrace[i,0],y,image) else: Itrace[i,0], sigtrace[i,0], powtrace[i,0], chi_vals[i,0] = np.nan, np.nan, np.nan, np.nan for i in range(1,len(yvals)): y = yvals[i] crsxn = image[:,y] ytrace[:,i] = y for j in range(num_fibers): if not np.isnan(xtrace[j,i-1]): #set boundaries lb = int(xtrace[j,i-1]-fiber_space/2) ub = int(xtrace[j,i-1]+fiber_space/2) #cutoff at edges if lb<0: lb = 0 if ub > xpix: ub = xpix #set subregion xregion = crsxn[lb:ub] #only look at if max is reasonably high (don't try to fit background) if np.max(xregion)>bg_cutoff: #estimate of trace position based on tallest peak xtrace[j,i] = np.argmax(xregion)+lb #quadratic fit for sub-pixel precision # print xtrace[j,i] xtrace[j,i], Itrace[j,i], sigtrace[j,i], powtrace[j,i], chi_vals[j,i] = fit_trace(xtrace[j,i],y,image) else: xtrace[j,i], Itrace[j,i], sigtrace[j,i], sigtrace[j,i], chi_vals[j,i] = np.nan, np.nan, np.nan, np.nan, np.nan else: xtrace[j,i], Itrace[j,i], sigtrace[j,i], sigtrace[j,i], chi_vals[j,i] = np.nan, np.nan, np.nan, np.nan, np.nan Itrace /= np.median(Itrace) #Rescale intensities #Finally fit x vs. y on traces. Start with quadratic for simple + close enough t_coeffs = np.zeros((3,num_fibers)) i_coeffs = np.zeros((3,num_fibers)) s_coeffs = np.zeros((3,num_fibers)) p_coeffs = np.zeros((3,num_fibers)) for i in range(num_fibers): #Given orientation makes more sense to swap x/y mask = ~np.isnan(xtrace[i,:]) profile = np.ones((len(ytrace[i,:][mask]),3)) #Quadratic fit profile[:,1] = (ytrace[i,:][mask]-ypix/2)/ypix #scale data to get better fit profile[:,2] = ((ytrace[i,:][mask]-ypix/2)/ypix)**2 noise = np.diag(chi_vals[i,:][mask]) if len(xtrace[i,:][mask])>3: tmp_coeffs, junk = sf.chi_fit(xtrace[i,:][mask],profile,noise) tmp_coeffs2, junk = sf.chi_fit(Itrace[i,:][mask],profile,noise) tmp_coeffs3, junk = sf.chi_fit(sigtrace[i,:][mask],profile,noise) tmp_coeffs4, junk = sf.chi_fit(powtrace[i,:][mask],profile,noise) else: tmp_coeffs = np.nan*np.ones((3)) tmp_coeffs2 = np.nan*np.ones((3)) tmp_coeffs3 = np.nan*np.ones((3)) tmp_coeffs4 = np.nan*np.ones((3)) t_coeffs[0,i] = tmp_coeffs[0] t_coeffs[1,i] = tmp_coeffs[1] t_coeffs[2,i] = tmp_coeffs[2] i_coeffs[0,i] = tmp_coeffs2[0] i_coeffs[1,i] = tmp_coeffs2[1] i_coeffs[2,i] = tmp_coeffs2[2] s_coeffs[0,i] = tmp_coeffs3[0] s_coeffs[1,i] = tmp_coeffs3[1] s_coeffs[2,i] = tmp_coeffs3[2] p_coeffs[0,i] = tmp_coeffs4[0] p_coeffs[1,i] = tmp_coeffs4[1] p_coeffs[2,i] = tmp_coeffs4[2] if return_all_coeffs: return t_coeffs, i_coeffs, s_coeffs, p_coeffs else: return t_coeffs