def parallel_chisquared(i, j, alpha, rv, res, snr, observation, host_models,
                        companion_models, output1, output2, x_memmap,
                        y_memmap):
    """Function for parallel processing in chisqr for resolution and snr uses numpy memap to store data.

    Inputs:
    i,j = matrix locations of this alpha and rv
    alpha =  flux ratio
    rv = radial velocity offset of companion
    snr = signal to noise ratio of observation

    Output:
    None: Changes data in output1 and output2 numpy memmaps.
    """
    # print("i, j, Resolution, snr, alpha, rv")
    # print(i, j, res, snr, alpha, rv)
    host_model = host_models[res]
    companion_model = companion_models[res]
    companion_model.doppler_shift(rv)
    combined_model = combine_spectra(host_model, companion_model, alpha)
    # model_new = combine_spectra(convolved_star_models[resolution],
    # convolved_planet_models[resolution].doppler_shift(rv), alpha)

    observation.wav_select(2100, 2200)
    # INTERPOLATE COMBINED MODEL VALUES TO OBSERVATION VALUES
    combined_model.spline_interpolate_to(observation)

    # print("i", i, "j", j, "chisqr",
    # scipy.stats.chisquare(observation.flux, combined_model.flux).statistic)

    output1[i, j] = scipy.stats.chisquare(observation.flux,
                                          combined_model.flux).statistic
    # output2[i, j] = chi_squared(observation.flux, combined_model.flux,
    #                            error=observation.flux/snr)
    output2[i, j] = chi_squared(observation.flux,
                                combined_model.flux,
                                error=None)
    x_memmap[i, j] = alpha
    y_memmap[i, j] = rv

    if False:
        # if i == j:
        spectrum_plotter(observation, label="Simulated obs", show=False)
        spectrum_plotter(combined_model, label="This model", show=False)
        plt.title("Parallel printing, alpha={},rv={}".format(alpha, rv))
        plt.show()
rv_grid_data = join_with_broadcast(mod1_spec_scaled, mod2_spec_scaled, rvs,
                                   [10], sample_x)
dual_grid_data = join_with_broadcast(mod1_spec_scaled, mod2_spec_scaled, rvs,
                                     gammas, sample_x)

for normalize in (True, False):
    print("normalizing", normalize)

    fake_data.remove_nans()

    print(fake_data.flux.shape)
    print(gamma_grid_data.shape)
    print(rv_grid_data.shape)
    print(dual_grid_data.shape)

    gamma_chi2 = chi_squared(fake_data.flux[:, np.newaxis], gamma_grid_data)
    rv_chi2 = chi_squared(fake_data.flux[:, np.newaxis], rv_grid_data)
    dual_chi2 = chi_squared(fake_data.flux[:, np.newaxis, np.newaxis],
                            dual_grid_data)

    plt.plot(gammas, gamma_chi2)
    plt.title("gamma chi2")
    plt.show()

    plt.plot(rvs, rv_chi2)
    plt.title("rv chi2")
    plt.show()

    gam, rv_grid = np.meshgrid(gammas, rvs)
    plt.contourf(gam, rv_grid, dual_chi2)
    plt.title("dual chi2 - gamma {}, rv {}".format(gamma, rv))
def iam_wrapper(num, params1, model2_pars, rvs, gammas, obs_spec, norm=False,
                verbose=False, save_only=True, chip=None, prefix=None, errors=None,
                area_scale=True, wav_scale=True, grid_slices=False, norm_method="scalar",
                fudge=None):
    """Wrapper for iteration loop of iam. params1 fixed, model2_pars are many.

    fudge is multiplicative on companion spectrum.
    """
    if prefix is None:
        sf = os.path.join(
            simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper(),
            "iam_{0}_{1}-{2}_part{6}_host_pars_[{3}_{4}_{5}].csv".format(
                obs_spec.header["OBJECT"].upper(), int(obs_spec.header["MJD-OBS"]), chip,
                params1[0], params1[1], params1[2], num))
        prefix = os.path.join(
            simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper())  # for fudge

    else:
        sf = "{0}_part{4}_host_pars_[{1}_{2}_{3}].csv".format(
            prefix, params1[0], params1[1], params1[2], num)
    save_filename = sf

    if os.path.exists(save_filename) and save_only:
        print("'{0}' exists, so not repeating calculation.".format(save_filename))
        return None
    else:
        if not save_only:
            iam_grid_chisqr_vals = np.empty(len(model2_pars))
        for jj, params2 in enumerate(model2_pars):
            if verbose:
                print(("Starting iteration with parameters: "
                       "{0}={1},{2}={3}").format(num, params1, jj, params2))

            # Main Part
            rv_limits = observation_rv_limits(obs_spec, rvs, gammas)

            obs_spec = obs_spec.remove_nans()
            assert ~np.any(np.isnan(obs_spec.flux)), "Observation has nan"

            # Load phoenix models and scale by area and wavelength limit
            mod1_spec, mod2_spec = \
                prepare_iam_model_spectra(params1, params2, limits=rv_limits,
                                          area_scale=area_scale, wav_scale=wav_scale)
            # Estimated flux ratio from models
            inherent_alpha = continuum_alpha(mod1_spec, mod2_spec, chip)

            # Combine model spectra with iam model
            mod1_spec.plot(label=params1)
            mod2_spec.plot(label=params2)
            plt.close()

            if fudge or (fudge is not None):
                fudge_factor = float(fudge)
                mod2_spec.flux *= fudge_factor  # fudge factor multiplication
                mod2_spec.plot(label="fudged {0}".format(params2))
                plt.title("fudges models")
                plt.legend()

                fudge_prefix = os.path.basename(os.path.normpath(prefix))
                fname = os.path.join(simulators.paths["output_dir"],
                                     obs_spec.header["OBJECT"].upper(), "iam", "fudgeplots",
                                     "{1}_fudged_model_spectra_factor={0}_num={2}_iter_{3}.png".format(fudge_factor,
                                                                                                       fudge_prefix,
                                                                                                       num, jj))
                plt.savefig(fname)
                plt.close()
                warnings.warn("Using a fudge factor = {0}".format(fudge_factor))

            iam_grid_func = inherent_alpha_model(mod1_spec.xaxis, mod1_spec.flux, mod2_spec.flux,
                                                 rvs=rvs, gammas=gammas)
            iam_grid_models = iam_grid_func(obs_spec.xaxis)

            # Continuum normalize all iam_gird_models
            def axis_continuum(flux):
                """Continuum to apply along axis with predefined variables parameters."""
                return continuum(obs_spec.xaxis, flux, splits=20, method="exponential", top=20)

            iam_grid_continuum = np.apply_along_axis(axis_continuum, 0, iam_grid_models)

            iam_grid_models = iam_grid_models / iam_grid_continuum

            # RE-NORMALIZATION
            if chip == 4:
                # Quadratically renormalize anyway
                obs_spec = renormalization(obs_spec, iam_grid_models, normalize=True, method="quadratic")
            obs_flux = renormalization(obs_spec, iam_grid_models, normalize=norm, method=norm_method)

            if grid_slices:
                # Long execution plotting.
                plot_iam_grid_slices(obs_spec.xaxis, rvs, gammas, iam_grid_models,
                                     star=obs_spec.header["OBJECT"].upper(),
                                     xlabel="wavelength", ylabel="rv", zlabel="gamma",
                                     suffix="iam_grid_models", chip=chip)

            old_shape = iam_grid_models.shape
            # Arbitrary_normalization of observation
            iam_grid_models, arb_norm = arbitrary_rescale(iam_grid_models,
                                                          *simulators.sim_grid["arb_norm"])
            # print("Arbitrary Normalized iam_grid_model shape.", iam_grid_models.shape)
            assert iam_grid_models.shape == (*old_shape, len(arb_norm))

            # Calculate Chi-squared
            obs_flux = np.expand_dims(obs_flux, -1)  # expand on last axis to match rescale
            iam_norm_grid_chisquare = chi_squared(obs_flux, iam_grid_models, error=errors)

            # Take minimum chi-squared value along Arbitrary normalization axis
            iam_grid_chisquare, arbitrary_norms = arbitrary_minimums(iam_norm_grid_chisquare, arb_norm)

            npix = obs_flux.shape[0]  # Number of pixels used

            if grid_slices:
                # Long execution plotting.
                plot_iam_grid_slices(rvs, gammas, arb_norm, iam_norm_grid_chisquare,
                                     star=obs_spec.header["OBJECT"].upper(),
                                     xlabel="rv", ylabel="gamma", zlabel="Arbitrary Normalization",
                                     suffix="iam_grid_chisquare", chip=chip)

            if not save_only:
                iam_grid_chisqr_vals[jj] = iam_grid_chisquare.ravel()[np.argmin(iam_grid_chisquare)]

            save_full_iam_chisqr(save_filename, params1, params2,
                                 inherent_alpha, rvs, gammas,
                                 iam_grid_chisquare, arbitrary_norms, npix, verbose=verbose)
        if save_only:
            return None
        else:
            return iam_grid_chisqr_vals
Example #4
0
def compare_spectra(table, params):
    """Plot the min chi2 result against the observations."""
    extractor = DBExtractor(table)
    gamma_df = extractor.simple_extraction(columns=["gamma"])
    extreme_gammas = [min(gamma_df.gamma.values), max(gamma_df.gamma.values)]
    for ii, chi2_val in enumerate(chi2_names[0:-2]):
        df = extractor.ordered_extraction(
            columns=["teff_1", "logg_1", "feh_1", "gamma", chi2_val],
            order_by=chi2_val,
            limit=1,
            asc=True)

        params1 = [
            df["teff_1"].values[0], df["logg_1"].values[0],
            df["feh_1"].values[0]
        ]

        params1 = [float(param1) for param1 in params1]

        gamma = df["gamma"].values

        obs_name, obs_params, output_prefix = bhm_helper_function(
            params["star"], params["obsnum"], ii + 1)
        print(obs_name)
        obs_spec = load_spectrum(obs_name)

        # Mask out bad portion of observed spectra
        # obs_spec = spectrum_masking(obs_spec, params["star"], params["obsnum"], ii + 1)

        # Barycentric correct spectrum
        _obs_spec = barycorr_crires_spectrum(obs_spec, extra_offset=None)
        normalization_limits = [obs_spec.xaxis[0] - 5, obs_spec.xaxis[-1] + 5]
        # models
        # print("params for models", params1)
        mod1 = load_starfish_spectrum(params1,
                                      limits=normalization_limits,
                                      hdr=True,
                                      normalize=False,
                                      area_scale=True,
                                      flux_rescale=True)

        bhm_grid_func = one_comp_model(mod1.xaxis, mod1.flux, gammas=gamma)
        bhm_upper_gamma = one_comp_model(mod1.xaxis,
                                         mod1.flux,
                                         gammas=extreme_gammas[1])
        bhm_lower_gamma = one_comp_model(mod1.xaxis,
                                         mod1.flux,
                                         gammas=extreme_gammas[0])

        bhm_grid_model = bhm_grid_func(obs_spec.xaxis).squeeze()
        bhm_grid_model_full = bhm_grid_func(mod1.xaxis).squeeze()
        bhm_upper_gamma = bhm_upper_gamma(obs_spec.xaxis).squeeze()
        bhm_lower_gamma = bhm_lower_gamma(obs_spec.xaxis).squeeze()

        model_spec_full = Spectrum(flux=bhm_grid_model_full, xaxis=mod1.xaxis)
        model_spec = Spectrum(flux=bhm_grid_model, xaxis=obs_spec.xaxis)
        bhm_upper_gamma = Spectrum(flux=bhm_upper_gamma, xaxis=obs_spec.xaxis)
        bhm_lower_gamma = Spectrum(flux=bhm_lower_gamma, xaxis=obs_spec.xaxis)

        model_spec = model_spec.remove_nans()
        model_spec = model_spec.normalize(method="exponential")
        model_spec_full = model_spec_full.remove_nans()
        model_spec_full = model_spec_full.normalize(method="exponential")
        bhm_lower_gamma = bhm_lower_gamma.remove_nans()
        bhm_lower_gamma = bhm_lower_gamma.normalize(method="exponential")
        bhm_upper_gamma = bhm_upper_gamma.remove_nans()
        bhm_upper_gamma = bhm_upper_gamma.normalize(method="exponential")

        from mingle.utilities.chisqr import chi_squared
        chisqr = chi_squared(obs_spec.flux, model_spec.flux)

        print("Recomputed chi^2 = {0}".format(chisqr))
        print("Database chi^2 = {0}".format(df[chi2_val]))
        fig, ax = plt.subplots(1, 1, figsize=(15, 8))
        plt.plot(obs_spec.xaxis,
                 obs_spec.flux + 0.01,
                 label="0.05 + Observation, {}".format(obs_name))
        plt.plot(model_spec.xaxis,
                 model_spec.flux,
                 label="Minimum \chi^2 model")
        plt.plot(model_spec_full.xaxis,
                 model_spec_full.flux,
                 "--",
                 label="Model_full_res")
        plt.plot(bhm_lower_gamma.xaxis,
                 bhm_lower_gamma.flux,
                 "-.",
                 label="gamma={}".format(extreme_gammas[0]))
        plt.plot(bhm_upper_gamma.xaxis,
                 bhm_upper_gamma.flux,
                 ":",
                 label="gamma={}".format(extreme_gammas[1]))
        plt.title("bhm spectrum")
        plt.legend()

        fig.tight_layout()
        name = "{0}-{1}_{2}_{3}_bhm_min_chi2_spectrum_comparison_{4}.png".format(
            params["star"], params["obsnum"], params["chip"], chi2_val,
            params["suffix"])
        plt.savefig(os.path.join(params["path"], "plots", name))
        plt.close()

        plt.plot(obs_spec.xaxis, obs_spec.flux, label="Observation")
        plt.show()
def test_spectrum_chi_squared(host):
    host2 = host.copy()
    host2.flux *= 2
    spec_chi2 = spectrum_chisqr(host, host2)
    chi2 = chi_squared(host.flux, host2.flux)
    assert np.allclose(spec_chi2, chi2)
Example #6
0
def tcm_analysis(obs_spec,
                 model1_pars,
                 model2_pars,
                 alphas=None,
                 rvs=None,
                 gammas=None,
                 verbose=False,
                 norm=False,
                 chip=None,
                 prefix=None):
    """Run two component model over all parameter cobinations in model1_pars and model2_pars."""
    if chip is None:
        chip = ""

    if alphas is None:
        alphas = np.array([0])
    elif isinstance(alphas, (float, int)):
        alphas = np.asarray(alphas, dtype=np.float32)
    if rvs is None:
        rvs = np.array([0])
    elif isinstance(rvs, (float, int)):
        rvs = np.asarray(rvs, dtype=np.float32)
    if gammas is None:
        gammas = np.array([0])
    elif isinstance(gammas, (float, int)):
        gammas = np.asarray(gammas, dtype=np.float32)

    if isinstance(model1_pars, list):
        logging.debug(
            __("Number of close model_pars returned {0}", len(model1_pars)))
    if isinstance(model2_pars, list):
        logging.debug(
            __("Number of close model_pars returned {0}", len(model2_pars)))

    print("host params", model1_pars)
    print("companion params", model2_pars)

    # Solution Grids to return
    # model_chisqr_vals = np.empty((len(model1_pars), len(model2_pars)))
    # model_xcorr_vals = np.empty(len(model1_pars), len(model2_pars))
    # model_xcorr_rv_vals = np.empty(len(model1_pars), len(model2_pars))
    broadcast_chisqr_vals = np.empty((len(model1_pars), len(model2_pars)))
    # broadcast_gamma = np.empty((len(model1_pars), len(model2_pars)))
    # full_broadcast_chisquare = np.empty((len(model1_pars), len(model2_pars), len(alphas), len(rvs), len(gammas)))

    normalization_limits = [2105, 2185]  # small as possible?
    # combined_params = itertools.product(model1_pars, model2_pars)
    for ii, params1 in enumerate(tqdm(model1_pars)):
        if prefix is None:
            sf = (
                "Analysis/{0}/tc_{0}_{1}-{2}_part{6}_host_pars_[{3}_{4}_{5}]_par"
                ".csv").format(obs_spec.header["OBJECT"].upper(),
                               int(obs_spec.header["MJD-OBS"]), chip,
                               params1[0], params1[1], params1[2], ii)

        else:
            sf = "{0}_part{4}_host_pars_[{1}_{2}_{3}]_par.csv".format(
                prefix, params1[0], params1[1], params1[2], ii)
        save_filename = sf

        for jj, params2 in enumerate(model2_pars):
            if verbose:
                print("Starting iteration with parameters:\n{0}={1},{2}={3}".
                      format(ii, params1, jj, params2))
            mod1_spec = load_starfish_spectrum(params1,
                                               limits=normalization_limits,
                                               hdr=True,
                                               normalize=True)
            mod2_spec = load_starfish_spectrum(params2,
                                               limits=normalization_limits,
                                               hdr=True,
                                               normalize=True)

            # TODO WHAT IS THE MAXIMUM (GAMMA + RV POSSIBLE? LIMIT IT TO THAT SHIFT?

            # Wavelength selection
            mod1_spec.wav_select(
                np.min(obs_spec.xaxis) - 5,
                np.max(obs_spec.xaxis) + 5)  # +- 5nm of obs for convolution
            mod2_spec.wav_select(
                np.min(obs_spec.xaxis) - 5,
                np.max(obs_spec.xaxis) + 5)
            obs_spec = obs_spec.remove_nans()

            # One component model with broadcasting over gammas
            # two_comp_model(wav, model1, model2, alphas, rvs, gammas)
            assert np.allclose(mod1_spec.xaxis, mod2_spec.xaxis)

            broadcast_result = two_comp_model(mod1_spec.xaxis,
                                              mod1_spec.flux,
                                              mod2_spec.flux,
                                              alphas=alphas,
                                              rvs=rvs,
                                              gammas=gammas)
            broadcast_values = broadcast_result(obs_spec.xaxis)

            assert ~np.any(np.isnan(obs_spec.flux)), "Observation is nan"

            #### NORMALIZATION NEEDED HERE
            if norm:
                obs_flux = chi2_model_norms(obs_spec.xaxis, obs_spec.flux,
                                            broadcast_values)

            else:
                obs_flux = obs_spec.flux[:, np.newaxis, np.newaxis, np.newaxis]
            #####

            broadcast_chisquare = chi_squared(obs_flux, broadcast_values)
            sp_chisquare = sp.stats.chisquare(obs_flux,
                                              broadcast_values,
                                              axis=0).statistic

            assert np.all(sp_chisquare == broadcast_chisquare)

            # print(broadcast_chisquare.shape)
            # print("chi squaresd = ", broadcast_chisquare.ravel()[np.argmin(broadcast_chisquare)])
            # New parameters to explore
            broadcast_chisqr_vals[ii, jj] = broadcast_chisquare.ravel()[
                np.argmin(broadcast_chisquare)]
            # broadcast_gamma[ii, jj] = gammas[np.argmin(broadcast_chisquare)]
            # full_broadcast_chisquare[ii, jj, :] = broadcast_chisquare

            save_full_chisqr(save_filename,
                             params1,
                             params2,
                             alphas,
                             rvs,
                             gammas,
                             broadcast_chisquare,
                             verbose=verbose)

    return broadcast_chisqr_vals  # Just output the best value for each model pair
def test_chi_squared_with_error_unequal_length(observed, expected, error):
    observed = np.asarray(observed)
    expected = np.asarray(expected)
    error = np.asarray(error)
    with pytest.raises(ValueError):
        chi_squared(observed, expected, error=error)
def test_model_equal_observed_chi_squared_returns_zero(observed):
    observed = np.asarray(observed)
    assert chi_squared(observed, observed) == 0
def test_chi_squared_without_errors_equals_scipy(observed, expected):
    observed = np.asarray(observed)
    expected = np.asarray(expected)
    assert chi_squared(observed, expected) == chisquare(observed,
                                                        expected).statistic
Example #10
0
def main():
    """Chisquare determinination to detect minimum alpha value."""
    print("Loading Data")

    path = "/home/jneal/Phd/Codes/Phd-codes/Simulations/saves"  # save path

    chip_limits = [2080, 2220]

    org_star_spec, org_bd_spec = load_PHOENIX_hd30501(limits=chip_limits,
                                                      normalize=True)

    # org_star_spec = Spectrum(xaxis=w_mod, flux=I_star, calibrated=True)
    # org_bd_spec = Spectrum(xaxis=w_mod, flux=I_bdmod, calibrated=True)

    resolutions = [50000]
    snrs = [100, 101, 110, 111]  # Signal to noise levels
    # alphas = 10**np.linspace(-5, -0.2, 200)
    # RVs = np.arange(10, 30, 0.1)

    # RV and alpha value of Simulations
    # rv_val = 0
    # Alpha = 0  # Vary this to determine detection limit
    # input_parameters = (rv_val, Alpha)

    convolved_star_model = store_convolutions(org_star_spec,
                                              resolutions,
                                              chip_limits=chip_limits)
    # convolved_planet_model = store_convolutions(org_bd_spec, resolutions, chip_limits=chip_limits)

    # print(type(convolved_star_model))
    # print(type(convolved_planet_model))
    noisey_obersvations = generate_noise_observations(convolved_star_model,
                                                      convolved_planet_model,
                                                      rv_val, Alpha,
                                                      resolutions, snrs)

    # Not used with gernerator function
    goal_planet_shifted = copy.copy(org_bd_spec)
    # RV shift BD spectra
    goal_planet_shifted.doppler_shift(rv_val)

    # These should be replaced by
    res_stored_chisquared = dict()
    res_error_stored_chisquared = dict()
    # This
    res_snr_storage_dict = defaultdict(dict)  # Dictionary of dictionaries
    error_res_snr_storage_dict = defaultdict(
        dict)  # Dictionary of dictionaries
    # Iterable over resolution and snr to process
    # res_snr_iter = itertools.product(resolutions, snrs)
    # Can then store to dict store_dict[res][snr]

    print("Starting loop")

    for resolution in tqdm(resolutions):
        chisqr_snr_dict = dict()  # store 2d array in dict of SNR
        error_chisqr_snr_dict = dict()
        print("\nSTARTING run of RESOLUTION={}\n".format(resolution))

        star_spec = apply_convolution(org_star_spec,
                                      R=resolution,
                                      chip_limits=chip_limits)
        goal_planet = apply_convolution(goal_planet_shifted,
                                        R=resolution,
                                        chip_limits=chip_limits)

        # if resolution is None:
        #    star_spec = copy.copy(org_star_spec)
        #    goal_planet = copy.copy(goal_planet_shifted)
        # else:
        #    ip_xaxis, ip_flux = IPconvolution(org_star_spec.xaxis,
        #             org_star_spec.flux, chip_limits, resolution,
        #            FWHM_lim=5.0, plot=False, verbose=True)

        #        star_spec = Spectrum(xaxis=ip_xaxis, flux=ip_flux,
        #                                 calibrated=True,
        #                                 header=org_star_spec.header)

        #        ip_xaxis, ip_flux = IPconvolution(goal_planet_shifted.xaxis,
        #            goal_planet_shifted.flux, chip_limits, resolution,
        #            FWHM_lim=5.0, plot=False, verbose=False)

        #        goal_planet = Spectrum(xaxis=ip_xaxis, flux=ip_flux,
        #                                     calibrated=True,
        #                                     header=goal_planet_shifted.header)

        print(
            "Starting SNR loop for resolution value of {}".format(resolution))
        for snr in snrs:
            loop_start = time.time()
            print("Calculation with snr level", snr)
            # This is the signal to try and recover
            alpha_combine = combine_spectra(star_spec, goal_planet, Alpha)
            alpha_combine.wav_select(2100, 2200)
            # alpha_combine.flux = add_noise2(alpha_combine.flux, snr)
            alpha_combine.add_noise(snr)

            # Test plot
            # plt.plot(alpha_combine.xaxis, alpha_combine.flux)
            sim_observation = simulated_obersvations[resolution][snr]
            # plt.plot(this_simulation.xaxis, this_simulation.flux, label="function generatred")
            # plt.legend()
            # plt.show()

            # chisqr_store = np.empty((len(alphas), len(RVs)))
            scipy_chisqr_store = np.empty((len(alphas), len(RVs)))
            error_chisqr_store = np.empty((len(alphas), len(RVs)))
            new_scipy_chisqr_store = np.empty((len(alphas), len(RVs)))
            new_error_chisqr_store = np.empty((len(alphas), len(RVs)))
            for i, alpha in enumerate(alphas):
                for j, RV in enumerate(RVs):
                    # print("RV", RV, "alpha", alpha, "snr", snr, "res", resolution)

                    # Generate model for this RV and alhpa
                    planet_shifted = copy.copy(org_bd_spec)
                    planet_shifted.doppler_shift(RV)
                    model = combine_spectra(star_spec, planet_shifted, alpha)
                    model.wav_select(2100, 2200)

                    # Try scipy chi_squared
                    scipy_chisquare = scipy.stats.chisquare(
                        alpha_combine.flux, model.flux)
                    error_chisquare = chi_squared(alpha_combine.flux,
                                                  model.flux,
                                                  error=alpha_combine.flux /
                                                  snr)

                    # print("Mine, scipy", chisqr, scipy_chisquare)
                    error_chisqr_store[i, j] = error_chisquare
                    scipy_chisqr_store[i, j] = scipy_chisquare.statistic

                    #########################
                    # using dictionary values
                    host_model = convolved_star_model[resolution]
                    companion_model = convolved_planet_model[resolution]
                    companion_model.doppler_shift(RV)
                    model_new = combine_spectra(host_model, companion_model,
                                                alpha)

                    # model_new = combine_spectra(convolved_star_model[resolution], convolved_planet_model[resolution].doppler_shift(RV), alpha)
                    model_new.wav_select(2100, 2200)
                    sim_observation.wav_select(2100, 2200)

                    new_scipy_chisquare = scipy.stats.chisquare(
                        sim_observation.flux, model_new.flux)
                    new_error_chisquare = chi_squared(
                        sim_observation.flux,
                        model_new.flux,
                        error=sim_observation.flux / snr)

                    new_error_chisqr_store[i, j] = new_error_chisquare
                    new_scipy_chisqr_store[i,
                                           j] = new_scipy_chisquare.statistic
                    ##############################

            chisqr_snr_dict[str(snr)] = scipy_chisqr_store
            error_chisqr_snr_dict[str(snr)] = error_chisqr_store

            res_snr_storage_dict[resolution][snr] = new_scipy_chisqr_store
            error_res_snr_storage_dict[resolution][
                snr] = new_error_chisqr_store

            # Save the results to a file to stop repeating loops

            for key, val in chisqr_snr_dict.items():
                np.save(
                    os.path.join(
                        path, "scipy_chisquare_data_snr_{0}_res{1}".format(
                            key, resolution)), val)
            for key, val in error_chisqr_snr_dict.items():
                np.save(
                    os.path.join(
                        path, "error_chisquare_data_snr_{0}_res{1}".format(
                            key, resolution)), val)
            # Store in dictionary
            res_stored_chisquared[resolution] = chisqr_snr_dict
            res_error_stored_chisquared[resolution] = error_chisqr_snr_dict

            print("SNR Loop time = {}".format(time.time() - loop_start))

    print("Finished Resolution {}".format(resolution))
def main():
    """Main function."""
    star = "HD30501"
    host_parameters = load_param_file(star)
    obsnum = 1
    chip = 1
    obs_name = select_observation(star, obsnum, chip)

    # Load observation
    # uncorrected_spectra = load_spectrum(obs_name)
    observed_spectra = load_spectrum(obs_name)
    _observed_spectra = barycorr_crires_spectrum(observed_spectra,
                                                 extra_offset=None)
    observed_spectra.flux /= 1.02

    obs_resolution = crires_resolution(observed_spectra.header)

    wav_model = fits.getdata(
        os.path.join(wav_dir, "WAVE_PHOENIX-ACES-AGSS-COND-2011.fits"))
    wav_model /= 10  # turn into nm
    logging.debug(__("Phoenix wav_model = {0}", wav_model))

    closest_model = phoenix_name_from_params(model_base_dir, host_parameters)
    original_model = "Z-0.0/lte05200-4.50-0.0.PHOENIX-ACES-AGSS-COND-2011-HiRes.fits"
    logging.debug(__("closest_model {0}", closest_model))
    logging.debug(__("original_model {0}", original_model))

    # Function to find the good models I need
    models = find_phoenix_model_names(model_base_dir, original_model)
    if isinstance(models, list):
        logging.debug(__("Number of close models returned {0}", len(models)))

    model_chisqr_vals = np.empty_like(models)
    model_xcorr_vals = np.empty_like(models)
    model_xcorr_rv_vals = np.empty_like(models)

    for ii, model_name in enumerate(models):
        mod_flux = fits.getdata(model_name)
        mod_header = fits.getheader(model_name)
        mod_spectrum = Spectrum(xaxis=wav_model,
                                flux=mod_flux,
                                header=mod_header,
                                calibrated=True)

        # Normalize Phoenix Spectrum
        # mod_spectrum.wav_select(2080, 2200)  # limits for simple normalization
        mod_spectrum.wav_select(2105, 2165)  # limits for simple normalization
        # norm_mod_spectrum = simple_normalization(mod_spectrum)
        norm_mod_spectrum = spec_local_norm(mod_spectrum, plot=False)

        # Wav select
        norm_mod_spectrum.wav_select(
            np.min(observed_spectra.xaxis) - 5,
            np.max(observed_spectra.xaxis) +
            5)  # +- 5nm of obs for convolution

        # Convolve to resolution of instrument
        conv_mod_spectrum = convolve_models([norm_mod_spectrum],
                                            obs_resolution,
                                            chip_limits=None)[0]

        # Find crosscorrelation RV
        # # Should run though all models and find best rv to apply uniformly
        rvoffset, cc_max = xcorr_peak(observed_spectra,
                                      conv_mod_spectrum,
                                      plot=False)

        # Interpolate to obs
        conv_mod_spectrum.spline_interpolate_to(observed_spectra)
        # conv_mod_spectrum.interpolate1d_to(observed_spectra)
        model_chi_val = chi_squared(observed_spectra.flux,
                                    conv_mod_spectrum.flux)

        # argmax = np.argmax(cc_max)
        model_chisqr_vals[ii] = model_chi_val
        model_xcorr_vals[ii] = cc_max
        model_xcorr_rv_vals[ii] = rvoffset

    logging.debug(pv("model_chisqr_vals"))
    logging.debug(pv("model_xcorr_vals"))
    chisqr_argmin_indx = np.argmin(model_chisqr_vals)
    xcorr_argmax_indx = np.argmax(model_xcorr_vals)

    logging.debug(pv("chisqr_argmin_indx"))
    logging.debug(pv("xcorr_argmax_indx"))

    logging.debug(pv("model_chisqr_vals"))
    print("Minimum  Chisqr value =",
          model_chisqr_vals[chisqr_argmin_indx])  # , min(model_chisqr_vals)
    print("Chisqr at max correlation value",
          model_chisqr_vals[chisqr_argmin_indx])

    print("model_xcorr_vals = {}".format(model_xcorr_vals))
    print("Maximum Xcorr value =",
          model_xcorr_vals[xcorr_argmax_indx])  # , max(model_xcorr_vals)
    print("Xcorr at min Chiqsr", model_xcorr_vals[chisqr_argmin_indx])

    logging.debug(pv("model_xcorr_rv_vals"))
    print("RV at max xcorr =", model_xcorr_rv_vals[xcorr_argmax_indx])
    # print("Median RV val =", np.median(model_xcorr_rv_vals))
    print(pv("model_xcorr_rv_vals[chisqr_argmin_indx]"))
    # print(pv("sp.stats.mode(np.around(model_xcorr_rv_vals))"))

    print("Max Correlation model = ",
          models[xcorr_argmax_indx].split("/")[-2:])
    print("Min Chisqr model = ", models[chisqr_argmin_indx].split("/")[-2:])

    limits = [2110, 2160]
    best_model = models[chisqr_argmin_indx]
    best_model_spec = load_phoenix_spectrum(best_model,
                                            limits=limits,
                                            normalize=True)
    best_model_spec = convolve_models([best_model_spec],
                                      obs_resolution,
                                      chip_limits=None)[0]

    best_xcorr_model = models[xcorr_argmax_indx]
    best_xcorr_model_spec = load_phoenix_spectrum(best_xcorr_model,
                                                  limits=limits,
                                                  normalize=True)
    best_xcorr_model_spec = convolve_models([best_xcorr_model_spec],
                                            obs_resolution,
                                            chip_limits=None)[0]

    close_model_spec = load_phoenix_spectrum(closest_model[0],
                                             limits=limits,
                                             normalize=True)
    close_model_spec = convolve_models([close_model_spec],
                                       obs_resolution,
                                       chip_limits=None)[0]

    plt.plot(observed_spectra.xaxis,
             observed_spectra.flux,
             label="Observations")
    plt.plot(best_model_spec.xaxis, best_model_spec.flux, label="Best Model")
    plt.plot(best_xcorr_model_spec.xaxis,
             best_xcorr_model_spec.flux,
             label="Best xcorr Model")
    plt.plot(close_model_spec.xaxis,
             close_model_spec.flux,
             label="Close Model")
    plt.legend()
    plt.xlim(*limits)
    plt.show()
    print("After plot")
def tcm_wrapper(num,
                params1,
                model2_pars,
                alphas,
                rvs,
                gammas,
                obs_spec,
                errors=None,
                norm=True,
                verbose=False,
                save_only=True,
                chip=None,
                prefix=None,
                wav_scale=True,
                norm_method="scalar"):
    """Wrapper for iteration loop of tcm. params1 fixed, model2_pars are many."""
    normalization_limits = [2105, 2185]  # small as possible?

    if prefix is None:
        sf = os.path.join(
            simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper(),
            "tc_{0}_{1}-{2}_part{6}_host_pars_[{3}_{4}_{5}].csv".format(
                obs_spec.header["OBJECT"].upper(),
                int(obs_spec.header["MJD-OBS"]), chip, params1[0], params1[1],
                params1[2], num))
    else:
        sf = "{0}_part{4}_host_pars_[{1}_{2}_{3}].csv".format(
            prefix, params1[0], params1[1], params1[2], num)
    save_filename = sf

    if os.path.exists(save_filename) and save_only:
        print("''{}' exists, so not repeating calculation.".format(
            save_filename))
        return None
    else:
        if not save_only:
            broadcast_chisqr_vals = np.empty(len(model2_pars))
        for jj, params2 in enumerate(model2_pars):
            if verbose:
                print("Starting iteration with parameters:\n {0}={1},{2}={3}".
                      format(num, params1, jj, params2))

            mod1_spec = load_starfish_spectrum(params1,
                                               limits=normalization_limits,
                                               hdr=True,
                                               normalize=True,
                                               wav_scale=wav_scale)
            mod2_spec = load_starfish_spectrum(params2,
                                               limits=normalization_limits,
                                               hdr=True,
                                               normalize=True,
                                               wav_scale=wav_scale)

            # Wavelength selection
            rv_limits = observation_rv_limits(obs_spec, rvs, gammas)
            mod1_spec.wav_select(*rv_limits)
            mod2_spec.wav_select(*rv_limits)

            obs_spec = obs_spec.remove_nans()

            # One component model with broadcasting over gammas
            # two_comp_model(wav, model1, model2, alphas, rvs, gammas)
            assert np.allclose(mod1_spec.xaxis, mod2_spec.xaxis)

            broadcast_result = two_comp_model(mod1_spec.xaxis,
                                              mod1_spec.flux,
                                              mod2_spec.flux,
                                              alphas=alphas,
                                              rvs=rvs,
                                              gammas=gammas)
            broadcast_values = broadcast_result(obs_spec.xaxis)

            assert ~np.any(np.isnan(obs_spec.flux)), "Observation is nan"

            # RE-NORMALIZATION
            if chip == 4:
                # Quadratically renormalize anyway
                obs_spec = renormalization(obs_spec,
                                           broadcast_values,
                                           normalize=True,
                                           method="quadratic")
            obs_flux = renormalization(obs_spec,
                                       broadcast_values,
                                       normalize=norm,
                                       method=norm_method)

            # sp_chisquare is much faster but don't think I can add masking.
            broadcast_chisquare = chi_squared(obs_flux,
                                              broadcast_values,
                                              error=errors)
            # sp_chisquare = stats.chisquare(obs_flux, broadcast_values, axis=0).statistic
            # broadcast_chisquare = sp_chisquare

            if not save_only:
                print(broadcast_chisquare.shape)
                print(broadcast_chisquare.ravel()[np.argmin(
                    broadcast_chisquare)])

                broadcast_chisqr_vals[jj] = broadcast_chisquare.ravel()[
                    np.argmin(broadcast_chisquare)]
            npix = obs_flux.shape[0]
            save_full_tcm_chisqr(save_filename,
                                 params1,
                                 params2,
                                 alphas,
                                 rvs,
                                 gammas,
                                 broadcast_chisquare,
                                 npix,
                                 verbose=verbose)

        if save_only:
            return None
        else:
            return broadcast_chisqr_vals
def bhm_analysis(obs_spec, model_pars, gammas=None, errors=None, prefix=None, verbose=False, chip=None, norm=False,
                 wav_scale=True, norm_method="scalar"):
    """Run one component model over all parameter combinations in model_pars."""
    # Gammas
    if gammas is None:
        gammas = np.array([0])
    elif isinstance(gammas, (float, int)):
        gammas = np.asarray(gammas, dtype=np.float32)

    if isinstance(model_pars, list):
        logging.debug(__("Number of close model_pars returned {0}", len(model_pars)))

    # Solution Grids to return
    model_chisqr_vals = np.empty(len(model_pars))
    model_xcorr_vals = np.empty(len(model_pars))
    model_xcorr_rv_vals = np.empty(len(model_pars))
    bhm_grid_chisqr_vals = np.empty(len(model_pars))
    bhm_grid_gamma = np.empty(len(model_pars))
    full_bhm_grid_chisquare = np.empty((len(model_pars), len(gammas)))

    normalization_limits = [2105, 2185]  # small as possible?

    for ii, params in enumerate(tqdm(model_pars)):
        if prefix is None:
            save_name = os.path.join(
                simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper(), "bhm",
                "bhm_{0}_{1}_{3}_part{2}.csv".format(
                    obs_spec.header["OBJECT"].upper(), obs_spec.header["MJD-OBS"], ii, chip))
        else:
            save_name = os.path.join("{0}_part{1}.csv".format(prefix, ii))

        if verbose:
            print("Starting iteration with parameter:s\n{}".format(params))

        mod_spec = load_starfish_spectrum(params, limits=normalization_limits, hdr=True,
                                          normalize=True, wav_scale=wav_scale)

        # Wavelength selection
        mod_spec.wav_select(np.min(obs_spec.xaxis) - 5,
                            np.max(obs_spec.xaxis) + 5)  # +- 5nm of obs

        obs_spec = obs_spec.remove_nans()

        # One component model with broadcasting over gammas
        bhm_grid_func = one_comp_model(mod_spec.xaxis, mod_spec.flux, gammas=gammas)
        bhm_grid_values = bhm_grid_func(obs_spec.xaxis)

        assert ~np.any(np.isnan(obs_spec.flux)), "Observation is nan"

        # RENORMALIZATION
        if chip == 4:
            # Quadratically renormalize anyway
            obs_spec = renormalization(obs_spec, bhm_grid_values, normalize=True, method="quadratic")
        obs_flux = renormalization(obs_spec, bhm_grid_values, normalize=norm, method=norm_method)

        # Simple chi2
        bhm_grid_chisquare_old = chi_squared(obs_flux, bhm_grid_values, error=errors)

        # Applying arbitrary scalar normalization to continuum
        bhm_norm_grid_values, arb_norm = arbitrary_rescale(bhm_grid_values,
                                                           *simulators.sim_grid["arb_norm"])

        # Calculate Chi-squared
        obs_flux = np.expand_dims(obs_flux, -1)  # expand on last axis to match rescale
        bhm_norm_grid_chisquare = chi_squared(obs_flux, bhm_norm_grid_values, error=errors)

        # Take minimum chi-squared value along Arbitrary normalization axis
        bhm_grid_chisquare, arbitrary_norms = arbitrary_minimums(bhm_norm_grid_chisquare, arb_norm)

        assert np.any(
            bhm_grid_chisquare_old >= bhm_grid_chisquare), "All chi2 values are not better or same with arbitrary scaling"

        # Interpolate to obs
        mod_spec.spline_interpolate_to(obs_spec)
        org_model_chi_val = chi_squared(obs_spec.flux, mod_spec.flux)

        model_chisqr_vals[ii] = org_model_chi_val  # This is gamma = 0 version

        # New parameters to explore
        bhm_grid_chisqr_vals[ii] = bhm_grid_chisquare[np.argmin(bhm_grid_chisquare)]
        bhm_grid_gamma[ii] = gammas[np.argmin(bhm_grid_chisquare)]
        full_bhm_grid_chisquare[ii, :] = bhm_grid_chisquare

        ################
        #  Find cross correlation RV
        # Should run though all models and find best rv to apply uniformly
        rvoffset, cc_max = xcorr_peak(obs_spec, mod_spec, plot=False)
        if verbose:
            print("Cross correlation RV = {}".format(rvoffset))
            print("Cross correlation max = {}".format(cc_max))

        model_xcorr_vals[ii] = cc_max
        model_xcorr_rv_vals[ii] = rvoffset
        ###################

        npix = obs_flux.shape[0]
        # print("bhm shape", bhm_grid_chisquare.shape)
        save_full_bhm_chisqr(save_name, params, gammas, bhm_grid_chisquare, arbitrary_norms,
                             npix, rvoffset)

    return (model_chisqr_vals, model_xcorr_vals, model_xcorr_rv_vals,
            bhm_grid_chisqr_vals, bhm_grid_gamma, full_bhm_grid_chisquare)