Esempio n. 1
0
    def telluric_correction_order(i):
        order = list_of_orders[i]
        order_cor = order * 0.0
        error = list_of_sigmas[i]
        error_cor = error * 0.0
        wl = copy.deepcopy(list_of_wls[i])  #input wl axis, either 1D or 2D.
        #If it is 1D, we make it 2D by tiling it vertically:
        if wl.ndim == 1: wl = np.tile(wl, (Nexp, 1))  #Tile it into a 2D thing.
        #If read (2D) or tiled (1D) correctly, wl and order should have the same shape:
        dimtest(wl, np.shape(order),
                f'Wl axis of order {i}/{No} in apply_telluric_correction()')
        dimtest(error, np.shape(order),
                f'errors {i}/{No} in apply_telluric_correction()')
        for j in range(Nexp):
            T_i = interp.interp1d(wlT[j], fxT[j],
                                  fill_value="extrapolate")(wl[j])
            postest(T_i,
                    f'T-spec of exposure {j} in apply_telluric_correction()')
            nantest(T_i,
                    f'T-spec of exposure {j} in apply_telluric_correction()')
            order_cor[j] = order[j] / T_i
            error_cor[j] = error[
                j] / T_i  #I checked that this works because the SNR before and after
            #telluric correction is identical.

        return (order_cor, error_cor)
Esempio n. 2
0
def envelope(wlm, fxm, binsize, selfrac=0.05, mode='top', threshold=''):
    """
    This program measures the top or bottom envelope of a spectrum (wl,fx), by
    chopping it up into bins of size binsze (unit of wl), and measuring the mean
    of the top n % of values in that bin. Setting the mode to 'bottom' will do the
    oppiste: The mean of the bottom n% of values. The output is the resulting wl
    and flux points of these bins.

    Example: wle,fxe = envelope(wl,fx,1.0,selfrac=3.0,mode='top')
    """
    import numpy as np
    import tayph.util as ut
    import tayph.functions as fun
    from tayph.vartests import typetest, dimtest, nantest, postest
    from matplotlib import pyplot as plt
    typetest(wlm, np.ndarray, 'wlm in ops.envelope()')
    typetest(fxm, np.ndarray, 'fxm in ops.envelope()')
    dimtest(wlm, [len(fxm)])
    typetest(binsize, [int, float], 'binsize in ops.envelope()')
    typetest(selfrac, float, 'percentage in ops.envelope()')
    typetest(mode, str, 'mode in envelope')
    nantest(fxm, 'fxm in ops.envelope()')
    nantest(wlm, 'wlm in ops.envelope()')
    postest(wlm, 'wlm in ops.envelope()')
    postest(binsize, 'binsize in ops.envelope()')

    if mode not in ['top', 'bottom']:
        raise Exception(
            f'RuntimeError in ops.envelope(): Mode should be set to "top" or "bottom" ({mode}).'
        )

    binsize = float(binsize)
    if mode == 'bottom':
        fxm *= -1.0

    wlcs = np.array([])  #Center wavelengths
    fxcs = np.array([])  #Flux at center wavelengths
    i_start = 0
    wlm_start = wlm[i_start]
    for i in range(0, len(wlm) - 1):
        if wlm[i] - wlm_start >= binsize:
            wlsel = wlm[i_start:i]
            fxsel = fxm[i_start:i]
            maxsel = fun.selmax(fxsel, selfrac)
            wlcs = np.append(wlcs, np.mean(wlsel[maxsel]))
            fxcs = np.append(fxcs, np.mean(fxsel[maxsel]))
            i_start = i + 1
            wlm_start = wlm[i + 1]

    if isinstance(threshold, float) == True:
        #This means that the threshold value is set, and we set all bins less than
        #that threshold to the threshold value:
        if mode == 'bottom':
            threshold *= -1.0
        fxcs[(fxcs < threshold)] = threshold

    if mode == 'bottom':
        fxcs *= -1.0
        fxm *= -1.0
    return wlcs, fxcs
Esempio 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)
Esempio n. 4
0
def apply_telluric_correction(inpath, list_of_wls, list_of_orders,
                              list_of_sigmas):
    """
    This applies a set of telluric spectra (computed by molecfit) for each exposure
    in our time series that were written to a pickle file by write_telluric_transmission_to_file.

    List of errors are provided to propagate telluric correction into the error array as well.

    Parameters
    ----------
    inpath : str, path like
        The path to the pickled transmission spectra.

    list_of_wls : list
        List of wavelength axes.

    list_of_orders :
        List of 2D spectral orders, matching to the wavelength axes in dimensions and in number.

    list_of_isgmas :
        List of 2D error matrices, matching dimensions and number of list_of_orders.

    Returns
    -------
    list_of_orders_corrected : list
        List of 2D spectral orders, telluric corrected.

    list_of_sigmas_corrected : list
        List of 2D error matrices, telluric corrected.

    """
    import scipy.interpolate as interp
    import numpy as np
    import tayph.util as ut
    import tayph.functions as fun
    from tayph.vartests import dimtest, postest, typetest, nantest
    import copy

    T = read_telluric_transmission_from_file(inpath)
    wlT = T[0]
    fxT = T[1]

    typetest(list_of_wls, list, 'list_of_wls in apply_telluric_correction()')
    typetest(list_of_orders, list,
             'list_of_orders in apply_telluric_correction()')
    typetest(list_of_sigmas, list,
             'list_of_sigmas in apply_telluric_correction()')
    typetest(wlT, list,
             'list of telluric wave-axes in apply_telluric_correction()')
    typetest(
        fxT, list,
        'list of telluric transmission spectra in apply_telluric_correction()')

    No = len(list_of_wls)  #Number of orders.
    x = fun.findgen(No)
    Nexp = len(wlT)

    #Test dimensions
    if No != len(list_of_orders):
        raise Exception(
            'Runtime error in telluric correction: List of wavelength axes and List '
            'of orders do not have the same length.')
    if Nexp != len(fxT):
        raise Exception(
            'Runtime error in telluric correction: List of telluric wls and telluric '
            'spectra read from file do not have the same length.')
    if Nexp != len(list_of_orders[0]):
        raise Exception(
            f'Runtime error in telluric correction: List of telluric spectra and data'
            f'spectra read from file do not have the same length ({Nexp} vs {len(list_of_orders[0])}).'
        )

    #Corrected orders will be stored here.
    list_of_orders_cor = []
    list_of_sigmas_cor = []

    #Do the correction order by order:
    for i in range(No):
        order = list_of_orders[i]
        order_cor = order * 0.0
        error = list_of_sigmas[i]
        error_cor = error * 0.0
        wl = copy.deepcopy(list_of_wls[i])  #input wl axis, either 1D or 2D.
        #If it is 1D, we make it 2D by tiling it vertically:
        if wl.ndim == 1: wl = np.tile(wl, (Nexp, 1))  #Tile it into a 2D thing.

        #If read (2D) or tiled (1D) correctly, wl and order should have the same shape:
        dimtest(wl, np.shape(order),
                f'Wl axis of order {i}/{No} in apply_telluric_correction()')
        dimtest(error, np.shape(order),
                f'errors {i}/{No} in apply_telluric_correction()')
        for j in range(Nexp):
            T_i = interp.interp1d(wlT[j], fxT[j],
                                  fill_value="extrapolate")(wl[j])
            postest(T_i,
                    f'T-spec of exposure {j} in apply_telluric_correction()')
            nantest(T_i,
                    f'T-spec of exposure {j} in apply_telluric_correction()')
            order_cor[j] = order[j] / T_i
            error_cor[j] = error[
                j] / T_i  #I checked that this works because the SNR before and after
            #telluric correction is identical.
        list_of_orders_cor.append(order_cor)
        list_of_sigmas_cor.append(error_cor)
        ut.statusbar(i, x)
    return (list_of_orders_cor, list_of_sigmas_cor)
Esempio n. 5
0
def blur_rotate(wl, order, dv, Rp, P, inclination, status=False, fast=False):
    """This function takes a spectrum and blurs it using a rotation x Gaussian
    kernel which has a FWHM width of dv km/s everywhere. Meaning that its width changes
    dynamically.
    Because the kernel needs to be recomputed on each element of the wavelength axis
    individually, this operation is much slower than convolution with a constant kernel,
    in which a simple shifting of the array, rather than a recomputation of the rotation
    profile is sufficient. By setting the fast keyword, the input array will first
    be oversampled onto a constant-velocity grid to enable the usage of a constant kernel,
    after which the result is interpolated back to the original grid.

    Input:
    The wavelength axis wl.
    The spectral axis order.
    The FHWM width of the resolution element in km/s.
    The Radius of the rigid body in Rj.
    The periodicity of the rigid body rotation in days.
    The inclination of the spin axis in degrees.

    Wavelength and order need to be numpy arrays and have the same number of elements.
    Rp, P and i need to be scalar floats.

    Output:
    The blurred spectral axis, with the same dimensions as wl and order.


    WARNING: THIS FUNCTION HANDLES NANS POORLY. I HAVE THEREFORE DECIDED CURRENTLY
    TO REQUIRE NON-NAN INPUT.




    This computes the simple numerical derivative of x by convolving with kernel [-1,0,1].

    Parameters
    ----------
    wl : list, np.ndarray
        The wavelength array.

    order : list, np.ndarray.
        The spectral axis.

    dv: float
        The FWHM of a resolution element in km/s.

    Rp: float
        The radius of the planet in jupiter radii.

    P: float
        The rotation period of the planet. For tidally locked planets, this is equal
        to the orbital period.

    inclination:
        The inclination of the spin axis in degrees. Presumed to be close to 90 degrees
        for transiting planets

    status: bool
        Output a statusbar, but only if fast == False.

    fast: bool
        Re-interpolate the input on a constant-v grid in order to speed up the computation
        of the convolution by eliminating the need to re-interpolate the kernel every step.



    Returns
    -------
    order_blurred : np.array
        The rotation-broadened spectrum on the same wavelength grid as the input.

    Example
    -------
    >>> import tayph.functions as fun
    >>> wl = fun.findgen(4000)*0.001+500.0
    >>> fx = wl*0.0
    >>> fx[2000] = 1.0
    >>> fx_blurred1 = blur_rotate(wl,fx,3.0,1.5,0.8,90.0,status=False,fast=False)
    >>> fx_blurred2 = blur_rotate(wl,fx,3.0,1.5,0.8,90.0,status=False,fast=True)
    """

    import numpy as np
    import tayph.util as ut
    import tayph.functions as fun
    from tayph.vartests import typetest, nantest, dimtest
    from matplotlib import pyplot as plt
    import astropy.constants as const
    import astropy.units as u
    import time
    import sys
    import pdb
    from scipy import interpolate
    typetest(dv, float, 'dv in blur_rotate()')
    typetest(wl, [list, np.ndarray], 'wl in blur_rotate()')
    typetest(order, [list, np.ndarray], 'order in blur_rotate()')
    typetest(P, float, 'P in blur_rotate()')
    typetest(Rp, float, 'Rp in blur_rotate()')
    typetest(inclination, float, 'inclination in blur_rotate()')
    typetest(status, bool, 'status in blur_rotate()')
    typetest(fast, bool, 'fast in blur_rotate()')
    nantest(wl, 'dv in blur_rotate()')
    nantest(order, 'order in blur_rotate()')
    dimtest(wl, [0], 'wl in blur_rotate()')
    dimtest(order, [len(wl)],
            'order in blur_rotate()')  #Test that wl and order are 1D, and that
    #they have the same length.

    if np.min(np.array([dv, P, Rp])) <= 0.0:
        raise Exception(
            "ERROR in blur_rotate: dv, P and Rp should be strictly positive.")

    #ut.typetest_array('wl',wl,np.float64)
    #ut.typetest_array('order',order,np.float64)
    #This is not possible because order may be 2D...
    #And besides, you can have floats, np.float32 and np.float64... All of these would
    #need to pass. Need to fix typetest_array some day.

    order_blurred = order * 0.0  #init the output.
    truncsize = 5.0  #The gaussian is truncated at 5 sigma from the extremest points of the RV amplitude.
    sig_dv = dv / (2 * np.sqrt(2.0 * np.log(2))
                   )  #Transform FWHM to Gaussian sigma. In km/s.
    deriv = derivative(wl)
    if max(deriv) < 0:
        raise Exception(
            "ERROR in ops.blur_rotate: WL derivative is smaller than 1.0. Sort wl in ascending order."
        )
    sig_wl = wl * sig_dv / (const.c.to('km/s').value)  #in nm
    sig_px = sig_wl / deriv

    n = 1000.0
    a = fun.findgen(n) / (n - 1) * np.pi
    rv = np.cos(a) * np.sin(
        np.radians(inclination)) * (2.0 * np.pi * Rp * const.R_jup /
                                    (P * u.day)).to('km/s').value  #in km/s
    trunc_dist = np.round(sig_px * truncsize + np.max(rv) * wl /
                          (const.c.to('km/s').value) / deriv).astype(int)
    # print('Maximum rotational rv: %s' % max(rv))
    # print('Sigma_px: %s' % np.nanmean(np.array(sig_px)))

    rvgrid_max = (np.max(trunc_dist) + 1.0) * sig_dv + np.max(rv)
    rvgrid_n = rvgrid_max / dv * 100.0  #100 samples per lsf fwhm.
    rvgrid = (
        fun.findgen(2 * rvgrid_n + 1) - rvgrid_n
    ) / rvgrid_n * rvgrid_max  #Need to make sure that this is wider than the truncation bin and more finely sampled than wl - everywhere.

    lsf = rvgrid * 0.0
    #We loop through velocities in the velocity grid to build up the sum of Gaussians
    #that is the LSF.
    for v in rv:
        lsf += fun.gaussian(
            rvgrid, 1.0, v, sig_dv
        )  #This defines the LSF on a velocity grid wih high fidelity.
    if fast:
        wlt, fxt, dv = constant_velocity_wl_grid(wl, order, 4)
        dv_grid = rvgrid[1] - rvgrid[0]

        len_rv_grid_low = int(max(rvgrid) / dv * 2 - 2)
        # print(len_rv_grid_low)
        # print(len(fun.findgen(len_rv_grid_low)))
        # print(len_rv_grid_low%2)
        if len_rv_grid_low % 2 == 0:
            len_rv_grid_low -= 1
        rvgrid_low = fun.findgen(
            len_rv_grid_low) * dv  #Slightly smaller than the original grid.
        rvgrid_low -= 0.5 * np.max(rvgrid_low)
        lsf_low = interpolate.interp1d(rvgrid, lsf)(rvgrid_low)
        lsf_low /= np.sum(
            lsf_low
        )  #This is now an LSF on a grid with the same spacing as the data has.
        #This means I can use it directly as a convolution kernel:
        fxt_blurred = convolve(fxt, lsf_low, edge_degree=1, fit_width=1)
        #And interpolate back to where it came from:
        order_blurred = interpolate.interp1d(wlt,
                                             fxt_blurred,
                                             bounds_error=False)(wl)
        #I can use interp1d because after blurring, we are now oversampled.
        # order_blurred2 = bin_avg(wlt,fxt_blurred,wl)
        return (order_blurred)

    #Now we loop through the wavelength grid to place this LSF at each wavelength position.
    for i in range(0, len(wl)):
        binstart = max([0, i - trunc_dist[i]])
        binend = i + trunc_dist[i]
        wlbin = wl[binstart:binend]

        wlgrid = wl[i] * rvgrid / (const.c.to('km/s').value) + wl[
            i]  #This converts the velocity grid to a d-wavelength grid centered on wk[i]
        #print([np.min(wlbin),np.min(wlgrid),np.max(wlbin),np.max(wlgrid)])

        i_wl = interpolate.interp1d(
            wlgrid, lsf, bounds_error=False, fill_value='extrapolate'
        )  #Extrapolate should not be necessary but sometimes there is a minute mismatch between the
        #start and end wavelengths of the constructed grid and the bin.
        try:
            lsf_wl = i_wl(wlbin)
        except:
            ut.tprint(
                'Error in interpolating LSF onto wlbin. Pausing to debug.')
            pdb.set_trace()
        k_n = lsf_wl / np.sum(
            lsf_wl
        )  #Normalize at each instance of the interpolation to make sure flux is conserved exactly.
        order_blurred[i] = np.sum(k_n * order[binstart:binend])
        if status == True:
            ut.statusbar(i, len(wl))
    return (order_blurred)
Esempio n. 6
0
def constant_velocity_wl_grid(wl, fx, oversampling=1.0):
    """This function will define a constant-velocity grid that is (optionally)
    sampled a number of times finer than the SMALLEST velocity difference that is
    currently in the grid.

    Example: wl_cv,fx_cv = constant_velocity_wl_grid(wl,fx,oversampling=1.5).

    This function is hardcoded to raise an exception if wl or fx contain NaNs,
    because interp1d does not handle NaNs.


    Parameters
    ----------
    wl : list, np.ndarray
        The wavelength array to be resampled.

    fx : list, np.ndarray
        The flux array to be resampled.

    oversampling : float
        The factor by which the wavelength array is *minimally* oversampled.


    Returns
    -------
    wl : np.array
        The new wavelength grid.

    fx : np.array
        The interpolated flux values.

    a : float
        The velocity step in km/s.


    """
    import astropy.constants as consts
    import numpy as np
    import tayph.functions as fun
    from tayph.vartests import typetest, nantest, dimtest, postest
    from scipy import interpolate
    import pdb
    import matplotlib.pyplot as plt
    typetest(
        oversampling,
        [int, float],
        'oversampling in constant_velocity_wl_grid()',
    )
    typetest(wl, [list, np.ndarray], 'wl in constant_velocity_wl_grid()')
    typetest(fx, [list, np.ndarray], 'fx in constant_velocity_wl_grid()')
    nantest(wl, 'wl in in constant_velocity_wl_grid()')
    nantest(fx, 'fx in constant_velocity_wl_grid()')
    dimtest(wl, [0], 'wl in constant_velocity_wl_grid()')
    dimtest(fx, [len(wl)], 'fx in constant_velocity_wl_grid()')
    postest(oversampling, 'oversampling in constant_velocity_wl_grid()')

    oversampling = float(oversampling)
    wl = np.array(wl)
    fx = np.array(fx)

    c = consts.c.to('km/s').value

    dl = derivative(wl)
    dv = dl / wl * c
    a = np.min(dv) / oversampling

    wl_new = 0.0
    #The following while loop will define the new pixel grid.
    #It starts trying 100,000 points, and if that's not enough to cover the entire
    #range from min(wl) to max(wl), it will add 100,000 more; until it's enough.
    n = len(wl)
    while np.max(wl_new) < np.max(wl):
        x = fun.findgen(n)
        wl_new = np.exp(a / c * x) * np.min(wl)
        n += len(wl)
    wl_new[0] = np.min(
        wl)  #Artificially set to zero to avoid making a small round
    #off error in that exponent.

    #Then at the end we crop the part that goes too far:
    wl_new_cropped = wl_new[(wl_new <= np.max(wl))]
    x_cropped = x[(wl_new <= np.max(wl))]
    i_fx = interpolate.interp1d(wl, fx)
    fx_new_cropped = i_fx(wl_new_cropped)
    return (wl_new_cropped, fx_new_cropped, a)
Esempio n. 7
0
def apply_telluric_correction(inpath, list_of_wls, list_of_orders,
                              list_of_sigmas):
    """
    This applies a set of telluric spectra (computed by molecfit) for each exposure
    in our time series that were written to a pickle file by write_telluric_transmission_to_file.

    List of errors are provided to propagate telluric correction into the error array as well.

    Parameters
    ----------
    inpath : str, path like
        The path to the pickled transmission spectra.

    list_of_wls : list
        List of wavelength axes.

    list_of_orders :
        List of 2D spectral orders, matching to the wavelength axes in dimensions and in number.

    list_of_isgmas :
        List of 2D error matrices, matching dimensions and number of list_of_orders.

    Returns
    -------
    list_of_orders_corrected : list
        List of 2D spectral orders, telluric corrected.

    list_of_sigmas_corrected : list
        List of 2D error matrices, telluric corrected.

    """
    import scipy.interpolate as interp
    import numpy as np
    import tayph.util as ut
    import tayph.functions as fun
    from tayph.vartests import dimtest, postest, typetest, nantest
    wlT, fxT = read_telluric_transmission_from_file(inpath)
    typetest(list_of_wls, list, 'list_of_wls in apply_telluric_correction()')
    typetest(list_of_orders, list,
             'list_of_orders in apply_telluric_correction()')
    typetest(list_of_sigmas, list,
             'list_of_sigmas in apply_telluric_correction()')
    typetest(wlT, list,
             'list of telluric wave-axes in apply_telluric_correction()')
    typetest(
        fxT, list,
        'list of telluric transmission spectra in apply_telluric_correction()')

    No = len(list_of_wls)
    x = fun.findgen(No)

    if No != len(list_of_orders):
        raise Exception(
            'Runtime error in telluric correction: List of data wls and List of orders do not have the same length.'
        )

    Nexp = len(wlT)

    if Nexp != len(fxT):
        raise Exception(
            'Runtime error in telluric correction: List of telluric wls and telluric spectra read from file do not have the same length.'
        )

    if Nexp != len(list_of_orders[0]):
        raise Exception(
            f'Runtime error in telluric correction: List of telluric spectra and data spectra read from file do not have the same length ({Nexp} vs {len(list_of_orders[0])}).'
        )
    list_of_orders_cor = []
    list_of_sigmas_cor = []
    # ut.save_stack('test.fits',list_of_orders)
    # pdb.set_trace()

    for i in range(No):  #Do the correction order by order.
        order = list_of_orders[i]
        order_cor = order * 0.0
        error = list_of_sigmas[i]
        error_cor = error * 0.0
        wl = list_of_wls[i]
        dimtest(order, [0, len(wl)],
                f'order {i}/{No} in apply_telluric_correction()')
        dimtest(error, np.shape(order),
                f'errors {i}/{No} in apply_telluric_correction()')

        for j in range(Nexp):
            T_i = interp.interp1d(wlT[j], fxT[j], fill_value="extrapolate")(wl)
            postest(T_i,
                    f'T-spec of exposure {j} in apply_telluric_correction()')
            nantest(T_i,
                    f'T-spec of exposure {j} in apply_telluric_correction()')
            order_cor[j] = order[j] / T_i
            error_cor[j] = error[
                j] / T_i  #I checked that this works because the SNR before and after telluric correction is identical.
        list_of_orders_cor.append(order_cor)
        list_of_sigmas_cor.append(error_cor)
        ut.statusbar(i, x)
    return (list_of_orders_cor, list_of_sigmas_cor)
Esempio n. 8
0
def xcor(list_of_wls,
         list_of_orders,
         wlm,
         fxm,
         drv,
         RVrange,
         plot=False,
         list_of_errors=None):
    """
    This routine takes a combined dataset (in the form of lists of wl spaces,
    spectral orders and possible a matching list of errors on those spectal orders),
    as well as a template (wlm,fxm) to cross-correlate with, and the cross-correlation
    parameters (drv,RVrange). The code takes on the order of ~10 minutes for an entire
    HARPS dataset, which appears to be superior to my old IDL pipe.

    The CCF used is the Geneva-style weighted average; not the Pearson CCF. Therefore
    it measures true 'average' planet lines, with flux on the y-axis of the CCF.
    The template must therefore be (something close to) a binary mask, with values
    inside spectral lines (the CCF is scale-invariant so their overall scaling
    doesn't matter),

    It returns the RV axis and the resulting CCF in a tuple.

    Thanks to Brett Morris (bmorris3), this code now implements a clever numpy broadcasting trick to
    instantly apply and interpolate the wavelength shifts of the model template onto
    the data grid in 2 dimensions. The matrix multiplication operator (originally
    recommended to me by Matteo Brogi) allowed this 2D template matrix to be multiplied
    with a 2D spectral order. np.hstack() is used to concatenate all orders end to end,
    effectively making a giant single spectral order (with NaNs in between due to masking).

    All these steps have eliminated ALL the forloops from the equation, and effectuated a
    speed gain of a factor between 2,000 and 3,000. The time to do cross correlations is now
    typically measured in 100s of milliseconds rather than minutes.

    This way of calculation does impose some strict rules on NaNs, though. To keep things fast,
    NaNs are now used to set the interpolated template matrix to zero wherever there are NaNs in the data.
    These NaNs are found by looking at the first spectrum in the stack, with the assumption that
    every NaN is in an all-NaN column. In the standard cross-correlation work-flow, isolated
    NaNs are interpolated over (healed), after all.

    The places where there are NaN columns in the data are therefore set to 0 in the template matrix.
    The NaN values themselves are then set to to an arbitrary value, since they will never
    weigh into the average by construction.


    Parameters
    ----------
    list_of_wls : list
        List of wavelength axes of the data.

    list_of_orders : list
        List of corresponding 2D orders.

    list_of_errors : list
        Optional, list of corresponding 2D error matrices.

    wlm : np.ndarray
        Wavelength axis of the template.

    fxm : np.ndarray
        Weight-axis of the template.

    drv : int,float
        The velocity step onto which the CCF is computed. Typically ~1 km/s.

    RVrange : int,float
        The velocity range in the positive and negative direction over which to
        evaluate the CCF. Typically >100 km/s.

    plot : bool
        Set to True for diagnostic plotting.

    Returns
    -------
    RV : np.ndarray
        The radial velocity grid over which the CCF is evaluated.

    CCF : np.ndarray
        The weighted average flux in the spectrum as a function of radial velocity.

    CCF_E : np.ndarray
        Optional. The error on each CCF point propagated from the error on the spectral values.

    Tsums : np.ndarray
        The sum of the template for each velocity step. Used for normalising the CCFs.
    """

    import tayph.functions as fun
    import astropy.constants as const
    import tayph.util as ut
    from tayph.vartests import typetest, dimtest, postest, nantest
    import numpy as np
    import scipy.interpolate
    import astropy.io.fits as fits
    import matplotlib.pyplot as plt
    import sys
    import pdb

    #===FIRST ALL SORTS OF TESTS ON THE INPUT===
    if len(list_of_wls) != len(list_of_orders):
        raise ValueError(
            f'In xcor(): List of wls and list of orders have different length ({len(list_of_wls)} & {len(list_of_orders)}).'
        )

    dimtest(wlm, [len(fxm)], 'wlm in ccf.xcor()')
    typetest(wlm, np.ndarray, 'wlm in ccf.xcor')
    typetest(fxm, np.ndarray, 'fxm in ccf.xcor')
    typetest(drv, [int, float], 'drv in ccf.xcor')
    typetest(
        RVrange,
        float,
        'RVrange in ccf.xcor()',
    )
    postest(RVrange, 'RVrange in ccf.xcor()')
    postest(drv, 'drv in ccf.xcor()')
    nantest(wlm, 'fxm in ccf.xcor()')
    nantest(fxm, 'fxm in ccf.xcor()')
    nantest(drv, 'drv in ccf.xcor()')
    nantest(RVrange, 'RVrange in ccf.xcor()')

    drv = float(drv)
    N = len(list_of_wls)  #Number of orders.

    if np.ndim(list_of_orders[0]) == 1.0:
        n_exp = 1
    else:
        n_exp = len(list_of_orders[0][:, 0])  #Number of exposures.

        #===Then check that all orders indeed have n_exp exposures===
        for i in range(N):
            if len(list_of_orders[i][:, 0]) != n_exp:
                raise ValueError(
                    f'In xcor(): Not all orders have {n_exp} exposures.')

#===END OF TESTS. NOW DEFINE CONSTANTS===
    c = const.c.to('km/s').value  #In km/s
    RV = fun.findgen(
        2.0 * RVrange / drv +
        1) * drv - RVrange  #..... CONTINUE TO DEFINE THE VELOCITY GRID
    beta = 1.0 - RV / c  #The doppler factor with which each wavelength is to be shifted.
    n_rv = len(RV)

    #===STACK THE ORDERS IN MASSIVE CONCATENATIONS===
    stack_of_orders = np.hstack(list_of_orders)
    stack_of_wls = np.concatenate(list_of_wls)
    if list_of_errors is not None:
        stack_of_errors = np.hstack(list_of_errors)  #Stack them horizontally

        #Check that the number of NaNs is the same in the orders as in the errors on the orders;
        #and that they are in the same place; meaning that if I add the errors to the orders, the number of
        #NaNs does not increase (NaN+value=NaN).
        if (np.sum(np.isnan(stack_of_orders)) != np.sum(
                np.isnan(stack_of_errors + stack_of_orders))) and (np.sum(
                    np.isnan(stack_of_orders)) != np.sum(
                        np.isnan(stack_of_errors))):
            raise ValueError(
                f"in CCF: The number of NaNs in list_of_orders and list_of_errors is not equal ({np.sum(np.isnan(list_of_orders))},{np.sum(np.isnan(list_of_errors))})"
            )

#===HERE IS THE JUICY BIT===
    shifted_wavelengths = stack_of_wls * beta[:, np.
                                              newaxis]  #2D broadcast of wl_data, each row shifted by beta[i].
    T = scipy.interpolate.interp1d(wlm, fxm, bounds_error=False, fill_value=0)(
        shifted_wavelengths)  #...making this a 2D thing.
    T[:, np.isnan(
        stack_of_orders[0]
    )] = 0.0  #All NaNs are assumed to be in all-NaN columns. If that is not true, the below nantest will fail.
    T_sums = np.sum(T, axis=1)

    #We check whether there are isolated NaNs:
    n_nans = np.sum(np.isnan(stack_of_orders),
                    axis=0)  #This is the total number of NaNs in each column.
    n_nans[n_nans == len(
        stack_of_orders
    )] = 0  #Whenever the number of NaNs equals the length of a column, set the flag to zero.
    if np.max(
            n_nans
    ) > 0:  #If there are any columns which still have NaNs in them, we need to crash.
        raise ValueError(
            f"in CCF: Not all NaN values are purely in columns. There are still isolated NaNs. Remove those."
        )

    stack_of_orders[np.isnan(
        stack_of_orders)] = 47e20  #Set NaNs to arbitrarily high values.
    CCF = stack_of_orders @ T.T / T_sums  #Here it the entire cross-correlation. Over all orders and velocity steps. No forloops.
    CCF_E = CCF * 0.0

    #If the errors were provided, we do the same to those:
    if list_of_errors is not None:
        stack_of_errors[np.isnan(
            stack_of_errors
        )] = 42e20  #we have already tested that these NaNs are in the same place.
        CCF_E = stack_of_errors**2 @ (
            T.T / T_sums)**2  #This has been mathematically proven.


#===THAT'S ALL. TEST INTEGRITY AND RETURN THE RESULT===
    nantest(
        CCF, 'CCF in ccf.xcor()'
    )  #If anything went wrong with NaNs in the data, these tests will fail because the matrix operation @ is non NaN-friendly.
    nantest(CCF_E, 'CCF_E in ccf.xcor()')

    if list_of_errors != None:
        return (RV, CCF, np.sqrt(CCF_E), T_sums)
    return (RV, CCF, T_sums)
Esempio n. 9
0
def clean_ccf(rv, ccf, ccf_e, dp):
    """
    This routine normalizes the CCF fluxes and subtracts the average out of
    transit CCF, using the transit lightcurve as a mask.


    Parameters
    ----------

    rv : np.ndarray
        The radial velocity axis

    ccf : np.ndarray
        The CCF with second dimension matching the length of rv.

    ccf_e : np.ndarray
        The error on ccf.

    dp : str or path-like
        The datapath of the present dataset, to establish which exposures in ccf
        are in or out of transit.

    Returns
    -------

    ccf_n : np.ndarray
        The transit-lightcurve normalised CCF.

    ccf_ne : np.ndarray
        The error on ccf_n

    ccf_nn : np.ndarray
        The CCF relative to the out-of-transit time-averaged, if sufficient (>25%
        of the time-series) out of transit exposures were available. Otherwise, the
        average over the entire time-series is used.

    ccf_ne : np.array
        The error on ccf_nn.


    """

    import numpy as np
    import tayph.functions as fun
    import tayph.util as ut
    from matplotlib import pyplot as plt
    import pdb
    import math
    import tayph.system_parameters as sp
    import tayph.operations as ops
    import astropy.io.fits as fits
    import sys
    import copy
    from tayph.vartests import typetest, dimtest, nantest

    typetest(rv, np.ndarray, 'rv in clean_ccf()')
    typetest(ccf, np.ndarray, 'ccf in clean_ccf')
    typetest(ccf_e, np.ndarray, 'ccf_e in clean_ccf')
    dp = ut.check_path(dp)
    dimtest(ccf, [0, len(rv)])
    dimtest(ccf_e, [0, len(rv)])
    nantest(rv, 'rv in clean_ccf()')
    nantest(ccf, 'ccf in clean_ccf()')
    nantest(ccf_e, 'ccf_e in clean_ccf()')
    #ADD PARAMGET DV HERE.

    transit = sp.transit(dp)
    # transitblock = fun.rebinreform(transit,len(rv))

    Nrv = int(math.floor(len(rv)))

    baseline_ccf = np.hstack((ccf[:,
                                  0:int(0.25 * Nrv)], ccf[:,
                                                          int(0.75 * Nrv):]))
    baseline_ccf_e = np.hstack(
        (ccf_e[:, 0:int(0.25 * Nrv)], ccf_e[:, int(0.75 * Nrv):]))
    baseline_rv = np.hstack((rv[0:int(0.25 * Nrv)], rv[int(0.75 * Nrv):]))
    meanflux = np.median(
        baseline_ccf, axis=1
    )  #Normalize the baseline flux, but away from the signal of the planet.
    meanflux_e = 1.0 / len(baseline_rv) * np.sqrt(
        np.nansum(baseline_ccf_e**2.0, axis=1))  #1/N times sum of squares.
    #I validated that this is approximately equal to ccf_e/sqrt(N).
    meanblock = fun.rebinreform(meanflux, len(rv))
    meanblock_e = fun.rebinreform(meanflux_e, len(rv))

    ccf_n = ccf / meanblock.T
    ccf_ne = np.abs(ccf_n) * np.sqrt(
        (ccf_e / ccf)**2.0 + (meanblock_e.T / meanblock.T)**
        2.0)  #R=X/Z -> dR = R*sqrt( (dX/X)^2+(dZ/Z)^2 )
    #I validated that this is essentially equal to ccf_e/meanblock.T; as expected because the error on the mean spectrum is small compared to ccf_e.

    if np.sum(transit == 1) == 0:
        print(
            '------WARNING in Cleaning: The data contains only in-transit exposures.'
        )
        print('------The mean ccf is taken over the entire time-series.')
        meanccf = np.nanmean(ccf_n, axis=0)
        meanccf_e = 1.0 / len(transit) * np.sqrt(
            np.nansum(ccf_ne**2.0,
                      axis=0))  #I validated that this is approximately equal
        #to sqrt(N)*ccf_ne, where N is the number of out-of-transit exposures.
    elif np.sum(transit == 1) <= 0.25 * len(transit):
        print(
            '------WARNING in Cleaning: The data contains very few (<25%) out of transit exposures.'
        )
        print('------The mean ccf is taken over the entire time-series.')
        meanccf = np.nanmean(ccf_n, axis=0)
        meanccf_e = 1.0 / len(transit) * np.sqrt(
            np.nansum(ccf_ne**2.0,
                      axis=0))  #I validated that this is approximately equal
        #to sqrt(N)*ccf_ne, where N is the number of out-of-transit exposures.
    if np.min(transit) == 1.0:
        print(
            '------WARNING in Cleaning: The data is not predicted to contain in-transit exposures.'
        )
        print(
            f'------If you expect to be dealing with transit-data, please check the ephemeris '
            f'at {dp}')
        print('------The mean ccf is taken over the entire time-series.')
        meanccf = np.nanmean(ccf_n, axis=0)
        meanccf_e = 1.0 / len(transit) * np.sqrt(
            np.nansum(ccf_ne**2.0,
                      axis=0))  #I validated that this is approximately equal
        #to sqrt(N)*ccf_ne, where N is the number of out-of-transit exposures.
    else:
        meanccf = np.nanmean(ccf_n[transit == 1.0, :], axis=0)
        meanccf_e = 1.0 / np.sum(transit == 1) * np.sqrt(
            np.nansum(ccf_ne[transit == 1.0, :]**2.0,
                      axis=0))  #I validated that this is approximately equal
        #to sqrt(N)*ccf_ne, where N is the number of out-of-transit exposures.

    meanblock2 = fun.rebinreform(meanccf, len(meanflux))
    meanblock2_e = fun.rebinreform(meanccf_e, len(meanflux))

    ccf_nn = ccf_n / meanblock2  #MAY NEED TO DO SUBTRACTION INSTEAD TOGETHER W. NORMALIZATION OF LIGHTCURVE. SEE ABOVE.
    ccf_nne = np.abs(
        ccf_n / meanblock2) * np.sqrt((ccf_ne / ccf_n)**2.0 +
                                      (meanblock2_e / meanblock2)**2.0)
    #I validated that this error is almost equal to ccf_ne/meanccf

    #ONLY WORKS IF LIGHTCURVE MODEL IS ACCURATE, i.e. if Euler observations are available.
    # print("---> WARNING IN CLEANING.CLEAN_CCF(): NEED TO ADD A FUNCTION THAT YOU CAN NORMALIZE BY THE LIGHTCURVE AND SUBTRACT INSTEAD OF DIVISION!")
    return (ccf_n, ccf_ne, ccf_nn - 1.0, ccf_nne)