working towards uniform analysis of five targets

bmorris3 · Jan 7, 2022 · a6067ad · a6067ad
1 parent fd0a1af
commit a6067ad
Show file tree

Hide file tree

Showing 10 changed files with 4,106 additions and 52 deletions.
diff --git a/jaxon/data/nea_transiting.csv b/jaxon/data/nea_transiting.csv
diff --git a/jaxon/data/rotation_table.pkl b/jaxon/data/rotation_table.pkl
diff --git a/jaxon/lightcurve.py b/jaxon/lightcurve.py
@@ -1,5 +1,6 @@
 import numpy as np
-from lightkurve import search_lightcurve
+from lightkurve import search_lightcurve, LightCurveCollection
+from lightkurve.correctors import CBVCorrector
 from astropy.stats import sigma_clip, mad_std
 import astropy.units as u
 from kelp import Filter
@@ -8,7 +9,8 @@
 import pymc3_ext as pmx
 
 from .utils import floatX
-from .hatp7 import get_planet_params
+# from .hatp7 import get_planet_params
+from .planets import get_planet_params
 
 __all__ = [
     'eclipse_model',
@@ -18,35 +20,37 @@
 
 cadence = "long"
 cadence_duration = 30 * u.min
+cbv_type = ['SingleScale']
+cbv_indices = [np.arange(1, 9)]
 
-
-def get_light_curve(quarter=None, cadence=cadence):
+def get_light_curve(planet_name, quarter=None, cadence=cadence):
     """
     Parameters
     ----------
     cadence : str {'long', 'short'}
         Kepler cadence mode
     """
     (planet_name, a_rs, a_rp, T_s, rprs, t0, period, eclipse_half_dur, b,
-        rstar, rho_star, rp_rstar, mstar, mass) = get_planet_params()
+        rstar, rho_star, rp_rstar, mstar, mass) = get_planet_params(
+        planet_name
+    )
 
     lcf = search_lightcurve(
         planet_name, mission="Kepler", cadence=cadence,
         quarter=quarter
-    ).download_all()
+    ).download_all(flux_column='sap_flux')
 
-    slc = lcf.stitch()
+    corrected_lcs = []
+    for each_quarter_lc in lcf:
+        cbvCorrector = CBVCorrector(each_quarter_lc)
+        # Perform the correction
+        cbvCorrector.correct_gaussian_prior(cbv_type=cbv_type,
+                                            cbv_indices=cbv_indices,
+                                            alpha=1e-4)
+        corrected_lcs.append(cbvCorrector.corrected_lc)
 
-    phases = ((slc.time.jd - t0) % period) / period
-    in_transit = (
-        (phases < 1.5 * eclipse_half_dur) |
-        (phases > 1 - 1.5 * eclipse_half_dur)
-    )
-    out_of_transit = np.logical_not(in_transit)
-
-    slc = slc.flatten(
-        polyorder=3, break_tolerance=10, window_length=1001, mask=~out_of_transit
-    ).remove_nans()
+    collection = LightCurveCollection(corrected_lcs)
+    slc = collection.stitch().remove_nans()
 
     phases = ((slc.time.jd - t0) % period) / period
     in_transit = (
@@ -57,7 +61,7 @@ def get_light_curve(quarter=None, cadence=cadence):
 
     sc = sigma_clip(
         np.ascontiguousarray(slc.flux[out_of_transit], dtype=floatX),
-        maxiters=100, sigma=8, stdfunc=mad_std
+        maxiters=100, sigma=5, stdfunc=mad_std
     )
 
     phase = np.ascontiguousarray(
@@ -94,7 +98,9 @@ def get_filter():
     return filt_wavelength, filt_trans
 
 
-def eclipse_model(quarter=None, cadence_duration=cadence_duration):
+def eclipse_model(
+        planet_name, quarter=None, cadence_duration=cadence_duration
+):
     """
     Compute the (static) eclipse model
 
@@ -110,8 +116,12 @@ def eclipse_model(quarter=None, cadence_duration=cadence_duration):
         in-eclipse.
     """
     (planet_name, a_rs, a_rp, T_s, rprs, t0, period, eclipse_half_dur, b,
-        rstar, rho_star, rp_rstar, mstar, mass) = get_planet_params()
-    phase, time, flux_normed, flux_normed_err = get_light_curve(quarter=quarter)
+        rstar, rho_star, rp_rstar, mstar, mass) = get_planet_params(
+        planet_name
+    )
+    phase, time, flux_normed, flux_normed_err = get_light_curve(
+        planet_name, quarter=quarter
+    )
 
     with pm.Model():
         # Define a Keplerian orbit using `exoplanet`:

diff --git a/jaxon/model.py b/jaxon/model.py
@@ -172,8 +172,7 @@ def model(
     theta = jnp.linspace(0, np.pi, n_theta, dtype=floatX)
     theta2d, phi2d = jnp.meshgrid(theta, phi)
 
-    # ln_C_11_kepler = -2.6
-    C_11_kepler = 0.35 # numpyro.sample('C_11', dist.Uniform(low=0, high=0.55))
+    C_11_kepler = 0.2 #numpyro.sample('C_11', dist.Uniform(low=0, high=0.55))
     # hml_eps = numpyro.sample('epsilon', dist.Uniform(low=0, high=8 / 5))
     hml_f = 0.73 #(2 / 3 - hml_eps * 5 / 12) ** 0.25
     delta_phi = 0 #numpyro.sample(
@@ -210,10 +209,18 @@ def model(
     sigma = numpyro.sample(
         "sigma", dist.TwoSidedTruncatedDistribution(
             dist.Normal(loc=y.std(), scale=y.std()/10),
-            low=0, high=1000 * y.ptp()
+            low=0, high=10 * y.std()
         )
     )
-    kernel = terms.Matern32Term(sigma=sigma, rho=22)
+
+    # rho = numpyro.sample(
+    #     "rho", dist.TwoSidedTruncatedDistribution(
+    #         dist.Normal(loc=22, scale=5),
+    #         low=6, high=50
+    #     )
+    # )
+
+    kernel = terms.Matern32Term(sigma=sigma, rho=30)
     jitter = numpyro.sample('jitter', dist.Uniform(low=0, high=100))
     gp = GaussianProcess(kernel, mean=flux_norm)
     gp.compute(time, yerr=jnp.sqrt(yerr ** 2 + jitter ** 2),
@@ -250,13 +257,13 @@ def model(
     numpyro.deterministic("interp_depths", interp_depths)
     numpyro.deterministic("Tarr", Tarr)
 
-    kepler_thermal_eclipse_depth_obs = numpyro.deterministic(
-        "kep_depth", jnp.interp(0, xi_grid, thermal_grid)
-    )
-    kepler_thermal_eclipse_depth_model = jnp.average(
-        fpfs_spectrum,
-        weights=(jnp.abs(kepler_mean_wl[0] - wav / 1000) < 0.3).astype(int)
-    )
+    # kepler_thermal_eclipse_depth_obs = numpyro.deterministic(
+    #     "kep_depth", jnp.interp(0, xi_grid, thermal_grid)
+    # )
+    # kepler_thermal_eclipse_depth_model = jnp.average(
+    #     fpfs_spectrum,
+    #     weights=(jnp.abs(kepler_mean_wl[0] - wav / 1000) < 0.15).astype(int)
+    # )
 
     numpyro.sample('phase_curve', gp.numpyro_dist(), obs=y)
     numpyro.factor('spectrum',
@@ -266,15 +273,15 @@ def model(
         ).log_prob(all_depths)
     )
 
-    kepler_thermal_eclipse_depth_err = numpyro.sample(
-        "kep_depth_err", dist.Uniform(low=1e-6, high=100e-6)
-    )
-    numpyro.factor(
-        "obs_spectrum_kepler", dist.Normal(
-            loc=kepler_thermal_eclipse_depth_model,
-            scale=kepler_thermal_eclipse_depth_err
-        ).log_prob(kepler_thermal_eclipse_depth_obs)
-    )
+    # kepler_thermal_eclipse_depth_err = 10e-6  # numpyro.sample(
+    # #     "kep_depth_err", dist.Uniform(low=1e-6, high=500e-6)
+    # # )
+    # numpyro.factor(
+    #     "obs_spectrum_kepler", dist.Normal(
+    #         loc=kepler_thermal_eclipse_depth_model,
+    #         scale=kepler_thermal_eclipse_depth_err
+    #     ).log_prob(kepler_thermal_eclipse_depth_obs)
+    # )
 
 
 def get_model_kwargs(quarter=None):