diff --git a/MOSFIRE/Fit.py b/MOSFIRE/Fit.py index d0431e5..bb9a0ed 100644 --- a/MOSFIRE/Fit.py +++ b/MOSFIRE/Fit.py @@ -19,37 +19,36 @@ #from MosfireDrpLog import debug, info, warning, error import unittest -def xcor(a,b,lags): - if len(a) != len(b): - error("cross correlation (xcor) requires a and b " - "to be of same length") - raise Exception( - "cross correlation (xcor) requires a and b " - "to be of same length") - cors = np.zeros(len(lags)) - - a_pad = np.zeros(len(a)+len(lags)) - b_pad = np.zeros(len(b)+len(lags)) - - st = np.argmin(np.abs(lags)) - a_pad[st:st+len(a)] = a - b_pad[st:st+len(b)] = b - - for i in range(len(lags)): - cors[i] = np.correlate(a_pad, np.roll(b_pad, lags[i]), 'valid') - - return cors - -def xcor_peak(a, b, lags): - '''Return the peak position in units of lags''' - - N = len(lags) - xcs = xcor(a, b, lags) +def find_shift(a, b, guess=0, maxshift=3): + '''Find the offset between two 1D arrays based on cross correlation. + + Intended to replace xcor and xcor_peak. Works out to be about 2x speed as + xcor_peak based on %timeit in ipython. + + Uses the following conditions for a point to be valid as consideration + for the shift result: + - second derivative of the correlation is negative (i.e. at local maxima) + - shift <= maxshift + + Given those conditions are met, the highest cross correlation value gives + the best value for the shift (or offset) between the two arrays. + ''' + a = a.filled(0) + b = b.filled(0) + assert len(a) == len(b) + shifts = np.arange(len(a)) - len(a)/2 + corr = np.correlate(a, b, mode='same') + first_derivative = np.gradient(corr) + second_derivative = np.gradient(first_derivative) - return lags[np.argmax(xcs)] + valid = np.where((second_derivative < 0) & (abs(shifts-guess) < maxshift))[0] + idx = np.argmax(corr[valid]) + bestcorrindex = valid[idx] + shift = shifts[bestcorrindex] + return shift, shifts, corr # TODO: Document mpfit_* functions diff --git a/MOSFIRE/Options.py b/MOSFIRE/Options.py index 413a37d..efa8f85 100644 --- a/MOSFIRE/Options.py +++ b/MOSFIRE/Options.py @@ -24,5 +24,7 @@ "version": 2, "fractional-wavelength-search": 0.99935, # used in determining oned wavelength solutions "chebyshev-degree": 5, # polynomial order for fitting wavelengths + "smooth": False, # smooth wavelength solution polynomial coefficients in Y direction + "smooth_scale": 15, # smoothing scale in pixels } diff --git a/MOSFIRE/Wavelength.py b/MOSFIRE/Wavelength.py index ed63ef5..6cad6a2 100644 --- a/MOSFIRE/Wavelength.py +++ b/MOSFIRE/Wavelength.py @@ -358,7 +358,7 @@ def fit_lambda(maskname, tock = time.time() - multicore = True + multicore = False if multicore: p = Pool() solutions = p.map(fit_lambda_helper, list(range(len(bs.ssl)))) @@ -1122,7 +1122,9 @@ def pick_linelist(header, neon=False, argon=False, short_exp = False): 13421.579]) if (band == 'K') and short_exp: - # remove: 19518.4784 , 19593.2626 , 19618.5719 ,19678.046 ,19839.7764 ,20193.1799 ,20499.237 ,21279.1406 ,21580.5093 ,21711.1235 , 21873.507 ,22460.4183 ,22690.1765 ,22985.9156,23914.55, 24041.62,22742.1907 + # remove: 19518.4784 , 19593.2626 , 19618.5719 ,19678.046 ,19839.7764 , + # 20193.1799 ,20499.237 ,21279.1406 ,21580.5093 ,21711.1235 , 21873.507 + #,22460.4183 ,22690.1765 ,22985.9156,23914.55, 24041.62,22742.1907 lines = np.array([ 19642.4493 , 19701.6455 , 19771.9063 , @@ -1336,7 +1338,7 @@ def refine_wavelength_guess(wave,spec,linelist): """ # find what is the average peak height to construct the reference - # spectrum template + # spectrum template peaks = signal.find_peaks_cwt(spec,np.array([2.0,4.0]),noise_perc=5.0,min_snr =1.5) avePeak = np.mean(spec[peaks]) @@ -1348,11 +1350,11 @@ def refine_wavelength_guess(wave,spec,linelist): corr = signal.correlate(spec,refSpec,mode='same') lags = np.arange(len(spec))-len(spec)/2 - # peak velocity corresponding to the pixel peak + # peak velocity corresponding to the pixel peak peakInd = np.argmax(corr) peakLag = lags[peakInd] - # return the number of pixels that need to be shifted + # return the number of pixels that need to be shifted return -peakLag def plot_mask_solution_ds9(fname, maskname, options): @@ -1419,49 +1421,95 @@ def estimate_half_power_points(slitno, header, bs): hpp = Filters.hpp[band] return [ np.argmin(np.abs(ll-hpp[0])), np.argmin(np.abs(ll-hpp[1])) ] - - -def xcor_known_lines(lines, ll, spec, spec0, options): +def find_known_lines(lines, wavs, spec, options, yhere=None, plot=False): """ lines[N]: list of lines in wavelength units - ll[2048]: lambda vector + wavs[2048]: lambda vector spec[2048]: spectrum vector (as function of lambda) options: wavelength options """ - inf = np.inf - dxs = [] - sigs = [] + from astropy.modeling import models, fitting - pix = np.arange(2048.) - - for lam in lines: - f = options["fractional-wavelength-search"] - roi = np.where((f*lam < ll) & (ll < lam/f))[0] - - - - if not roi.any(): - dxs.append(inf) - sigs.append(inf) - continue - - lags = np.arange(-len(roi)/2, len(roi)/2) - cors = Fit.xcor(spec[roi], spec0[roi], lags) - - fit = Fit.mpfitpeak(lags, cors) - - if (fit.perror is None) or (fit.status < 0): - dxs.append(inf) - sigs.append(inf) - continue - - dxs.append(fit.params[1]) - sigs.append(fit.params[2]) - - return list(map(np.array, [dxs, sigs])) + pix = np.arange(len(spec)) + xinits = [] + yinits = [] + xs = [] + sxs = [] + sigmas = [] + model = None + for line in lines: + lambda_frac = options["fractional-wavelength-search"] + roi = (lambda_frac*line < wavs) & (wavs < line/lambda_frac) + if roi.any(): + linepix = (min(pix[roi]) + max(pix[roi]))/2 + peak = max(spec[roi]) + g_init = models.Gaussian1D(amplitude=peak, mean=linepix, stddev=1) + g_init.bounds['amplitude'] = (0.5*peak, None) + g_init.bounds['mean'] = (linepix*lambda_frac, linepix/lambda_frac) + g_init.bounds['stddev'] = (1,3) + fit_g = fitting.LevMarLSQFitter() + g = fit_g(g_init, pix, spec) + + if fit_g.fit_info['ierr'] not in [1,2,3,4]: + xs.append(0.0) + xinits.append(linepix) + yinits.append(peak) + sxs.append(np.inf) + else: + try: + pcov = fit_g.fit_info['cov_x'] + perr = np.sqrt(np.diag(pcov)) + except ValueError: + xs.append(0.0) + xinits.append(linepix) + yinits.append(peak) + sxs.append(np.inf) + else: + xs.append(g.mean.value) + xinits.append(linepix) + yinits.append(peak) + sxs.append(perr[1]) + sigmas.append(g.stddev.value) + if model is None: + model = g + else: + model += g + else: + xs.append(0.0) + sxs.append(np.inf) + + residuals = spec-model(pix) + tot = np.ma.sum(np.ma.abs(residuals)) + if plot is True: + from matplotlib import pyplot as pl + roi_plot = (min(lines) < wavs) & (wavs < max(lines)) + + pl.figure(figsize=(18,9)) + + pl.subplot(2,1,1) + if yhere is not None: + pl.title(f"Line {yhere:d}") + pl.plot(pix, spec, 'k-') + pl.plot(pix, model(pix), 'g-', alpha=0.7) + pl.xlim(min(pix[roi_plot]), max(pix[roi_plot])) + for i,xi in enumerate(xinits): + pl.plot([xi, xi], [yinits[i]+100, yinits[i]+1000], 'b-') + + pl.subplot(2,1,2) + xmin = min(pix[roi_plot]) + xmax = max(pix[roi_plot]) + pl.title(f"Residuals = {tot:.0f}") + pl.plot(pix, residuals, 'b-') + pl.xlim(xmin, xmax) + pl.axvspan(xmin, xmax, ymin=1000, ymax=10000, facecolor='r', alpha=0.5) + pl.axvspan(xmin, xmax, ymin=-10000, ymax=-1000, facecolor='r', alpha=0.5) + pl.ylim(min([-100, residuals.min()-50]), max([900, residuals.max()+50])) + pl.show() + + return list(map(np.array, [xs, sxs, sigmas])) -def find_known_lines(lines, ll, spec, options): +def find_known_lines_old(lines, ll, spec, options): """ lines[N]: list of lines in wavelength units ll[2048]: lambda vector @@ -1539,6 +1587,8 @@ def fit_chebyshev_to_lines(xs, sxs, lines, options): raise Exception("Units fed to this function are likely in micron but " "should be in angstrom") +# import pdb; pdb.set_trace() + cfit = CV.chebfit(xs[ok], lines[ok], options["chebyshev-degree"]) delt = CV.chebval(xs[ok], cfit) - lines[ok] @@ -2014,7 +2064,8 @@ def __call__(self, event): actions_mouseless = {".": self.fastforward, "n": self.nextobject, "p": self.prevobject, "q": self.quit, "r": self.reset, "f": - self.fit_event, "k": self.toggle_sigma_clip, "\\": self.fit_event, "b": self.toggle_noninteractive} + self.fit_event, "k": self.toggle_sigma_clip, "\\": self.fit_event, + "b": self.toggle_noninteractive} actions = { "c": self.shift, "d": self.drop_point, "z": self.zoom, "x": self.unzoom, "s": self.savefig} @@ -2203,6 +2254,64 @@ def construct_model(slitno): # # Two dimensional wavelength fitting # +def fit_smooth_solution(sol_2d, order=3, top=None, bottom=None): + positions = sol_2d['positions'] + if top is None: + top = max(positions) + if bottom is None: + bottom = min(positions) + all_positions = np.arange(bottom+1, top, 1) + coeffs = sol_2d['coeffs'] + + s_sol_2d = {'coeffs': np.zeros((len(all_positions), sol_2d['coeffs'].shape[1])), + 'positions': all_positions, + 'delts': sol_2d['delts'], + 'lambdaRMS': sol_2d['lambdaRMS'], + 'lambdaMAD': sol_2d['lambdaMAD']} + + from astropy.modeling import models, fitting + p_init = models.Polynomial1D(order) + fitter = fitting.LinearLSQFitter() + p0 = fitter(p_init, positions, coeffs[:,0]) + p1 = fitter(p_init, positions, coeffs[:,1]) + p2 = fitter(p_init, positions, coeffs[:,2]) + p3 = fitter(p_init, positions, coeffs[:,3]) + p4 = fitter(p_init, positions, coeffs[:,4]) + p5 = fitter(p_init, positions, coeffs[:,5]) + + + s_sol_2d['coeffs'][:,0] = p0(all_positions) + s_sol_2d['coeffs'][:,1] = p1(all_positions) + s_sol_2d['coeffs'][:,2] = p2(all_positions) + s_sol_2d['coeffs'][:,3] = p3(all_positions) + s_sol_2d['coeffs'][:,4] = p4(all_positions) + s_sol_2d['coeffs'][:,5] = p5(all_positions) + + return s_sol_2d + + +def smooth_solution(sol_2d, filter_size=15): + from scipy.signal import medfilt + + s_coeffs = np.zeros(sol_2d['coeffs'].shape) + sorted_order = np.argsort(sol_2d['positions']) + + # for each of the coefficients, pass a median boxcar over it + info(f"Smoothing fit coefficients in Y direction") + for j in range(sol_2d['coeffs'].shape[1]): + coeff = sol_2d['coeffs'][:,j][sorted_order] + smoothed_coeff = medfilt(coeff, kernel_size=filter_size) + s_coeffs[:,j] = smoothed_coeff[sol_2d['positions']-min(sol_2d['positions'])] + + sol_2d_smoothed = {'coeffs': s_coeffs, + 'delts': sol_2d['delts'], + 'lambdaRMS': sol_2d['lambdaRMS'], + 'lambdaMAD': sol_2d['lambdaMAD'], + "positions": sol_2d['positions']} + + return sol_2d_smoothed + + def fit_outwards_refit(data, bs, sol_1d, lines, options, start, bottom, top, slitno, linelist2=None, data2=None, sol_1d2=None): '''Fit Chebyshev polynomials across a slit bounded by bottom to top. @@ -2253,7 +2362,7 @@ def fit_outwards_refit(data, bs, sol_1d, lines, options, start, bottom, top, pix = np.arange(2048.) linelist = lines - def fit_parameters(yhere): + def fit_parameters(yhere, guess=0, pmlines=2): """ Return chebyshev fit to a pixel column 2014 June 17 MK- Added a second set a variables to indicate that there @@ -2261,19 +2370,22 @@ def fit_parameters(yhere): for the arc line data. We want to fit both Ne and Argon simultaneously. - """ + """ cfit = sol_1d[1] - spec_here = np.ma.median(data[int(yhere)-2:int(yhere)+2, :], axis=0) - shift = Fit.xcor_peak(spec_here, spec0, lags) + spec_here = np.ma.median(data[int(yhere)-pmlines:int(yhere)+pmlines, :], axis=0) + shift, shifts, corr = Fit.find_shift(spec_here, spec0, guess=guess, maxshift=4*pmlines) + ll_here = CV.chebval(pix - shift, cfit) [xs, sxs, sigmas] = find_known_lines(linelist, - ll_here, spec_here, options) + ll_here, spec_here, options, yhere=yhere) +# [xs2, sxs2, sigmas2] = find_known_lines_old(linelist, +# ll_here, spec_here, options) if data2 is not None: cfit2 = sol_1d2[1] - spec_here2 = np.ma.median(data2[yhere-2:yhere+2, :], axis=0) - shift2 = Fit.xcor_peak(spec_here2, spec2, lags) + spec_here2 = np.ma.median(data2[yhere-pmlines:yhere+pmlines, :], axis=0) + shift2 = Fit.find_shift(spec_here2, spec2, guess=guess, maxshift=4*pmlines) ll_here2 = CV.chebval(pix - shift2, cfit2) [xs2, sxs2, sigmas2] = find_known_lines(linelist2, @@ -2296,20 +2408,30 @@ def fit_parameters(yhere): linelist, options) #if np.std(delt) < .01: pdb.set_trace() - debug("resid ang S%2.2i @ p%4.0i: %1.2f rms %1.2f mad [shift%2.0f]" % \ - (slitno+1, yhere, np.std(delt), np.median(np.abs(delt)), - shift)) +# debug("resid ang S%2.2i @ p%4.0i: %1.2f rms %1.2f mad [shift%2.0f]" % \ +# (slitno+1, yhere, np.std(delt), np.median(np.abs(delt)), +# shift)) + debug(f"resid ang S{slitno+1:02d} @ p{yhere:4d}: {np.std(delt):4.2f} rms"\ + +f" {np.median(np.abs(delt)):5.2f} mad [shift {shift:.1f} pix]") - return cfit, delt + return cfit, delt, shift - def sweep(positions): + def sweep(positions, step=1, pmlines=1): ret = [] cfits = [] sds = [] mads = [] + lastshift = 0 + lastposition = positions[0] for position in positions: - cfit, delt = fit_parameters(position) + if abs(position-lastposition) < step*3: + guess = lastshift + else: + guess = 0 + cfit, delt, shift = fit_parameters(position, guess=guess, pmlines=pmlines) + lastposition = position + lastshift = shift cfits.append(cfit) sds.append(np.std(delt)) @@ -2326,16 +2448,40 @@ def sweep(positions): """ Start of main section of fit_outwards """ pix = np.arange(2048.) - positions = np.concatenate((np.arange(start, top, 1), - np.arange(start-1,bottom,-1))) - positions = np.arange(bottom, top, 1) + step = 1 + positions = np.concatenate((np.arange(start, top, step), + np.arange(start-1,bottom,-step))) +# positions = np.concatenate((np.arange(start, top, 1), +# np.arange(start-1,bottom,-1))) +# positions = np.arange(bottom, top, 1) + info("Computing 0 spectrum at %i" % start) spec0 = np.ma.median(data[start-1:start+1, :], axis=0) if data2 is not None: - spec2 = np.ma.median(data2[start-1:start+1, :], axis=0) - params = sweep(positions) + spec2 = np.ma.median(data2[start-1:start+1, :], axis=0) + params = sweep(positions, step=step, pmlines=1) + + ## Smooth behavior of coefficients in Y direction + if options["smooth"] == True: +# params_s = smooth_solution(params, filter_size=options['smooth_scale']) + params_s = fit_smooth_solution(params, order=2, top=top, bottom=bottom) + params_s['smoothed'] = True + + pl.figure(figsize=(15,10)) + for j in range(params_s['coeffs'].shape[1]): + pl.subplot(3,2,j+1) + pl.title(f'Wavelength fit coefficient {j}') + pl.plot(params['positions'], params['coeffs'][:,j], 'ko') + pl.ylim(np.percentile(params['coeffs'][:,j], 1), + np.percentile(params['coeffs'][:,j], 99)) + s = np.argsort(params_s['positions']) + pl.plot(params_s['positions'][s], params_s['coeffs'][:,j][s], 'g-', alpha=0.5) + pl.show() - return params + if options["smooth"] == True: + return params_s + else: + return params class NoSuchFit(Exception): pass