Exemplo n.º 1
0
    def __init__(self, fig, ax, rv, ccf, primer, dp, outname):
        """We initialize with a figure object, three axis objects (in a list)
        an rv axis, the CCF, the user-generated prior on the doppler shadow (the two)
        points, and the pointer to the data in order to load physics."""
        import tayph.system_parameters as sp
        import numpy as np
        import pdb
        import scipy.interpolate as interpol
        import tayph.functions as fun
        import tayph.util as ut
        import sys
        #Upon initialization, we pass the keywords onto self.
        nexp = np.shape(ccf)[0]
        nrv = np.shape(ccf)[1]
        self.rv = rv
        self.ccf = ccf
        self.p = sp.phase(dp)
        if len(self.p) != nexp:
            print('ERROR IN FIT_DOPPLER_MODEL __INIT__:')
            print('The height of the CCF does not match nexp.')
            sys.exit()
        transit = sp.transit(dp) - 1.0
        self.T = abs(transit / max(abs(transit)))
        self.ax = ax
        self.aRstar = sp.paramget('aRstar', dp)
        self.vsys = sp.paramget('vsys', dp)
        self.RVp = sp.RV(dp) + self.vsys
        self.inclination = sp.paramget('inclination', dp)
        self.n_components = 1
        self.maskHW = 10.0  #Default masking half-width
        self.offset = 0.0
        self.D = 0  #Initial degree of polynomial fitting.
        self.S = 0
        self.dp = ut.check_path(dp, exists=True)
        self.outpath = self.dp / (outname + '.pkl')
        #Translate the pivot to RV-phase points.
        #Need to interpolate on phase only, because the primer was already defined in RV space.
        p_i = interpol.interp1d(np.arange(nexp, dtype=float), self.p)
        p1 = float(p_i(primer[0][1]))
        p2 = float(p_i(primer[1][1]))
        v1 = primer[0][0]
        v2 = primer[1][0]

        self.primer = [[v1, p1], [v2, p2]]
        self.fit_spinorbit()
        # self.primer = a*self.rv+b
        self.v_star_primer = fun.local_v_star(self.p, self.aRstar,
                                              self.inclination, self.vsini_fit,
                                              self.l_fit) + self.vsys
        # ax[0].plot(v_star_primer,fun.findgen(nexp),'--',color='black',label='Spin-orbit fit to primer')
        self.mask_ccf()
        self.fit_model()
Exemplo n.º 2
0
def start_run(configfile):
    """
    This is the main command-line initializer of the cross-correlation routine provided by Tayph.
    It parses a configuration file located at configfile which should contain predefined keywords.
    These keywords are passed on to the run_instance routine, which executes the analysis cascade.
    A call to this function is called a "run" of Tayph. A run has this configuration file as input,
    as well as a dataset, a (list of) cross-correlation template(s) and a (list of) models for injection
    purposes.
    """
    import tayph.system_parameters as sp
    cf = configfile

    print('')
    print('')
    print('')
    print('')
    print('')
    print('')
    print('')
    print('')
    print('')
    print('')
    print('')
    print(' = = = = = = = = = = = = = = = = =')
    print(' = = = = WELCOME TO TAYPH! = = = =')
    print(' = = = = = = = = = = = = = = = = =')
    print('')
    print(f'    Running {cf}')
    print('')
    print(' = = = = = = = = = = = = = = = = =')
    print('')
    print('')
    print('')

    print('---Start')
    print('---Load parameters from config file')
    modellist = sp.paramget('model', cf, full_path=True).split(',')
    templatelist = sp.paramget('template', cf, full_path=True).split(',')

    params = {
        'dp':
        sp.paramget('datapath', cf, full_path=True),
        'shadowname':
        sp.paramget('shadowname', cf, full_path=True),
        'maskname':
        sp.paramget('maskname', cf, full_path=True),
        'RVrange':
        sp.paramget('RVrange', cf, full_path=True),
        'drv':
        sp.paramget('drv', cf, full_path=True),
        'f_w':
        sp.paramget('f_w', cf, full_path=True),
        'do_colour_correction':
        sp.paramget('do_colour_correction', cf, full_path=True),
        'do_telluric_correction':
        sp.paramget('do_telluric_correction', cf, full_path=True),
        'do_xcor':
        sp.paramget('do_xcor', cf, full_path=True),
        'plot_xcor':
        sp.paramget('plot_xcor', cf, full_path=True),
        'make_mask':
        sp.paramget('make_mask', cf, full_path=True),
        'apply_mask':
        sp.paramget('apply_mask', cf, full_path=True),
        'do_berv_correction':
        sp.paramget('do_berv_correction', cf, full_path=True),
        'do_keplerian_correction':
        sp.paramget('do_keplerian_correction', cf, full_path=True),
        'make_doppler_model':
        sp.paramget('make_doppler_model', cf, full_path=True),
        'skip_doppler_model':
        sp.paramget('skip_doppler_model', cf, full_path=True),
        'modellist':
        modellist,
        'templatelist':
        templatelist,
        'c_subtract':
        sp.paramget('c_subtract', cf, full_path=True),
        'template_library':
        sp.paramget('template_library', cf, full_path=True),
        'model_library':
        sp.paramget('model_library', cf, full_path=True),
    }
    run_instance(params)
Exemplo n.º 3
0
def run_instance(p):
    """This runs the entire cross correlation analysis cascade."""
    import numpy as np
    from astropy.io import fits
    import astropy.constants as const
    import astropy.units as u
    from matplotlib import pyplot as plt
    import os.path
    import scipy.interpolate as interp
    import pylab
    import pdb
    import os.path
    import os
    import sys
    import glob
    import distutils.util
    import pickle
    import copy
    from pathlib import Path

    import tayph.util as ut
    import tayph.operations as ops
    import tayph.functions as fun
    import tayph.system_parameters as sp
    import tayph.tellurics as telcor
    import tayph.masking as masking
    import tayph.models as models
    from tayph.ccf import xcor, clean_ccf, filter_ccf, construct_KpVsys
    from tayph.vartests import typetest, notnegativetest, nantest, postest, typetest_array, dimtest
    import tayph.shadow as shadow
    # from lib import analysis
    # from lib import cleaning
    # from lib import masking as masking
    # from lib import shadow as shadow
    # from lib import molecfit as telcor

    #First parse the parameter dictionary into required variables and test them.
    typetest(p, dict, 'params in run_instance()')

    dp = Path(p['dp'])
    ut.check_path(dp, exists=True)

    modellist = p['modellist']
    templatelist = p['templatelist']
    model_library = p['model_library']
    template_library = p['template_library']
    typetest(modellist, [str, list], 'modellist in run_instance()')
    typetest(templatelist, [str, list], 'templatelist in run_instance()')
    typetest(model_library, str, 'model_library in run_instance()')
    typetest(template_library, str, 'template_library in run_instance()')
    ut.check_path(model_library, exists=True)
    ut.check_path(template_library, exists=True)
    if type(modellist) == str:
        modellist = [modellist]  #Force this to be a list
    if type(templatelist) == str:
        templatelist = [templatelist]  #Force this to be a list
    typetest_array(modellist, str, 'modellist in run_instance()')
    typetest_array(templatelist, str, 'modellist in run_instance()')

    shadowname = p['shadowname']
    maskname = p['maskname']
    typetest(shadowname, str, 'shadowname in run_instance()')
    typetest(maskname, str, 'shadowname in run_instance()')

    RVrange = p['RVrange']
    drv = p['drv']
    f_w = p['f_w']
    resolution = sp.paramget('resolution', dp)
    typetest(RVrange, [int, float], 'RVrange in run_instance()')
    typetest(drv, [int, float], 'drv in run_instance()')
    typetest(f_w, [int, float], 'f_w in run_instance()')
    typetest(resolution, [int, float], 'resolution in run_instance()')
    nantest(RVrange, 'RVrange in run_instance()')
    nantest(drv, 'drv in run_instance()')
    nantest(f_w, 'f_w in run_instance()')
    nantest(resolution, 'resolution in run_instance()')
    postest(RVrange, 'RVrange in run_instance()')
    postest(drv, 'drv in run_instance()')
    postest(resolution, 'resolution in run_instance()')
    notnegativetest(f_w, 'f_w in run_instance()')

    do_colour_correction = p['do_colour_correction']
    do_telluric_correction = p['do_telluric_correction']
    do_xcor = p['do_xcor']
    plot_xcor = p['plot_xcor']
    make_mask = p['make_mask']
    apply_mask = p['apply_mask']
    c_subtract = p['c_subtract']
    do_berv_correction = p['do_berv_correction']
    do_keplerian_correction = p['do_keplerian_correction']
    make_doppler_model = p['make_doppler_model']
    skip_doppler_model = p['skip_doppler_model']
    typetest(do_colour_correction, bool,
             'do_colour_correction in run_instance()')
    typetest(do_telluric_correction, bool,
             'do_telluric_correction in run_instance()')
    typetest(do_xcor, bool, 'do_xcor in run_instance()')
    typetest(plot_xcor, bool, 'plot_xcor in run_instance()')
    typetest(make_mask, bool, 'make_mask in run_instance()')
    typetest(apply_mask, bool, 'apply_mask in run_instance()')
    typetest(c_subtract, bool, 'c_subtract in run_instance()')
    typetest(do_berv_correction, bool, 'do_berv_correction in run_instance()')
    typetest(do_keplerian_correction, bool,
             'do_keplerian_correction in run_instance()')
    typetest(make_doppler_model, bool, 'make_doppler_model in run_instance()')
    typetest(skip_doppler_model, bool, 'skip_doppler_model in run_instance()')

    #We start by defining constants and preparing for generating output.
    c = const.c.value / 1000.0  #in km/s
    colourdeg = 3  #A fitting degree for the colour correction.

    print(
        f'---Passed parameter input tests. Initiating output folder tree in {Path("output")/dp}.'
    )
    libraryname = str(template_library).split('/')[-1]
    if str(dp).split('/')[0] == 'data':
        dataname = str(dp).replace('data/', '')
        print(
            f'------Data is located in data/ folder. Assuming output name for this dataset as {dataname}'
        )
    else:
        dataname = dp
        print(
            f'------Data is NOT located in data/ folder. Assuming output name for this dataset as {dataname}'
        )

    list_of_wls = []  #This will store all the data.
    list_of_orders = []  #All of it needs to be loaded into your memory.
    list_of_sigmas = []

    trigger2 = 0  #These triggers are used to limit the generation of output in the forloop.
    trigger3 = 0
    n_negative_total = 0  #This will hold the total number of pixels that were set to NaN because they were zero when reading in the data.
    air = sp.paramget('air', dp)  #Read bool from str in config file.
    typetest(air, bool, 'air in run_instance()')

    filelist_orders = [str(i) for i in Path(dp).glob('order_*.fits')]
    if len(filelist_orders) == 0:
        raise Exception(
            f'Runtime error: No orders_*.fits files were found in {dp}.')
    try:
        order_numbers = [
            int(i.split('order_')[1].split('.')[0]) for i in filelist_orders
        ]
    except:
        raise Exception(
            'Runtime error: Failed casting fits filename numerals to ints. Are the filenames of the spectral orders correctly formatted?'
        )
    order_numbers.sort()  #This is the ordered list of numerical order IDs.
    n_orders = len(order_numbers)
    if n_orders == 0:
        raise Exception(
            f'Runtime error: n_orders may not have ended up as zero. ({n_orders})'
        )

#Loading the data from the datafolder.
    if do_xcor == True or plot_xcor == True or make_mask == True:
        print(f'---Loading orders from {dp}.')

        # for i in range(startorder,endorder+1):
        for i in order_numbers:
            wavepath = dp / f'wave_{i}.fits'
            orderpath = dp / f'order_{i}.fits'
            sigmapath = dp / f'sigma_{i}.fits'
            ut.check_path(wavepath, exists=True)
            ut.check_path(orderpath, exists=True)
            ut.check_path(sigmapath, exists=False)
            wave_axis = fits.getdata(wavepath)
            dimtest(wave_axis, [0], 'wavelength grid in run_instance()')
            n_px = len(wave_axis)  #Pixel width of the spectral order.
            if air == False:
                if i == np.min(order_numbers):
                    print("------Assuming wavelengths are in vaccuum.")
                list_of_wls.append(1.0 * wave_axis)
            else:
                if i == np.min(order_numbers):
                    print("------Applying airtovac correction.")
                list_of_wls.append(ops.airtovac(wave_axis))

            order_i = fits.getdata(orderpath)
            if i == np.min(order_numbers):
                dimtest(
                    order_i, [0, n_px], f'order {i} in run_instance()'
                )  #For the first order, check that it is 2D and that is has a width equal to n_px.
                n_exp = np.shape(
                    order_i
                )[0]  #then fix n_exp. All other orders should have the same n_exp.
                print(f'------{n_exp} exposures recognised.')
            else:
                dimtest(order_i, [n_exp, n_px], f'order {i} in run_instance()')

            #Now test for negatives, set them to NaN and track them.
            n_negative = len(order_i[order_i <= 0])
            if trigger3 == 0 and n_negative > 0:
                print("------Setting negative values to NaN.")
                trigger3 = -1
            n_negative_total += n_negative
            order_i[order_i <= 0] = np.nan
            postest(order_i, f'order {i} in run_instance().'
                    )  #make sure whatever comes out here is strictly positive.
            list_of_orders.append(order_i)

            try:  #Try to get a sigma file. If it doesn't exist, we raise a warning. If it does, we test its dimensions and append it.
                sigma_i = fits.getdata(sigmapath)
                dimtest(sigma_i, [n_exp, n_px],
                        f'order {i} in run_instance().')
                list_of_sigmas.append(sigma_i)
            except FileNotFoundError:
                if trigger2 == 0:
                    print(
                        '------WARNING: Sigma (flux error) files not provided. Assuming sigma = sqrt(flux). This is standard practise for HARPS data, but e.g. ESPRESSO has a pipeline that computes standard errors on each pixel for you.'
                    )
                    trigger2 = -1
                list_of_sigmas.append(np.sqrt(order_i))
        print(
            f"------{n_negative_total} negative values set to NaN ({np.round(100.0*n_negative_total/n_exp/n_px/len(order_numbers),2)}% of total spectral pixels in dataset.)"
        )

    if len(list_of_orders) != n_orders:
        raise Exception(
            'Runtime error: n_orders is not equal to the length of list_of_orders. Something went wrong when reading them in?'
        )

    print('---Finished loading dataset to memory.')

    #Apply telluric correction file or not.
    # plt.plot(list_of_wls[60],list_of_orders[60][10],color='red')
    # plt.plot(list_of_wls[60],list_of_orders[60][10]+list_of_sigmas[60][10],color='red',alpha=0.5)#plot corrected spectra
    # plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='red',alpha=0.5)#plot SNR
    if do_telluric_correction == True and n_orders > 0:
        print('---Applying telluric correction')
        telpath = dp / 'telluric_transmission_spectra.pkl'
        list_of_orders, list_of_sigmas = telcor.apply_telluric_correction(
            telpath, list_of_wls, list_of_orders, list_of_sigmas)

    # plt.plot(list_of_wls[60],list_of_orders[60][10],color='blue')
    # plt.plot(list_of_wls[60],list_of_orders[60][10]+list_of_sigmas[60][10],color='blue',alpha=0.5)#plot corrected spectra

    # plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='blue',alpha=0.5) #plot SNR
    # plt.show()
    # pdb.set_trace()

#Do velocity correction of wl-solution. Explicitly after telluric correction
#but before masking. Because the cross-correlation relies on columns being masked.
#Then if you start to move the CCFs around before removing the time-average,
#each masked column becomes slanted. Bad deal.
    rv_cor = 0
    if do_berv_correction == True:
        rv_cor += sp.berv(dp)
    if do_keplerian_correction == True:
        rv_cor -= sp.RV_star(dp) * (1.0)

    if type(rv_cor) != int and len(list_of_orders) > 0:
        print('---Reinterpolating data to correct velocities')
        list_of_orders_cor = []
        list_of_sigmas_cor = []
        for i in range(len(list_of_wls)):
            order = list_of_orders[i]
            sigma = list_of_sigmas[i]
            order_cor = order * 0.0
            sigma_cor = sigma * 0.0
            for j in range(len(list_of_orders[0])):
                wl_i = interp.interp1d(list_of_wls[i],
                                       order[j],
                                       bounds_error=False)
                si_i = interp.interp1d(list_of_wls[i],
                                       sigma[j],
                                       bounds_error=False)
                wl_cor = list_of_wls[i] * (
                    1.0 - (rv_cor[j] * u.km / u.s / const.c)
                )  #The minus sign was tested on a slow-rotator.
                order_cor[j] = wl_i(wl_cor)
                sigma_cor[j] = si_i(
                    wl_cor
                )  #I checked that this works because it doesn't affect the SNR, apart from wavelength-shifting it.
            list_of_orders_cor.append(order_cor)
            list_of_sigmas_cor.append(sigma_cor)
            ut.statusbar(i, fun.findgen(len(list_of_wls)))
        # plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='blue')
        # plt.plot(list_of_wls[60],list_of_orders_cor[60][10]/list_of_sigmas_cor[60][10],color='red')
        # plt.show()
        # sys.exit()
        list_of_orders = list_of_orders_cor
        list_of_sigmas = list_of_sigmas_cor

    if len(list_of_orders) != n_orders:
        raise RuntimeError(
            'n_orders is no longer equal to the length of list_of_orders, though it was before. Something went wrong during telluric correction or velocity correction.'
        )

#Compute / create a mask and save it to file (or not)
    if make_mask == True and len(list_of_orders) > 0:
        if do_colour_correction == True:
            print(
                '---Constructing mask with intra-order colour correction applied'
            )
            masking.mask_orders(list_of_wls,
                                ops.normalize_orders(list_of_orders,
                                                     list_of_sigmas,
                                                     colourdeg)[0],
                                dp,
                                maskname,
                                40.0,
                                5.0,
                                manual=True)
        else:
            print(
                '---Constructing mask WITHOUT intra-order colour correction applied.'
            )
            print(
                '---Switch on colour correction if you see colour variations in the 2D spectra.'
            )
            masking.mask_orders(list_of_wls,
                                list_of_orders,
                                dp,
                                maskname,
                                40.0,
                                5.0,
                                manual=True)
        if apply_mask == False:
            print(
                '---WARNING in run_instance: Mask was made but is not applied to data (apply_mask == False)'
            )

#Apply the mask that was previously created and saved to file.
    if apply_mask == True:
        print('---Applying mask')
        list_of_orders = masking.apply_mask_from_file(dp, maskname,
                                                      list_of_orders)
        list_of_sigmas = masking.apply_mask_from_file(dp, maskname,
                                                      list_of_sigmas)
#Interpolate over all isolated NaNs and set bad columns to NaN (so that they are ignored in the CCF)
    if do_xcor == True:
        print('---Healing NaNs')
        list_of_orders = masking.interpolate_over_NaNs(
            list_of_orders
        )  #THERE IS AN ISSUE HERE: INTERPOLATION SHOULD ALSO HAPPEN ON THE SIGMAS ARRAY!
        list_of_sigmas = masking.interpolate_over_NaNs(list_of_sigmas)

#Normalize the orders to their average flux in order to effectively apply a broad-band colour correction (colour is typically a function of airmass and seeing).
    if do_colour_correction == True:
        print('---Normalizing orders to common flux level')
        # plt.plot(list_of_wls[60],list_of_orders[60][10]/list_of_sigmas[60][10],color='blue',alpha=0.4)
        list_of_orders_normalised, list_of_sigmas_normalised, meanfluxes = ops.normalize_orders(
            list_of_orders, list_of_sigmas, colourdeg
        )  #I tested that this works because it doesn't alter the SNR.

        meanfluxes_norm = meanfluxes / np.nanmean(meanfluxes)
    else:
        meanfluxes_norm = fun.findgen(len(
            list_of_orders[0])) * 0.0 + 1.0  #All unity.
        # plt.plot(list_of_wls[60],list_of_orders_normalised[60][10]/list_of_sigmas[60][10],color='red',alpha=0.4)
        # plt.show()
        # sys.exit()

    if len(list_of_orders) != n_orders:
        raise RuntimeError(
            'n_orders is no longer equal to the length of list_of_orders, though it was before. Something went wrong during masking or colour correction.'
        )

#Construct the cross-correlation templates in case we will be computing or plotting the CCF.
    if do_xcor == True or plot_xcor == True:

        list_of_wlts = []
        list_of_templates = []
        outpaths = []

        for templatename in templatelist:
            print(f'---Building template {templatename}')
            wlt, T = models.build_template(templatename,
                                           binsize=0.5,
                                           maxfrac=0.01,
                                           resolution=resolution,
                                           template_library=template_library,
                                           c_subtract=c_subtract)
            T *= (-1.0)
            if np.mean(wlt) < 50.0:  #This is likely in microns:
                print(
                    '------WARNING: The loaded template has a mean wavelength less than 50.0, meaning that it is very likely not in nm, but in microns. I have divided by 1,000 now and hope for the best...'
                )
                wlt *= 1000.0
            list_of_wlts.append(wlt)
            list_of_templates.append(T)

            outpath = Path('output') / Path(dataname) / Path(
                libraryname) / Path(templatename)

            if not os.path.exists(outpath):
                print(
                    f"------The output location ({outpath}) didn't exist, I made it now."
                )
                os.makedirs(outpath)
            outpaths.append(outpath)

#Perform the cross-correlation on the entire list of orders.
    for i in range(len(list_of_wlts)):
        templatename = templatelist[i]
        wlt = list_of_wlts[i]
        T = list_of_templates[i]
        outpath = outpaths[i]
        if do_xcor == True:
            print(
                f'---Cross-correlating spectra with template {templatename}.')
            t1 = ut.start()
            rv, ccf, ccf_e, Tsums = xcor(
                list_of_wls,
                list_of_orders_normalised,
                np.flipud(np.flipud(wlt)),
                T,
                drv,
                RVrange,
                list_of_errors=list_of_sigmas_normalised)
            ut.end(t1)
            print(f'------Writing CCFs to {str(outpath)}')
            ut.writefits(outpath / 'ccf.fits', ccf)
            ut.writefits(outpath / 'ccf_e.fits', ccf_e)
            ut.writefits(outpath / 'RV.fits', rv)
            ut.writefits(outpath / 'Tsum.fits', Tsums)
        else:
            print(
                f'---Reading CCFs with template {templatename} from {str(outpath)}.'
            )
            if os.path.isfile(outpath / 'ccf.fits') == False:
                raise FileNotFoundError(
                    f'CCF output not located at {outpath}. Rerun with do_xcor=True to create these files?'
                )
        rv = fits.getdata(outpath / 'rv.fits')
        ccf = fits.getdata(outpath / 'ccf.fits')
        ccf_e = fits.getdata(outpath / 'ccf_e.fits')
        Tsums = fits.getdata(outpath / 'Tsum.fits')

        ccf_cor = ccf * 1.0
        ccf_e_cor = ccf_e * 1.0

        print('---Cleaning CCFs')
        ccf_n, ccf_ne, ccf_nn, ccf_nne = clean_ccf(rv, ccf_cor, ccf_e_cor, dp)

        if make_doppler_model == True and skip_doppler_model == False:
            shadow.construct_doppler_model(rv,
                                           ccf_nn,
                                           dp,
                                           shadowname,
                                           xrange=[-200, 200],
                                           Nxticks=20.0,
                                           Nyticks=10.0)
            make_doppler_model = False  # This sets it to False after it's been run once, for the first template.
        if skip_doppler_model == False:
            print('---Reading doppler shadow model from ' + shadowname)
            doppler_model, dsmask = shadow.read_shadow(
                dp, shadowname, rv, ccf
            )  #This returns both the model evaluated on the rv,ccf grid, as well as the mask that blocks the planet trace.
            ccf_clean, matched_ds_model = shadow.match_shadow(
                rv, ccf_nn, dsmask, dp, doppler_model
            )  #THIS IS AN ADDITIVE CORRECTION, SO CCF_NNE DOES NOT NEED TO BE ALTERED AND IS STILL VALID VOOR CCF_CLEAN
        else:
            print('---Not performing shadow correction')
            ccf_clean = ccf_nn * 1.0
            matched_ds_model = ccf_clean * 0.0

        if f_w > 0.0:
            print('---Performing high-pass filter on the CCF')
            ccf_clean_filtered, wiggles = filter_ccf(
                rv, ccf_clean, v_width=f_w
            )  #THIS IS ALSO AN ADDITIVE CORRECTION, SO CCF_NNE IS STILL VALID.
        else:
            print('---Skipping high-pass filter')
            ccf_clean_filtered = ccf_clean * 1.0
            wiggles = ccf_clean * 0.0  #This filtering is additive so setting to zero is accurate.

        print('---Weighing CCF rows by mean fluxes that were normalised out')
        ccf_clean_weighted = np.transpose(
            np.transpose(ccf_clean_filtered) * meanfluxes_norm
        )  #MULTIPLYING THE AVERAGE FLUXES BACK IN! NEED TO CHECK THAT THIS ALSO GOES PROPERLY WITH THE ERRORS!
        ccf_nne = np.transpose(np.transpose(ccf_nne) * meanfluxes_norm)

        ut.save_stack(outpath / 'cleaning_steps.fits', [
            ccf, ccf_cor, ccf_nn, ccf_clean, matched_ds_model,
            ccf_clean_filtered, wiggles, ccf_clean_weighted
        ])
        ut.writefits(outpath / 'ccf_cleaned.fits', ccf_clean_weighted)
        ut.writefits(outpath / 'ccf_cleaned_error.fits', ccf_nne)

        print('---Constructing KpVsys')
        Kp, KpVsys, KpVsys_e = construct_KpVsys(rv, ccf_clean_weighted,
                                                ccf_nne, dp)
        ut.writefits(outpath / 'KpVsys.fits', KpVsys)
        ut.writefits(outpath / 'KpVsys_e.fits', KpVsys_e)
        ut.writefits(outpath / 'Kp.fits', Kp)

    return
    sys.exit()

    if plot_xcor == True and inject_model == False:
        print('---Plotting KpVsys')
        analysis.plot_KpVsys(rv, Kp, KpVsys, dp)

    #Now repeat it all for the model injection.
    if inject_model == True:
        for modelname in modellist:
            outpath_i = outpath + modelname + '/'
            if do_xcor == True:
                print('---Injecting model ' + modelname)
                list_of_orders_injected = models.inject_model(
                    list_of_wls,
                    list_of_orders,
                    dp,
                    modelname,
                    model_library=model_library
                )  #Start with the unnormalised orders from before.
                #Normalize the orders to their average flux in order to effectively apply
                #a broad-band colour correction (colour is a function of airmass and seeing).
                if do_colour_correction == True:
                    print(
                        '------Normalizing injected orders to common flux level'
                    )
                    list_of_orders_injected, list_of_sigmas_injected, meanfluxes_injected = ops.normalize_orders(
                        list_of_orders_injected, list_of_sigmas, colourdeg)
                    meanfluxes_norm_injected = meanfluxes_injected / np.mean(
                        meanfluxes_injected)
                else:
                    meanfluxes_norm_injected = fun.findgen(
                        len(list_of_orders_injected[0])
                    ) * 0.0 + 1.0  #All unity.

                print('------Cross-correlating injected orders')
                rv_i, ccf_i, ccf_e_i, Tsums_i = analysis.xcor(
                    list_of_wls,
                    list_of_orders_injected,
                    np.flipud(np.flipud(wlt)),
                    T,
                    drv,
                    RVrange,
                    list_of_errors=list_of_sigmas_injected)
                print('------Writing injected CCFs to ' + outpath_i)
                if not os.path.exists(outpath_i):
                    print("---------That path didn't exist, I made it now.")
                    os.makedirs(outpath_i)
                ut.writefits(outpath_i + '/' + 'ccf_i_' + modelname + '.fits',
                             ccf_i)
                ut.writefits(
                    outpath_i + '/' + 'ccf_e_i_' + modelname + '.fits',
                    ccf_e_i)
            else:
                print('---Reading injected CCFs from ' + outpath_i)
                if os.path.isfile(outpath_i + 'ccf_i_' + modelname +
                                  '.fits') == False:
                    print('------ERROR: Injected CCF not located at ' +
                          outpath_i + 'ccf_i_' + modelname + '.fits' +
                          '. Set do_xcor and inject_model to True?')
                    sys.exit()
                if os.path.isfile(outpath_i + 'ccf_e_i_' + modelname +
                                  '.fits') == False:
                    print('------ERROR: Injected CCF error not located at ' +
                          outpath_i + 'ccf_e_i_' + modelname + '.fits' +
                          '. Set do_xcor and inject_model to True?')
                    sys.exit()
                # f.close()
                # f2.close()
                ccf_i = fits.getdata(outpath_i + 'ccf_i_' + modelname +
                                     '.fits')
                ccf_e_i = fits.getdata(outpath_i + 'ccf_e_i_' + modelname +
                                       '.fits')

            print('---Cleaning injected CCFs')
            ccf_n_i, ccf_ne_i, ccf_nn_i, ccf_nne_i = cleaning.clean_ccf(
                rv, ccf_i, ccf_e_i, dp)
            ut.writefits(outpath_i + 'ccf_normalized_i.fits', ccf_nn_i)
            ut.writefits(outpath_i + 'ccf_ne_i.fits', ccf_ne_i)

            # if make_doppler_model == True and skip_doppler_model == False:
            # shadow.construct_doppler_model(rv,ccf_nn,dp,shadowname,xrange=[-200,200],Nxticks=20.0,Nyticks=10.0)
            if skip_doppler_model == False:
                # print('---Reading doppler shadow model from '+shadowname)
                # doppler_model,maskHW = shadow.read_shadow(dp,shadowname,rv,ccf)
                ccf_clean_i, matched_ds_model_i = shadow.match_shadow(
                    rv, ccf_nn_i, dp, doppler_model, maskHW)
            else:
                print(
                    '---Not performing shadow correction on injected spectra either.'
                )
                ccf_clean_i = ccf_nn_i * 1.0
                matched_ds_model_i = ccf_clean_i * 0.0

            if f_w > 0.0:
                ccf_clean_i_filtered, wiggles_i = cleaning.filter_ccf(
                    rv, ccf_clean_i, v_width=f_w)
            else:
                ccf_clean_i_filtered = ccf_clean_i * 1.0

            ut.writefits(outpath_i + 'ccf_cleaned_i.fits',
                         ccf_clean_i_filtered)
            ut.writefits(outpath + 'ccf_cleaned_i_error.fits', ccf_nne)

            print(
                '---Weighing injected CCF rows by mean fluxes that were normalised out'
            )
            ccf_clean_i_filtered = np.transpose(
                np.transpose(ccf_clean_i_filtered) * meanfluxes_norm_injected
            )  #MULTIPLYING THE AVERAGE FLUXES BACK IN! NEED TO CHECK THAT THIS ALSO GOES PROPERLY WITH THE ERRORS!
            ccf_nne_i = np.transpose(
                np.transpose(ccf_nne_i) * meanfluxes_norm_injected)

            print('---Constructing injected KpVsys')
            Kp, KpVsys_i, KpVsys_e_i = analysis.construct_KpVsys(
                rv, ccf_clean_i_filtered, ccf_nne_i, dp)
            ut.writefits(outpath_i + 'KpVsys_i.fits', KpVsys_i)
            # ut.writefits(outpath+'KpVsys_e_i.fits',KpVsys_e_i)
            if plot_xcor == True:
                print('---Plotting KpVsys with ' + modelname + ' injected.')
                analysis.plot_KpVsys(rv, Kp, KpVsys, dp, injected=KpVsys_i)
Exemplo n.º 4
0
def inject_model(list_of_wls,
                 list_of_orders,
                 dp,
                 modelname,
                 model_library='library/models'):
    """This function takes a list of spectral orders and injects a model with library
    identifier modelname, and system parameters as defined in dp. The model is blurred taking into
    account spectral resolution and rotation broadening (with an LSF as per Brogi et al.) and
    finite-exposure broadening (with a box LSF).

    It returns a copy of the list of orders with the model injected."""

    import tayph.util as ut
    import tayph.system_parameters as sp
    import tayph.models
    import astropy.constants as const
    import numpy as np
    import scipy
    import tayph.operations as ops
    from tayph.vartests import typetest, dimtest
    import pdb
    import copy
    import matplotlib.pyplot as plt

    # dimtest(order,[0,len(wld)])
    dp = ut.check_path(dp)
    typetest(modelname, str, 'modelname in models.inject_model()')
    typetest(model_library, str, 'model_library in models.inject_model()')

    c = const.c.to('km/s').value  #In km/s
    Rd = sp.paramget('resolution', dp)
    planet_radius = sp.paramget('Rp', dp)
    inclination = sp.paramget('inclination', dp)
    P = sp.paramget('P', dp)
    transit = sp.transit(dp)
    n_exp = len(transit)
    vsys = sp.paramget('vsys', dp)
    rv = sp.RV(dp) + vsys
    dRV = sp.dRV(dp)
    phi = sp.phase(dp)
    dimtest(transit, [n_exp])
    dimtest(rv, [n_exp])
    dimtest(phi, [n_exp])
    dimtest(dRV, [n_exp])

    mask = (transit - 1.0) / (np.min(transit - 1.0))

    wlm, fxm = get_model(modelname, library=model_library)
    if wlm[-1] <= wlm[0]:  #Reverse the wl axis if its sorted the wrong way.
        wlm = np.flipud(wlm)
        fxm = np.flipud(fxm)

    #With the model and the revelant parameters in hand, now only select that
    #part of the model that covers the wavelengths of the order provided.
    #A larger wavelength range would take much extra time because the convolution
    #is a slow operation.

    N_orders = len(list_of_wls)
    if N_orders != len(list_of_orders):
        raise RuntimeError(
            f'in models.inject_model: List_of_wls and list_of_orders are not of the '
            f'same length ({N_orders} vs {len(list_of_orders)})')

    if np.min(wlm) > np.min(list_of_wls) - 1.0 or np.max(
            wlm) < np.max(list_of_wls) + 1.0:
        ut.tprint(
            'WARNING in model injection: Data grid falls (partly) outside of model range. '
            'Setting missing area to 1.0. (meaning, no planet absorption.)')

    modelsel = [
        (wlm >= np.min(list_of_wls) - 1.0) & (wlm <= np.max(list_of_wls) + 1.0)
    ]

    wlm = wlm[tuple(modelsel)]
    fxm = fxm[tuple(modelsel)]

    fxm_b = ops.blur_rotate(wlm, fxm, c / Rd, planet_radius, P,
                            inclination)  #Only do this once per dataset.

    oversampling = 2.5
    wlm_cv, fxm_bcv, vstep = ops.constant_velocity_wl_grid(
        wlm, fxm_b, oversampling=oversampling)

    if np.min(dRV) < c / Rd / 10.0:

        dRV_min = c / Rd / 10.0  #If the minimum dRV is less than 10% of the spectral
        #resolution, we introduce a lower limit to when we are going to blur, because the effect
        #becomes insignificant.
    else:
        dRV_min = np.min(dRV)

    if dRV_min / vstep < 3:  #Meaning, if the smoothing is going to be undersampled by this choice
        #in v_step, it means that the oversampling parameter in ops.constant_velocity_wl_grid was
        #not high enough. Then we adjust it. I choose a value of 3 here to be safe, even though
        #ops.smooth below accepts values as low as 2.
        oversampling_new = 3.0 / (
            dRV_min / vstep) * oversampling  #scale up the oversampling factor.
        wlm_cv, fxm_bcv, vstep = ops.constant_velocity_wl_grid(
            wlm, fxm_b, oversampling=oversampling_new)

    list_of_orders_injected = copy.deepcopy(list_of_orders)

    for i in range(n_exp):
        if dRV[i] >= c / Rd / 10.0:
            fxm_b2 = ops.smooth(fxm_bcv, dRV[i] / vstep, mode='box')
        else:
            fxm_b2 = copy.deepcopy(fxm_bcv)
        shift = 1.0 + rv[i] / c
        fxm_i = scipy.interpolate.interp1d(
            wlm_cv * shift, fxm_b2, fill_value=1.0,
            bounds_error=False)  #This is a class that can be called.
        #Fill_value = 1 because if the model does not fully cover the order, it will be padded with 1.0s,
        #assuming that we are dealing with a model that is in transmission.
        for j in range(len(list_of_orders)):
            list_of_orders_injected[j][i] *= (
                1.0 + mask[i] * (fxm_i(list_of_wls[j]) - 1.0))  #This assumes
            #that the model is in transit radii. This can definitely be vectorised!

        #These are for checking that the broadening worked as expected:
        # injection_total[i,:]= scipy.interpolate.interp1d(wlm_cv*shift,fxm_b2)(wld)
        # injection_rot_only[i,:]=scipy.interpolate.interp1d(wlm*shift,fxm_b)(wld)
        # injection_pure[i,:]=scipy.interpolate.interp1d(wlm*shift,fxm)(wld)

    # ut.save_stack('test.fits',[injection_pure,injection_rot_only,injection_total])
    # pdb.set_trace()
    # ut.writefits('test.fits',injection)
    # pdb.set_trace()

    return (list_of_orders_injected)
Exemplo n.º 5
0
def test_molecfit_config(molecfit_config):
    """This tests the existence and integrity of the system-wide molecfit configuration folder."""
    import tayph.util as ut
    import tayph.system_parameters as sp
    from pathlib import Path
    import sys

    try:
        molecfit_input_folder = Path(
            sp.paramget('molecfit_input_folder',
                        molecfit_config,
                        full_path=True))
        molecfit_prog_folder = Path(
            sp.paramget('molecfit_prog_folder',
                        molecfit_config,
                        full_path=True))
        python_alias = sp.paramget('python_alias',
                                   molecfit_config,
                                   full_path=True)
    except:
        err_msg = (
            f'ERROR in initialising Molecfit. The molecfit configuration file '
            f'({molecfit_config}) exists, but it does not contain the right keywords. The required '
            'parameters are molecfit_input_folder, molecfit_prog_folder and python_alias'
        )
        ut.tprint(err_msg)
        sys.exit()

    if molecfit_input_folder.exists() == False:
        err_msg = (
            f"ERROR in initialising Molecfit. The molecfit configuration file "
            f"({molecfit_config}) exists and it has the correct parameter keywords "
            f"(molecfit_input_folder, molecfit_prog_folder and python_alias), but the "
            f"molecfit_input_folder path ({molecfit_input_folder}) does not exist. Please run "
            f"tayph.tellurics.set_molecfit_config() to resolve this.")
        ut.tprint(err_msg)
        sys.exit()

    if molecfit_prog_folder.exists() == False:
        err_msg = (
            f"ERROR in initialising Molecfit. The molecfit configuration file "
            f"({molecfit_config}) exists and it has the correct parameter keywords "
            f"(molecfit_input_folder, molecfit_prog_folder and python_alias), but the "
            f"molecfit_prog_folder path ({molecfit_prog_folder}) does not exist. Please run "
            f"tayph.tellurics.set_molecfit_config() to resolve this.")
        ut.tprint(err_msg)
        sys.exit()
    binarypath = molecfit_prog_folder / 'molecfit'
    guipath = molecfit_prog_folder / 'molecfit_gui'

    if binarypath.exists() == False:
        err_msg = (
            f"ERROR in initialising Molecfit. The molecfit configuration file "
            f"({molecfit_config}) exists and it has the correct parameter keywords "
            f"(molecfit_input_folder, molecfit_prog_folder and python_alias), but the molecfit "
            f"binary ({binarypath}) does not exist. Please run "
            f"tayph.tellurics.set_molecfit_config() to resolve this.")
        ut.tprint(err_msg)
        sys.exit()

    if guipath.exists() == False:
        err_msg = (
            f"ERROR in initialising Molecfit. The molecfit configuration file "
            f"({molecfit_config}) exists and it has the correct parameter keywords "
            f"(molecfit_input_folder, molecfit_prog_folder and python_alias), but the molecfit "
            f"gui binary ({guipath}) does not exist. Please run "
            f"tayph.tellurics.set_molecfit_config() to resolve this.")
        ut.tprint(err_msg)
        sys.exit()

    if ut.test_alias(python_alias) == False:
        err_msg = (
            f'ERROR in initialising Molecfit. The molecfit configuration file '
            f'({molecfit_config}) exists and it has the correct parameter keywords '
            f'(molecfit_input_folder, molecfit_prog_folder and python_alias), but the python '
            f'alias ({python_alias}) does not exist. Please run '
            f'tayph.tellurics.set_molecfit_config() to resolve this.')
        ut.tprint(err_msg)
        sys.exit()
Exemplo n.º 6
0
def set_molecfit_config(configpath):
    import pkg_resources
    from pathlib import Path
    import os
    import subprocess
    import tayph.system_parameters as sp
    import tayph.util as ut

    #Prepare for making formatted output.
    # terminal_height,terminal_width = subprocess.check_output(['stty', 'size']).split()

    Q1 = (
        'In what folder are parameter files defined and should (intermediate) molecfit output be '
        'written to?')
    Q2 = 'In what folder is the molecfit binary located?'
    Q3 = 'What is your python 3.x alias?'

    # configpath=get_molecfit_config()
    configpath = Path(configpath)
    if configpath.exists():
        ut.tprint(
            f'Molecfit configuration file already exists at {configpath}.')
        print('Overwriting existing values.')
        current_molecfit_input_folder = sp.paramget('molecfit_input_folder',
                                                    configpath,
                                                    full_path=True)
        current_molecfit_prog_folder = sp.paramget('molecfit_prog_folder',
                                                   configpath,
                                                   full_path=True)
        current_python_alias = sp.paramget('python_alias',
                                           configpath,
                                           full_path=True)

        ut.tprint(Q1)
        ut.tprint(
            f'Currently: {current_molecfit_input_folder} (leave empty to keep current '
            'value).')

        new_input_folder_input = str(input())
        if len(new_input_folder_input) == 0:
            new_molecfit_input_folder = ut.check_path(
                current_molecfit_input_folder, exists=True)
        else:
            new_molecfit_input_folder = ut.check_path(new_input_folder_input,
                                                      exists=True)
        print('')
        ut.tprint(Q2)
        ut.tprint(f'Currently: {current_molecfit_prog_folder}')
        new_prog_folder_input = str(input())
        if len(new_prog_folder_input) == 0:
            new_molecfit_prog_folder = ut.check_path(
                current_molecfit_prog_folder, exists=True)
        else:
            new_molecfit_prog_folder = ut.check_path(new_prog_folder_input,
                                                     exists=True)
        print('')
        ut.tprint(Q3)
        ut.tprint(f'Currently: {current_python_alias}')
        new_python_alias_input = str(input())
        if len(new_python_alias_input) == 0:
            new_python_alias = current_python_alias
        else:
            new_python_alias = new_python_alias_input
    else:  #This is actually the default mode of using this, because this function is generally
        #only called when tel.molecfit() is run for the first time and the config file doesn't exist yet.
        ut.tprint(Q1)
        new_molecfit_input_folder = ut.check_path(str(input()), exists=True)
        print('')
        ut.tprint(Q2)
        new_molecfit_prog_folder = ut.check_path(str(input()), exists=True)
        print('')
        ut.tprint(Q3)
        new_python_alias = str(input())

    with open(configpath, "w") as f:
        f.write(f'molecfit_input_folder   {str(new_molecfit_input_folder)}\n')
        f.write(f'molecfit_prog_folder   {str(new_molecfit_prog_folder)}\n')
        f.write(f'python_alias   {str(new_python_alias)}\n')

    ut.tprint(
        f'New molecfit configation file successfully written to {configpath}')
Exemplo n.º 7
0
def plot_ccf(rv,
             ccf,
             dp,
             xrange=[-200, 200],
             yrange=[0, 0],
             Nxticks=10.0,
             Nyticks=10.0,
             title='',
             doppler_model=[],
             i_legend=True,
             show=True):
    """
    This is a routine that does all the plotting of the cleaned 2D CCF. It overplots
    the expected planet velocity, modified by the systemic velocity given in the config file.
    Optionally, a trace of the doppler model is added as a list (of y-values) in the
    doppler_model parameter. Set i_legend to False if you wish to remove the interactive legend.
    """

    import numpy as np
    import matplotlib.pyplot as plt
    import pdb
    import tayph.drag_colour as dcb
    import tayph.functions as fun
    import pylab as pl
    import tayph.system_parameters as sp

    #Load necessary physics for overplotting planet velocity.
    vsys = sp.paramget('vsys', dp)
    RVp = sp.RV(dp) + vsys
    nexp = np.shape(ccf)[0]

    #Default to the entire y-axis if yrange = [0,0]
    if all(v == 0 for v in yrange):
        yrange = [0, nexp - 1]

    x2, y2, z, rv_sel, y_sel, xticks, yticks, vmin, vmax = plotting_scales_2D(
        rv,
        fun.findgen(nexp),
        ccf,
        xrange,
        yrange,
        Nxticks=Nxticks,
        Nyticks=Nyticks,
        nsigma=3.0)
    #The plotting
    fig, ax = plt.subplots(figsize=(12, 6))
    img = ax.pcolormesh(x2, y2, z, vmin=vmin, vmax=vmax, cmap='hot')
    ax.axis([x2.min(), x2.max(), y2.min(), y2.max()])
    line1, = ax.plot(RVp,
                     fun.findgen(nexp),
                     '--',
                     color='black',
                     label='Planet rest-frame')
    if len(doppler_model) > 0:
        line2, = ax.plot(doppler_model + vsys,
                         fun.findgen(nexp),
                         '--',
                         color='black',
                         label='Doppler shadow')
    ax.set_xticks(xticks)
    ax.set_yticks(yticks)
    ax.set_title(title)
    ax.set_xlabel('Radial velocity (km/s)')
    ax.set_ylabel('Exposure')
    #The colourbar
    cbar = plt.colorbar(img, format='%05.4f', aspect=15)
    # cbar.set_norm(dcb.MyNormalize(vmin=vmin,vmax=vmax,stretch='linear'))
    cbar = dcb.DraggableColorbar_fits(cbar, img, 'hot')
    cbar.connect()

    #The clickable legend.
    if len(doppler_model) > 0:
        lines = [line1, line2]
    else:
        lines = [line1]

    if i_legend == True:
        interactive_legend(fig, ax, lines)
    if show == True:
        plt.show()
    return (fig, ax, cbar)
Exemplo n.º 8
0
def construct_doppler_model(rv,
                            ccf,
                            dp,
                            shadowname,
                            xrange=[-200, 200],
                            Nxticks=20.0,
                            Nyticks=10.0):
    """This is the the main function to construct a doppler model. The above are mostly dependencies."""
    import numpy as np
    import matplotlib.pyplot as plt
    import tayph.drag_colour as dcb
    import tayph.functions as fun
    import tayph.system_parameters as sp
    import tayph.plotting as fancyplots
    import sys
    from matplotlib.widgets import Slider, Button, RadioButtons, CheckButtons

    #This is for setting plot axes in the call to plotting_scales_2D below.
    nexp = np.shape(ccf)[0]
    yrange = [0, nexp - 1]
    y_axis = np.arange(nexp, dtype=float)  #fun.findgen(nexp)
    #And for adding the planet line:
    vsys = sp.paramget('vsys', dp)
    vsini = sp.paramget('vsini', dp)
    RVp = sp.RV(dp) + vsys
    transit = sp.transit(dp)
    sel_transit = y_axis[(transit < 1.0)]
    transit_start = min(sel_transit)
    transit_end = max(sel_transit)
    fig, ax, cbar = fancyplots.plot_ccf(
        rv,
        ccf,
        dp,
        xrange=xrange,
        Nxticks=Nxticks,
        Nyticks=Nyticks,
        i_legend=False,
        show=False)  #Initiates the plot for the primer.
    primer = prime_doppler_model(fig, ax, cbar)  #We use the active Figure
    #to let the user indicate where the doppler shadow is located (the primer). This calls the plt.show()
    #which was not called by plot_ccf. To proceed, we would like to model the shadow using the
    #primer, along with fancy visualization of the ccf.
    #We first re-instate a plot. This plot is a bit more complex than the primer so we don't
    #use plot_ccf anymore; though the philosophy is similar.

    x2, y2, z, rv_sel, y_sel, xticks, yticks, vmin, vmax = fancyplots.plotting_scales_2D(
        rv,
        y_axis,
        ccf,
        xrange,
        yrange,
        Nxticks=Nxticks,
        Nyticks=Nyticks,
        nsigma=3.0)

    #Create empty plots.
    fig, ax = plt.subplots(3, 1, sharex=True, figsize=(13, 6))
    plt.subplots_adjust(right=0.75)

    #Here we initiate the model instance that does the fitting and handles the GUI.
    #This does an initial fit based on the primer.
    model_callback = fit_doppler_model(fig, ax, rv_sel, z, primer, dp,
                                       shadowname)
    s_init = model_callback.maskHW  #Masking half-width around the planet RV to start the slider with.

    #We continue by overplotting the velocty traces and transit markers onto the 3 subplots, saving the references to the lines.
    #These lines can be set to be visible/invisible using the checkbuttons below, and are set to be invisible
    #when the plot opens.
    l_planet = []
    t_start = []
    t_end = []
    l_primer = []
    l_vfit = []
    for sub_ax in ax:
        sub_ax.axis([x2.min(), x2.max(), y2.min(), y2.max()])
        sub_ax.set_xticks(xticks)
        sub_ax.set_yticks(yticks)
        sub_ax.set_ylabel('Exposure')
        l1 = sub_ax.plot(RVp,
                         y_axis,
                         '--',
                         color='black',
                         label='Planet rest-frame',
                         visible=False)[0]
        l2 = sub_ax.plot(rv,
                         rv * 0.0 + transit_start,
                         '--',
                         color='white',
                         label='Transit start',
                         visible=False)[0]
        l3 = sub_ax.plot(rv,
                         rv * 0.0 + transit_end,
                         '--',
                         color='white',
                         label='Transit end',
                         visible=False)[0]
        l4 = sub_ax.plot(model_callback.v_star_primer,
                         y_axis,
                         '--',
                         color='black',
                         label='Local velocity (primer)',
                         visible=False)[0]
        l5 = sub_ax.plot(model_callback.v_star_fit,
                         y_axis,
                         '--',
                         color='black',
                         label='Local velocity (fit)',
                         visible=False)[0]
        l_planet.append(l1)
        t_start.append(l2)
        t_end.append(l3)
        l_primer.append(l4)
        l_vfit.append(l5)

    ax[0].set_title('Data')
    ax[1].set_title('Model shadow')
    ax[2].set_title('Residual')
    ax[2].set_xlabel('Radial velocity (km/s)')

    #Here we actually plot the initial fit, which will be modified each time the parameters are changed
    #using the GUI buttons/sliders.
    img1 = ax[0].pcolormesh(x2, y2, z, vmin=vmin, vmax=vmax, cmap='hot')
    img2 = ax[1].pcolormesh(x2,
                            y2,
                            model_callback.model,
                            vmin=vmin,
                            vmax=vmax,
                            cmap='hot')
    img3 = ax[2].pcolormesh(x2,
                            y2,
                            z - model_callback.model,
                            vmax=vmax,
                            cmap='hot')
    #This trick to associate a single CB to multiple axes comes from
    #https://stackoverflow.com/questions/13784201/matplotlib-2-subplots-1-colorbar
    cbar = fig.colorbar(img1,
                        ax=ax.ravel().tolist(),
                        format='%05.4f',
                        aspect=15)
    cbar = dcb.DraggableColorbar_fits(cbar, [img1, img2, img3], 'hot')
    cbar.connect()

    #We define the interface and the bahaviour of button/slider callbacks.
    #First the check buttons for showing the lines defined above.
    rax_top = plt.axes([0.8, 0.65, 0.15, 0.25])
    rax_top.set_title('Plot:')
    labels = [
        'Planet velocity', 'Transit start/end', 'Shadow v$_c$ primer',
        'Shadow $v_c$ fit', 'Masked area'
    ]
    start = [False, False, False, False, False,
             False]  #Start with none of these actually visible.
    check = CheckButtons(rax_top, labels, start)

    def func(label):
        index = labels.index(label)
        lines = [l_planet, np.append(t_end, t_start), l_primer, l_vfit]
        if index < len(lines):
            for l in lines[index]:
                l.set_visible(not l.get_visible())
        if index == len(
                lines
        ):  #I.e. if we are on the last element, which is not a line an option for SHOWING the masked area:
            status = check.get_status()[-1]
            if status == True:  #If true, mask the image.
                data = z * model_callback.ccf_mask
                data[np.isnan(
                    data
                )] = np.inf  #The colobar doesn't eat NaNs, set them to inf instead for the plot. Makes them white, too.
                img1.set_array((data).ravel())
                img3.set_array((data - model_callback.model).ravel())
            if status == False:  #If false (unclicked), then just the data w/o mask.
                img1.set_array(z.ravel())
                img3.set_array((z - model_callback.model).ravel())
        plt.draw()

    check.on_clicked(func)

    #Then the choice for 1 or 2 Gaussian fitting components:
    rax_middle = plt.axes([0.8, 0.53, 0.15, 0.10])
    clabels = ['1 component', '2 components']
    radio = RadioButtons(rax_middle, clabels)

    def cfunc(label):
        index = clabels.index(label)
        model_callback.n_components = index + 1
        model_callback.fit_model()  #Each time we change the choice, refit.
        status = check.get_status()[-1]
        if status == True:  #If true, mask the image.
            data = z * model_callback.ccf_mask
            data[np.isnan(
                data
            )] = np.inf  #The colobar doesn't eat NaNs, set them to inf instead for the plot.
            img2.set_array(model_callback.model.ravel())
            img3.set_array((data - model_callback.model).ravel())
        if status == False:  #If false (unclicked), then just the data w/o mask.
            img2.set_array(model_callback.model.ravel())
            img3.set_array((z - model_callback.model).ravel())
        plt.draw()

    radio.on_clicked(cfunc)

    rax_deg = plt.axes([0.8, 0.35, 0.07, 0.13])
    plt.title('Polynomial degree', fontsize=8)
    rax_poly = plt.axes([0.88, 0.35, 0.07, 0.13])
    dlabels = ['0 (off)', '2', '4', '6']
    plabels = ['Single', 'Even', 'Full']
    dradio = RadioButtons(rax_deg, dlabels)
    pradio = RadioButtons(rax_poly, plabels)

    def update_degree(label):
        model_callback.D = dlabels.index(label) * 2
        model_callback.fit_model()  #Each time we change the choice, refit.
        status = check.get_status()[-1]
        if status == True:  #If true, mask the image.
            data = z * model_callback.ccf_mask
            data[np.isnan(
                data
            )] = np.inf  #The colobar doesn't eat NaNs, set them to inf instead for the plot.
            img2.set_array(model_callback.model.ravel())
            img3.set_array((data - model_callback.model).ravel())
        if status == False:  #If false (unclicked), then just the data w/o mask.
            img2.set_array(model_callback.model.ravel())
            img3.set_array((z - model_callback.model).ravel())
        plt.draw()

    def update_poly(label):
        model_callback.S = plabels.index(label)
        model_callback.fit_model()  #Each time we change the choice, refit.
        status = check.get_status()[-1]
        if status == True:  #If true, mask the image.
            data = z * model_callback.ccf_mask
            data[np.isnan(
                data
            )] = np.inf  #The colobar doesn't eat NaNs, set them to inf instead for the plot.
            img2.set_array(model_callback.model.ravel())
            img3.set_array((data - model_callback.model).ravel())
        if status == False:  #If false (unclicked), then just the data w/o mask.
            img2.set_array(model_callback.model.ravel())
            img3.set_array((z - model_callback.model).ravel())
        plt.draw()

    dradio.on_clicked(update_degree)
    pradio.on_clicked(update_poly)

    #Then the offset slider:
    rax_slider2 = plt.axes([0.8, 0.25, 0.15, 0.02])
    rax_slider2.set_title('Offset 2nd component')
    offset_slider = Slider(rax_slider2,
                           '',
                           -1.0 * np.ceil(vsini),
                           np.ceil(vsini),
                           valinit=0.0,
                           valstep=1.0)

    def update_offset(val):
        model_callback.offset = offset_slider.val
        status = radio.value_selected
        if status == clabels[
                1]:  #Only update the fit if we are actually asked to do 2 components.
            model_callback.fit_model()
            # data = z*model_callback.ccf_mask
            # data[np.isnan(data)] = np.inf#The colobar doesn't eat NaNs...
            # img1.set_array((data).ravel())
            img2.set_array(model_callback.model.ravel())
            img3.set_array((z - model_callback.model).ravel())
        # if status == False:#If false (unclicked), then just the data w/o mask.
        #     img1.set_array(z.ravel())
        #     img2.set_array(model_callback.model.ravel())
        #     img3.set_array((z-model_callback.model).ravel())
        plt.draw()

    offset_slider.on_changed(update_offset)

    #Then the slider:
    rax_slider = plt.axes([0.8, 0.18, 0.15, 0.02])
    rax_slider.set_title('Mask width (km/s)')
    mask_slider = Slider(rax_slider,
                         '',
                         0.0,
                         30.0,
                         valinit=s_init,
                         valstep=1.0)

    def update(val):
        model_callback.maskHW = mask_slider.val
        model_callback.mask_ccf()
        model_callback.fit_model()

        status = check.get_status()[-1]
        if status == True:  #If true, mask the image.
            data = z * model_callback.ccf_mask
            data[np.isnan(data)] = np.inf  #The colobar doesn't eat NaNs...
            img1.set_array((data).ravel())
            img2.set_array(model_callback.model.ravel())
            img3.set_array((data - model_callback.model).ravel())
        if status == False:  #If false (unclicked), then just the data w/o mask.
            img1.set_array(z.ravel())
            img2.set_array(model_callback.model.ravel())
            img3.set_array((z - model_callback.model).ravel())
        plt.draw()

    mask_slider.on_changed(update)

    #And finally the save button.
    rax_save = plt.axes([0.875, 0.1, 0.06, 0.05])
    bsave = Button(rax_save, 'Save')
    bsave.on_clicked(model_callback.save)

    rax_cancel = plt.axes([0.8, 0.1, 0.06, 0.05])
    bcancel = Button(rax_cancel, 'Cancel')
    bcancel.on_clicked(model_callback.cancel)
    plt.show()  #All fitting is done before this line through event handling.
Exemplo n.º 9
0
def do_molecfit(headers,
                spectra,
                configfile,
                wave=[],
                mode='HARPS',
                load_previous=False):
    """This is the main wrapper for molecfit that pipes a list of s1d spectra and
    executes it. It first launces the molecfit gui on the middle spectrum of the
    sequence, and then loops through the entire list, returning the transmission
    spectra of the Earths atmosphere in the same order as the list provided.
    These can then be used to correct the s1d spectra or the e2ds spectra.
    Note that the s1d spectra are assumed to be in the barycentric frame in vaccuum,
    but that the output transmission spectrum is in the observers frame, and e2ds files
    are in air wavelengths by default.

    If you have run do_molecfit before, and want to reuse the output of the previous run
    for whatever reason, set the load_previous keyword to True. This will reload the
    list of transmission spectra created last time, if available.
    """

    import pdb
    import numpy as np
    import matplotlib.pyplot as plt
    import sys
    import os.path
    import pickle
    import copy
    from pathlib import Path
    import tayph.util as ut
    import tayph.system_parameters as sp
    configfile = ut.check_path(configfile, exists=True)

    molecfit_input_folder = sp.paramget('molecfit_input_folder',
                                        configfile,
                                        full_path=True)
    molecfit_prog_folder = sp.paramget('molecfit_prog_folder',
                                       configfile,
                                       full_path=True)
    temp_specname = copy.deepcopy(
        mode)  #The name of the temporary file used (without extension).
    #The spectrum will be named like this.fits There should be a this.par file as well,
    #that contains a line pointing molecfit to this.fits:
    parname = temp_specname + '.par'

    #====== ||  START OF PROGRAM   ||======#
    N = len(headers)
    if N != len(spectra):
        raise RuntimeError(
            f' in prep_for_molecfit: Length of list of headers is not equal to length of list of spectra ({N},{len(spectra)})'
        )

    #Test that the input root and molecfit roots exist; that the molecfit root contains the molecfit executables.
    #that the input root contains the desired parfile and later fitsfile.
    molecfit_input_root = Path(molecfit_input_folder)
    molecfit_prog_root = Path(molecfit_prog_folder)

    ut.check_path(molecfit_input_root, exists=True)
    ut.check_path(molecfit_prog_root, exists=True)
    ut.check_path(
        molecfit_input_root / parname, exists=True
    )  #Test that the molecfit parameter file for this mode exists.
    ut.check_path(molecfit_prog_root / 'molecfit', exists=True)
    ut.check_path(molecfit_prog_root / 'molecfit_gui', exists=True)

    pickle_outpath = molecfit_input_root / 'previous_run_of_do_molecfit.pkl'

    if load_previous == True:
        if os.path.isfile(pickle_outpath) == False:
            print(
                'WARNING in do_molecfit(): Previously saved run was asked for but is not available.'
            )
            print(
                'The program will proceed to re-fit. That run will then be saved.'
            )
            load_previous = False
        else:
            pickle_in = open(pickle_outpath, "rb")
            list_of_wls, list_of_fxc, list_of_trans = pickle.load(pickle_in)

    if load_previous == False:
        list_of_wls = []
        list_of_fxc = []
        list_of_trans = []

        middle_i = int(
            round(0.5 * N)
        )  #We initialize molecfit on the middle spectrum of the time series.
        write_file_to_molecfit(molecfit_input_root,
                               temp_specname + '.fits',
                               headers,
                               spectra,
                               middle_i,
                               mode=mode,
                               wave=wave)
        print(molecfit_input_root)
        print(temp_specname + '.fits')
        execute_molecfit(molecfit_prog_root,
                         molecfit_input_root / parname,
                         gui=True)
        wl, fx, trans = retrieve_output_molecfit(molecfit_input_root /
                                                 temp_specname)
        remove_output_molecfit(molecfit_input_root, temp_specname)
        for i in range(N):  #range(len(spectra)):
            print('Fitting spectrum %s from %s' % (i + 1, len(spectra)))
            t1 = ut.start()
            write_file_to_molecfit(molecfit_input_root,
                                   temp_specname + '.fits',
                                   headers,
                                   spectra,
                                   i,
                                   mode=mode,
                                   wave=wave)
            execute_molecfit(molecfit_prog_root,
                             molecfit_input_root / parname,
                             gui=False)
            wl, fx, trans = retrieve_output_molecfit(molecfit_input_root /
                                                     temp_specname)
            remove_output_molecfit(molecfit_input_root, temp_specname)
            list_of_wls.append(wl * 1000.0)  #Convert to nm.
            list_of_fxc.append(fx / trans)
            list_of_trans.append(trans)
            ut.end(t1)

        pickle_outpath = molecfit_input_root / 'previous_run_of_do_molecfit.pkl'
        with open(pickle_outpath, 'wb') as f:
            pickle.dump((list_of_wls, list_of_fxc, list_of_trans), f)

    to_do_manually = check_fit_gui(list_of_wls, list_of_fxc, list_of_trans)
    if len(to_do_manually) > 0:
        print('The following spectra were selected to redo manually:')
        print(to_do_manually)
        #CHECK THAT THIS FUNCIONALITY WORKS:
        for i in to_do_manually:
            write_file_to_molecfit(molecfit_input_root,
                                   temp_specname + '.fits',
                                   headers,
                                   spectra,
                                   int(i),
                                   mode=mode,
                                   wave=wave)
            execute_molecfit(molecfit_prog_root,
                             molecfit_input_root / parname,
                             gui=True)
            wl, fx, trans = retrieve_output_molecfit(molecfit_input_root /
                                                     temp_specname)
            list_of_wls[int(i)] = wl * 1000.0  #Convert to nm.
            list_of_fxc[int(i)] = fxc
            list_of_trans[int(i)] = trans
    return (list_of_wls, list_of_trans)