Fit a transit of an oblate planet#
Here we’ll attempt to recover the parameters describing planetary oblateness from a realistic-ish simulated light curve. We’ll focus on Kepler-167 e, a cold Jupiter-like planet that will be observed by JWST NIRSpec in Fall 2024.
[1]:
import jax
jax.config.update("jax_enable_x64", True)
import warnings
import astropy.units as u
import corner
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from exotic_ld import StellarLimbDarkening
from exotic_ld.ld_laws import nonlinear_4param_ld_law, quadratic_ld_law
from nautilus import Prior, Sampler
from scipy import stats
from scipy.interpolate import CubicSpline
from scipy.optimize import minimize
from tqdm import tqdm
from squishyplanet import OblateSystem
np.random.seed(13)
/Users/cassese/.cache/uv/archive-v0/XcIReEEbwKltJmbRLt3mX/lib/python3.13/site-packages/exotic_ld/ld_grids.py:4: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
import pkg_resources
Simulating the transit#
Limb darkening#
Since limb darkening varies most during ingress/egress, the same time our expected oblateness signal will be strongest, it’s important that we try to simulate it as accurately as possible in this injection/recovery exercise. There are many ways to approach fitting limb darkening, and we’ll get to that later: here our goal is just simulation.
We’ll use the ExoTiC-LD (Grand and Wakeford 2022) package’s ability to convert stellar grids into limb darkening profiles. This is a great package and is capable of quickly fitting these profiles to some of the commonly used approximate limb darkening laws, like quadratic and four-parameter non-linear. However, we won’t actually use this last fitting step: instead,
we’ll take advantage of squishyplanet’s ability to model high-order polynomials for limb darkening and directly fit the profiles ourselves.
We’ll use the ATLAS library of model stellar atmosphere (Kostogryz et al. 2022) which ExoTiC-LD calls “mps1”. ExoTiC-LD will download the relevant models automatically.
First, let’s take a look at the spectrum and expected limb darkening profile for Kepler-167 e. ExoTiC-LD works in units of \(\mu=\sqrt(1-r^2)\), where \(r\) is the projected distance to the center of the star.
[2]:
sld = StellarLimbDarkening(
M_H=0.020,
Teff=4884,
logg=4.579,
ld_model="mps1",
ld_data_path="exotic_ld_data",
interpolate_type="trilinear",
custom_wavelengths=None,
custom_mus=None,
custom_stellar_model=None,
ld_data_version="3.2.0",
verbose=2,
)
quad_coeffs = sld.compute_quadratic_ld_coeffs(
wavelength_range=[6_000, 53_000], # the bounds in angstroms of the Prism throughput
mode="JWST_NIRSpec_Prism",
custom_wavelengths=None,
custom_throughput=None,
mu_min=0.1,
)
print()
four_coeffs = sld.compute_4_parameter_non_linear_ld_coeffs(
wavelength_range=[6_000, 53_000],
mode="JWST_NIRSpec_Prism",
custom_wavelengths=None,
custom_throughput=None,
mu_min=0.1,
)
# from the ExoTIC-LD tutorials:
fig, ax = plt.subplots()
for mu_idx in np.arange(0, sld.stellar_intensities.shape[1], 5):
ax.plot(
sld.stellar_wavelengths,
sld.stellar_intensities[:, mu_idx],
label=rf"$\mu$={sld.mus[mu_idx]:.2f}",
)
ax.legend()
ax.set(
xlim=(0, 53_000),
xlabel=r"wavelength [$\AA$]",
ylabel=r"intensity [/ $n_{\gamma} s^{-1} cm^{-2} \AA{-1} sr^{-1}$]",
)
mu_grid = np.linspace(0.0, 1.0, 1_000) # for plotting the LD laws
quad_prediction = quadratic_ld_law(mu_grid, *quad_coeffs)
four_prediction = nonlinear_4param_ld_law(mu_grid, *four_coeffs)
fig, ax = plt.subplots()
ax.plot(mu_grid, quad_prediction, label="quadratic model")
ax.plot(mu_grid, four_prediction, label="4-parameter model")
ax.scatter(sld.mus, sld.I_mu, label="grid evaluation", color="k")
ax.set(xlabel=r"$\mu$", ylabel="normalized intensity")
ax.legend();
Input stellar parameters are M_H=0.02, Teff=4884, logg=4.579.
Loading stellar model from mps1 grid.
Using interpolation type = trilinear.
Trilinear interpolation within M_H=0.0-0.05, Teff=4800.0-4900.0, logg=4.5-4.6.
Stellar model loaded.
Loading instrument mode=JWST_NIRSpec_Prism with wavelength range 6000.0-53000.0 A.
Integrating I(mu) for wavelength limits of 6000-53000 A.
Integral done for I(mu).
Fitting limb-darkening law to 16 I(mu) data points where 0.1 <= mu <= 1, with the Levenberg-Marquardt algorithm.
Fit done, resulting coefficients are [0.21401188 0.23248826].
Loading instrument mode=JWST_NIRSpec_Prism with wavelength range 6000.0-53000.0 A.
Integrating I(mu) for wavelength limits of 6000-53000 A.
Integral done for I(mu).
Fitting limb-darkening law to 16 I(mu) data points where 0.1 <= mu <= 1, with the Levenberg-Marquardt algorithm.
Fit done, resulting coefficients are [ 0.50297679 0.12608929 -0.06546206 -0.00304464].
It’s clear that the non-linear law does a much better job than the standard quadratic approximation: unfortunately, squishyplanet can’t model limb darkening in this form. What it can do is model limb darkening as a high-order polynomial, which should be able to capture the non-linear behavior. Since we have 24 points in the profile, we could fit up to 23rd-order, but we’ll limit ourselves to something lower to avoid overfitting issues.
We can perform the fit with squishyplanet’s OblateSystem.fit_limb_darkening_profile method. We can then visualize the fit with squishyplanet’s OblateSystem.limb_darkening_profile method.
[3]:
# compute the best-fit polynomial model
u_coeffs = OblateSystem.fit_limb_darkening_profile(
intensities=sld.I_mu, mus=sld.mus, order=8
)
# evaluate that model on the mu grid
poly_model = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs, mu=mu_grid)
# ExoTIC-LD scales the intensity profile by the intensity at mu=0,
# while squishyplanet scales it so the integrated intensity is 1.0. rescale
# ours to match theirs:
norm = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs, r=0)
poly_model /= norm
# same game, higher order polynomial
u_coeffs2 = OblateSystem.fit_limb_darkening_profile(
intensities=sld.I_mu, mus=sld.mus, order=12
)
poly_model2 = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs2, mu=mu_grid)
norm2 = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs2, r=0)
poly_model2 /= norm2
fig, ax = plt.subplots()
ax.plot(mu_grid, quad_prediction, label="quadratic model")
ax.plot(mu_grid, four_prediction, label="4-parameter model")
ax.plot(mu_grid, poly_model, label="8th-order poly model")
ax.plot(mu_grid, poly_model2, label="12th-order poly model")
ax.scatter(sld.mus, sld.I_mu, label="grid evaluation", color="k")
ax.set(xlabel=r"$\mu$", ylabel="normalized intensity")
ax.legend();
The 8th-order is a much better than the quadratic approximation! But, as you can see with the 12th order fit, polynomials are really sensitive to overfitting and doing strange things between the evaluated grid points. To avoid this, let’s fit a cubic spline to the profile, evaluate it on a finer grid, then fit a polynomial to that instead:
[4]:
# interpolate between the grid points
order = np.argsort(sld.mus)
spline_approx = CubicSpline(x=sld.mus[order], y=sld.I_mu[order])
interpolated_vals = spline_approx(mu_grid)
# exact same as above, but now fitting to the interpolated values
u_coeffs = OblateSystem.fit_limb_darkening_profile(
intensities=interpolated_vals, mus=mu_grid, order=8
)
poly_model = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs, mu=mu_grid)
norm = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs, r=0)
poly_model /= norm
u_coeffs2 = OblateSystem.fit_limb_darkening_profile(
intensities=interpolated_vals, mus=mu_grid, order=12
)
poly_model2 = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs2, mu=mu_grid)
norm2 = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs2, r=0)
poly_model2 /= norm2
fig, ax = plt.subplots()
ax.plot(mu_grid, quad_prediction, label="quadratic model")
ax.plot(mu_grid, four_prediction, label="4-parameter model")
ax.plot(mu_grid, poly_model, label="8th-order poly model")
ax.plot(mu_grid, poly_model2, label="12th-order poly model")
ax.scatter(sld.mus, sld.I_mu, label="grid evaluation", color="k")
ax.set(xlabel=r"$\mu$", ylabel="normalized intensity")
ax.legend();
Looking much better!
As an aside, although it looks as thought the quadratic is a pretty bad fit, it’s worth remembering that \(\mu\) is defined as \(\sqrt(1-r^2)\), where \(r\) is the projected distance to the center of the star. A planet’s motion will be more linear in \(r\) than \(\mu\) (exactly linear if it has a non-zero impact parameter), so let’s see how the fits look when reparameterized:
[5]:
# convert mu->r
r_grid = np.sqrt(1 - mu_grid**2)
# evaluate our models on the r grid
poly_model1 = OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs, r=r_grid) / norm
poly_model2 = (
OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs2, r=r_grid) / norm2
)
fig, ax = plt.subplots(nrows=2, sharex=True, figsize=(8, 6))
ax[0].plot(r_grid, quad_prediction, label="quadratic")
ax[0].plot(r_grid, four_prediction, label="4-parameter")
ax[0].plot(r_grid, poly_model, label="8th-order poly")
ax[0].plot(r_grid, poly_model2, label="12th-order poly")
ax[0].scatter(
np.sqrt(1 - sld.mus**2), sld.I_mu, label="grid evaluation", color="k", zorder=10
)
ax[0].set(ylabel="normalized intensity")
ax[0].legend()
rs = np.sqrt(1 - sld.mus**2)
ax[1].scatter(rs, sld.I_mu - quadratic_ld_law(sld.mus, *quad_coeffs), label="quadratic")
ax[1].scatter(
rs, sld.I_mu - nonlinear_4param_ld_law(sld.mus, *four_coeffs), label="4-parameter"
)
ax[1].scatter(
rs,
sld.I_mu - OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs, r=rs) / norm,
label="8th-order poly",
)
ax[1].scatter(
rs,
sld.I_mu - OblateSystem.limb_darkening_profile(ld_u_coeffs=u_coeffs2, r=rs) / norm2,
label="12th-order poly",
)
ax[1].set(xlabel="r", ylabel="residuals")
ax[1].legend()
plt.subplots_adjust(hspace=0);
So, quadratic really only breaks down at the very edge of the star and mostly agrees with the others for \(\sim98\)% of the radial positions, which for most use cases is fine. Still, that can have an impact when you’re looking for subtle effects that manifest at those exact moments. We can quantify the effect of slightly mis-specified limb darkening on a transit by creating light curves with the different order polynomial approximations:
[6]:
t_exp = 5 * u.min
times = jnp.arange(-15, 15, t_exp.to(u.hour).value) * u.hour.to(u.day)
# comparison curve with high-order polynomial
u_coeffs = OblateSystem.fit_limb_darkening_profile(
intensities=interpolated_vals, mus=mu_grid, order=14
)
state = {
"t_peri": -250.0,
"times": times,
"a": 540.0,
"period": 1_000.0,
"r": (1 * u.R_jup).to(u.R_sun).value,
"ld_u_coeffs": u_coeffs,
"tidally_locked": False,
}
limb_darkening_demo_system = OblateSystem(**state)
lc_base = limb_darkening_demo_system.lightcurve()
lc_base2 = limb_darkening_demo_system.lightcurve(
{"r": (1 * u.R_earth).to(u.R_sun).value}
)
# loop through different orders of the polynomial
fig, ax = plt.subplots(ncols=2, figsize=(12, 4))
for order in tqdm(np.arange(2, 14, 1)):
u_coeffs = OblateSystem.fit_limb_darkening_profile(
intensities=interpolated_vals, mus=mu_grid, order=order
)
# need to create a new system each time, since the order of the limb darkening
# law sets the Green's basis transformation matrix which is only computed on
# initialization
state = {
"t_peri": -250.0,
"times": times,
"a": 540.0,
"period": 1_000.0,
"r": (1 * u.R_jup).to(u.R_sun).value,
"ld_u_coeffs": u_coeffs,
"tidally_locked": False,
}
limb_darkening_demo_system = OblateSystem(**state)
lc = limb_darkening_demo_system.lightcurve()
ax[0].plot(times * u.day.to(u.hour), (lc_base - lc) * 1e6, label=f"{order}th order")
lc2 = limb_darkening_demo_system.lightcurve(
{"r": (1 * u.R_earth).to(u.R_sun).value}
)
ax[1].plot(
times * u.day.to(u.hour), (lc_base2 - lc2) * 1e6, label=f"{order}th order"
)
ax[0].set(
xlabel="time [hours]",
ylabel="difference w/ 13th order [ppm]",
title="effect on Jupiter-sized planet",
)
ax[0].legend(loc="lower center", ncol=3, prop={"size": 10})
ax[1].set(
xlabel="time [hours]",
ylabel="difference w/ 13th order [ppm]",
title="effect on Earth-sized planet",
)
ax[1].legend(loc="lower center", ncol=3, prop={"size": 10})
plt.tight_layout();
100%|██████████| 12/12 [00:10<00:00, 1.18it/s]
Unsurprisingly, the larger the planet, the more subtle changes to the limb darkening law matters. A Jupiter-sized planet covers more radial positions than an Earth-sized planet, so changes between laws start to stand out more. Keep in mind, the 14th order polynomial isn’t “correct” either, it’s just our comparison curve. It’s a bummer that most oblate planets will be closer to Jupiter-sized than Earth-sized, but there’s not much we can do about that. We’ll just to use a fitting procedure that appropriately handles trade-offs between limb darkening and oblateness.
Moving ahead, we’ll use the 14th-order coefficients to simulate the transit. With those set, let’s go ahead and simulate the “ideal” transit. For this, we’ll use parameters reported in Chachan et al. 2020, then assume the planet is as oblate as and has the same obliquity of Saturn, and a randomly-chosen precession angle of 60 degrees.
The JWST observations will have a cadence of ~1.6s, a white light precision of ~540 ppm, and a baseline of ~60 hours. Since that’s an enormous amount of data, we’ll only simulate 40 hours and bin the data to a 5-minute cadence, but will be sure to add the finite exposure time correction to model the resulting transit blurring.
[7]:
t_exp = 5 * u.min
times = jnp.arange(-20, 20, t_exp.to(u.hour).value) * u.hour.to(u.day)
# generate the stellar intensities
sld = StellarLimbDarkening(
M_H=0.020,
Teff=4884,
logg=4.579,
ld_model="mps1",
ld_data_path="exotic_ld_data",
interpolate_type="trilinear",
custom_wavelengths=None,
custom_mus=None,
custom_stellar_model=None,
ld_data_version="3.2.0",
verbose=0,
)
sld._integrate_I_mu(
wavelength_range=[6_000, 53_000],
mode="JWST_NIRSpec_Prism",
custom_wavelengths=None,
custom_throughput=None,
)
# interpolate the stellar intensities
mu_grid = np.linspace(0.0, 1.0, 1_000)
order = np.argsort(sld.mus)
f = CubicSpline(x=sld.mus[order], y=sld.I_mu[order])
interpolated_vals = f(mu_grid)
# fit the limb darkening profile
u_coeffs = OblateSystem.fit_limb_darkening_profile(
intensities=interpolated_vals, mus=mu_grid, order=14
)
# create the planet
injected_state = {
"t_peri": -1071.23205 / 4, # going to make it circular with mid-transit at t=0
"times": times,
"exposure_time": t_exp.to(u.day).value,
"oversample": 3, # 3x more samples under-the-hood, then binned back down
"oversample_correction_order": 2,
"a": ((1.883 * u.au) / (0.749 * u.R_sun)).to(u.dimensionless_unscaled),
"period": 1071.23205,
"r": ((0.9064 * u.R_jup) / (0.749 * u.R_sun)).to(u.dimensionless_unscaled),
"i": 89.9720 * jnp.pi / 180,
"ld_u_coeffs": u_coeffs,
"f1": 0.1, # ~Saturn oblateness
"obliq": 26.7 * jnp.pi / 180, # Saturn obliquity
"prec": 60.0 * jnp.pi / 180, # randomly chosen
"tidally_locked": False,
}
injected_planet = OblateSystem(**injected_state)
# create a spherical planet with the same projected area and orbital parameters for comparison
spherical_planet_state = injected_state.copy()
spherical_planet_state["r"] = injected_planet.state["projected_effective_r"]
spherical_planet_state["f1"] = 0.0
spherical_planet = OblateSystem(**spherical_planet_state)
# generate the lightcurves
injected_transit = injected_planet.lightcurve()
spherical_transit = spherical_planet.lightcurve()
Let’s inspect what we just created. We can plot the projected outline of the planet using injected_planet.illustrate, then plot the light curve and its comparison to the one created by the spherical planet:
[8]:
print(f"Projected flattening: {injected_planet.state['projected_f']}")
injected_planet.illustrate(star_fill=False) # suppress the star fill for clarity
fig, ax = plt.subplots(nrows=2, sharex=True, figsize=(8, 7))
ax[0].plot(times * u.day.to(u.hour), injected_transit)
ax[0].set(ylabel="flux")
ax[1].plot(times * u.day.to(u.hour), (injected_transit - spherical_transit) * 1e6)
ax[1].set(ylabel="diff. w/ \nsph. planet [ppm]", xlabel="time [hours]")
plt.subplots_adjust(hspace=0);
Projected flattening: 0.08412143970386853
So, we’re looking for an effect that’s ~100 ppm, and asymmetric.
The projected flattening value, f, is less than the true injected f value: that’s because we’ve given the planet some obliquity, which leaves us with a more circular projected ellipse when we only flatten along the \(z\) axis with \(f_1\) (if the obliquity was a full 90 degrees, we’d see a perfect circle). See Berardo and de Wit 2022 for more details.
Lastly, let’s double check that our oversampling correction is sufficient to account for smearing caused by the finite exposure time binning:
[9]:
no_oversample_planet_state = injected_state.copy()
no_oversample_planet_state["oversample"] = 1
no_oversample_planet = OblateSystem(**no_oversample_planet_state)
no_oversample_transit = no_oversample_planet.lightcurve()
supersampled_planet_state = injected_state.copy()
supersampled_planet_state["oversample"] = 10
supersampled_planet = OblateSystem(**supersampled_planet_state)
supersampled_transit = supersampled_planet.lightcurve()
fig, ax = plt.subplots()
ax.plot(
times * u.day.to(u.hour),
(no_oversample_transit - supersampled_transit) * 1e6,
label="no oversample",
)
ax.plot(
times * u.day.to(u.hour),
(injected_transit - supersampled_transit) * 1e6,
label="3x oversample (current sim)",
)
ax.legend()
ax.set(ylabel="diff. w/ 10x\noversampling [ppm]", xlabel="time [hours]");
We’re safely in the <2ppm regime here, so we’ll stick with the 3x oversample to avoid unnecessary computation.
Noise addition#
Now we need to add noise to this injected model. This will include white noise from photon counting statistics and from instrument readout (which here we’ll also assume is white). The JWST Exposure Time Calculator suggests that combined, this will be about ~540 ppm for each 1.6s integration, taken with 6 groups per up-the-ramp sample. We’ll scale that to our 5-minute cadence and add it to the light curve.
In principle the host star could also contribute correlated (non-white) noise, but Kepler-167 is a quiet K dwarf. Inspecting its archival Kepler photometry (for example with a Lomb-Scargle periodogram or an FFT of the long- and short-cadence data) reveals a nearly flat power spectrum at the frequencies relevant to our ~60-hour observation, with no significant p-mode or granulation peaks and only a mild 1/f rise at the very lowest frequencies. A pure white-noise model is therefore a reasonable approximation here, so we simply inject Gaussian white noise at the amplitude set by the ETC.
[10]:
# scale the shot noise/integration to our longer binned light curve
shot_noise_amplitude = (540 * 1e-6) / jnp.sqrt(
((t_exp) / (1.6 * u.s)).to(u.dimensionless_unscaled)
)
shot_noise = (
jax.random.normal(jax.random.PRNGKey(0), times.shape) * shot_noise_amplitude
)
noised_data = injected_transit + shot_noise
fig, ax = plt.subplots()
ax.plot(times * u.day.to(u.hour), noised_data, marker="o", markersize=2, ls="none")
ax.set(ylabel="flux", xlabel="time [hours]");
And that’s our simulated data!
Fitting the transit#
Initial best-fit solution#
Now that we have some simulated data, let’s use squishyplanet to fit it. We’ll start by finding the maximum likelihood solution for the transit parameters.
[11]:
# slightly perturbed parameters compared to the injected_state
init_state = {
"times": times,
"exposure_time": t_exp.to(u.day).value,
"oversample": 3,
"oversample_correction_order": 2,
"period": 1071.23205,
"t_peri": -1071.23205 / 4 + 0.01,
"tidally_locked": False,
"parameterize_with_projected_ellipse": True,
"a": 540.0,
"i": 90 * jnp.pi / 180,
"ld_u_coeffs": jnp.array([0.3, 0.2]),
"projected_effective_r": 0.12,
"projected_f": 0.0,
"projected_theta": 0.0,
}
planet = OblateSystem(**init_state) # the system we'll use for all our fits!
# the observed data and per-point uncertainties are no longer stored on the
# OblateSystem, so we keep them here and build our own Gaussian log-likelihood
uncertainties = jnp.ones_like(noised_data) * shot_noise_amplitude
@jax.jit
def loglike(params):
# `log_jitter` (if present) is added in quadrature to the uncertainties; every
# other key is passed straight through to the lightcurve model
params = dict(params)
log_jitter = params.pop("log_jitter")
resids = noised_data - planet.lightcurve(params)
var = jnp.exp(log_jitter) ** 2 + uncertainties**2
return jnp.sum(-0.5 * (resids**2 / var + jnp.log(2 * jnp.pi * var)))
# evaluate this likelihood (log_jitter=-inf recovers the pure-uncertainty case)
print(f"initial loglike: {loglike({'log_jitter': -jnp.inf})}")
init_lc = planet.lightcurve()
fig, ax = plt.subplots()
ax.errorbar(
times * u.day.to(u.hour),
(init_lc - noised_data) * 1e6,
yerr=shot_noise_amplitude,
marker="o",
markersize=2,
ls="none",
)
ax.set(ylabel="residuals [ppm]", xlabel="time [hours]");
initial loglike: -115736.42788386013
That’s obviously not a great fit (the errorbars are there, just smaller than the points), but it’s just our starting point. Though the above cell takes a few seconds to run, thanks to the jit-compilation, each subsequent likelihood evaluation will be faster than the initial one:
[12]:
%timeit loglike({'log_jitter': -jnp.inf}).block_until_ready()
5.29 ms ± 199 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
To get an initial estimate of the best fit, we can use scipy’s minimize routines. Though we could leave the limb darkening fixed to the values we calculated above, to demonstrate how you might fit for the limb darkening parameters, we’ll use the Kipping 2013 parameterization and fit for \(q_1\) and \(q_2\), which have uniform priors between [0,1] and can be directly converted into the \(u_1\) and \(u_2\) coefficients that define a quadratic limb darkening law. We’ll also put wide bounds on the other parameters that roughly match the priors we’ll use later, just to keep anything from running away.
[13]:
np.random.seed(10)
# scipy works in arrays, squishyplanet in dictionaries
# this function translates between the two
@jax.jit
def translate_to_dict(x):
# we're using t_peri here instead of t0, but the posterior shapes will look the same
# for circular orbits
t_peri = x[0]
a = x[1]
impact_param = x[2]
i = jnp.arccos(impact_param / a)
q1 = x[3] # using the Kipping 2013 uninformative quadratic limb darkening setup
q2 = x[4]
u1 = 2 * jnp.sqrt(q1) * q2 # convert to u1, u2
u2 = jnp.sqrt(q1) * (1 - 2 * q2)
projected_effective_r = x[5]
projected_f = x[6]
projected_theta = x[7]
log_jitter = x[8]
param_dict = {
"t_peri": t_peri,
"a": a,
"i": i,
"ld_u_coeffs": jnp.array([u1, u2]),
"projected_effective_r": projected_effective_r,
"projected_f": projected_f,
"projected_theta": projected_theta,
"log_jitter": log_jitter,
}
return param_dict
# our objective
def minimizer_fn(x):
param_dict = translate_to_dict(x)
return -loglike(param_dict)
# starting point
init_guess = jnp.array(
[
init_state["t_peri"],
init_state["a"],
init_state["a"] * jnp.cos(init_state["i"]) + 0.05,
init_state["ld_u_coeffs"][0] ** 2,
init_state["ld_u_coeffs"][0] * init_state["ld_u_coeffs"][1] / 2,
init_state["projected_effective_r"],
init_state["projected_f"],
init_state["projected_theta"],
jnp.log(1e-7), # start log_jitter at its lower prior bound
]
)
# run the minimizer with wide bounds on each parameter
res1 = minimize(
minimizer_fn,
init_guess,
bounds=[
(-1071.23205 / 4 - 0.15, -1071.23205 / 4 + 0.15), # t_peri, +/- 0.15 days
(530.0, 550.0), # a, 530-550 days
(
0.0,
1.0,
), # impact parameter (b), 0-1 (so excludes extreme grazing but guarantees a transit)
(0.0, 1.0), # q1, 0-1
(0.0, 1.0), # q2, 0-1
(0.10, 0.13), # projected_effective_r, 0.10-0.13
(0.0, 0.3), # projected_f, 0-0.3 (0.3 is a pretty extreme oblateness)
(0.0, jnp.pi), # projected_theta, 0-180 degrees
(jnp.log(1e-7), jnp.log(1e-4)), # log_jitter, meaning an jitter of 1e-7 to 1e-4
],
)
oblate_resids = planet.lightcurve(translate_to_dict(res1.x)) - noised_data
fig, ax = plt.subplots(figsize=(8, 6))
ax.errorbar(
times * u.day.to(u.hour),
oblate_resids * 1e6,
yerr=shot_noise_amplitude * 1e6,
marker="o",
markersize=2,
ls="none",
alpha=0.5,
)
ax.set(
ylabel="residuals [ppm]", xlabel="time [hours]", title="oblate planet optimization"
)
print(res1)
print("\nOptimized parameters:")
for key in translate_to_dict(res1.x):
print(f"{key}: {translate_to_dict(res1.x)[key]}")
message: CONVERGENCE: RELATIVE REDUCTION OF F <= FACTR*EPSMCH
success: True
status: 0
fun: -4182.210301988128
x: [-2.678e+02 5.420e+02 2.592e-01 1.644e-01 3.134e-01
1.191e-01 9.315e-02 2.552e-01 -1.574e+01]
nit: 98
jac: [-2.614e+01 9.277e-03 4.634e+00 1.553e+00 4.155e-01
-3.644e+01 1.992e-02 -1.919e-02 0.000e+00]
nfev: 1130
njev: 113
hess_inv: <9x9 LbfgsInvHessProduct with dtype=float64>
Optimized parameters:
a: 542.0104839570314
i: 1.570318060793753
ld_u_coeffs: [0.25416645 0.15132339]
log_jitter: -15.73629190710488
projected_effective_r: 0.11910978330188558
projected_f: 0.09315439133040007
projected_theta: 0.2551704572498602
t_peri: -267.80804498071956
That’s much better!
To get a sense of how important oblateness is in the fit, let’s do the same optimization while fixing the oblateness to zero:
[14]:
# exact same as above, but with the f and theta terms left out
np.random.seed(10)
@jax.jit
def translate_to_dict(x):
t_peri = x[0]
a = x[1]
impact_param = x[2]
i = jnp.arccos(impact_param / a)
q1 = x[3] # using the Kipping 2013 uninformative quadratic limb darkening setup
q2 = x[4]
u1 = 2 * jnp.sqrt(q1) * q2
u2 = jnp.sqrt(q1) * (1 - 2 * q2)
projected_effective_r = x[5]
# projected_f = x[6]
# projected_theta = x[7]
log_jitter = x[6]
param_dict = {
"t_peri": t_peri,
"a": a,
"i": i,
"ld_u_coeffs": jnp.array([u1, u2]),
"projected_effective_r": projected_effective_r,
# "projected_f": projected_f,
# "projected_theta": projected_theta,
"log_jitter": log_jitter,
}
return param_dict
def minimizer_fn(x):
param_dict = translate_to_dict(x)
return -loglike(param_dict)
init_guess = jnp.array(
[
init_state["t_peri"],
init_state["a"],
init_state["a"] * jnp.cos(init_state["i"]) + 0.05,
init_state["ld_u_coeffs"][0] ** 2,
init_state["ld_u_coeffs"][0] * init_state["ld_u_coeffs"][1] / 2,
init_state["projected_effective_r"],
# init_state["projected_f"],
# init_state["projected_theta"],
jnp.log(1e-7), # start log_jitter at its lower prior bound
]
)
res2 = minimize(
minimizer_fn,
init_guess,
bounds=[
(-1071.23205 / 4 - 0.15, -1071.23205 / 4 + 0.15),
(530.0, 550.0),
(0.0, 1.0),
(0.0, 1.0),
(0.0, 1.0),
(0.10, 0.13),
# (0.0, 0.3),
# (0.0, jnp.pi),
(jnp.log(1e-7), jnp.log(1e-4)),
],
)
sphere_resids = planet.lightcurve(translate_to_dict(res2.x)) - noised_data
fig, ax = plt.subplots(figsize=(8, 6))
ax.errorbar(
times * u.day.to(u.hour),
sphere_resids * 1e6,
yerr=shot_noise_amplitude * 1e6,
marker="o",
markersize=2,
ls="none",
alpha=0.5,
)
ax.set(
ylabel="residuals [ppm]",
xlabel="time [hours]",
title="spherical planet optimization",
)
print(res2)
print("\nOptimized parameters:")
for key in translate_to_dict(res2.x):
print(f"{key}: {translate_to_dict(res2.x)[key]}")
message: CONVERGENCE: RELATIVE REDUCTION OF F <= FACTR*EPSMCH
success: True
status: 0
fun: -4169.089109780205
x: [-2.678e+02 5.324e+02 3.154e-01 1.665e-01 3.090e-01
1.194e-01 -1.609e+01]
nit: 76
jac: [-5.288e+00 1.610e-02 4.878e+00 5.076e-01 -1.701e-02
-8.546e+01 -3.638e-04]
nfev: 728
njev: 91
hess_inv: <7x7 LbfgsInvHessProduct with dtype=float64>
Optimized parameters:
a: 532.3967297775962
i: 1.570203892186498
ld_u_coeffs: [0.25217657 0.1558381 ]
log_jitter: -16.08890731387245
projected_effective_r: 0.1193751281193505
t_peri: -267.80804855068305
Even though we’re forcing it to fit a mis-specified model, by eye at least the fit is still pretty good. To more rigorously compare the two models, let’s use two statistical tests: the Bayesian Information Criterion (BIC) and the Akaike Information Criterion (AIC). These tests penalize models with more free parameters, so we can use them to determine if the likelihood improvement we get using the oblate model is justified by its extra terms:
[15]:
# plot the residual differences
fig, ax = plt.subplots(figsize=(6, 4))
ax.errorbar(
times * u.day.to(u.hour),
oblate_resids * 1e6,
color="r",
yerr=shot_noise_amplitude * 1e6,
marker="o",
markersize=2,
ls="none",
label="oblate model",
alpha=0.3,
)
ax.errorbar(
times * u.day.to(u.hour),
sphere_resids * 1e6,
color="g",
yerr=shot_noise_amplitude * 1e6,
marker="o",
markersize=2,
ls="none",
label="spherical model",
alpha=0.3,
)
ax.set(
ylabel="residuals [ppm]",
xlabel="time [hours]",
title="comparison of oblate and spherical models",
)
ax.legend(prop={"size": 12})
# plot the model differences
fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(
times * u.day.to(u.hour),
(oblate_resids - sphere_resids) * 1e6,
color="k",
marker="o",
markersize=2,
ls="none",
)
ax.set(
ylabel="difference in residuals [ppm]",
xlabel="time [hours]",
title="difference between oblate and spherical models",
)
# do the model comparison:
# definitions:
# bic = -2 * log_likelihood + num_params * np.log(num_data_points)
# aic = -2 * log_likelihood + 2 * num_params
bic_oblate = -2 * (-res1.fun) + 8 * np.log(len(noised_data))
bic_sphere = -2 * (-res2.fun) + 6 * np.log(len(noised_data))
aic_oblate = -2 * (-res1.fun) + 16
aic_sphere = -2 * (-res2.fun) + 12
print(f"BIC oblate: {bic_oblate}")
print(f"BIC spherical: {bic_sphere}")
print(f"AIC oblate: {aic_oblate}")
print(f"AIC spherical: {aic_sphere}\n")
if bic_oblate < bic_sphere:
print(
f"The oblate model is preferred by BIC, delta BIC = {bic_sphere - bic_oblate}"
)
else:
print(
f"The spherical model is preferred by BIC, delta BIC = {bic_oblate - bic_sphere}"
)
if aic_oblate < aic_sphere:
print(
f"The oblate model is preferred by AIC, delta AIC = {aic_sphere - aic_oblate}"
)
else:
print(
f"The spherical model is preferred by AIC, delta AIC = {aic_oblate - aic_sphere}"
)
BIC oblate: -8315.03031514504
BIC spherical: -8301.135502936999
AIC oblate: -8348.420603976256
AIC spherical: -8326.17821956041
The oblate model is preferred by BIC, delta BIC = 13.894812208041913
The oblate model is preferred by AIC, delta AIC = 22.242384415845663
So, under these tests at least, it looks like the oblate model is preferred. In a real analysis situation we’d want to play with different limb darkening parameterizations and make sure that our recovery is robust, but hopefully this gives a starting point for how you might use squishyplanet to fit a model of an oblate planet.
Nested Sampling#
If we wanted to take another approach to model comparison and compute the Bayesian evidence for each fit, we’ll have to run a nested sampling algorithm. To show how you might do that with squishyplanet, we’ll use the nautilus library. nautilus, which is described in full in Lange 2023, uses neural nets to help select new live points, which can
dramatically speed up sampling. If we were doing a real analysis project, it’d be important to run convergence checks on the final results, but here we’re just going to use conservative settings and hope for the best. Each of the following fits takes ~2 hours my laptop with no parallelization.
[16]:
# construct the prior
np.random.seed(10)
prior = Prior()
# estimate uncertainty on mid-transit time: assume the uncertainties on the period and
# t_c are Gaussian and uncorrelated
tc_uncertainty = 0.00039 # u.day, from the 2010 Kepler epoch, Chachan+ 2020
period_uncertainty = 0.00059 # u.day, Chachan+ 2020
tc_draw = np.random.normal(loc=0, scale=tc_uncertainty, size=10_000)
period_deviation_draws = np.random.normal(loc=0, scale=period_uncertainty, size=10_000)
sims = period_deviation_draws * 5 + tc_uncertainty
t_peri_uncertainty = np.std(sims)
prior.add_parameter(
"t_peri", dist=stats.norm(loc=-1071.23205 / 4, scale=t_peri_uncertainty)
)
a_uncertainty = ((0.027 * u.au) / (0.749 * u.R_sun)).to(u.dimensionless_unscaled).value
prior.add_parameter(
"a",
dist=stats.norm(
loc=((1.883 * u.au) / (0.749 * u.R_sun)).to(u.dimensionless_unscaled),
scale=a_uncertainty,
),
)
# straight from Chachan+ 2020
prior.add_parameter("b", dist=stats.norm(0.271, 0.073))
# using the Kipping 2013 uninformative quadratic limb darkening setup
prior.add_parameter("q1", dist=(0.0, 1.0))
prior.add_parameter("q2", dist=(0.0, 1.0))
# log-uniform prior on the jitter. between 0.1 and 100 ppm extra noise seems reasonable
prior.add_parameter("log_jitter", dist=(jnp.log(1e-7), jnp.log(1e-4)))
# to account for different bandpass, we'll scale the uncertainty by 3x
# (sort of accidentally simulated this by not adjusting the planet radius: the projected
# effective radius of the injected planet is smaller than the true radius, will be
# ~1sigma from this mean)
r_draws = (
(
np.random.normal(loc=10.16, scale=0.42, size=10_000)
* u.R_earth
/ (0.749 * u.R_sun)
)
.to(u.dimensionless_unscaled)
.value
)
r_mean = np.mean(r_draws)
r_std = np.std(r_draws) * 3
prior.add_parameter("projected_effective_r", dist=stats.norm(loc=r_mean, scale=r_std))
# large uniform prior on this
prior.add_parameter("projected_f", dist=(0.0, 0.15))
# completely unconstrained
prior.add_parameter("projected_theta", dist=(0.0, jnp.pi))
[17]:
# since we're fitting for impact parameter b (not inclination i) and the limb darkening
# parameters q1 and q2 (not u1 and u2), we need to add a transformation before feeding
# the dictionary into the loglike helper defined above
@jax.jit
def nautilus_loglike(x):
inc = jnp.arccos(x["b"] / x["a"])
u1 = 2 * jnp.sqrt(x["q1"]) * x["q2"]
u2 = jnp.sqrt(x["q1"]) * (1 - 2 * x["q2"])
param_dict = {
"t_peri": x["t_peri"],
"a": x["a"],
"i": inc,
"ld_u_coeffs": jnp.array([u1, u2]),
"projected_effective_r": x["projected_effective_r"],
"projected_f": x["projected_f"],
"projected_theta": x["projected_theta"],
"log_jitter": x["log_jitter"],
}
return loglike(param_dict)
# create the sampler
sampler1 = Sampler(
prior,
nautilus_loglike,
n_live=3000,
seed=0,
enlarge_per_dim=1.2, # larger than 1.1 default, but still manually set (i.e. no bootstrapping like UltraNest)
filepath="nautilus_oblate_checkpoint.hdf5",
resume=True,
)
# run it
sampler1.run(
f_live=0.001, # default is 0.01, but we want to be more aggressive
n_eff=50_000, # pretty large
n_like_max=np.inf,
discard_exploration=True, # see the nautlius paper, if False results will be faster but biased
timeout=np.inf,
verbose=True,
)
Resuming nautilus run...
Please report issues at github.com/johannesulf/nautilus.
Status | Bounds | Ellipses | Networks | Calls | f_live | N_eff | log Z
Finished | 55 | 2 | 4 | 275600 | N/A | 50083 | +4142.34
[17]:
np.True_
And now we can do the same for the spherical model:
[18]:
# construct the prior
np.random.seed(10)
prior = Prior()
# estimate uncertainty on mid-transit time: assume the uncertainties on the period and
# t_c are Gaussian and uncorrelated
tc_uncertainty = 0.00039 # u.day, from the 2010 Kepler epoch, Chachan+ 2020
period_uncertainty = 0.00059 # u.day, Chachan+ 2020
tc_draw = np.random.normal(loc=0, scale=tc_uncertainty, size=10_000)
period_deviation_draws = np.random.normal(loc=0, scale=period_uncertainty, size=10_000)
sims = period_deviation_draws * 5 + tc_uncertainty
t_peri_uncertainty = np.std(sims)
prior.add_parameter(
"t_peri", dist=stats.norm(loc=-1071.23205 / 4, scale=t_peri_uncertainty)
)
a_uncertainty = ((0.027 * u.au) / (0.749 * u.R_sun)).to(u.dimensionless_unscaled).value
prior.add_parameter(
"a",
dist=stats.norm(
loc=((1.883 * u.au) / (0.749 * u.R_sun)).to(u.dimensionless_unscaled),
scale=a_uncertainty,
),
)
# straight from Chachan+ 2020
prior.add_parameter("b", dist=stats.norm(0.271, 0.073))
# using the Kipping 2013 uninformative quadratic limb darkening setup
prior.add_parameter("q1", dist=(0.0, 1.0))
prior.add_parameter("q2", dist=(0.0, 1.0))
# log-uniform prior on the jitter. between 0.1 and 100 ppm extra noise seems reasonable
prior.add_parameter("log_jitter", dist=(jnp.log(1e-7), jnp.log(1e-4)))
# to account for different bandpass, we'll scale the uncertainty by 3x
# (sort of accidentally simulated this by not adjusting the planet radius: the projected
# effective radius of the injected planet is smaller than the true radius, will be
# ~1sigma from this mean)
r_draws = (
(
np.random.normal(loc=10.16, scale=0.42, size=10_000)
* u.R_earth
/ (0.749 * u.R_sun)
)
.to(u.dimensionless_unscaled)
.value
)
r_mean = np.mean(r_draws)
r_std = np.std(r_draws) * 3
prior.add_parameter("projected_effective_r", dist=stats.norm(loc=r_mean, scale=r_std))
# # large uniform prior on this
# prior.add_parameter("projected_f", dist=(0.0, 0.15))
# # completely unconstrained
# prior.add_parameter("projected_theta", dist=(0.0, jnp.pi))
[19]:
# same as the oblate case
@jax.jit
def nautilus_loglike(x):
inc = jnp.arccos(x["b"] / x["a"])
u1 = 2 * jnp.sqrt(x["q1"]) * x["q2"]
u2 = jnp.sqrt(x["q1"]) * (1 - 2 * x["q2"])
param_dict = {
"t_peri": x["t_peri"],
"a": x["a"],
"i": inc,
"ld_u_coeffs": jnp.array([u1, u2]),
"projected_effective_r": x["projected_effective_r"],
"projected_f": 0.0,
"projected_theta": 0.0,
"log_jitter": x["log_jitter"],
}
return loglike(param_dict)
sampler2 = Sampler(
prior,
nautilus_loglike,
n_live=3000,
seed=1,
enlarge_per_dim=1.2,
filepath="nautilus_sphere_checkpoint.hdf5",
resume=True,
)
sampler2.run(
f_live=0.001,
n_eff=50_000,
n_like_max=np.inf,
discard_exploration=True,
timeout=np.inf,
verbose=True,
)
Resuming nautilus run...
Please report issues at github.com/johannesulf/nautilus.
Status | Bounds | Ellipses | Networks | Calls | f_live | N_eff | log Z
Finished | 53 | 1 | 4 | 235300 | N/A | 50078 | +4132.57
[19]:
np.True_
[20]:
# the bayes factor is the ratio between the evidence of the two models:
print(
f"the oblate model is preffered over the spherical model by a Bayes factor of: {jnp.exp(sampler1.log_z - sampler2.log_z)}"
)
the oblate model is preffered over the spherical model by a Bayes factor of: 17592.312534727695
So, with some caveats about this being conditional to the choice of prior and limb darkening parameterization, it looks like we successfully recovered an oblateness signal! Let’s take a look at the posteriors:
[21]:
o_points, o_log_w, o_log_l = sampler1.posterior()
_s_points, _s_log_w, _s_log_l = sampler2.posterior()
s_points = np.zeros((len(_s_points), o_points.shape[1]))
s_points[:, :-2] = _s_points
truth = [
injected_state["t_peri"],
injected_state["a"].value,
(injected_state["a"] * jnp.cos(injected_state["i"])).value,
None,
None,
None,
injected_planet.state["projected_effective_r"][0],
injected_planet.state["projected_f"],
injected_planet.state["projected_theta"][0],
]
warnings.filterwarnings("ignore")
f = corner.corner(
s_points,
weights=np.exp(_s_log_w),
bins=20,
labels=[
"t_peri",
"a",
"b",
"q1",
"q2",
"log_jitter",
"projected_effective_r",
"projected_f",
"projected_theta",
],
color="#1285ff",
plot_datapoints=False,
range=np.repeat(0.999, 9),
)
corner.corner(
o_points,
fig=f,
weights=np.exp(o_log_w),
bins=20,
labels=[
"t_peri",
"a",
"b",
"q1",
"q2",
"log_jitter",
"projected_effective_r",
"projected_f",
"projected_theta",
],
color="#c10029",
plot_datapoints=False,
range=np.repeat(0.999, 9),
truths=truth,
)
plt.savefig("corner_plot.png")
WARNING:root:Too few points to create valid contours
Above, the spherical model is in blue while the oblate is in red. Interestingly, the a and b posterior for the spherical model appear significantly biased away from the injected values: to handle the slightly different ingress/egress times, the spherical model was forced to change the velocity of the planet and the length of its transit chord.
We can focus in on the marginalized posterior of the projected oblateness:
[22]:
fig, ax = plt.subplots()
ax.hist(
o_points[:, -2],
bins=50,
histtype="step",
color="purple",
weights=np.exp(o_log_w),
)
med = np.percentile(o_points[:, -2], 50)
low = np.percentile(o_points[:, -2], 16)
high = np.percentile(o_points[:, -2], 84)
ax.axvline(med, color="purple", ls="--", label="median")
ax.axvline(low, color="purple", ls=":")
ax.axvline(high, color="purple", ls=":")
ax.axvline(injected_planet.state["projected_f"], color="k", ls="--", label="truth")
ax.annotate(
f"projected f:\n{med:.3f} +{high-med:.3f} -{med-low:.3f}",
(0.95, 0.95),
xycoords="axes fraction",
ha="right",
va="top",
)
ax.legend(loc="lower right")
ax.set(xlabel="projected f", ylabel="posterior mass");
We can see that the true value is nicely within the posterior, and that zero flattening is confidently ruled out.
Hopefully this gives a sense of how you might go about using squishyplanet. We used nautlius for our inference here, but we just as easily could have used other nested sampling frameworks like dynesty or ultranest, or even MCMC frameworks like emcee. We also could have used the squishyplanet-generated lightcurve as the mean for a more complicated systematics model, for example if we wanted to include a GP to model the stellar noise, or some kind of polynomial regression
against engineering telemetry to model observatory systematics. squishyplanet is designed to leave the inference choices to the user, and just provides the tools to do the transit modeling and likelihood evaluation.