Merge pull request #55 from arjunsavel/test_rm_effect

test emcee hires a bit

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"]

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,

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,

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,
+    )

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