Esempio n. 1
0
def test_calc_AV_RV():
    # get the location of the data files
    data_path = pkg_resources.resource_filename("measure_extinction", "data/")

    # read in the observed data of the stars
    redstar = StarData("hd229238.dat", path=data_path)
    compstar = StarData("hd204172.dat", path=data_path)

    # calculate the extinction curve
    ext = ExtData()
    ext.calc_elx(redstar, compstar)

    # calculate A(V)
    ext.calc_AV()
    np.testing.assert_almost_equal(ext.columns["AV"], 2.5626900237367805)

    # calculate R(V)
    ext.calc_RV()
    np.testing.assert_almost_equal(ext.columns["RV"], 2.614989769244703)
def fit_spex_ext(
    starpair,
    path,
    functype="pow",
    dense=False,
    profile="drude_asym",
    exclude=None,
    bootstrap=False,
    fixed=False,
):
    """
    Fit the observed SpeX NIR extinction curve

    Parameters
    ----------
    starpair : string
        Name of the star pair for which to fit the extinction curve, in the format "reddenedstarname_comparisonstarname" (no spaces), or "average" to fit the average extinction curve

    path : string
        Path to the data files

    functype : string [default="pow"]
        Fitting function type ("pow" for powerlaw or "pol" for polynomial)

    dense : boolean [default=False]
        Whether or not to fit the features around 3 and 3.4 micron

    profile : string [default="drude_asym"]
        Profile to use for the features if dense = True (options are "gauss", "drude", "lorentz", "gauss_asym", "drude_asym", "lorentz_asym")

    exclude : list of tuples [default=None]
        list of tuples (min,max) with wavelength regions (in micron) that need to be excluded from the fitting, e.g. [(0.8,1.2),(2.2,5)]

    bootstrap : boolean [default=False]
        Whether or not to do a quick bootstrap fitting to get more realistic uncertainties on the fitting results

    fixed : boolean [default=False]
        Whether or not to add a fixed feature around 3 micron (for diffuse sightlines)

    Returns
    -------
    Updates extdata.model["type", "waves", "exts", "residuals", "chi2", "params"] and extdata.columns["AV"] with the fitting results:
        - type: string with the type of model (e.g. "pow_elx_Drude")
        - waves: np.ndarray with the SpeX wavelengths
        - exts: np.ndarray with the fitted model to the extinction curve at "waves" wavelengths
        - residuals: np.ndarray with the residuals, i.e. data-fit, at "waves" wavelengths
        - chi2 : float with the chi square of the fitting
        - params: list with output Parameter objects
    """
    # retrieve the SpeX data to be fitted, and sort the curve from short to long wavelengths
    filename = "%s%s_ext.fits" % (path, starpair.lower())
    if fixed:
        filename = filename.replace(".", "_ice.")
    extdata = ExtData(filename)
    (waves, exts, exts_unc) = extdata.get_fitdata(["SpeX_SXD", "SpeX_LXD"])
    indx = np.argsort(waves)
    waves = waves[indx].value
    exts = exts[indx]
    exts_unc = exts_unc[indx]

    # exclude wavelength regions if requested
    if exclude:
        mask = np.full_like(waves, False, dtype=bool)
        for region in exclude:
            mask += (waves > region[0]) & (waves < region[1])
        waves = waves[~mask]
        exts = exts[~mask]
        exts_unc = exts_unc[~mask]

    # get a quick estimate of A(V)
    if extdata.type == "elx":
        extdata.calc_AV()
        AV_guess = extdata.columns["AV"]
    else:
        AV_guess = None

    # convert to A(lambda)/A(1 micron)
    # ind1 = np.abs(waves - 1).argmin()
    # exts = exts / exts[ind1]
    # exts_unc = exts_unc / exts[ind1]

    # obtain the function to fit
    if "SpeX_LXD" not in extdata.waves.keys():
        dense = False
        fixed = False
    func = fit_function(
        dattype=extdata.type,
        functype=functype,
        dense=dense,
        profile=profile,
        AV_guess=AV_guess,
        fixed=fixed,
    )

    # for dense sightlines, add more weight to the feature region
    weights = 1 / exts_unc
    if dense:
        mask_ice = (waves > 2.88) & (waves < 3.19)
        mask_tail = (waves > 3.4) & (waves < 4)
        weights[mask_ice + mask_tail] *= 2

    # use the Levenberg-Marquardt algorithm to fit the data with the model
    fit = LevMarLSQFitter()
    fit_result_lev = fit(func, waves, exts, weights=weights, maxiter=10000)

    # set up the backend to save the samples for the emcee runs
    emcee_samples_file = path + "Fitting_results/" + starpair + "_emcee_samples.h5"

    # do the fitting again, with MCMC, using the results from the first fitting as input
    fit2 = EmceeFitter(nsteps=10000, burnfrac=0.1, save_samples=emcee_samples_file)

    # add parameter bounds
    for param in fit_result_lev.param_names:
        if "amplitude" in param:
            getattr(fit_result_lev, param).bounds = (0, 2)
        elif "alpha" in param:
            getattr(fit_result_lev, param).bounds = (0, 4)
        elif "Av" in param:
            getattr(fit_result_lev, param).bounds = (0, 10)

    fit_result_mcmc = fit2(fit_result_lev, waves, exts, weights=weights)

    # create standard MCMC plots
    fit2.plot_emcee_results(
        fit_result_mcmc, filebase=path + "Fitting_results/" + starpair
    )

    # choose the fit result to save
    fit_result = fit_result_mcmc
    # fit_result = fit_result_lev
    print(fit_result)

    # determine the wavelengths at which to evaluate and save the fitted model curve: all SpeX wavelengths, sorted from short to long (to avoid problems with overlap between SXD and LXD), and shortest and longest wavelength should have data
    if "SpeX_LXD" not in extdata.waves.keys():
        full_waves = extdata.waves["SpeX_SXD"].value
        full_npts = extdata.npts["SpeX_SXD"]
    else:
        full_waves = np.concatenate(
            (extdata.waves["SpeX_SXD"].value, extdata.waves["SpeX_LXD"].value)
        )
        full_npts = np.concatenate((extdata.npts["SpeX_SXD"], extdata.npts["SpeX_LXD"]))
    # sort the wavelengths
    indxs_sort = np.argsort(full_waves)
    full_waves = full_waves[indxs_sort]
    full_npts = full_npts[indxs_sort]
    # cut the wavelength region
    indxs = np.logical_and(full_waves >= np.min(waves), full_waves <= np.max(waves))
    full_waves = full_waves[indxs]
    full_npts = full_npts[indxs]

    # calculate the residuals and put them in an array of the same length as "full_waves" for plotting
    residuals = exts - fit_result(waves)
    full_res = np.full_like(full_npts, np.nan)
    if exclude:
        mask = np.full_like(full_waves, False, dtype=bool)
        for region in exclude:
            mask += (full_waves > region[0]) & (full_waves < region[1])
        full_res[(full_npts > 0) * ~mask] = residuals

    else:
        full_res[(full_npts > 0)] = residuals

    # bootstrap to get more realistic uncertainties on the parameter results
    if bootstrap:
        red_star = StarData(extdata.red_file, path=path, use_corfac=True)
        comp_star = StarData(extdata.comp_file, path=path, use_corfac=True)
        red_V_unc = red_star.data["BAND"].get_band_mag("V")[1]
        comp_V_unc = comp_star.data["BAND"].get_band_mag("V")[1]
        unc_V = np.sqrt(red_V_unc ** 2 + comp_V_unc ** 2)
        fit_result_mcmc_low = fit2(fit_result_lev, waves, exts - unc_V, weights=weights)
        fit_result_mcmc_high = fit2(
            fit_result_lev, waves, exts + unc_V, weights=weights
        )

    # save the fitting results to the fits file
    if dense:
        functype += "_" + profile
    extdata.model["type"] = functype + "_" + extdata.type
    extdata.model["waves"] = full_waves
    extdata.model["exts"] = fit_result(full_waves)
    extdata.model["residuals"] = full_res
    extdata.model["chi2"] = np.sum((residuals / exts_unc) ** 2)
    print("Chi2", extdata.model["chi2"])
    extdata.model["params"] = []
    for param in fit_result.param_names:
        # update the uncertainties when bootstrapping
        if bootstrap:
            min_val = min(
                getattr(fit_result_mcmc, param).value,
                getattr(fit_result_mcmc_low, param).value,
                getattr(fit_result_mcmc_high, param).value,
            )
            max_val = max(
                getattr(fit_result_mcmc, param).value,
                getattr(fit_result_mcmc_low, param).value,
                getattr(fit_result_mcmc_high, param).value,
            )
            sys_unc = (max_val - min_val) / 2
            getattr(fit_result, param).unc_minus = np.sqrt(
                getattr(fit_result, param).unc_minus ** 2 + sys_unc ** 2
            )
            getattr(fit_result, param).unc_plus = np.sqrt(
                getattr(fit_result, param).unc_plus ** 2 + sys_unc ** 2
            )

        extdata.model["params"].append(getattr(fit_result, param))

        # save the column information (A(V), E(B-V) and R(V))
        if "Av" in param:
            extdata.columns["AV"] = (
                getattr(fit_result, param).value,
                getattr(fit_result, param).unc_minus,
                getattr(fit_result, param).unc_plus,
            )
            # calculate the distrubtion of R(V) and 1/R(V) from the distributions of A(V) and E(B-V)
            nsamples = getattr(fit_result, param).posterior.n_samples
            av_dist = unc.normal(
                extdata.columns["AV"][0],
                std=(extdata.columns["AV"][1] + extdata.columns["AV"][2]) / 2,
                n_samples=nsamples,
            )
            b_indx = np.abs(extdata.waves["BAND"] - 0.438 * u.micron).argmin()
            ebv_dist = unc.normal(
                extdata.exts["BAND"][b_indx],
                std=extdata.uncs["BAND"][b_indx],
                n_samples=nsamples,
            )
            ebv_per = ebv_dist.pdf_percentiles([16.0, 50.0, 84.0])
            extdata.columns["EBV"] = (
                ebv_per[1],
                ebv_per[1] - ebv_per[0],
                ebv_per[2] - ebv_per[1],
            )
            rv_dist = av_dist / ebv_dist
            rv_per = rv_dist.pdf_percentiles([16.0, 50.0, 84.0])
            extdata.columns["RV"] = (
                rv_per[1],
                rv_per[1] - rv_per[0],
                rv_per[2] - rv_per[1],
            )
            inv_rv_dist = ebv_dist / av_dist
            inv_rv_per = inv_rv_dist.pdf_percentiles([16.0, 50.0, 84.0])
            extdata.columns["IRV"] = (
                inv_rv_per[1],
                inv_rv_per[1] - inv_rv_per[0],
                inv_rv_per[2] - inv_rv_per[1],
            )
            print(extdata.columns)

    # save the fits file
    extdata.save(filename)

    # print information about the ice feature
    if fixed:
        print(
            "Ice feature strength: ",
            extdata.model["params"][3].value,
            extdata.model["params"][3].unc_minus,
            extdata.model["params"][3].unc_plus,
        )