diff --git a/pyproject.toml b/pyproject.toml index 70fb145..307fa61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,12 +39,14 @@ dependencies = [ "jax", "exoplanet", "tqdm", + "emcee", "numba", "scikit-learn>=1.3.0", "matplotlib", "pandas", "exoplanet-core", - "pymc>=4" + "pymc>=4", + "schwimmbad", ] dynamic = ["version"] diff --git a/src/scope/emcee_fit_hires.py b/src/scope/emcee_fit_hires.py index 09fa6c6..af2b14f 100644 --- a/src/scope/emcee_fit_hires.py +++ b/src/scope/emcee_fit_hires.py @@ -3,15 +3,35 @@ import emcee from schwimmbad import MPIPool # todo: add these as dependencies...? -from utils import * from scope.run_simulation import * do_pca = True np.random.seed(42) +# load the data +test_data_path = os.path.join(os.path.dirname(__file__), "../data") -def log_prob(x): + +def log_prob( + x, + best_kp, + wl_cube_model, + Fp_conv, + n_order, + n_exposure, + n_pixel, + A_noplanet, + star, + n_princ_comp, + flux_cube, + wl_model, + Fstar_conv, + Rp_solar, + Rstar, + phases, + do_pca, +): """ just add the log likelihood and the log prob. @@ -23,33 +43,40 @@ def log_prob(x): ------- :log_prob: (float) log probability. """ + rv_semiamp_orbit = 0.0 Kp, Vsys, log_scale = x scale = np.power(10, log_scale) - prior_val = prior(x) + prior_val = prior(x, best_kp) + if not np.isfinite(prior_val): return -np.inf - return ( - prior_val - + calc_log_likelihood( - Vsys, - Kp, - scale, - wl_cube_model, - fTemp, - Ndet, - n_exposure, - n_pixel, - A_noplanet=A_noplanet, - do_pca=do_pca, - NPC=n_princ_comp, - star=star, - )[0] - ) + ll = calc_log_likelihood( + Vsys, + Kp, + scale, + wl_cube_model, + wl_model, + Fp_conv, + Fstar_conv, + flux_cube, + n_order, + n_exposure, + n_pixel, + phases, + Rp_solar, + Rstar, + rv_semiamp_orbit, + A_noplanet, + do_pca=do_pca, + n_princ_comp=n_princ_comp, + star=star, + observation="transmission", + )[0] + return prior_val + ll -# van Sluijs+22 fit for log10 a. so do Brogi et al. # @numba.njit -def prior(x): +def prior(x, best_kp): """ Prior on the parameters. Only uniform! @@ -63,7 +90,11 @@ def prior(x): """ Kp, Vsys, log_scale = x # do I sample in log_scale? - if 146.0 < Kp < 246.0 and -50 < Vsys < 50 and -1 < log_scale < 1: + if ( + best_kp - 50.0 < Kp < best_kp + 50.0 + and -50.0 < Vsys < 50.0 + and -1 < log_scale < 1 + ): return 0 return -np.inf @@ -72,7 +103,19 @@ def sample( nchains, nsample, A_noplanet, - fTemp, + Fp_conv, + wl_cube_model, + n_order, + n_exposure, + n_pixel, + star, + n_princ_comp, + flux_cube, + wl_model, + Fstar_conv, + Rp_solar, + Rstar, + phases, do_pca=True, best_kp=192.06, best_vsys=0.0, @@ -111,13 +154,38 @@ def sample( nwalkers, ndim = pos.shape - sampler = emcee.EnsembleSampler(nwalkers, ndim, log_prob, pool=pool) + sampler = emcee.EnsembleSampler( + nwalkers, + ndim, + log_prob, + args=( + best_kp, + wl_cube_model, + Fp_conv, + n_order, + n_exposure, + n_pixel, + A_noplanet, + star, + n_princ_comp, + flux_cube, + wl_model, + Fstar_conv, + Rp_solar, + Rstar, + phases, + do_pca, + ), + pool=pool, + ) sampler.run_mcmc(pos, nsample, progress=True) return sampler if __name__ == "__main__": + # example below! + # ind = eval(sys.argv[1]) ind = 111 # SNR = 60, blaze, tell, star, 4 PCA components param_dict = parameter_list[ind] @@ -132,7 +200,6 @@ def sample( lls, ccfs = np.zeros((50, 50)), np.zeros((50, 50)) - # redoing the grid. how close does PCA get to a tellurics-free signal detection? A_noplanet, fTemp, fTemp_nopca, just_tellurics = make_data( 1.0, wl_cube_model, diff --git a/src/scope/run_simulation.py b/src/scope/run_simulation.py index 5463ca2..9f00c18 100644 --- a/src/scope/run_simulation.py +++ b/src/scope/run_simulation.py @@ -278,7 +278,7 @@ def make_data( ) # returning CCF and logL values -@njit +# @njit def calc_log_likelihood( v_sys, Kp, diff --git a/src/scope/tests/conftest.py b/src/scope/tests/conftest.py new file mode 100644 index 0000000..c66d69b --- /dev/null +++ b/src/scope/tests/conftest.py @@ -0,0 +1,100 @@ +import os +import sys +import pytest +import numpy as np +import pandas as pd + +from scope.run_simulation import * +from scope.utils import * + +# Define the path to the test data +test_data_path = os.path.join(os.path.dirname(__file__), "../data") + + +@pytest.fixture +def test_inputs(): + """ + sets up the spectral things. + """ + data_cube_path = os.path.join(test_data_path, "data_RAW_20201214_wl_algn_03.pic") + planet_spectrum_path = os.path.join( + test_data_path, "best_fit_spectrum_tran_gj1214_steam_nir.pic" + ) + wl_model, Fp, Fstar = np.load(planet_spectrum_path, allow_pickle=True) + mike_wave, mike_cube = pickle.load(open(data_cube_path, "rb"), encoding="latin1") + wl_cube_model = mike_wave.copy().astype(np.float64) + + v_rot = 4.33 # km/s + rot_ker = get_rot_ker(v_rot, wl_model) + Fp_conv_rot = np.convolve(Fp, rot_ker, mode="same") + # instrument profile convolution + instrument_kernel = get_instrument_kernel() + Fp_conv = np.convolve(Fp_conv_rot, instrument_kernel, mode="same") + star_spectrum_path = os.path.join(test_data_path, "PHOENIX_5605_4.33.txt") + star_wave, star_flux = np.loadtxt( + star_spectrum_path + ).T # Phoenix stellar model packing + Fstar_conv = get_star_spline( + star_wave, star_flux, wl_model, instrument_kernel, smooth=False + ) + + return Fp_conv, Fstar_conv, wl_cube_model, wl_model + + +@pytest.fixture +def test_baseline_outouts(test_inputs): + Fp_conv, Fstar_conv, wl_cube_model, wl_model = test_inputs + + scale = 1 + + n_exposure = 10 + n_order, n_pixel = (44, 1848) # todo: fix. + phases = np.linspace(-0.01, 0.01, n_exposure) + Rp_solar = 0.1 + Rstar = 1.0 + + # Test the function + A_noplanet, flux_cube, flux_cube_nopca, just_tellurics = make_data( + scale, + wl_cube_model, + wl_model, + Fp_conv, + Fstar_conv, + n_order, + n_exposure, + n_pixel, + phases, + Rp_solar, + Rstar, + do_pca=True, + blaze=True, + do_airmass_detrending=False, + tellurics=True, + n_princ_comp=4, + v_sys=0, + Kp=192.06, + star=True, + SNR=0, + rv_semiamp_orbit=0.3229, + observation="transmission", + tell_type="data-driven", + time_dep_tell=False, + wav_error=False, + order_dep_throughput=True, + a=1, # AU + u1=0.3, + u2=0.3, + LD=True, + b=0.0, # impact parameter + divide_out_of_transit=False, + out_of_transit_dur=0.1, + ) + return ( + A_noplanet, + flux_cube, + flux_cube_nopca, + just_tellurics, + n_exposure, + n_order, + n_pixel, + ) diff --git a/src/scope/tests/test_emcee_fit_hires.py b/src/scope/tests/test_emcee_fit_hires.py new file mode 100644 index 0000000..083f8ee --- /dev/null +++ b/src/scope/tests/test_emcee_fit_hires.py @@ -0,0 +1,90 @@ +import pytest + +from scope.emcee_fit_hires import * +from scope.tests.conftest import test_baseline_outouts, test_inputs + +test_data_path = os.path.join(os.path.dirname(__file__), "../data") + + +@pytest.mark.parametrize( + "values, output", + [ + ([150, 0, 0], 0.0), + ([2, 0, 0], -np.inf), + ([150, -200, 0], -np.inf), + ([150, 0, -100], -np.inf), + ([150, -200, -100], -np.inf), + ([2, 0, -100], -np.inf), + ([150, -200, 0], -np.inf), + ([2, -200, -100], -np.inf), + ], +) +def test_prior(values, output): + """ + Test the prior function. + """ + Kp, Vsys, log_scale = values + assert prior(values, best_kp=160.0) == output + + +def test_likelihood(test_baseline_outouts, test_inputs): + ( + A_noplanet, + flux_cube, + flux_cube_nopca, + just_tellurics, + n_exposure, + n_order, + n_pixel, + ) = test_baseline_outouts + star = True + Rp_solar = 0.1 + Rstar = 1.0 + phases = np.linspace(-0.01, 0.01, n_exposure) + Fp_conv, Fstar_conv, wl_cube_model, wl_model = test_inputs + best_kp = 192.06 + + x_good = [best_kp, 0, 0] + x_bad = [144, -50, -0.9] + n_princ_comp = 4 + log_prob_good = log_prob( + x_good, + best_kp, + wl_cube_model, + Fp_conv, + n_order, + n_exposure, + n_pixel, + A_noplanet, + star, + n_princ_comp, + flux_cube, + wl_model, + Fstar_conv, + Rp_solar, + Rstar, + phases, + do_pca, + ) + + log_prob_bad = log_prob( + x_bad, + best_kp, + wl_cube_model, + Fp_conv, + n_order, + n_exposure, + n_pixel, + A_noplanet, + star, + n_princ_comp, + flux_cube, + wl_model, + Fstar_conv, + Rp_solar, + Rstar, + phases, + do_pca, + ) + + assert log_prob_good > log_prob_bad diff --git a/src/scope/tests/test_run_simulation.py b/src/scope/tests/test_run_simulation.py index 1d0e3c3..74b8886 100644 --- a/src/scope/tests/test_run_simulation.py +++ b/src/scope/tests/test_run_simulation.py @@ -10,99 +10,13 @@ from scope.run_simulation import * from scope.utils import * +from scope.tests.conftest import test_inputs, test_baseline_outouts # Define the path to the test data test_data_path = os.path.join(os.path.dirname(__file__), "../data") # make a pytest fixture to store the test data -@pytest.fixture -def test_inputs(): - """ - sets up the spectral things. - """ - data_cube_path = os.path.join(test_data_path, "data_RAW_20201214_wl_algn_03.pic") - planet_spectrum_path = os.path.join( - test_data_path, "best_fit_spectrum_tran_gj1214_steam_nir.pic" - ) - wl_model, Fp, Fstar = np.load(planet_spectrum_path, allow_pickle=True) - mike_wave, mike_cube = pickle.load(open(data_cube_path, "rb"), encoding="latin1") - wl_cube_model = mike_wave.copy().astype(np.float64) - - v_rot = 4.33 # km/s - rot_ker = get_rot_ker(v_rot, wl_model) - Fp_conv_rot = np.convolve(Fp, rot_ker, mode="same") - # instrument profile convolution - instrument_kernel = get_instrument_kernel() - Fp_conv = np.convolve(Fp_conv_rot, instrument_kernel, mode="same") - star_spectrum_path = os.path.join(test_data_path, "PHOENIX_5605_4.33.txt") - star_wave, star_flux = np.loadtxt( - star_spectrum_path - ).T # Phoenix stellar model packing - Fstar_conv = get_star_spline( - star_wave, star_flux, wl_model, instrument_kernel, smooth=False - ) - - return Fp_conv, Fstar_conv, wl_cube_model, wl_model - - -@pytest.fixture -def test_baseline_outouts(test_inputs): - Fp_conv, Fstar_conv, wl_cube_model, wl_model = test_inputs - - scale = 1 - - n_exposure = 10 - n_order, n_pixel = (44, 1848) # todo: fix. - phases = np.linspace(-0.01, 0.01, n_exposure) - Rp_solar = 0.1 - Rstar = 1.0 - - # Test the function - A_noplanet, flux_cube, flux_cube_nopca, just_tellurics = make_data( - scale, - wl_cube_model, - wl_model, - Fp_conv, - Fstar_conv, - n_order, - n_exposure, - n_pixel, - phases, - Rp_solar, - Rstar, - do_pca=True, - blaze=True, - do_airmass_detrending=False, - tellurics=True, - n_princ_comp=4, - v_sys=0, - Kp=192.06, - star=True, - SNR=0, - rv_semiamp_orbit=0.3229, - observation="emission", - tell_type="data-driven", - time_dep_tell=False, - wav_error=False, - order_dep_throughput=True, - a=1, # AU - u1=0.3, - u2=0.3, - LD=True, - b=0.0, # impact parameter - divide_out_of_transit=False, - out_of_transit_dur=0.1, - ) - return ( - A_noplanet, - flux_cube, - flux_cube_nopca, - just_tellurics, - n_exposure, - n_order, - n_pixel, - ) @pytest.fixture @@ -317,7 +231,7 @@ def test_baseline_outputs_take2(test_inputs): star=True, SNR=0, rv_semiamp_orbit=0.3229, - observation="emission", + observation="transmission", tell_type="data-driven", time_dep_tell=False, wav_error=False,