def womcho(hop): """choose region of spectrum with wavelengths or mouse""" import matplotlib.pyplot as plt from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element wave = hop[0].wave.copy() flux = hop[0].flux.copy() var = hop[0].var.copy() print('Current A/pix is {}'.format(wave[1] - wave[0])) wave, flux, mode = womwaverange(wave, flux, 'none') indexblue = womget_element(hop[0].wave, wave[0]) indexred = womget_element(hop[0].wave, wave[-1]) var = var[indexblue:indexred + 1].copy() print('Overplotting chosen section') plt.cla() plt.plot(hop[0].wave, hop[0].flux, drawstyle='steps-mid') plt.plot(wave, flux, drawstyle='steps-mid') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.title(hop[0].obname) hop[0].wave = wave.copy() hop[0].flux = flux.copy() hop[0].var = var.copy() ##FIX header for fits return hop
def womvelcho(hop): """select region of spectrum, put onto velocity scale""" import matplotlib.pyplot as plt import numpy as np import logging from tmath.wombat.inputter import inputter from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element from tmath.wombat.womscipyrebin import womscipyrebin wave=hop[0].wave.copy() flux=hop[0].flux.copy() var=hop[0].var.copy() nwave,nflux,mode=womwaverange(wave,flux,'none') wavebindex=womget_element(wave,nwave[0]) waverindex=womget_element(wave,nwave[-1]) nvar=var[wavebindex:waverindex+1].copy() binvec=np.arange(len(nwave)) wavelog=np.log(nwave[0])+((np.log(nwave[-1])-np.log(nwave[0]))/len(nwave))*binvec wavelog=np.exp(wavelog) fluxlog=womscipyrebin(nwave,nflux,wavelog) varlog=womscipyrebin(nwave,nvar,wavelog) zp=wavelog[0]-1.0 while (zp < wavelog[0]) or (zp > wavelog[-1]): zp=inputter('Enter the zero point for velocity (in angstroms): ','float',False) indexzp=womget_element(wavelog, zp) print(' ') logging.info('Zero at bin {}'.format(indexzp)) logging.info('with lambda {}'.format(wavelog[indexzp])) print(' ') z=(wavelog - zp)/zp square=(z+1)*(z+1) wavekmsrel=((square-1.)/(square+1.))*299792.458 kmsperbin=np.zeros(len(nwave)) for i in range(1,len(nwave)): kmsperbin[i]=2.0*2.99792458e5*(wavelog[i]-wavelog[i-1])/\ (wavelog[i]+wavelog[i-1]) kmsmean=np.mean(kmsperbin[1:]) logging.info('Average km/s per bin: {}'.format(kmsmean)) logging.info('km/s at bin 1: {}'.format(kmsperbin[1])) logging.info('km/s at bin n: {}'.format(kmsperbin[-1])) print(' ') wavekms=kmsmean*binvec+kmsmean/2.0 indexzp=womget_element(wavekms,zp) wavekms=wavekms-wavekms[indexzp] hop[0].wave=wavekmsrel.copy() hop[0].flux=fluxlog.copy() hop[0].var=varlog.copy() return hop
def womwaverange(wave, flux, mode): import matplotlib.pyplot as plt from tmath.wombat.waveparse import waveparse from tmath.wombat.womget_element import womget_element from tmath.wombat.inputter_single import inputter_single from tmath.wombat.yesno import yesno from tmath.wombat.wshow import wshow if (mode == 'none'): print('Enter method to select wavelength range') mode = inputter_single( 'Enter (w)avelengths or mark with the (m)ouse? (w/m) ', 'wm') print(mode) print('Spectrum runs from {} to {}.'.format(wave[0], wave[-1])) plt.cla() plt.plot(wave, flux, drawstyle='steps-mid') plt.ylabel('Flux') plt.xlabel('Wavelength') wshow() done = False while (not done): if (mode == 'w'): waveb, waver = waveparse() else: pickpoints = plt.ginput(2, timeout=-1) waveb = pickpoints[0][0] waver = pickpoints[1][0] if (waveb > waver): waveb, waver = waver, waveb indexblue = womget_element(wave, waveb) indexred = womget_element(wave, waver) print('Range selected: {} to {}'.format(wave[indexblue], wave[indexred])) plt.plot(wave[indexblue:indexred+1],flux[indexblue:indexred+1], \ drawstyle='steps-mid',color='r') plt.pause(0.01) print('Is this range correct?') answer = yesno('y') if (answer == 'n'): plt.cla() plt.plot(wave, flux, drawstyle='steps-mid') plt.ylabel('Flux') plt.xlabel('Wavelength') wshow() else: done = True return wave[indexblue:indexred + 1], flux[indexblue:indexred + 1], mode
def womexam(hop): """print spectrum bin values""" from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element from tmath.wombat.yesno import yesno done = False while (not done): answer = 'y' wave, flux, mode = womwaverange(hop[0].wave, hop[0].flux, 'none') indexblue = womget_element(hop[0].wave, wave[0]) indexred = womget_element(hop[0].wave, wave[-1]) var = hop[0].var[indexblue:indexred + 1] if (len(wave) > 30): print('This will print out {} values'.format(len(wave))) print('Do you really want to do that?') answer = yesno('n') if (answer == 'y'): done = True print('\nWave Flux') for i, _ in enumerate(wave): print('{} {} {}'.format(wave[i], flux[i], var[i])) return hop
def womint(hop): """calculate intensity in given wavelength range""" import logging import numpy as np from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element print(' ') logging.info('Object is {}'.format(hop[0].obname)) print(' ') print('Spectrum runs from {} to {}'.format(hop[0].wave[0], hop[0].wave[-1])) print(' ') print('This routine expects the spectrum to be in flambda units.') print('It also expects a linear wavelength scale.') print(' ') print('Choose general region of spectrum\n') nwave, nflux, mode = womwaverange(hop[0].wave, hop[0].flux, 'none') print('\nNow pick the exact range for the intensity calculation') wavint, fluxint, mode = womwaverange(nwave, nflux, mode) indexblue = womget_element(nwave, wavint[0]) indexred = womget_element(nwave, wavint[-1]) wdelt = nwave[1] - nwave[0] lineflux = np.sum(nflux[indexblue:indexred + 1]) * wdelt linefluxin = np.sum(nflux[indexblue + 1:indexred]) * wdelt linefluxout = np.sum(nflux[indexblue - 1:indexred + 2]) * wdelt print(' ') logging.info('FWZI (approximate): {}'.format(nwave[indexred] - nwave[indexblue])) logging.info('Line flux (ergs/sec/cm^2): {}'.format(lineflux)) logging.info('Line flux one pixel in: {}'.format(linefluxin)) logging.info('Line flux one pixel out: {}'.format(linefluxout)) logging.info('Note that flux may need to be scaled by 1e-15') logging.info('Average difference (between line flux and one pixel') avgdiff = (np.abs(linefluxin - lineflux) + np.abs(linefluxout - lineflux)) / 2.0 logging.info('in or out): {}'.format(avgdiff)) logging.info('As a percentage of line flux: {}'.format(100.0 * avgdiff / lineflux)) return hop
def womblo(hop, fig): """blotch out bad data""" import logging import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import splrep, splev from tmath.wombat.inputter_single import inputter_single from tmath.wombat.waveparse import waveparse from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element from tmath.wombat.yesno import yesno global nsplinepoints, tmpsplptsx, tmpsplptsy, pflag print('Select regions to blotch') wave = hop[0].wave.copy() flux = hop[0].flux.copy() done = False while (not done): plt.cla() plt.plot(wave, flux, drawstyle='steps-mid') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.title(hop[0].obname) wavesub, fluxsub, mode = womwaverange(wave, flux, 'none') wavebind = womget_element(wave, wavesub[0]) waverind = womget_element(wave, wavesub[-1]) plt.cla() plt.plot(wave[wavebind:waverind+1],flux[wavebind:waverind+1], \ drawstyle='steps-mid') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.title(hop[0].obname) plt.pause(0.01) print('Do you want to enter blotch wavelengths by hand (w),') print('mark points (m), fit a spline (s), or quit (q)?') choice = inputter_single('(w/m/s/q): ', 'wmsq') if (choice == 'w') or (choice == 'm'): blotchgood = False while (not blotchgood): wavechoicedone = False while (not wavechoicedone): if (choice == 'w'): waveselb, waveselr = waveparse() else: print('Mark the two endpoints of the blotch region') endpoints = plt.ginput(2, timeout=-1) waveselb = endpoints[0][0] waveselr = endpoints[1][0] if (waveselb > waveselr): waveselb, waveselr = waveselr, waveselb waveselbind = womget_element(wave, waveselb) waveselrind = womget_element(wave, waveselr) print(waveselb, waveselr, waveselbind, waveselrind) if (waveselbind == 0) or (waveselrind == (len(wave) - 1)): print('Wavelengths incorrect--too close to endpoints') else: wavechoicedone = True contblue = flux[waveselbind - 1] contred = flux[waveselrind + 1] delta = (contred - contblue) / (waveselrind - waveselbind + 1) fluxcor = flux.copy() for i in range(waveselbind, waveselrind + 1): fluxcor[i] = contblue + (i - waveselbind + 1) * delta plt.plot(wave[wavebind:waverind+1],fluxcor[wavebind:waverind+1], \ drawstyle='steps-mid') plt.pause(0.01) print('Is this acceptable') answer = yesno('y') if (answer == 'y'): flux = fluxcor.copy() blotchgood = True logging.info('File {} blotched from {} to {}'.format( hop[0].obname, wave[waveselbind], wave[waveselrind])) elif (choice == 's'): xmin, xmax = plt.xlim() ymin, ymax = plt.ylim() plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) nsplinepoints = 0 tmpsplptsx = [] tmpsplptsy = [] spldone = False while (not spldone): plt.cla() plt.plot(wave[wavebind:waverind+1],flux[wavebind:waverind+1], \ drawstyle='steps-mid') if (len(tmpsplptsx) > 0): plt.plot(tmpsplptsx, tmpsplptsy, 'ro') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.title(hop[0].obname) plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) cid = fig.canvas.mpl_connect('button_press_event', onclick) print('\nClick on continuum points for spline fit.') print( 'Spline will replace values between first and last point') print('Left button = add point') print('Middle button = delete point') print('Right button = done\n') pflag = '' while (pflag != 'done'): plt.pause(0.01) fig.canvas.mpl_disconnect(cid) splptsy = [z for _, z in sorted(zip(tmpsplptsx, tmpsplptsy))] splptsx = sorted(tmpsplptsx) spline = splrep(splptsx, splptsy, k=3) splblueindex = womget_element(wave, splptsx[0]) splredindex = womget_element(wave, splptsx[-1]) splwave = wave[splblueindex:splredindex + 1].copy() splineresult = splev(splwave, spline) fluxcor = flux.copy() fluxcor[splblueindex:splredindex + 1] = splineresult.copy() plt.plot(splwave, splineresult, drawstyle='steps-mid') print('Is this acceptable') answer = yesno('y') if (answer == 'y'): flux = fluxcor.copy() spldone = True logging.info( 'File {} blotched with spline from {} to {}'.format( hop[0].obname, wave[splblueindex], wave[splredindex])) else: done = True print('Do another region?') another = yesno('n') if (another == 'n'): done = True hop[0].flux = flux.copy() return hop
def secondcat(wave, obj1, obj2, sig1, sig2, secondtimeflag, wavebeg, waveend, brscale): import numpy as np import matplotlib.pyplot as plt from tmath.wombat.womget_element import womget_element from tmath.wombat.inputter_single import inputter_single from tmath.wombat.yesno import yesno print('\nCombining two calibrated spectra for second-order correction') waveblue = wave wavered = wave fluxblue = obj1 fluxred = obj2 fluxtemp = fluxred if (not secondtimeflag): print('Plotting blue side as blue, red side as red') indexblue = womget_element(waveblue, wavered[0]) indexred = womget_element(wavered, waveblue[-1]) meanblue = np.mean(fluxblue[indexblue:]) meanred = np.mean(fluxred[0:indexred]) if ((meanblue / meanred) > 1.2) or (meanblue / meanred < 0.8): print('Averages very different, scaling red to blue for plot') scale = meanblue / meanred print('Red multiplied by {}'.format(scale)) fluxtemp = fluxred * scale plt.clf() plt.plot(waveblue[indexblue:], fluxblue[indexblue:], drawstyle='steps-mid', color='b') plt.plot(wavered[0:indexred], fluxred[:indexred], drawstyle='steps-mid', color='r') plt.xlabel('Wavelength') plt.ylabel('Flux') donescale = False while (not donescale): print('Change y-scale?') answer = yesno('n') if (answer == 'y'): ymin, ymax = scaleparse() plt.ylim([ymin, ymax]) plt.pause(0.01) print('Do again?') answer = yesno('n') if (answer == 'n'): donescale = True regiondone = False while (not regiondone): print( '\nOK, mark the two end points of the region to compute average' ) endpoints = plt.ginput(2, timeout=-1) element1 = endpoints[0][0] element2 = endpoints[1][0] plt.plot(endpoints, 'ro') if (element1 > element2): element1, element2 = element2, element1 binend = womget_element(waveblue, element2) binbeg = womget_element(waveblue, element1) if (binend > len(waveblue) - 1): binend = len(waveblue) - 1 if (binbeg < indexblue): binbeg = indexblue print('\nAre these points OK?') answer = yesno('y') if (answer == 'y'): regiondone = False else: plt.plot(endpoints, 'wo') wavebeg = waveblue[binbeg] waveend = waveblue[binend] binbegred = womget_element(wavered, wavebeg) binendred = womget_element(wavered, waveend) meanblue = np.mean(fluxblue[binbeg:binend + 1]) meanred = np.mean(fluxred[binbegred:binendred + 1]) print('Average for {}:{}'.format(wavebeg, waveend)) print('Blue side: {}'.format(meanblue)) print('Red side: {}'.format(meanred)) brscale = inputter_single('Scale to blue or red? (b/r) ', 'br') if (brscale == 'b'): print('Scaling to blue by {}'.format(meanblue / meanred)) fluxred = fluxred * meanblue / meanred else: print('Scaling to red by {}'.format(meanred / meanblue)) fluxblue = fluxblue * meanred / meanblue print('\nPlotting blue side as blue, red side as red') plt.clf() plt.plot(waveblue[binbeg:binend + 1], fluxblue[binbeg:binend + 1], drawstyle='steps-mid', color='b') plt.plot(wavered[binbegred:binendred + 1], fluxred[binbegred:binendred + 1], drawstyle='steps-mid', color='r') plt.xlabel('Wavelength') plt.ylabel('Flux') donescale = False while (not donescale): print('Change y-scale?') answer = yesno('n') if (answer == 'y'): ymin, ymax = scaleparse() plt.ylim([ymin, ymax]) plt.pause(0.01) print('Do again?') answer = yesno('n') if (answer == 'n'): donescale = True regiondone = False while (not regiondone): print('\nOK, mark the two end points of the region to combine') endpoints = plt.ginput(2, timeout=-1) element1 = endpoints[0][0] element2 = endpoints[1][0] plt.plot(endpoints, 'ro') if (element1 > element2): element1, element2 = element2, element1 binend = womget_element(waveblue, element2) binbeg = womget_element(waveblue, element1) if (binend > len(waveblue) - 1): binend = len(waveblue) - 1 if (binbeg < indexblue): binbeg = indexblue print('\nAre these points OK?') answer = yesno('y') if (answer == 'y'): regiondone = False else: plt.plot(endpoints, 'wo') wavebeg = waveblue[binbeg] waveend = waveblue[binend] binbegred = womget_element(wavered, wavebeg) binendred = womget_element(wavered, waveend) else: binbeg = womget_element(waveblue, wavebeg) binend = womget_element(waveblue, waveend) binbegred = womget_element(wavered, wavebeg) binendred = womget_element(wavered, waveend) meanblue = np.mean(fluxblue[binbeg:binend + 1]) meanred = np.mean(fluxred[binbegred:binendred + 1]) if (brscale == 'b'): scale = meanblue / meanred fluxred = fluxred * scale sig2 = sig2 * scale else: scale = meanred / meanblue fluxblue = fluxblue * scale sig1 = sig1 * scale overflux = (fluxblue[binbeg:binend + 1] + fluxred[binbegred:binendred + 1]) / 2.0 oversig = np.sqrt(sig1[binbeg:binend + 1]**2 + sig2[binbegred:binendred + 1]**2) newwave = waveblue.copy() newflux = fluxblue.copy() newsig = sig1.copy() newsig[binbeg:binend + 1] = oversig newsig[binend + 1:] = sig2[binendred + 1:] newflux[binbeg:binend + 1] = overflux newflux[binend + 1:] = fluxred[binendred + 1:] plt.clf() plt.pause(0.01) axarr = fig.subplots(2, sharex=True) fig.subplots_adjust(hspace=0) axarr[0].plot(waveblue[binbeg:binend + 1], fluxblue[binbeg:binend + 1], drawstyle='steps-mid', color='b') axarr[0].plot(wavered[binbegred:binendred + 1], fluxred[binbegred:binendred + 1], drawstyle='steps-mid', color='r') axarr[0].plot(waveblue[binbeg:binend + 1], overflux, drawstyle='steps-mid', color='k') axarr[0].set_title('Overlap region--with average') axarr[0].set_ylabel('Flux') plt.pause(0.01) axarr[1].plot(newwave, newflux, drawstyle='steps-mid', color='k') axarr[1].plot(newwave[binbeg:binend + 1], newflux[binbeg:binend + 1], drawstyle='steps-mid', color='r') axarr[1].set_xlabel('Wavelength') axarr[1].set_ylabel('Flux') plt.pause(0.01) return newflux, newsig, wavebeg, waveend, brscale
def womscale2match(hop): """scale one spectrum to match another""" import matplotlib.pyplot as plt import numpy as np import logging from tmath.wombat.inputter import inputter from tmath.wombat.inputter_single import inputter_single from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element from tmath.wombat import HOPSIZE print('This will scale one hopper to match another') done = False while (not done): hopnum1 = inputter('Enter first hopper: ', 'int', False) hopnum2 = inputter('Enter second hopper: ', 'int', False) if (hopnum1 < 1) or (hopnum1 > HOPSIZE) or (hopnum2 < 1) or \ (hopnum2 > HOPSIZE): print('Hopper numbers must be in range 1-{}'.format(HOPSIZE)) else: done = True if (hop[hopnum1].wave[0] != hop[hopnum2].wave[0]) or \ (hop[hopnum1].wave[1] != hop[hopnum2].wave[1]) or \ (hop[hopnum1].wave[-1] != hop[hopnum2].wave[-1]): print('Hoppers to not have the same wavelength scale!') return hop print('\nSpectra run from {} to {}'.format(hop[hopnum1].wave[0], \ hop[hopnum1].wave[-1])) plt.cla() plt.plot(hop[hopnum1].wave,hop[hopnum1].flux,drawstyle='steps-mid', \ color='k') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.title(hop[hopnum1].obname) plt.plot(hop[hopnum2].wave,hop[hopnum2].flux,drawstyle='steps-mid', \ color='r') print('\nHopper A in black, hopper B in red') print('\nChoose region to compute average') wave, flux, mode = womwaverange(hop[hopnum1].wave, hop[hopnum1].flux, 'none') indexblue = womget_element(hop[hopnum1].wave, wave[0]) indexred = womget_element(hop[hopnum1].wave, wave[-1]) avg1 = np.mean(hop[hopnum1].flux[indexblue:indexred + 1]) avg2 = np.mean(hop[hopnum2].flux[indexblue:indexred + 1]) print('\n Hopper A: {}'.format(avg1)) print('\n Hopper B: {}'.format(avg2)) choice = inputter_single('Scale to (A) or (B) ', 'ab') print(' ') done = False while (not done): hopnum3 = inputter('Store scaled spec in which hopper? ', 'int', False) if (hopnum3 < 1) or (hopnum3 > HOPSIZE): print('Hopper numbers must be in range 1-{}'.format(HOPSIZE)) else: done = True if (choice == 'a'): hop[hopnum3].flux = hop[hopnum2].flux * (avg1 / avg2).copy() print('Scale: {}'.format(avg1 / avg2)) hop[hopnum3].obname = hop[hopnum2].obname hop[hopnum3].header = hop[hopnum2].header hop[hopnum3].var = hop[hopnum2].var * (avg1 / avg2).copy() lstring='File: {} scaled by {} to match {}'.format(hop[hopnum2], \ avg1/avg2, \ hop[hopnum1]) else: hop[hopnum3].flux = hop[hopnum1].flux * (avg2 / avg1).copy() print('Scale: {}'.format(avg1 / avg2)) hop[hopnum3].obname = hop[hopnum1].obname hop[hopnum3].header = hop[hopnum1].header hop[hopnum3].var = hop[hopnum1].var * (avg2 / avg1).copy() lstring='File: {} scaled by {} to match {}'.format(hop[hopnum1].obname, \ avg2/avg1, \ hop[hopnum2].obname) hop[hopnum3].wave = hop[hopnum1].wave.copy() logging.debug(lstring) return hop
def womcat(hop): """concatenate data with overlapping wavelength regions""" import numpy as np import logging import matplotlib.pyplot as plt from tmath.wombat.inputter import inputter from tmath.wombat.inputter_single import inputter_single from tmath.wombat.womget_element import womget_element from tmath.wombat.yesno import yesno from tmath.wombat.get_screen_size import get_screen_size from tmath.wombat.womwaverange2 import womwaverange2 from tmath.wombat import HOPSIZE from matplotlib.widgets import Cursor plt.ion() screen_width, screen_height = get_screen_size() screenpos = '+{}+{}'.format(int(screen_width * 0.2), int(screen_height * 0.05)) fig = plt.figure() ax = fig.add_subplot(111) cursor = Cursor(ax, useblit=True, color='k', linewidth=1) fig.canvas.manager.window.wm_geometry(screenpos) fig.canvas.set_window_title('Cat') fig.set_size_inches(9, 6) # turns off key stroke interaction fig.canvas.mpl_disconnect(fig.canvas.manager.key_press_handler_id) print("\nThis will combine blue and red pieces from two hoppers\n") hopnum1 = 0 hopnum2 = 0 while (hopnum1 < 1) or (hopnum1 > HOPSIZE): hopnum1 = inputter('Enter first hopper: ', 'int', False) while (hopnum2 < 1) or (hopnum2 > HOPSIZE): hopnum2 = inputter('Enter second hopper: ', 'int', False) if (hop[hopnum1].wave[0] > hop[hopnum2].wave[0]): hopnum1, hopnum2 = hopnum2, hopnum1 wdelt1 = hop[hopnum1].wave[1] - hop[hopnum1].wave[0] wdelt2 = hop[hopnum2].wave[1] - hop[hopnum2].wave[0] # check if wavelength dispersion same if (abs(wdelt1 - wdelt2) > 0.00001): print('Spectra do not have same Angstrom/pixel') print('Blue side: {}'.format(wdelt1)) print('Red side: {}'.format(wdelt2)) return hop if hop[hopnum1].wave[-1] < hop[hopnum2].wave[0]: print('Spectra do not overlap\n') return hop print("\nOverlap range is {} to {}".format(hop[hopnum2].wave[0], hop[hopnum1].wave[-1])) print("\nPlotting blue side as blue, red side as red\n") waveblue = hop[hopnum1].wave.copy() fluxblue = hop[hopnum1].flux.copy() varblue = hop[hopnum1].var.copy() wavered = hop[hopnum2].wave.copy() fluxred = hop[hopnum2].flux.copy() varred = hop[hopnum2].var.copy() indexblue = womget_element(waveblue, wavered[0]) indexred = womget_element(wavered, waveblue[-1]) fluxcor = 1.0 blue_mean = np.mean(fluxblue[indexblue:]) red_mean = np.mean(fluxred[0:indexred + 1]) if (blue_mean / red_mean < 0.8) or (blue_mean / red_mean > 1.2): fluxcor = blue_mean / red_mean print("Averages very different, scaling red to blue for plot") print("Red multiplied by {}".format(fluxcor)) plt.cla() plt.plot(waveblue[indexblue:], fluxblue[indexblue:], drawstyle='steps-mid', color='b') plt.plot(wavered[0:indexred], fluxred[0:indexred] * fluxcor, drawstyle='steps-mid', color='r') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.pause(0.01) print('Change scale?') answer = yesno('n') if (answer == 'y'): xmin_old, xmax_old = plt.xlim() ymin_old, ymax_old = plt.ylim() done = False while (not done): plt.xlim([xmin_old, xmax_old]) plt.ylim([ymin_old, ymax_old]) print('Click corners of box to change plot scale') newlims = plt.ginput(2, timeout=-1) xmin = newlims[0][0] ymin = newlims[0][1] xmax = newlims[1][0] ymax = newlims[1][1] plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) print('Is this OK?') loopanswer = yesno('y') if (loopanswer == 'y'): done = True print('\nEnter method to select wavelength ranges\n') mode = inputter_single( 'Enter (w)avelengths or mark with the (m)ouse? (w/m) ', 'wm') print('\nChoose end points of region to compute average\n') waveb, waver, mode = womwaverange2(waveblue[indexblue:], fluxblue[indexblue:], wavered[0:indexred], fluxred[0:indexred] * fluxcor, mode) indexblueb = womget_element(waveblue, waveb) indexbluer = womget_element(waveblue, waver) indexredb = womget_element(wavered, waveb) indexredr = womget_element(wavered, waver) mean_blue = np.mean(fluxblue[indexblueb:indexbluer + 1]) mean_red = np.mean(fluxred[indexredb:indexredr + 1]) print("\nAverage for {}:{}".format(waveb, waver)) print("Blue side: {}".format(mean_blue)) print("Red side: {}\n".format(mean_red)) brscale = inputter_single('Scale to blue or red (b/r)? ', 'br') if (brscale == 'b'): brscalefac = mean_blue / mean_red logging.info('Cat scaling to blue by {}'.format(brscalefac)) fluxred = fluxred * brscalefac varred = varred * brscalefac**2 else: brscalefac = mean_red / mean_blue logging.info('Cat scaling to red by {}'.format(brscalefac)) fluxblue = fluxblue * brscalefac varblue = varblue * brscalefac**2 # print("\nPlotting blue side as blue, red side as red\n") # plt.cla() # plt.plot(waveblue[indexblueb:indexbluer+1],fluxblue[indexblueb:indexbluer+1],drawstyle='steps-mid',color='b') # plt.plot(wavered[indexredb:indexredr+1],fluxred[indexredb:indexredr+1],drawstyle='steps-mid',color='r') # plt.xlabel('Wavelength') # plt.ylabel('Flux') # plt.pause(0.01) # print('Change scale?') # answer=yesno('n') # if (answer == 'y'): # xmin_old,xmax_old=plt.xlim() # ymin_old,ymax_old=plt.ylim() # done=False # while (not done): # plt.xlim([xmin_old,xmax_old]) # plt.ylim([ymin_old,ymax_old]) # print('Click corners of box to change plot scale') # newlims=plt.ginput(2,timeout=-1) # xmin=newlims[0][0] # ymin=newlims[0][1] # xmax=newlims[1][0] # ymax=newlims[1][1] # plt.xlim([xmin,xmax]) # plt.ylim([ymin,ymax]) # print('Is this OK?') # loopanswer=yesno('y') # if (loopanswer == 'y'): # done=True # print('\nChoose end points of region to compute average\n') # waveb,waver,mode=womwaverange2(waveblue[indexblueb:indexbluer+1], # fluxblue[indexblueb:indexbluer+1], # wavered[indexredb:indexredr+1], # fluxred[indexredb:indexredr+1],mode) # indexblueb=womget_element(waveblue,waveb) # indexbluer=womget_element(waveblue,waver) # indexredb=womget_element(wavered,waveb) # indexredr=womget_element(wavered,waver) # ewadd=inputter_single('Add overlap region (e)qually or with (w)eights (e/w)?','ew') # if (ewadd == 'e'): # overflux=(fluxblue[indexblueb:indexbluer+1]+fluxred[indexredb:indexredr+1])/2. # overvar=(varblue[indexblueb:indexbluer+1]+varred[indexredb:indexredr+1]) #replacing with inverse variance weighted average overflux = np.average([ fluxblue[indexblueb:indexbluer + 1], fluxred[indexredb:indexredr + 1] ], weights=[ 1. / varblue[indexblueb:indexbluer + 1], 1. / varred[indexredb:indexredr + 1] ], axis=0) overvar = np.sum( [varblue[indexblueb:indexbluer + 1], varred[indexredb:indexredr + 1]], axis=0) logging.info('Cat combined sides with inverse variance weighted average') # else: # wei_done = False # while (not wei_done): # weiblue=inputter('Enter fractional weight for blue side: ','float',False) # weired=inputter('Enter fractional weight for red side: ','float',False) # if (np.abs((weiblue+weired)-1.0) > 0.000001): # print('Weights do not add to 1.0') # else: # wei_done = True # overflux=(fluxblue[indexblueb:indexbluer+1]*weiblue+ # fluxred[indexredb:indexredr+1]*weired) # overvar=(varblue[indexblueb:indexbluer+1]*weiblue**2+ # varred[indexredb:indexredr+1]*weired**2) # logging.info('Cat adds blue side with weight {} and red side with weight {}'.format(weiblue, weired)) newbluewave = waveblue[:indexblueb] newblueflux = fluxblue[:indexblueb] newbluevar = varblue[:indexblueb] overwave = waveblue[indexblueb:indexbluer + 1] newredwave = wavered[indexredr + 1:] newredflux = fluxred[indexredr + 1:] newredvar = varred[indexredr + 1:] newwave = np.concatenate([newbluewave, overwave, newredwave]) newflux = np.concatenate([newblueflux, overflux, newredflux]) newvar = np.concatenate([newbluevar, overvar, newredvar]) logging.info('File {} and'.format(hop[hopnum1].obname)) logging.info('File {} concatenated'.format(hop[hopnum2].obname)) logging.info('over wavelength range {} to {}'.format( waveblue[indexblueb], waveblue[indexbluer])) plt.clf() axarr = fig.subplots(2) axarr[0].plot(overwave, overflux, drawstyle='steps-mid', color='k') axarr[0].plot(waveblue[indexblueb:indexbluer + 1], fluxblue[indexblueb:indexbluer + 1], drawstyle='steps-mid', color='b') axarr[0].plot(wavered[indexredb:indexredr + 1], fluxred[indexredb:indexredr + 1], drawstyle='steps-mid', color='r') axarr[0].set_title('Overlap region, with inputs and combination') axarr[1].set_ylabel('Flux') axarr[1].plot(newwave, newflux, drawstyle='steps-mid', color='k') axarr[1].plot(newwave[indexblueb:indexbluer + 1], newflux[indexblueb:indexbluer + 1], drawstyle='steps-mid', color='r') plt.pause(0.01) hopout = 0 while (hopout < 1) or (hopout > HOPSIZE): hopout = inputter('Enter hopper to store combined spectrum: ', 'int', False) hop[hopout].wave = newwave hop[hopout].flux = newflux hop[hopout].var = newvar hop[hopout].obname = hop[hopnum1].obname hop[hopout].header = hop[hopnum1].header plt.close() #fig.clf() #plt.cla() return hop
def womgau(hop): """fit gaussian to line""" import numpy as np import logging import matplotlib.pyplot as plt from scipy.optimize import curve_fit from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element from tmath.wombat.inputter import inputter from tmath.wombat.inputter_single import inputter_single from tmath.wombat.gauss import gauss from tmath.wombat.gauss_cont import gauss_cont from tmath.wombat.yesno import yesno print(' ') logging.info('Object is {}'.format(hop[0].obname)) print(' ') print('Spectrum runs from {} to {}'.format(hop[0].wave[0], hop[0].wave[-1])) print(' ') print('This routine expects the spectrum to be in flambda units.') print('It also expects a linear wavelength scale.') print(' ') print('Choose general region of spectrum\n') nwave, nflux, mode = womwaverange(hop[0].wave, hop[0].flux, 'none') print('\nNow pick the exact range for the fit') waveint, fluxint, mode = womwaverange(nwave, nflux, mode) indexblue = womget_element(nwave, waveint[0]) indexred = womget_element(nwave, waveint[-1]) if (mode == 'w'): done = False while (not done): print(' ') wavecenter = inputter('Enter approximate center of Gaussian : ', 'float', False) indexcenter = womget_element(waveint, wavecenter) if (indexcenter <= 0) or (wavecenter > waveint[-1]): print('Bad central wavelength, try again') else: done = True else: done = False while (not done): print('Mark the approximate center of the Gaussian') pickcent = plt.ginput(1, timeout=-1) indexcenter = womget_element(waveint, pickcent[0][0]) print('\nApproximate center at {}'.format(waveint[indexcenter])) print('\nIs this OK?') answer = yesno('y') if (answer == 'y'): done = True weights = np.sqrt(hop[0].var[indexblue:indexred + 1]) print(' ') continuum = inputter_single( 'Do you want to fit gaussian with (c)ontinuum, or (n)o continuum? ', 'cn') if (continuum == 'c'): p = [fluxint[indexcenter], waveint[indexcenter], 3.0, 1.0, waveint[0]] result = curve_fit(gauss_cont, waveint, fluxint, sigma=weights, p0=p, absolute_sigma=True, full_output=True) else: p = [fluxint[indexcenter], waveint[indexcenter], 3.0] result = curve_fit(gauss, waveint, fluxint, sigma=weights, p0=p, absolute_sigma=True, full_output=True) coefferr = np.sqrt(np.diag(result[1])) coeff = result[0] # make 'finer-grained' version of fit, 0.2A/pix for calculations wavecalc = np.arange(2 * 5 * 50 * abs( coeff[2])) * 0.2 + coeff[1] - 0.2 * 5 * 50 * abs(coeff[2]) calccenter = womget_element(wavecalc, coeff[1]) if (continuum == 'c'): fluxcalc = gauss_cont(wavecalc, *coeff) fluxcont = wavecalc * coeff[3] + coeff[4] fluxgaussian = fluxcalc - fluxcont linecont = fluxcont[calccenter] else: fluxcalc = gauss(wavecalc, *coeff) deltafit = wavecalc[1] - wavecalc[0] calcindexblue = womget_element(wavecalc, waveint[0]) calcindexred = womget_element(wavecalc, waveint[-1]) sumfluxcalc = np.sum(fluxcalc[calcindexblue:calcindexred + 1] * deltafit) sumallfluxcalc = np.sum(fluxcalc * deltafit) chi = (result[2]['fvec']**2).sum() redchi = chi / (len(waveint) - len(coeff)) if (continuum == 'c'): sumfluxgaussian = np.sum(fluxgaussian[calcindexblue:calcindexred + 1] * deltafit) sumallfluxgaussian = np.sum(fluxgaussian * deltafit) sumfluxcont = np.sum(fluxcont[calcindexblue:calcindexred + 1] * deltafit) sumallfluxcont = np.sum(fluxcont * deltafit) # propagate uncertainty (from old version) not sure this is correct height_pct = coefferr[0] / coeff[0] sigma_pct = coefferr[2] / coeff[2] flux_pct = np.sqrt(height_pct**2 + sigma_pct**2) sumfluxgaussiansig = sumfluxgaussian * flux_pct sumallfluxgaussiansig = sumallfluxgaussian * flux_pct plt.cla() plt.plot(nwave, nflux, drawstyle='steps-mid', color='k') plt.ylabel('Flux') plt.xlabel('Wavelength') xmin, xmax = plt.xlim() ymin, ymax = plt.ylim() plt.plot(wavecalc, fluxcalc, drawstyle='steps-mid', color='b') if (continuum == 'c'): plt.plot(wavecalc, fluxgaussian, drawstyle='steps-mid', color='r') plt.plot(wavecalc, fluxcont, drawstyle='steps-mid', color='g') plt.plot([waveint[0], waveint[0]], [ymin, ymax], color='k', linestyle='--') plt.plot([waveint[-1], waveint[-1]], [ymin, ymax], color='k', linestyle='--') plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) logging.info('For object {} Gaussian fit'.format(hop[0].obname)) if (continuum == 'c'): print( '\nData = Black, Fit = Blue, Continuum = Green, Fit-Continuum = Red\n' ) else: print('\nData = Black, Fit = Blue\n') logging.info('Height {:16.8f}+/-{:16.8f}'.format( coeff[0], coefferr[0])) logging.info('Center {:16.8f}+/-{:16.8f}'.format( coeff[1], coefferr[1])) logging.info('Sigma {:16.8f}+/-{:16.8f}'.format( coeff[2], coefferr[2])) if (continuum == 'c'): logging.info('Slope {:16.8f}+/-{:16.8f}'.format( coeff[3], coefferr[3])) logging.info('Y-intercept {:16.8f}+/-{:16.8f}'.format( coeff[4], coefferr[4])) logging.info('FWHM {:16.8f}+/-{:16.8f}'.format( 2.35482 * np.abs(coeff[2]), 2.35482 * coefferr[2])) logging.info( 'Flux between dotted lines (Gaussian): {:16.8f}+/-{:16.8f}'.format( sumfluxgaussian, sumfluxgaussiansig)) logging.info('EW between dotted lines (Gaussian): {:16.8f}'.format( sumfluxgaussian / linecont)) logging.info('Flux for full (Gaussian): {:16.8f}+/-{:16.8f}'.format( sumallfluxgaussian, sumallfluxgaussiansig)) logging.info('EW for full (Gaussian): {:16.8f}'.format( sumallfluxgaussian / linecont)) logging.info( 'Continuum flux at line center: {:16.8f}'.format(linecont)) logging.info('Chi^2: {}'.format(chi)) logging.info('Reduced chi^2: {}'.format(redchi)) logging.info('All fluxes might need to be scaled by 1e-15') print(' ') return hop
def womscalemany(hop): """scale many hoppers to one""" import matplotlib.pyplot as plt import numpy as np import logging from tmath.wombat.inputter import inputter from tmath.wombat.womwaverange import womwaverange from tmath.wombat.womget_element import womget_element from tmath.wombat import HOPSIZE print('This will scale many hoppers to one') done = False while (not done): hopnum1 = inputter('Enter fiducial hopper: ', 'int', False) if (hopnum1 < 1) or (hopnum1 > HOPSIZE): print('Hopper number must be in range 1-{}'.format(HOPSIZE)) else: done = True hoplist = [] hopnum2 = 0 while (hopnum2 != 99): done = False while (not done): hopnum2 = inputter('Enter another hopper: (99 to end) ', 'int', False) if (hopnum2 != 99): if (hopnum2 < 1) or (hopnum2 > HOPSIZE): print( 'Hopper numbers must be in range 1-{}'.format(HOPSIZE)) elif (len(hop[hopnum2].wave) == 0): print('Nothing in that hopper') else: hoplist.append(hopnum2) done = True if (hopnum2 > 0) and (hopnum2 < 21) and \ (len(hop[hopnum2].wave) != 0): if (hop[hopnum1].wave[0] != hop[hopnum2].wave[0]) \ or (hop[hopnum1].wave[1] != hop[hopnum2].wave[1]) \ or (hop[hopnum1].wave[-1] != hop[hopnum2].wave[-1]): print('Hoppers to not have the same wavelength scale!') return hop else: done = True print('\nSpectra run from {} to {}'.format(hop[hopnum1].wave[0], \ hop[hopnum1].wave[-1])) plt.cla() plt.plot(hop[hopnum1].wave,hop[hopnum1].flux,drawstyle='steps-mid', \ color='k') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.title(hop[hopnum1].obname) print('\nPlotting fiducial spectrum') print('\nChoose region to compute average') wave, flux, mode = womwaverange(hop[hopnum1].wave, hop[hopnum1].flux, 'none') indexblue = womget_element(hop[hopnum1].wave, wave[0]) indexred = womget_element(hop[hopnum1].wave, wave[-1]) avg1 = np.mean(hop[hopnum1].flux[indexblue:indexred + 1]) print('Fiducial spectrum mean in range: {}'.format(avg1)) for i, _ in enumerate(hoplist): print('{} {}'.format(i, hoplist[i])) avgloop = np.mean(hop[hoplist[i]].flux[indexblue:indexred + 1]) print('Hopper {} mean in range: {}'.format(hoplist[i], avgloop)) hop[hoplist[i]].flux = hop[hoplist[i]].flux * avg1 / avgloop hop[hoplist[i]].var = hop[hoplist[i]].var * avg1 / avgloop logging.debug('File: {} scaled by {} to match {}'.format( hop[hoplist[i]].obname, avg1 / avgloop, hop[hopnum1].obname)) return hop
def telluric_remove(bstarwave, bstar, bairmass, wave, object, airmass, variance, spectrum, shift=None): import numpy as np import pdb import matplotlib.pyplot as plt from tmath.wombat.inputter import inputter from tmath.wombat.yesno import yesno from tmath.wombat.womscipyrebin import womscipyrebin from tmath.wombat.womget_element import womget_element from tmath.pydux.xcor import xcor from tmath.pydux.finalscaler import finalscaler bstartmp = womscipyrebin(bstarwave, bstar, wave) # plt.cla() # plt.plot(bstarwave,bstartmp) # plt.pause(0.01) # answer=yesno('y') print('\nThe ratio of airmasses (object/B-star) is {}'.format(airmass / bairmass)) if (airmass / bairmass > 3.0) or (airmass / bairmass < 0.33): print('\nWARNING: OBJECT AND B-STAR HAVE WILDLY DIFFERENT') print('AIRMASSES: ATMOSPHERIC BAND DIVISION MAY BE LOUSY\n') wmin = wave[0] wmax = wave[-1] npix = len(object) wdelt = wave[1] - wave[0] print('wdelt', wdelt) lag = np.zeros(3) lagflag = [False] * 3 xfactor = 10 maxlag = 200 if not shift: print('\nCross-correlating object with B-star spectrum\n') fig = plt.figure() axarr = fig.subplots(2, 1) if (wmin < 6200) and (wmax > 6400) and (wmax < 6900): indblue = womget_element(wave, 6200) indred = womget_element(wave, 6400) lag[0] = xcor(object[indblue:indred + 1], bstartmp[indblue:indred + 1], xfactor, maxlag) lagflag[0] = True print('The shift at the 6250A band is {} angstroms'.format(lag[0] * wdelt)) if (wmin < 6800) and (wmax > 6500): indblue = womget_element(wave, 6800) indred = womget_element(wave, 6950) scale = 1. / np.max(object[indblue:indred + 1]) obb = scale * object[indblue:indred + 1] bb = bstartmp[indblue:indred + 1] lag[1] = xcor(obb, bb, xfactor, maxlag) lagflag[1] = True print('The shift at the B band is {} angstroms'.format(lag[1] * wdelt)) plt.cla() # ymin,ymax=finalscaler(object) # plt.plot(wave,object,drawstyle='steps-mid',color='r') # plt.plot(wave,newobject,drawstyle='steps-mid',color='k') ymin, ymax = finalscaler(bstartmp[indblue:indred + 1]) axarr[0].plot(wave[indblue:indred + 1], scale * object[indblue:indred + 1], drawstyle='steps-mid', color='r') axarr[0].plot(wave[indblue:indred + 1], bstartmp[indblue:indred + 1], drawstyle='steps-mid', color='k') axarr[0].plot(wave[indblue:indred + 1] + lag[1] * wdelt, bstartmp[indblue:indred + 1], drawstyle='steps-mid', color='g') plt.pause(0.01) if (wmin < 7500) and (wmax > 8000): indblue = womget_element(wave, 7500) indred = womget_element(wave, 8000) scale = 1. / np.max(object[indblue:indred + 1]) lag[2] = xcor(scale * object[indblue:indred + 1], bstartmp[indblue:indred + 1], xfactor, maxlag) print('The shift at the A band is {} angstroms'.format(lag[2] * wdelt)) lagflag[2] = True # ymin,ymax=finalscaler(object) # plt.plot(wave,object,drawstyle='steps-mid',color='r') # plt.plot(wave,newobject,drawstyle='steps-mid',color='k') ymin, ymax = finalscaler(bstartmp[indblue:indred + 1]) axarr[1].plot(wave[indblue:indred + 1], scale * object[indblue:indred + 1], drawstyle='steps-mid', color='r') axarr[1].plot(wave[indblue:indred + 1], bstartmp[indblue:indred + 1], drawstyle='steps-mid', color='k') axarr[1].plot(wave[indblue:indred + 1] + lag[2] * wdelt, bstartmp[indblue:indred + 1], drawstyle='steps-mid', color='g') plt.pause(0.01) check = inputter('Check plot [enter when done]: ', 'string', False) if (sum(lagflag) > 0): avglag = np.sum(lag) / sum(lagflag) angshift = avglag * wdelt print('The mean shift is {} Angstroms'.format(angshift)) else: angshift = 0.0 plt.close() else: angshift = shift fig = plt.figure(figsize=[9, 5]) telluric_done = False bstartmpcopy = bstartmp.copy() while (not telluric_done): print('Applying a shift of {} Angstroms'.format(angshift)) bstartmp = bstartmpcopy.copy() tmp = womscipyrebin(wave + angshift, bstartmp, wave) bstartmp = tmp.copy() bstartmp = bstartmp**((airmass / bairmass)**0.55) # newobject=object/bstartmp newobject = spectrum / bstartmp bvar = variance / bstartmp print('\nPlotting before and after atmospheric band correction\n') plt.cla() # ymin,ymax=finalscaler(object) # plt.plot(wave,object,drawstyle='steps-mid',color='r') # plt.plot(wave,newobject,drawstyle='steps-mid',color='k') ymin, ymax = finalscaler(spectrum) plt.plot(wave, spectrum, drawstyle='steps-mid', color='r') plt.plot(wave, newobject, drawstyle='steps-mid', color='k') plt.ylim([ymin, ymax]) plt.pause(0.01) if not shift: print('Is this OK?') answer = yesno('y') if (answer == 'n'): angshift = inputter('Enter B-star shift in Angstroms: ', 'float', False) else: telluric_done = True else: check = inputter('Check plot [enter when done]: ', 'string', False) telluric_done = True plt.close() return newobject, bvar, angshift
def fitspl_dev(wave, flux, airlimit, fig, cal=None): import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import splrep, splev import tmath.wombat.womconfig as womconfig from tmath.wombat.womget_element import womget_element from tmath.pydux.wave_telluric import wave_telluric from tmath.wombat.yesno import yesno from tmath.wombat.onclick import onclick from tmath.wombat.onkeypress import onkeypress import glob """fit spline to spectrum""" # starting points for spline # bandpts = np.array([3000, 3050, 3090, 3200, 3430, 3450, 3500, 3550, 3600, \ # 3650, 3700, 3767, 3863, 3945, 4025, 4144, 4200, 4250, \ # 4280, 4390, 4450, 4500, 4600, 4655, 4717, 4750, 4908, \ # 4950, 5000, 5050, 5100, 5150, 5200, 5250, 5280, 5350, \ # 5387, 5439, 5500, 5550, 6100, 6150, 6400, 6430, 6650, \ # 6700, 6750, 6800, 7450, 7500, 7550, 8420, 8460, 8520, \ # 8570, 8600, 8725, 8770, 9910, 10000, 10200, 10300, \ # 10400, 10500, 10600, 10700]) binWidth = 60 bandpts = np.arange(3000, 11000, binWidth) locsinrange = np.logical_and((bandpts > wave[10]), (bandpts < wave[-10])) #useband=bandpts[locsinrange] # now mask based on spectral features # you can add mask features here, and the feature name can # be anything, they aren't used explicitly featureMask = {} # empirical masks from Dimitriadis featureMask['feature1'] = [3722.56 - 10.0 / 2., 3722.56 + 10.0 / 2.] featureMask['feature2'] = [3736.90 - 15.0 / 2., 3736.90 + 15.0 / 2.] featureMask['feature3'] = [3752.00 - 15.0 / 2., 3752.00 + 15.0 / 2.] featureMask['feature4'] = [3772.00 - 17.0 / 2., 3772.00 + 17.0 / 2.] featureMask['feature5'] = [3800.00 - 18.0 / 2., 3800.00 + 18.0 / 2.] featureMask['feature6'] = [3835.38 - 20.0 / 2., 3835.38 + 20.0 / 2.] featureMask['feature7'] = [3889.05 - 24.0 / 2., 3889.05 + 24.0 / 2.] featureMask['feature8'] = [3933.66 - 16.0 / 2., 3933.66 + 16.0 / 2.] featureMask['feature9'] = [3970.07 - 18.0 / 2., 3970.07 + 18.0 / 2.] featureMask['feature10'] = [4101.74 - 28.0 / 2., 4101.74 + 28.0 / 2.] featureMask['feature11'] = [4340.46 - 30.0 / 2., 4340.46 + 30.0 / 2.] featureMask['feature12'] = [4471.48 - 20.0 / 2., 4471.48 + 20.0 / 2.] featureMask['feature13'] = [4685.70 - 30.0 / 2., 4685.70 + 30.0 / 2.] featureMask['feature14'] = [4861.36 - 35.0 / 2., 4861.36 + 35.0 / 2.] featureMask['feature15'] = [5411.52 - 35.0 / 2., 5411.52 + 35.0 / 2.] featureMask['feature16'] = [5889.95 - 32.0 / 2., 5889.95 + 32.0 / 2.] featureMask['feature17'] = [6562.85 - 40.0 / 2., 6562.85 + 40.0 / 2.] featureMask['feature18'] = [8498.02 - 25.0 / 2., 8498.02 + 25.0 / 2.] featureMask['feature19'] = [8542.09 - 25.0 / 2., 8542.09 + 25.0 / 2.] featureMask['feature20'] = [8662.14 - 25.0 / 2., 8662.14 + 25.0 / 2.] featureMask['feature21'] = [8763.96 - 20.0 / 2., 8763.96 + 20.0 / 2.] featureMask['feature22'] = [8865.75 - 25.0 / 2., 8865.75 + 25.0 / 2.] featureMask['feature23'] = [9010.00 - 28.0 / 2., 9010.00 + 28.0 / 2.] featureMask['feature24'] = [9213.90 - 30.0 / 2., 9213.90 + 30.0 / 2.] featureMask['feature25'] = [9545.97 - 34.0 / 2., 9545.97 + 34.0 / 2.] featureMask['telluric1'] = [3216.0 - binWidth / 2., 3420.0 + binWidth / 2.] #featureMask['telluric2'] = [5600.0-binWidth/2., 6050.0+binWidth/2.] featureMask['telluric3'] = [6250.0 - binWidth / 2., 6360.0 + binWidth / 2.] featureMask['telluric4'] = [6450.0 - binWidth / 2., 6530.0 + binWidth / 2.] featureMask['telluric5'] = [6840.0 - binWidth / 2., 7410.0 + binWidth / 2.] featureMask['telluric6'] = [7560.0 - binWidth / 2., 8410.0 + binWidth / 2.] featureMask['telluric7'] = [8925.0 - binWidth / 2., 9900.0 + binWidth / 2.] # inital testing masks from XIDL # featureMask['balmer1'] = [3714.0-binWidth/2., 3723.0+binWidth/2.] # featureMask['balmer2'] = [3650.0-binWidth/2., 3820.0+binWidth/2.] # featureMask['balmer3'] = [3725.0-binWidth/2., 3735.0+binWidth/2.] # featureMask['balmer4'] = [3740.0-binWidth/2., 3755.0+binWidth/2.] # featureMask['balmer5'] = [3760.0-binWidth/2., 3775.0+binWidth/2.] # featureMask['balmer6'] = [3785.0-binWidth/2., 3806.0+binWidth/2.] # featureMask['balmer7'] = [3810.0-binWidth/2., 3820.0+binWidth/2.] # featureMask['balmer8'] = [3824.0-binWidth/2., 3841.0+binWidth/2.] # featureMask['balmer9'] = [3880.0-binWidth/2., 3895.0+binWidth/2.] # featureMask['balmer10'] = [3957.0-binWidth/2., 3979.0+binWidth/2.] # featureMask['balmer11'] = [4000.0-binWidth/2., 4030.0+binWidth/2.] # featureMask['balmer12'] = [4087.0-binWidth/2., 4120.0+binWidth/2.] # featureMask['balmer13'] = [4135.0-binWidth/2., 4145.0+binWidth/2.] # featureMask['balmer14'] = [4328.0-binWidth/2., 4355.0+binWidth/2.] # featureMask['balmer15'] = [4677.0-binWidth/2., 4692.0+binWidth/2.] # featureMask['balmer16'] = [4830.0-binWidth/2., 4931.0+binWidth/2.] # featureMask['balmer17'] = [5402.0-binWidth/2., 5417.0+binWidth/2.] # featureMask['balmer18'] = [6535.0-binWidth/2., 6590.0+binWidth/2.] # featureMask['telluric1'] = [3216.0-binWidth/2., 3420.0+binWidth/2.] # #featureMask['telluric2'] = [5600.0-binWidth/2., 6050.0+binWidth/2.] # featureMask['telluric3'] = [6250.0-binWidth/2., 6360.0+binWidth/2.] # featureMask['telluric4'] = [6450.0-binWidth/2., 6530.0+binWidth/2.] # featureMask['telluric5'] = [6840.0-binWidth/2., 7410.0+binWidth/2.] # featureMask['telluric6'] = [7560.0-binWidth/2., 8410.0+binWidth/2.] # featureMask['telluric7'] = [8925.0-binWidth/2., 9900.0+binWidth/2.] # generate a mask via a boolean array that excludes regions in stellar features maskInds = np.full(len(bandpts), True, dtype=bool) for i, key in enumerate(featureMask.keys()): featureMaskInds = np.logical_or((bandpts < featureMask[key][0]), (bandpts > featureMask[key][1])) maskInds *= featureMaskInds fullMask = locsinrange * maskInds # apply the mask useband = bandpts[fullMask] for i, _ in enumerate(useband): index = womget_element(wave, useband[i]) useband[i] = index # useband now has indices of wavelength positions if (min(flux) < 0): flux[np.where(flux < 0)] = 0.0 plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') xmin, xmax = plt.xlim() ymin, ymax = plt.ylim() plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) if (airlimit): loc = wave_telluric(wave, 'high') else: loc = wave_telluric(wave, 'low') #invert loc and convert all good sections to np.nan so they won't plot loc = np.invert(loc) wavetell = wave.copy() fluxtell = flux.copy() wavetell[loc] = np.nan fluxtell[loc] = np.nan plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if '../../master_files/' + cal + '_splpts_master.txt' not in glob.glob( '../../master_files/*'): womconfig.nsplinepoints = len(useband) womconfig.tmpsplptsx = wave[useband].copy().tolist() womconfig.tmpsplptsy = [] for i, _ in enumerate(useband): womconfig.tmpsplptsy.append( np.median(flux[useband[i] - 2:useband[i] + 3])) spline = splrep(womconfig.tmpsplptsx, womconfig.tmpsplptsy, k=3) else: masterx, mastery = np.genfromtxt('../../master_files/' + cal + '_splpts_master.txt') womconfig.nsplinepoints = len(masterx) womconfig.tmpsplptsx = list(masterx) womconfig.tmpsplptsy = list(mastery) # print (type(womconfig.tmpsplptsx)) spline = splrep(womconfig.tmpsplptsx, womconfig.tmpsplptsy, k=3) splineresult = splev(wave, spline) plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) plt.pause(0.01) done = False print('Is this OK? ') answer = yesno('n') if (answer == 'y'): done = True splptsy = [ z for _, z in sorted(zip(womconfig.tmpsplptsx, womconfig.tmpsplptsy)) ] splptsx = sorted(womconfig.tmpsplptsx) while (not done): plotdone = False while (not plotdone): plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) plt.pause(0.01) print('Change scale? ') answer = yesno('n') if (answer == 'y'): plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') print('Click corners of box to change plot scale') newlims = plt.ginput(2, timeout=-1) xmin = newlims[0][0] ymin = newlims[0][1] xmax = newlims[1][0] ymax = newlims[1][1] plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) plotdone = True else: plotdone = True cid = fig.canvas.mpl_connect('button_press_event', onclick) cid2 = fig.canvas.mpl_connect('key_press_event', onkeypress) print('\nClick on continuum points for spline fit.') print('Left button or (a) = add point') print('Middle button or (s) = delete point') print('Right button or (d) = done\n') womconfig.pflag = '' while (womconfig.pflag != 'done'): plt.pause(0.01) fig.canvas.mpl_disconnect(cid) fig.canvas.mpl_disconnect(cid2) splptsy = [ z for _, z in sorted(zip(womconfig.tmpsplptsx, womconfig.tmpsplptsy)) ] splptsx = sorted(womconfig.tmpsplptsx) spline = splrep(splptsx, splptsy, k=3) splineresult = splev(wave, spline) plt.plot(wave, splineresult, drawstyle='steps-mid', color='g') plt.pause(0.01) print('Is this fit OK? ') answer = yesno('y') if (answer == 'y'): done = True if cal != None: np.savetxt('../../master_files/' + cal + '_splpts_master.txt', [splptsx, splptsy]) return splineresult
def fitspl(wave, flux, airlimit, fig, cal=None): """fit spline to spectrum""" import numpy as np import matplotlib.pyplot as plt from scipy.interpolate import splrep, splev import tmath.wombat.womconfig as womconfig from tmath.wombat.womget_element import womget_element from tmath.pydux.wave_telluric import wave_telluric from tmath.wombat.yesno import yesno from tmath.wombat.onclick import onclick from tmath.wombat.onkeypress import onkeypress import glob # global nsplinepoints, tmpsplptsx, tmpsplptsy, pflag # starting points for spline bandpts = np.array([3000, 3050, 3090, 3200, 3430, 3450, 3500, 3550, 3600, \ 3650, 3700, 3767, 3863, 3945, 4025, 4144, 4200, 4250, \ 4280, 4390, 4450, 4500, 4600, 4655, 4717, 4750, 4908, \ 4950, 5000, 5050, 5100, 5150, 5200, 5250, 5280, 5350, \ 5387, 5439, 5500, 5550, 6100, 6150, 6400, 6430, 6650, \ 6700, 6750, 6800, 7450, 7500, 7550, 8420, 8460, 8520, \ 8570, 8600, 8725, 8770, 9910, 10000, 10200, 10300, \ 10400, 10500, 10600, 10700]) locsinrange = np.logical_and((bandpts > wave[10]), (bandpts < wave[-10])) useband = bandpts[locsinrange] for i, _ in enumerate(useband): index = womget_element(wave, useband[i]) useband[i] = index # useband now has indices of wavelength positions if (min(flux) < 0): flux[np.where(flux < 0)] = 0.0 plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') xmin, xmax = plt.xlim() ymin, ymax = plt.ylim() plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) if (airlimit): loc = wave_telluric(wave, 'high') else: loc = wave_telluric(wave, 'low') #invert loc and convert all good sections to np.nan so they won't plot loc = np.invert(loc) wavetell = wave.copy() fluxtell = flux.copy() wavetell[loc] = np.nan fluxtell[loc] = np.nan plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if '../../master_files/' + cal + '_splpts_master.txt' not in glob.glob( '../../master_files/*'): womconfig.nsplinepoints = len(useband) womconfig.tmpsplptsx = wave[useband].copy().tolist() womconfig.tmpsplptsy = [] for i, _ in enumerate(useband): womconfig.tmpsplptsy.append( np.median(flux[useband[i] - 2:useband[i] + 3])) spline = splrep(womconfig.tmpsplptsx, womconfig.tmpsplptsy, k=3) else: masterx, mastery = np.genfromtxt('../../master_files/' + cal + '_splpts_master.txt') womconfig.nsplinepoints = len(masterx) womconfig.tmpsplptsx = masterx womconfig.tmpsplptsy = mastery spline = splrep(womconfig.tmpsplptsx, womconfig.tmpsplptsy, k=3) splineresult = splev(wave, spline) plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) plt.pause(0.01) done = False print('Is this OK? ') answer = yesno('n') if (answer == 'y'): done = True splptsy = [ z for _, z in sorted(zip(womconfig.tmpsplptsx, womconfig.tmpsplptsy)) ] splptsx = sorted(womconfig.tmpsplptsx) while (not done): plotdone = False while (not plotdone): plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) plt.pause(0.01) print('Change scale? ') answer = yesno('n') if (answer == 'y'): plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') print('Click corners of box to change plot scale') newlims = plt.ginput(2, timeout=-1) xmin = newlims[0][0] ymin = newlims[0][1] xmax = newlims[1][0] ymax = newlims[1][1] plt.cla() plt.plot(wave, flux, drawstyle='steps-mid', color='k') plt.plot(wavetell, fluxtell, drawstyle='steps-mid', color='violet') if (len(womconfig.tmpsplptsx) > 0): plt.plot(womconfig.tmpsplptsx, womconfig.tmpsplptsy, 'ro') plt.plot(wave, splineresult, color='g') plt.xlabel('Wavelength') plt.ylabel('Flux') plt.xlim([xmin, xmax]) plt.ylim([ymin, ymax]) plotdone = True else: plotdone = True cid = fig.canvas.mpl_connect('button_press_event', onclick) cid2 = fig.canvas.mpl_connect('key_press_event', onkeypress) print('\nClick on continuum points for spline fit.') print('Left button or (a) = add point') print('Middle button or (s) = delete point') print('Right button or (d) = done\n') womconfig.pflag = '' while (womconfig.pflag != 'done'): plt.pause(0.01) fig.canvas.mpl_disconnect(cid) fig.canvas.mpl_disconnect(cid2) splptsy = [ z for _, z in sorted(zip(womconfig.tmpsplptsx, womconfig.tmpsplptsy)) ] splptsx = sorted(womconfig.tmpsplptsx) spline = splrep(splptsx, splptsy, k=3) splineresult = splev(wave, spline) plt.plot(wave, splineresult, drawstyle='steps-mid', color='g') plt.pause(0.01) print('Is this fit OK? ') answer = yesno('y') if (answer == 'y'): done = True if cal != None: np.savetxt('../../master_files/' + cal + '_splpts_master.txt', [splptsx, splptsy]) return splineresult