<jupyter_start><jupyter_text>TESS atlas fit for TOI 144In this notebook, we fit...[Open in Google Colab](https://colab.research.google.com/github/dfm/tess-atlas/blob/master/notebooks/toi-144.ipynb)<jupyter_code>!pip uninstall -y lightkurve
!pip install -q https://github.com/nksaunders/lightkurve/archive/pld-pca-method.zip
!pip install -q fbpca exoplanet corner
%matplotlib inline
import matplotlib.pyplot as plt
plt.style.use('default')

from matplotlib import rcParams
rcParams["savefig.dpi"] = 100
rcParams["figure.dpi"] = 100
rcParams["font.size"] = 16
rcParams["text.usetex"] = False
rcParams["font.family"] = ["sans-serif"]
rcParams["font.sans-serif"] = ["cmss10"]
rcParams["axes.unicode_minus"] = False

import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=FutureWarning)

import corner
import numpy as np
import pandas as pd
import lightkurve as lk
import matplotlib.pyplot as plt

import pymc3 as pm
import exoplanet as xo
import theano.tensor as tt
toi_num = 144
tois = pd.read_csv("https://exofop.ipac.caltech.edu/tess/download_toi.php?sort=toi&output=csv")

toi = tois[tois["TOI"] == toi_num + 0.01].iloc[0]
tic = toi['TIC ID']
tois = tois[tois["TIC ID"] == tic]
periods = np.array(tois["Period (days)"], dtype=float)
epochs = np.array(tois["Epoch (BJD)"], dtype=float)
depths = np.array(tois["Depth (ppm)"], dtype=float)
durations = np.array(tois["Duration (hours)"], dtype=float) / 24.0

# Stellar properties
toi_r_star = toi['Stellar Radius (R_Sun)']
toi_r_star_err = toi['Stellar Radius (R_Sun) err']
has_r_star = True
if not (np.isfinite(toi_r_star) and np.isfinite(toi_r_star_err)):
    has_r_star = False
    
letters = "bcdefghijklmnopqrstuvwxyz"[:len(periods)]
# Download fits files
sr = lk.search_targetpixelfile('TIC %i' % tic)
tpf_collection = sr.download_all()

def get_transit_mask(t, transit_time, period, duration=0.2):
    hp = 0.5*period
    return np.abs((t-transit_time+hp) % period - hp) < 0.5*duration

# Run PLD on each TPF
lc_collection = []
for tpf in tpf_collection:
    mask = np.ones_like(tpf.time, dtype=bool)
    for i in range(len(periods)):
        mask &= get_transit_mask(tpf.time, epochs[i], periods[i], duration=1.5*durations[i])
    pld = tpf.to_corrector('pld')
    lc = pld.correct(cadence_mask=~mask, use_gp=False, pld_order=2)
    lc_collection.append(lc)

# Normalize and stitch
lc = lc_collection[0].normalize()
if len(lc_collection) > 1:
    lc = lc.append([next_lc.normalize() for next_lc in lc_collection[1:]])
lc = lc.remove_outliers()

lc = lc.flatten(window_length=901)
    
lc.scatter();<jupyter_output><empty_output><jupyter_text>Some more info about the modeling...<jupyter_code>def build_model(x, y, yerr, periods, t0s, depths, mask=None, start=None):
    """Build an exoplanet model for a dataset and set of planets
    
    Args:
        x: The time series (in days); this should probably be centered
        y: The relative fluxes (in parts per thousand)
        yerr: The uncertainties on ``y``
        periods: The periods of the planets (in days)
        t0s: The phases of the planets in the same coordinates as ``x``
        depths: The depths of the transits in parts per thousand
        
    """
    if mask is None:
        mask = np.ones(len(x), dtype=bool)
    
    periods = np.atleast_1d(periods)
    t0s = np.atleast_1d(t0s)
    depths = np.atleast_1d(depths)
    n_planets = len(periods)
    
    with pm.Model() as model:
        
        model.x = x[mask]
        model.y = y[mask]
        model.yerr = (yerr + np.zeros_like(x))[mask]
        model.mask = mask

        # The baseline flux
        mean = pm.Normal("mean", mu=0.0, sd=10.0)

        # The time of a reference transit for each planet
        t0 = pm.Normal("t0", mu=t0s, sd=1.0, shape=n_planets)

        # The log period; also tracking the period itself
        logP = pm.Normal("logP", mu=np.log(periods), sd=0.1, shape=n_planets)
        period = pm.Deterministic("period", tt.exp(logP))

        # The Kipping (2013) parameterization for quadratic limb darkening paramters
        u = xo.distributions.QuadLimbDark("u")

        # The Espinoza (2018) parameterization for the joint radius ratio and
        # impact parameter distribution
        r, b = xo.distributions.get_joint_radius_impact(
            min_radius=0.001, max_radius=0.5,
            testval_r=np.sqrt(1e-3*np.array(depths)),
            testval_b=0.5+np.zeros(n_planets)
        )

        # This shouldn't make a huge difference, but I like to put a uniform
        # prior on the *log* of the radius ratio instead of the value. This
        # can be implemented by adding a custom "potential" (log probability).
        pm.Potential("r_prior", -pm.math.log(r))

        # Set up a Keplerian orbit for the planets
        model.orbit = xo.orbits.KeplerianOrbit(
            period=period, t0=t0, b=b)
        
        # Compute the model light curve using starry
        model.light_curves = xo.StarryLightCurve(u).get_light_curve(
            orbit=model.orbit, r=r, t=model.x)
        model.light_curve = pm.math.sum(model.light_curves, axis=-1) * 1e3 + mean

        # Jitter & GP parameters
        logs2 = pm.Normal("logs2", mu=np.log(np.var(model.y)), sd=10)
        pm.Normal("obs", mu=model.light_curve,
                  sd=tt.sqrt(model.yerr**2 + tt.exp(logs2)),
                  observed=model.y)

        # Fit for the maximum a posteriori parameters, I've found that I can get
        # a better solution by trying different combinations of parameters in turn
        if start is None:
            start = model.test_point
        map_soln = start
        map_soln = xo.optimize(start=map_soln, vars=[logs2, mean])
        map_soln = xo.optimize(start=map_soln, vars=[model.rb, mean])
        map_soln = xo.optimize(start=map_soln, vars=[logP, t0, mean])
        map_soln = xo.optimize(start=map_soln)
        model.map_soln = map_soln
        
    return model

def build_model_sigma_clip(x, y, yerr, periods, t0s, depths, sigma=5.0, maxiter=10, start=None):
    ntot = len(x)
    for i in range(maxiter):
        # Build the model
        model = build_model(x, y, yerr, periods, t0s, depths, start=start)
        start = model.map_soln

        # Compute the map prediction
        with model:
            mod = xo.utils.eval_in_model(model.light_curve, model.map_soln)
            
        # Do sigma clipping
        resid = y - mod
        rms = np.sqrt(np.median(resid**2))
        mask = np.abs(resid) < sigma * rms
        print(ntot, mask.sum())
        if ntot == mask.sum():
            break
        ntot = mask.sum()

    return model<jupyter_output><empty_output><jupyter_text>Setup and fit for map model<jupyter_code># periods = periods
t0s = epochs - 2457000
depths = depths * 1e-3

x = lc.time
y = (lc.flux - 1.0) * 1e3
yerr = lc.flux_err * 1e3

model = build_model_sigma_clip(x, y, yerr, periods, t0s, depths)
with model:
    mean = model.map_soln["mean"]
    light_curves = xo.utils.eval_in_model(model.light_curves, model.map_soln)

plt.plot(model.x, model.y - mean, "k", label="data")
for n, l in enumerate(letters):
    plt.plot(model.x, 1e3 * light_curves[:, n], label="planet {0}".format(l))

plt.xlabel("time [days]")
plt.ylabel("flux [ppt]")
plt.title("initial fit")
plt.xlim(model.x.min(), model.x.max())
plt.legend(fontsize=10);
np.random.seed(123)
sampler = xo.PyMC3Sampler(window=50, start=50, finish=500)
with model:
    burnin = sampler.tune(tune=3000, start=model.map_soln,
                          step_kwargs=dict(target_accept=0.9),
                          chains=2)
with model:
    trace = sampler.sample(draws=1000, chains=2)
pm.summary(trace, varnames=["mean", "u", "period", "t0", "r", "b"])
with model:
    light_curves = np.empty((500, len(model.x), len(periods)))
    func = xo.utils.get_theano_function_for_var(model.light_curves)
    for i, sample in enumerate(xo.utils.get_samples_from_trace(
            trace, size=len(light_curves))):
        light_curves[i] = func(*xo.utils.get_args_for_theano_function(sample))

for n, letter in enumerate(letters):
    plt.figure()

    # Compute the GP prediction
    mean_mod = np.median(trace["mean"][:, None])

    # Get the posterior median orbital parameters
    p = np.median(trace["period"][:, n])
    t0 = np.median(trace["t0"][:, n])

    # Compute the median of posterior estimate of the contribution from
    # the other planet. Then we can remove this from the data to plot
    # just the planet we care about.
    inds = np.arange(len(periods)) != n
    others = np.median(1e3*np.sum(light_curves[:, :, inds], axis=-1), axis=0)

    # Plot the folded data
    x_fold = (model.x - t0 + 0.5*p) % p - 0.5*p
    plt.plot(x_fold, model.y - mean_mod - others, ".k", label="data", zorder=-1000)

    # Plot the folded model
    inds = np.argsort(x_fold)
    inds = inds[np.abs(x_fold)[inds] < 0.3]
    pred = 1e3 * light_curves[:, inds, n]
    pred = np.percentile(pred, [16, 50, 84], axis=0)
    plt.plot(x_fold[inds], pred[1], color="C1", label="model")
    art = plt.fill_between(x_fold[inds], pred[0], pred[2], color="C1", alpha=0.5,
                           zorder=1000)
    art.set_edgecolor("none")

    # Annotate the plot with the planet's period
    txt = "period = {0:.4f} +/- {1:.4f} d".format(
        np.mean(trace["period"][:, n]), np.std(trace["period"][:, n]))
    plt.annotate(txt, (0, 0), xycoords="axes fraction",
                 xytext=(5, 5), textcoords="offset points",
                 ha="left", va="bottom", fontsize=12)

    plt.legend(fontsize=10, loc=4)
    plt.xlim(-0.5*p, 0.5*p)
    plt.xlabel("time since transit [days]")
    plt.ylabel("de-trended flux")
    plt.title("TOI {0}{1}".format(toi_num, letter));
    plt.xlim(-0.3, 0.3)
ror_samps = trace["r"]
r_star_samps = toi_r_star + toi_r_star_err * np.random.randn(len(ror_samps))
r_pl = ror_samps * r_star_samps[:, None] * 109.07637070600963
samples = np.concatenate((r_pl, trace["b"]), axis=-1)

labels = ["$R_{{\mathrm{{Pl}},{0}}}$ [$R_\oplus$]".format(i) for i in letters]
labels += ["impact param {0}".format(i) for i in letters]

corner.corner(samples, labels=labels);
labels = ["$P_{{{0}}}$ [days]".format(i) for i in letters]
labels += ["$t0_{{{0}}}$ [TBJD]".format(i) for i in letters]
samples = np.concatenate((trace["period"], trace["t0"]), axis=-1)
corner.corner(samples, labels=labels);
samples = trace["u"]
labels = ["$u_1$", "$u_2$"]
corner.corner(samples, labels=labels);<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1433599695585997, 'fraction_numerical': 0.022355403348554032, 'mean_word_length': 3.305077805077805, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 14, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '416302', 'license': 'MIT', 'n_tokens_mistral': 4021, 'n_tokens_neox': 3807, 'n_words': 891, 'path': 'notebooks/toi-144.ipynb', 'repo_name': 'barentsen/tess-atlas', 'size': '10512'}