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
def arc_peaks(data,wvln,invar,ts,sampling_est=3,pad=4):
""" Finds the position, wavelength, amplitude, width, etc. of distinct
peaks along an extracted arc frame.
INPUTS:
data - extracted arc frame (a x b x c) a=telescope, b=fiber,
c=pixel
wvln - wavelength corresponding to each point in 'data'
invar - inverse variance of each point in 'data'
ts - telescope number to use ('a' in data)
sampling_est - estimated FWHM in pixels (must be integer)
pad - width to fit to either side of gaussian max
OUTPUTS:
pos_d - position in pixels of each peak
wl_d, mx_it_d, stddev_d, chi_d, err_d
"""
#use dictionaries since number of peaks per fiber varies
mx_it_d = dict() #max intensities of each peak
stddev_d = dict() #FWHM of each peak
wl_d = dict() #position of each peak in wavelength space
# herm3_d = dict() #amplitude of third order hermite polynomial
# herm4_d = dict() #amplitude of fourth order hermite polynomial
pos_d = dict() #position of each peak in pixel space
chi_d = dict() #reduced chi squared
err_d = dict() #stddev parameter error
for i in range(len(data[0,:,0])):
##Optional - use scipy to get initial guesses of peak locations
##Problem is this misses many of the low amplitude peaks.
#pos_est = np.array(sig.find_peaks_cwt(data[ts,i,:],np.arange(3,4)))
#Since spectrum has ~no background, can use my own peak finder.
pos_est = np.zeros((len(data[ts,i,:])),dtype=int)
for j in range(2*sampling_est,len(data[ts,i,:])-2*sampling_est):
#if point is above both its neighbors, call it a peak
if data[ts,i,j]>data[ts,i,j-1] and data[ts,i,j]>data[ts,i,j+1]:
pos_est[j] = j
#Then remove extra elements from pos_est
pos_est = pos_est[np.nonzero(pos_est)[0]]
#Cut out any that are within 2*sampling of each other (won't be able to fit well)
pos_diff = ediff1d(pos_est)
if np.count_nonzero(pos_diff<(2*sampling_est))>0:
close_inds = np.nonzero(pos_diff<(2*sampling_est))[0]
### Try 1x sampling and see if that gives any more peaks in the end...
# if np.count_nonzero(pos_diff<(1*sampling_est))>0:
# close_inds = np.nonzero(pos_diff<(1*sampling_est))[0]
close_inds = np.concatenate((close_inds,close_inds+1))
close_inds = np.unique(close_inds)
close_inds = np.sort(close_inds)
pos_est = np.delete(pos_est,close_inds)
#Also cut out any with a zero 1 pixels or less to either side
# tspl = pos_est-2
ospl = pos_est-1
ospr = pos_est+1
# tspr = pos_est+2
zero_inds = np.zeros((len(pos_est)))
for tt in range(len(zero_inds)):
# if tt < 6:
# print tt, ":"
# print data[ts,i,tspl[tt]]==0
# print data[ts,i,ospl[tt]]==0
# print data[ts,i,ospr[tt]]==0
# print data[ts,i,tspr[tt]]==0
if data[ts,i,ospl[tt]]==0 or data[ts,i,ospr[tt]]==0:# or data[ts,i,tspl[tt]]==0 or data[ts,i,tspr[tt]]==0:
zero_inds[tt]=1
# print zero_inds
# print pos_est
# plt.plot(data[0,0,:])
# plt.figure()
# plt.plot(pos_est)
# plt.show()
pos_est = pos_est[zero_inds==0]
# if i == 0:
# print pos_est
#variable length arrays to dump into dictionary
num_pks = len(pos_est)
mx_it = zeros((num_pks))
stddev = zeros((num_pks))
# herm3 = zeros((num_pks))
# herm4 = zeros((num_pks))
pos = zeros((num_pks))
chi = zeros((num_pks))
err = zeros((num_pks))
pos_idx = zeros((num_pks),dtype=int)
# slp = zeros((num_pks))
#Now fit gaussian with background to each (can improve function later)
for j in range(num_pks):
pos_idx[j] = pos_est[j]
xarr = pos_est[j] + np.arange(-pad,pad,1)
xarr = xarr[(xarr>0)*(xarr<2048)]
yarr = data[ts,i,:][xarr]
wlarr = wvln[ts,i,:][xarr]
invarr = invar[ts,i,:][xarr]
try:
params, errarr = sf.gauss_fit(wlarr,yarr,invr=invarr,fit_background='n')
# params = sf.fit_gauss_herm1d(wlarr,yarr,invarr)
# errarr = np.diag(np.ones(len(params)))
except RuntimeError:
params = np.zeros(3)
# params = np.zeros(5)
pos_idx[j] = 0
tot = sf.gaussian(wlarr,abs(params[0]),params[1],params[2])#,params[3],params[4])
# tot = sf.gauss_herm1d(wlarr,abs(params[0]),params[1],params[2],params[3],params[4])
chi_sq = sum((yarr-tot)**2*invarr)
chi[j] = chi_sq/len(yarr)
if chi_sq/len(yarr) > 10: #arbitrary cutoff
params = np.zeros(5)
pos_idx[j] = 0
mx_it[j] = params[2] #height
stddev[j] = params[0]#*2*sqrt(2*log(2)) #converted from std dev
# herm3[j] = params[3]
# herm4[j] = params[4]
err[j] = np.sqrt(errarr[0,0])
pos[j] = params[1] #center
# slp[j] = params[4] #bg_slope
mx_it_d[i] = mx_it[np.nonzero(pos)[0]] #Remove zero value points
stddev_d[i] = stddev[np.nonzero(pos)[0]]
# herm3_d[i] = herm3[np.nonzero(pos)[0]]
# herm4_d[i] = herm4[np.nonzero(pos)[0]]
wl_d[i] = pos[np.nonzero(pos)[0]]
pos_d[i] = pos_idx[np.nonzero(pos)[0]]
plt.show()
chi_d[i] = chi[np.nonzero(pos)[0]]
err_d[i] = err[np.nonzero(pos)[0]]
# if i == 0:
# plt.plot(pos_idx,data[ts,i,:][pos_idx],'ks')
# plt.plot(data[ts,i,:])
# plt.show()
return pos_d, wl_d, mx_it_d, stddev_d, chi_d, err_d
# plt.show()
xmx = -coeffsn[1]/(2*coeffsn[2])
zmx = coeffsn[2]*xmx**2+coeffsn[1]*xmx+coeffsn[0]
line_amps[j,k] = zmx[0]#/Is[j,k]
'''
#Lines to find sum of 10px square around the peak (total intensity)
intensity = sum(arc[-5+xs[j,k]:6+xs[j,k],-5+k:6+k])
line_amps[j,k] = intensity/Is[j,k]
'''
#print zmx[0]/Is
pad2 = 7
xsub = np.arange(-pad2,pad2+1)+int(xs[j,k])
xsub = xsub[(xsub>=0)*(xsub<xpix)]
zvals = arc[xsub,k]
try:
params, errarr = sf.gauss_fit(xsub,zvals)
except RuntimeError:
print("No fit for fiber {fib} at pixel {px}".format(fib=j,px=k))
line_gauss[j,k] = 0
else:
line_gauss[j,k] = params[2] #-background?
# if j==0 and k==332:
# fit = sf.gaussian(xsub,abs(params[0]),params[1],params[2],params[3],params[4])
# plt.plot(xsub,zvals,xsub,fit)
# print xs[j,k]
# plt.show()
nonzero_count+=1
else:
line_amps[j,k] = 0
line_gauss[j,k] = 0
zero_count+=1
