def fit_ttv(self, n, run_MCMC=False, ttv_start=None, verbose=True):
"""
Fit a single transit with a transit timing variation, using the shape given by best-fit orbital parameters.
"""
if verbose: print("Fitting ttv for transit number", n)
# Get the transit lightcurve
transit = self.lightcurve.get_transit(n, self.p_ref, self.t0_ref)
t = transit.time * self.p_ref
y = transit.flux
sd = transit.flux_err
if ttv_start is None:
ttv_start = np.median(self.pars['ttvs'])
with pm.Model() as model:
ttv = pm.Normal("ttv", mu=ttv_start, sd=0.025) # sd = 36 minutes
orbit = xo.orbits.KeplerianOrbit(period=self.p_ref,
t0=ttv,
b=self.pars['b'])
light_curves = xo.LimbDarkLightCurve(
self.pars['u']).get_light_curve(orbit=orbit,
r=self.pars['r'],
t=t)
light_curve = pm.math.sum(light_curves, axis=-1) + 1
pm.Deterministic("transit_" + str(n), light_curve)
pm.Normal("obs", mu=light_curve, sd=sd, observed=y)
map_soln = xo.optimize(start=model.test_point,
verbose=False,
progress_bar=False)
self.pars['ttvs'][n] = float(map_soln['ttv'])
if verbose: print(f"\t ttv {n} = {self.pars['ttvs'][n]}")
if run_MCMC:
np.random.seed(42)
with model:
trace = pm.sample(
tune=500,
draws=500,
start=map_soln,
cores=1,
chains=2,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
self.pars['ttvs'][n] = np.median(trace['ttv'])
self.pars['e_ttvs'][n] = self.pars['ttvs'][n] - np.percentile(
trace['ttv'], 16, axis=0)
self.pars['E_ttvs'][n] = -self.pars['ttvs'][n] + np.percentile(
trace['ttv'], 84, axis=0)
if verbose:
print(
f"\t ttv {n} = {self.pars['ttvs'][n]} /+ {self.pars['E_ttvs'][n]} /- {self.pars['e_ttvs'][n]}"
)
def get_transit_curve(self, t, t_offset=0):
orbit = xo.orbits.KeplerianOrbit(period=self.p_ref,
t0=t_offset,
b=self.pars['b'])
lightcurve = xo.LimbDarkLightCurve(self.pars['u']).get_light_curve(
orbit=orbit, r=self.pars['r'], t=t).eval() + self.pars['mean']
return lightcurve
def xo_lightcurve(time, period=3, r=0.1, t0=0, plot=False):
import exoplanet as xo
orbit = xo.orbits.KeplerianOrbit(period=0.7, t0=0.1)
light_curve = xo.LimbDarkLightCurve([0.1, 0.4]).get_light_curve(
orbit=orbit, r=r, t=time).eval() + 1
if plot:
plt.plot(time, light_curve, color="C0", lw=2)
plt.ylabel("relative flux")
plt.xlabel("time [days]")
return light_curve.flatten()
def transit_model(params, t, texp=30 / (60 * 24), mstar=1, rstar=1):
period = params[0]
t0 = params[1]
r = params[2]
b = params[3]
u = params[4]
mean = params[5]
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
mstar=mstar,
rstar=rstar)
return (mean + xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=t, texp=texp).eval().flatten())
def get_model_transit(paramd, time_eval, t_exp=2 / (60 * 24)):
"""
you know the paramters, and just want to evaluate the median lightcurve.
"""
import exoplanet as xo
period = paramd['period']
t0 = paramd['t0']
try:
r = paramd['r']
except KeyError:
r = np.exp(paramd['log_r'])
b = paramd['b']
u0 = paramd['u[0]']
u1 = paramd['u[1]']
r_star = paramd['r_star']
logg_star = paramd['logg_star']
try:
mean = paramd['mean']
except KeyError:
mean_key = [k for k in list(paramd.keys()) if 'mean' in k]
assert len(mean_key) == 1
mean_key = mean_key[0]
mean = paramd[mean_key]
# factor * 10**logg / r_star = rho
factor = 5.141596357654149e-05
rho_star = factor * 10**logg_star / r_star
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
rho_star=rho_star)
u = [u0, u1]
mu_transit = xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=time_eval, texp=t_exp).T.flatten()
return mu_transit.eval() + mean
def test_exoplanet_transit():
# The light curve calculation requires an orbit
orbit = xo.orbits.KeplerianOrbit(period=3.456)
texp = 30/(60*24)
# Compute a limb-darkened light curve using starry
t = np.linspace(-0.1, 0.1, 1000)
u = [0.3, 0.2]
light_curve = (
xo.LimbDarkLightCurve(u)
.get_light_curve(orbit=orbit, r=0.1, t=t, texp=texp)
.eval()
)
# Note: the `eval` is needed because this is using Theano in
# the background
plt.plot(t, light_curve, color="C0", lw=2)
plt.ylabel("relative flux")
plt.xlabel("time [days]")
_ = plt.xlim(t.min(), t.max())
plt.savefig('../results/test_results/test_exoplanet_transit.png',
bbox_inches='tight')
plt.close('all')
def build_model(mask=None, start=None):
with pm.Model() as model:
# The baseline flux
mean = pm.Normal("mean", mu=0.0, sd=0.00001)
# The time of a reference transit for each planet
t0 = pm.Normal("t0", mu=t0s, sd=1.0, shape=1)
# The log period; also tracking the period itself
logP = pm.Normal("logP", mu=np.log(periods), sd=0.01, shape=1)
rho_star = pm.Normal("rho_star", mu=0.14, sd=0.01, shape=1)
r_star = pm.Normal("r_star", mu=2.7, sd=0.01, shape=1)
period = pm.Deterministic("period", pm.math.exp(logP))
# The Kipping (2013) parameterization for quadratic limb darkening paramters
u = xo.distributions.QuadLimbDark("u", testval=np.array([0.3, 0.2]))
r = pm.Uniform(
"r", lower=0.01, upper=0.3, shape=1, testval=0.15)
b = xo.distributions.ImpactParameter(
"b", ror=r, shape=1, testval=0.5)
# Transit jitter & GP parameters
logs2 = pm.Normal("logs2", mu=np.log(np.var(y)), sd=10)
logw0 = pm.Normal("logw0", mu=0, sd=10)
logSw4 = pm.Normal("logSw4", mu=np.log(np.var(y)), sd=10)
# Set up a Keplerian orbit for the planets
orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, b=b, rho_star=rho_star,r_star=r_star)
# Compute the model light curve using starry
light_curves = xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=t
)
light_curve = pm.math.sum(light_curves, axis=-1) + mean
# Here we track the value of the model light curve for plotting
# purposes
pm.Deterministic("light_curves", light_curves)
kernel = xo.gp.terms.SHOTerm(log_Sw4=logSw4, log_w0=logw0, Q=1 / np.sqrt(2))
gp = xo.gp.GP(kernel, t, tt.exp(logs2) + tt.zeros(len(t)), mean=light_curve)
gp.marginal("gp", observed=y)
pm.Deterministic("gp_pred", gp.predict())
# The likelihood function assuming known Gaussian uncertainty
pm.Normal("obs", mu=light_curve, sd=yerr, observed=y)
# Fit for the maximum a posteriori parameters given the simuated
# dataset
map_soln = xo.optimize(start=model.test_point)
return model, map_soln
model, map_soln = build_model()
gp_mod = map_soln["gp_pred"] + map_soln["mean"]
plt.clf()
plt.plot(t, y, ".k", ms=4, label="data")
plt.plot(t, gp_mod, lw=1,label="gp model")
plt.plot(t, map_soln["light_curves"], lw=1,label="transit model")
plt.xlim(t.min(), t.max())
plt.ylabel("relative flux")
plt.xlabel("time [days]")
plt.legend(fontsize=10)
_ = plt.title("map model")
np.random.seed(42)
with model:
trace = pm.sample(
tune=3000,
draws=3000,
start=map_soln,
cores=2,
chains=2,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
pm.summary(trace, varnames=["period", "t0", "r", "b", "u", "mean", "rho_star","logw0","logSw4","logs2"])
import corner
samples = pm.trace_to_dataframe(trace, varnames=["period", "r"])
truth = np.concatenate(
xo.eval_in_model([period, r], model.test_point, model=model)
)
_ = corner.corner(
samples,
truths=truth,
labels=["period 1", "radius 1"],
)
# Compute the GP prediction
gp_mod = np.median(trace["gp_pred"] + trace["mean"][:, None], axis=0)
# Get the posterior median orbital parameters
p = np.median(trace["period"])
t0 = np.median(trace["t0"])
# Plot the folded data
x_fold = (t - t0 + 0.5 * p) % p - 0.5 * p
plt.plot(x_fold, y - gp_mod, ".k", label="data", zorder=-1000)
# Overplot the phase binned light curve
bins = np.linspace(-0.41, 0.41, 50)
denom, _ = np.histogram(x_fold, bins)
num, _ = np.histogram(x_fold, bins, weights=y)
denom[num == 0] = 1.0
plt.plot(0.5 * (bins[1:] + bins[:-1]), num / denom, "o", color="C2", label="binned")
# Plot the folded model
inds = np.argsort(x_fold)
inds = inds[np.abs(x_fold)[inds] < 0.3]
pred = trace["light_curves"][:, inds, 0]
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:.5f} +/- {1:.5f} d".format(
np.mean(trace["period"]), np.std(trace["period"])
)
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.xlim(-0.3, 0.3);
def run_onetransit_inference(self, prior_d, pklpath, make_threadsafe=True):
"""
Similar to "run_transit_inference", but with more restrictive priors on
ephemeris. Also, it simultaneously fits for quadratic trend.
"""
# if the model has already been run, pull the result from the
# pickle. otherwise, run it.
if os.path.exists(pklpath):
d = pickle.load(open(pklpath, 'rb'))
self.model = d['model']
self.trace = d['trace']
self.map_estimate = d['map_estimate']
return 1
with pm.Model() as model:
assert len(self.data.keys()) == 1
name = list(self.data.keys())[0]
x_obs = list(self.data.values())[0][0]
y_obs = list(self.data.values())[0][1]
y_err = list(self.data.values())[0][2]
t_exp = list(self.data.values())[0][3]
# Fixed data errors.
sigma = y_err
# Define priors and PyMC3 random variables to sample over.
# Stellar parameters. (Following tess.world notebooks).
logg_star = pm.Normal("logg_star", mu=LOGG, sd=LOGG_STDEV)
r_star = pm.Bound(pm.Normal, lower=0.0)("r_star",
mu=RSTAR,
sd=RSTAR_STDEV)
rho_star = pm.Deterministic("rho_star",
factor * 10**logg_star / r_star)
# Transit parameters.
t0 = pm.Normal("t0",
mu=prior_d['t0'],
sd=1e-3,
testval=prior_d['t0'])
period = pm.Normal('period',
mu=prior_d['period'],
sd=3e-4,
testval=prior_d['period'])
# NOTE: might want to implement kwarg for flexibility
# u = xo.distributions.QuadLimbDark(
# "u", testval=prior_d['u']
# )
u0 = pm.Uniform('u[0]',
lower=prior_d['u[0]'] - 0.15,
upper=prior_d['u[0]'] + 0.15,
testval=prior_d['u[0]'])
u1 = pm.Uniform('u[1]',
lower=prior_d['u[1]'] - 0.15,
upper=prior_d['u[1]'] + 0.15,
testval=prior_d['u[1]'])
u = [u0, u1]
# # 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=1.0,
# testval_r=prior_d['r'],
# testval_b=prior_d['b']
# )
# # NOTE: apparently, it's been deprecated. DFM's manuscript notes
# that it leads to Rp/Rs values biased high
log_r = pm.Uniform('log_r',
lower=np.log(1e-2),
upper=np.log(1),
testval=prior_d['log_r'])
r = pm.Deterministic('r', tt.exp(log_r))
b = xo.distributions.ImpactParameter("b",
ror=r,
testval=prior_d['b'])
# the transit
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
rho_star=rho_star)
transit_lc = pm.Deterministic(
'transit_lc',
xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=x_obs, texp=t_exp).T.flatten())
# quadratic trend parameters
mean = pm.Normal(f"{name}_mean",
mu=prior_d[f'{name}_mean'],
sd=1e-2,
testval=prior_d[f'{name}_mean'])
a1 = pm.Normal(f"{name}_a1",
mu=prior_d[f'{name}_a1'],
sd=1,
testval=prior_d[f'{name}_a1'])
a2 = pm.Normal(f"{name}_a2",
mu=prior_d[f'{name}_a2'],
sd=1,
testval=prior_d[f'{name}_a2'])
_tmid = np.nanmedian(x_obs)
lc_model = pm.Deterministic(
'mu_transit', mean + a1 * (x_obs - _tmid) + a2 *
(x_obs - _tmid)**2 + transit_lc)
roughdepth = pm.Deterministic(f'roughdepth',
pm.math.abs_(transit_lc).max())
#
# Derived parameters
#
# planet radius in jupiter radii
r_planet = pm.Deterministic(
"r_planet",
(r * r_star) * (1 * units.Rsun / (1 * units.Rjup)).cgs.value)
#
# eq 30 of winn+2010, ignoring planet density.
#
a_Rs = pm.Deterministic("a_Rs", (rho_star * period**2)**(1 / 3) *
(((1 * units.gram / (1 * units.cm)**3) *
(1 * units.day**2) * const.G /
(3 * np.pi))**(1 / 3)).cgs.value)
#
# cosi. assumes e=0 (e.g., Winn+2010 eq 7)
#
cosi = pm.Deterministic("cosi", b / a_Rs)
# safer than tt.arccos(cosi)
sini = pm.Deterministic("sini", pm.math.sqrt(1 - cosi**2))
#
# transit durations (T_14, T_13) for circular orbits. Winn+2010 Eq 14, 15.
# units: hours.
#
T_14 = pm.Deterministic('T_14', (period / np.pi) * tt.arcsin(
(1 / a_Rs) * pm.math.sqrt((1 + r)**2 - b**2) * (1 / sini)) *
24)
T_13 = pm.Deterministic('T_13', (period / np.pi) * tt.arcsin(
(1 / a_Rs) * pm.math.sqrt((1 - r)**2 - b**2) * (1 / sini)) *
24)
#
# mean model and likelihood
#
# mean_model = mu_transit + mean
# mu_model = pm.Deterministic('mu_model', lc_model)
likelihood = pm.Normal('obs',
mu=lc_model,
sigma=sigma,
observed=y_obs)
# Optimizing
map_estimate = pm.find_MAP(model=model)
# start = model.test_point
# if 'transit' in self.modelcomponents:
# map_estimate = xo.optimize(start=start,
# vars=[r, b, period, t0])
# map_estimate = xo.optimize(start=map_estimate)
if make_threadsafe:
pass
else:
# as described in
# https://github.com/matplotlib/matplotlib/issues/15410
# matplotlib is not threadsafe. so do not make plots before
# sampling, because some child processes tries to close a
# cached file, and crashes the sampler.
print(map_estimate)
# sample from the posterior defined by this model.
trace = pm.sample(
tune=self.N_samples,
draws=self.N_samples,
start=map_estimate,
cores=self.N_cores,
chains=self.N_chains,
step=xo.get_dense_nuts_step(target_accept=0.8),
)
with open(pklpath, 'wb') as buff:
pickle.dump(
{
'model': model,
'trace': trace,
'map_estimate': map_estimate
}, buff)
self.model = model
self.trace = trace
self.map_estimate = map_estimate
def build_model_and_get_map_soln(lc,
pri_t0,
pri_p,
pri_rprs,
pri_m_star,
pri_r_star,
verbose=False):
"""
Build the transit light curve model.
Parameters
----------
lc : `~lightkurve.LightCurve`
A light curve object with the data.
pri_t0 : float
Initial guess for mid-transit time.
pri_p : float
Initial guess for period.
pri_rprs : float
Initial guess for planet-to-star radius ratio.
pri_m_star : ndarray
Mean and standard deviation of the stellar mass estimate
in solar masses.
pri_r_star : ndarray
Mean and standard deviation of the stellar radius estimate
in solar radii.
verbose : bool, optional
Print details of optimization.
Returns
-------
model : `~pymc3.model`
A model object.
map_soln : dict
A dictionary with the maximum a posteriori estimates of the variables.
"""
# Ignore warnings from theano, unless specified elsewise
if not verbose:
warnings.filterwarnings(action='ignore',
category=FutureWarning,
module='theano')
# Ensure right data type for theano
dtype = np.dtype('float64')
dts = [arr.dtype for arr in [lc.time, lc.flux, lc.flux_err]]
if not all(dt is dtype for dt in dts):
lc.time = np.array(lc.time, dtype=dtype)
lc.flux = np.array(lc.flux, dtype=dtype)
lc.flux_err = np.array(lc.flux_err, dtype=dtype)
# Estimate flux uncertainties if not given
idx_nan = np.isnan(lc.flux_err)
if idx_nan.any():
mad = median_abs_deviation(lc.flux, scale='normal')
lc.flux_err[idx_nan] = mad
# Define the model for the light curve
with pm.Model() as model:
# Stellar mass
m_star = pm.Normal('m_star', mu=pri_m_star[0], sigma=pri_m_star[1])
# Stellar radius
r_star = pm.Normal('r_star', mu=pri_r_star[0], sigma=pri_r_star[1])
# Quadratic limb-darkening parameters
u = xo.distributions.QuadLimbDark('u', testval=np.array([0.3, 0.2]))
# Radius ratio
r = pm.Uniform('r', lower=0., upper=1., testval=pri_rprs)
# Impact parameter
b = xo.distributions.ImpactParameter(
'b',
ror=r,
)
# Period
logperiod = pm.Uniform(
'logperiod',
lower=-2.3, # 0.1 d
upper=3.4, # 30 d
testval=np.log(pri_p))
period = pm.Deterministic('period', tt.exp(logperiod))
# Mid-transit time
t0 = pm.Uniform('t0',
lower=lc.time.min(),
upper=lc.time.max(),
testval=pri_t0)
# Keplerian orbit
orbit = xo.orbits.KeplerianOrbit(m_star=m_star,
r_star=r_star,
period=period,
t0=t0,
b=b)
# Model transit light curve
light_curves = xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit,
r=r * r_star,
t=lc.time)
transit_model = pm.math.sum(light_curves, axis=-1)
# The baseline flux
f0 = pm.Normal('f0',
mu=np.median(lc.flux),
sigma=median_abs_deviation(lc.flux, scale='normal'))
# The full model
lc_model = transit_model + f0
#######################
# Track some parameters
#######################
# Track transit depth
pm.Deterministic('depth', r**2)
# Track planet radius (in Earth radii)
pm.Deterministic('rearth',
r * r_star * (units.solRad / units.earthRad).si.scale)
# Track semimajor axis (in AU)
au_per_rsun = (units.solRad / units.AU).si.scale
pm.Deterministic('a', orbit.a * au_per_rsun)
# Track system scale
pm.Deterministic('aRs',
orbit.a / r_star) # normalize by stellar radius
# Track inclination
pm.Deterministic('incl', np.rad2deg(orbit.incl))
# Track transit duration
# Seager and Mallen-Ornelas (2003) Eq. 3
sini = np.sin(orbit.incl)
t14 = ((period / np.pi) * np.arcsin(
(r_star / orbit.a * sini) * np.sqrt(
(1. + r)**2 - b**2))) * 24. * 60. # min
t14 = pm.Deterministic('t14', t14)
# Track stellar density (in cgs units)
rho_star = pm.Deterministic('rho_star', orbit.rho_star)
# Track stellar density (in units of solar density)
rho_sol = (units.solMass /
(4. / 3. * np.pi * units.solRad**3)).cgs.value
pm.Deterministic('rho_star_sol', orbit.rho_star / rho_sol)
# Track x2
x2 = pm.math.sum(((lc.flux - lc_model) / lc.flux_err)**2)
x2 = pm.Deterministic('x2', x2)
# # Fit for variance
# logs2 = pm.Normal('logs2', mu=np.log(np.var(lc.flux)), sigma=1)
# sigma = pm.Deterministic('sigma', pm.math.sqrt(pm.math.exp(logs2)))
# The likelihood function
pm.Normal('obs', mu=lc_model, sigma=lc.flux_err, observed=lc.flux)
# Fit for the maximum a posteriori parameters
map_soln = xo.optimize(start=model.test_point, verbose=verbose)
map_soln = xo.optimize(start=map_soln,
vars=[f0, period, t0, r],
verbose=verbose)
map_soln = xo.optimize(start=map_soln, vars=rho_star, verbose=verbose)
map_soln = xo.optimize(start=map_soln, vars=t14, verbose=verbose)
map_soln = xo.optimize(start=map_soln, verbose=verbose)
# Reset warning filter
warnings.resetwarnings()
return model, map_soln
def run_allindivtransit_inference(self,
prior_d,
pklpath,
make_threadsafe=True,
target_accept=0.8):
# if the model has already been run, pull the result from the
# pickle. otherwise, run it.
if os.path.exists(pklpath):
d = pickle.load(open(pklpath, 'rb'))
self.model = d['model']
self.trace = d['trace']
self.map_estimate = d['map_estimate']
return 1
with pm.Model() as model:
# Shared parameters
# Stellar parameters. (Following tess.world notebooks).
logg_star = pm.Normal("logg_star", mu=LOGG, sd=LOGG_STDEV)
r_star = pm.Bound(pm.Normal, lower=0.0)("r_star",
mu=RSTAR,
sd=RSTAR_STDEV)
rho_star = pm.Deterministic("rho_star",
factor * 10**logg_star / r_star)
# fix Rp/Rs across bandpasses, b/c you're assuming it's a planet
log_r = pm.Uniform('log_r',
lower=np.log(1e-2),
upper=np.log(1),
testval=prior_d['log_r'])
r = pm.Deterministic('r', tt.exp(log_r))
# Some orbital parameters
t0 = pm.Normal("t0",
mu=prior_d['t0'],
sd=1e-1,
testval=prior_d['t0'])
period = pm.Normal('period',
mu=prior_d['period'],
sd=1e-1,
testval=prior_d['period'])
b = xo.distributions.ImpactParameter("b",
ror=r,
testval=prior_d['b'])
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
rho_star=rho_star)
# NOTE: limb-darkening should be bandpass specific, but we don't
# have the SNR to justify that, so go with TESS-dominated
# u = xo.QuadLimbDark("u")
# NOTE: deprecated(?)
delta_u = 0.15
u0 = pm.Uniform('u[0]',
lower=prior_d['u[0]'] - delta_u,
upper=prior_d['u[0]'] + delta_u,
testval=prior_d['u[0]'])
u1 = pm.Uniform('u[1]',
lower=prior_d['u[1]'] - delta_u,
upper=prior_d['u[1]'] + delta_u,
testval=prior_d['u[1]'])
u = [u0, u1]
star = xo.LimbDarkLightCurve(u)
# Loop over "instruments" (TESS, then each ground-based lightcurve)
parameters = dict()
lc_models = dict()
roughdepths = dict()
for n, (name, (x, y, yerr, texp)) in enumerate(self.data.items()):
if 'tess' in name:
delta_trend = 0.100
else:
delta_trend = 0.050
# Define per-instrument parameters in a submodel, to not need
# to prefix the names. Yields e.g., "TESS_0_mean",
# "elsauce_0_mean", "elsauce_2_a2"
mean = pm.Normal(f'{name}_mean',
mu=prior_d[f'{name}_mean'],
sd=1e-2,
testval=prior_d[f'{name}_mean'])
a1 = pm.Uniform(f'{name}_a1',
lower=-delta_trend,
upper=delta_trend,
testval=prior_d[f'{name}_a1'])
a2 = pm.Uniform(f'{name}_a2',
lower=-delta_trend,
upper=delta_trend,
testval=prior_d[f'{name}_a2'])
# midpoint for this definition of the quadratic trend
_tmid = np.nanmedian(x)
transit_lc = star.get_light_curve(orbit=orbit,
r=r,
t=x,
texp=texp).T.flatten()
lc_models[name] = pm.Deterministic(
f'{name}_mu_transit',
mean + a1 * (x - _tmid) + a2 * (x - _tmid)**2 + transit_lc)
roughdepths[name] = pm.Deterministic(
f'{name}_roughdepth',
pm.math.abs_(transit_lc).max())
# NOTE: might want error bar fudge.
likelihood = pm.Normal(f'{name}_obs',
mu=lc_models[name],
sigma=yerr,
observed=y)
#
# Derived parameters
#
# planet radius in jupiter radii
r_planet = pm.Deterministic(
"r_planet",
(r * r_star) * (1 * units.Rsun / (1 * units.Rjup)).cgs.value)
#
# eq 30 of winn+2010, ignoring planet density.
#
a_Rs = pm.Deterministic("a_Rs", (rho_star * period**2)**(1 / 3) *
(((1 * units.gram / (1 * units.cm)**3) *
(1 * units.day**2) * const.G /
(3 * np.pi))**(1 / 3)).cgs.value)
#
# cosi. assumes e=0 (e.g., Winn+2010 eq 7)
#
cosi = pm.Deterministic("cosi", b / a_Rs)
# probably safer than tt.arccos(cosi)
sini = pm.Deterministic("sini", pm.math.sqrt(1 - cosi**2))
#
# transit durations (T_14, T_13) for circular orbits. Winn+2010 Eq 14, 15.
# units: hours.
#
T_14 = pm.Deterministic('T_14', (period / np.pi) * tt.arcsin(
(1 / a_Rs) * pm.math.sqrt((1 + r)**2 - b**2) * (1 / sini)) *
24)
T_13 = pm.Deterministic('T_13', (period / np.pi) * tt.arcsin(
(1 / a_Rs) * pm.math.sqrt((1 - r)**2 - b**2) * (1 / sini)) *
24)
map_estimate = pm.find_MAP(model=model)
# if make_threadsafe:
# pass
# else:
# # NOTE: would usually plot MAP estimate here, but really
# # there's not a huge need.
# print(map_estimate)
# for k,v in map_estimate.items():
# if 'transit' not in k:
# print(k, v)
# NOTE: could start at map_estimate, which currently is not being
# used for anything.
start = model.test_point
trace = pm.sample(
tune=self.N_samples,
draws=self.N_samples,
start=start,
cores=self.N_cores,
chains=self.N_chains,
step=xo.get_dense_nuts_step(target_accept=target_accept),
)
with open(pklpath, 'wb') as buff:
pickle.dump(
{
'model': model,
'trace': trace,
'map_estimate': map_estimate
}, buff)
self.model = model
self.trace = trace
self.map_estimate = map_estimate
def run_alltransit_inference(self, prior_d, pklpath, make_threadsafe=True):
# if the model has already been run, pull the result from the
# pickle. otherwise, run it.
if os.path.exists(pklpath):
d = pickle.load(open(pklpath, 'rb'))
self.model = d['model']
self.trace = d['trace']
self.map_estimate = d['map_estimate']
return 1
with pm.Model() as model:
# Shared parameters
# Stellar parameters. (Following tess.world notebooks).
logg_star = pm.Normal("logg_star", mu=LOGG, sd=LOGG_STDEV)
r_star = pm.Bound(pm.Normal, lower=0.0)("r_star",
mu=RSTAR,
sd=RSTAR_STDEV)
rho_star = pm.Deterministic("rho_star",
factor * 10**logg_star / r_star)
# fix Rp/Rs across bandpasses, b/c you're assuming it's a planet
if 'quaddepthvar' not in self.modelid:
log_r = pm.Uniform('log_r',
lower=np.log(1e-2),
upper=np.log(1),
testval=prior_d['log_r'])
r = pm.Deterministic('r', tt.exp(log_r))
else:
log_r_Tband = pm.Uniform('log_r_Tband',
lower=np.log(1e-2),
upper=np.log(1),
testval=prior_d['log_r_Tband'])
r_Tband = pm.Deterministic('r_Tband', tt.exp(log_r_Tband))
log_r_Rband = pm.Uniform('log_r_Rband',
lower=np.log(1e-2),
upper=np.log(1),
testval=prior_d['log_r_Rband'])
r_Rband = pm.Deterministic('r_Rband', tt.exp(log_r_Rband))
log_r_Bband = pm.Uniform('log_r_Bband',
lower=np.log(1e-2),
upper=np.log(1),
testval=prior_d['log_r_Bband'])
r_Bband = pm.Deterministic('r_Bband', tt.exp(log_r_Bband))
r = r_Tband
# Some orbital parameters
t0 = pm.Normal("t0",
mu=prior_d['t0'],
sd=5e-3,
testval=prior_d['t0'])
period = pm.Normal('period',
mu=prior_d['period'],
sd=5e-3,
testval=prior_d['period'])
b = xo.distributions.ImpactParameter("b",
ror=r,
testval=prior_d['b'])
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
rho_star=rho_star)
# NOTE: limb-darkening should be bandpass specific, but we don't
# have the SNR to justify that, so go with TESS-dominated
u0 = pm.Uniform('u[0]',
lower=prior_d['u[0]'] - 0.15,
upper=prior_d['u[0]'] + 0.15,
testval=prior_d['u[0]'])
u1 = pm.Uniform('u[1]',
lower=prior_d['u[1]'] - 0.15,
upper=prior_d['u[1]'] + 0.15,
testval=prior_d['u[1]'])
u = [u0, u1]
star = xo.LimbDarkLightCurve(u)
# Loop over "instruments" (TESS, then each ground-based lightcurve)
parameters = dict()
lc_models = dict()
roughdepths = dict()
for n, (name, (x, y, yerr, texp)) in enumerate(self.data.items()):
# Define per-instrument parameters in a submodel, to not need
# to prefix the names. Yields e.g., "TESS_mean",
# "elsauce_0_mean", "elsauce_2_a2"
with pm.Model(name=name, model=model):
# Transit parameters.
mean = pm.Normal("mean",
mu=prior_d[f'{name}_mean'],
sd=1e-2,
testval=prior_d[f'{name}_mean'])
if 'quad' in self.modelid:
if name != 'tess':
# units: rel flux per day.
a1 = pm.Normal("a1",
mu=prior_d[f'{name}_a1'],
sd=1,
testval=prior_d[f'{name}_a1'])
# units: rel flux per day^2.
a2 = pm.Normal("a2",
mu=prior_d[f'{name}_a2'],
sd=1,
testval=prior_d[f'{name}_a2'])
if self.modelid == 'alltransit':
lc_models[name] = pm.Deterministic(
f'{name}_mu_transit', mean + star.get_light_curve(
orbit=orbit, r=r, t=x, texp=texp).T.flatten())
elif self.modelid == 'alltransit_quad':
if name != 'tess':
# midpoint for this definition of the quadratic trend
_tmid = np.nanmedian(x)
lc_models[name] = pm.Deterministic(
f'{name}_mu_transit', mean + a1 * (x - _tmid) +
a2 * (x - _tmid)**2 + star.get_light_curve(
orbit=orbit, r=r, t=x, texp=texp).T.flatten())
elif name == 'tess':
lc_models[name] = pm.Deterministic(
f'{name}_mu_transit', mean + star.get_light_curve(
orbit=orbit, r=r, t=x, texp=texp).T.flatten())
elif self.modelid == 'alltransit_quaddepthvar':
if name != 'tess':
# midpoint for this definition of the quadratic trend
_tmid = np.nanmedian(x)
# do custom depth-to-
if (name == 'elsauce_20200401'
or name == 'elsauce_20200426'):
r = r_Rband
elif name == 'elsauce_20200521':
r = r_Tband
elif name == 'elsauce_20200614':
r = r_Bband
transit_lc = star.get_light_curve(
orbit=orbit, r=r, t=x, texp=texp).T.flatten()
lc_models[name] = pm.Deterministic(
f'{name}_mu_transit', mean + a1 * (x - _tmid) +
a2 * (x - _tmid)**2 + transit_lc)
roughdepths[name] = pm.Deterministic(
f'{name}_roughdepth',
pm.math.abs_(transit_lc).max())
elif name == 'tess':
r = r_Tband
transit_lc = star.get_light_curve(
orbit=orbit, r=r, t=x, texp=texp).T.flatten()
lc_models[name] = pm.Deterministic(
f'{name}_mu_transit', mean + transit_lc)
roughdepths[name] = pm.Deterministic(
f'{name}_roughdepth',
pm.math.abs_(transit_lc).max())
# TODO: add error bar fudge
likelihood = pm.Normal(f'{name}_obs',
mu=lc_models[name],
sigma=yerr,
observed=y)
#
# Derived parameters
#
if self.modelid == 'alltransit_quaddepthvar':
r = r_Tband
# planet radius in jupiter radii
r_planet = pm.Deterministic(
"r_planet",
(r * r_star) * (1 * units.Rsun / (1 * units.Rjup)).cgs.value)
#
# eq 30 of winn+2010, ignoring planet density.
#
a_Rs = pm.Deterministic("a_Rs", (rho_star * period**2)**(1 / 3) *
(((1 * units.gram / (1 * units.cm)**3) *
(1 * units.day**2) * const.G /
(3 * np.pi))**(1 / 3)).cgs.value)
#
# cosi. assumes e=0 (e.g., Winn+2010 eq 7)
#
cosi = pm.Deterministic("cosi", b / a_Rs)
# probably safer than tt.arccos(cosi)
sini = pm.Deterministic("sini", pm.math.sqrt(1 - cosi**2))
#
# transit durations (T_14, T_13) for circular orbits. Winn+2010 Eq 14, 15.
# units: hours.
#
T_14 = pm.Deterministic('T_14', (period / np.pi) * tt.arcsin(
(1 / a_Rs) * pm.math.sqrt((1 + r)**2 - b**2) * (1 / sini)) *
24)
T_13 = pm.Deterministic('T_13', (period / np.pi) * tt.arcsin(
(1 / a_Rs) * pm.math.sqrt((1 - r)**2 - b**2) * (1 / sini)) *
24)
# Optimizing
map_estimate = pm.find_MAP(model=model)
# start = model.test_point
# if 'transit' in self.modelcomponents:
# map_estimate = xo.optimize(start=start,
# vars=[r, b, period, t0])
# map_estimate = xo.optimize(start=map_estimate)
if make_threadsafe:
pass
else:
# NOTE: would usually plot MAP estimate here, but really
# there's not a huge need.
print(map_estimate)
pass
# sample from the posterior defined by this model.
trace = pm.sample(
tune=self.N_samples,
draws=self.N_samples,
start=map_estimate,
cores=self.N_cores,
chains=self.N_chains,
step=xo.get_dense_nuts_step(target_accept=0.8),
)
with open(pklpath, 'wb') as buff:
pickle.dump(
{
'model': model,
'trace': trace,
'map_estimate': map_estimate
}, buff)
self.model = model
self.trace = trace
self.map_estimate = map_estimate
def run_fitting():
with pm.Model() as model:
# Stellar parameters
mean = pm.Normal("mean", mu=0.0, sigma=10.0 * 1e-3)
u = xo.distributions.QuadLimbDark("u")
star_params = [mean, u]
# Gaussian process noise model
sigma = pm.InverseGamma("sigma",
alpha=3.0,
beta=2 *
np.median(self_.relative_flux_errors))
log_Sw4 = pm.Normal("log_Sw4", mu=0.0, sigma=10.0)
log_w0 = pm.Normal("log_w0",
mu=np.log(2 * np.pi / 10.0),
sigma=10.0)
kernel = xo.gp.terms.SHOTerm(log_Sw4=log_Sw4,
log_w0=log_w0,
Q=1.0 / 3)
noise_params = [sigma, log_Sw4, log_w0]
# Planet parameters
log_ror = pm.Normal("log_ror",
mu=0.5 * np.log(self_.depth),
sigma=10.0 * 1e-3)
ror = pm.Deterministic("ror", tt.exp(log_ror))
depth = pm.Deterministic("depth", tt.square(ror))
# Orbital parameters
log_period = pm.Normal("log_period",
mu=np.log(self_.period),
sigma=1.0)
t0 = pm.Normal("t0", mu=self_.transit_epoch, sigma=1.0)
log_dur = pm.Normal("log_dur", mu=np.log(0.1), sigma=10.0)
b = xo.distributions.ImpactParameter("b", ror=ror)
period = pm.Deterministic("period", tt.exp(log_period))
dur = pm.Deterministic("dur", tt.exp(log_dur))
# Set up the orbit
orbit = xo.orbits.KeplerianOrbit(period=period,
duration=dur,
t0=t0,
b=b,
r_star=self.star_radius)
# We're going to track the implied density for reasons that will become clear later
pm.Deterministic("rho_circ", orbit.rho_star)
# Set up the mean transit model
star = xo.LimbDarkLightCurve(u)
def lc_model(t):
return mean + tt.sum(star.get_light_curve(
orbit=orbit, r=ror * self.star_radius, t=t),
axis=-1)
# Finally the GP observation model
gp = xo.gp.GP(kernel,
self_.times,
(self_.relative_flux_errors**2) + (sigma**2),
mean=lc_model)
gp.marginal("obs", observed=self_.relative_fluxes)
# Double check that everything looks good - we shouldn't see any NaNs!
print(model.check_test_point())
# Optimize the model
map_soln = model.test_point
map_soln = xo.optimize(map_soln, [sigma])
map_soln = xo.optimize(map_soln, [log_ror, b, log_dur])
map_soln = xo.optimize(map_soln, noise_params)
map_soln = xo.optimize(map_soln, star_params)
map_soln = xo.optimize(map_soln)
with model:
gp_pred, lc_pred = xo.eval_in_model(
[gp.predict(), lc_model(self_.times)], map_soln)
x_fold = (self_.times - map_soln["t0"] + 0.5 * map_soln["period"]
) % map_soln["period"] - 0.5 * map_soln["period"]
inds = np.argsort(x_fold)
initial_fit_data_source.data['Folded time (days)'] = x_fold
initial_fit_data_source.data[
'Relative flux'] = self_.relative_fluxes - gp_pred - map_soln[
"mean"]
initial_fit_data_source.data[
'Fit'] = lc_pred[inds] - map_soln["mean"]
initial_fit_data_source.data['Fit time'] = x_fold[
inds] # TODO: This is terrible, you should be able to line them up *afterward* to not make a duplicate time column
with model:
trace = pm.sample(
tune=2000,
draws=2000,
start=map_soln,
chains=4,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
trace_summary = pm.summary(
trace, round_to='none'
) # Not a typo. PyMC3 wants 'none' as a string here.
epoch = round(
trace_summary['mean']['t0'],
3) # Round the epoch differently, as BTJD needs more digits.
trace_summary['mean'] = self_.round_series_to_significant_figures(
trace_summary['mean'], 5)
trace_summary['mean']['t0'] = epoch
parameters_table_data_source.data = trace_summary
parameters_table_data_source.data[
'parameter'] = trace_summary.index
with pd.option_context('display.max_columns', None,
'display.max_rows', None):
print(trace_summary)
print(f'Star radius: {self.star_radius}')
(assuming you haven't already put together a ~/.theanorc file)
"""
import numpy as np
import matplotlib.pyplot as plt
import theano
theano.config.gcc.cxxflags = "-Wno-c++11-narrowing"
import exoplanet as xo
# The light curve calculation requires an orbit
orbit = xo.orbits.KeplerianOrbit(period=3.456)
# Compute a limb-darkened light curve using starry
# Note: the `eval` is needed because this is using Theano in
# the background
t = np.linspace(-0.1, 0.1, 1000)
u = [0.3, 0.2]
light_curve = (xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit,
r=0.1,
t=t,
texp=0.02).eval())
plt.plot(t, light_curve, color="C0", lw=2)
plt.ylabel("relative flux")
plt.xlabel("time [days]")
_ = plt.xlim(t.min(), t.max())
outpath = '../results/test_results/theano_import_check.png'
plt.savefig(outpath)
print(f'saved {outpath}')
# %% [markdown]
# The *exoplanet* package is mostly just glue that connects many other ideas and software.
# In a situation like this, it can be easy to forget about the important infrastructure upon which our science is built.
# In order to make sure that you can easily give credit where credit is due, we have tried to make it as painless as possible to work out which citations are expected for a model fit using *exoplanet* by including a :func:`exoplanet.citations.get_citations_for_model` function that introspects the current PyMC3 model and constructs a list of citations for the functions used in that model.
#
# For example, you might compute a quadratically limb darkened light curve using `starry` (via the :class:`exoplanet.light_curve.StarryLightCurve` class):
# %%
import pymc3 as pm
import exoplanet as xo
with pm.Model() as model:
u = xo.distributions.QuadLimbDark("u")
orbit = xo.orbits.KeplerianOrbit(period=10.0)
light_curve = xo.LimbDarkLightCurve(u)
transit = light_curve.get_light_curve(r=0.1, orbit=orbit, t=[0.0, 0.1])
txt, bib = xo.citations.get_citations_for_model()
# %% [markdown]
# The :func:`exoplanet.citations.get_citations_for_model` function would generate an acknowledgement that cites:
#
# * [PyMC3](https://docs.pymc.io/#citing-pymc3): for the inference engine and modeling framework,
# * [Theano](http://deeplearning.net/software/theano/citation.html): for the numerical infrastructure,
# * [AstroPy](http://www.astropy.org/acknowledging.html): for units and constants,
# * [Kipping (2013)](https://arxiv.org/abs/1308.0009): for the reparameterization of the limb darkening parameters for a quadratic law, and
# * [Luger, et al. (2018)](https://arxiv.org/abs/1810.06559): for the light curve calculation.
#
# The first output from :func:`exoplanet.citations.get_citations_for_model` gives the acknowledgement text:
def build_model(mask=None, start=None):
with pm.Model() as model:
# The baseline flux
mean = pm.Normal("mean", mu=0.0, sd=0.00001)
# The time of a reference transit for each planet
t0 = pm.Normal("t0", mu=t0s, sd=1.0, shape=1)
# The log period; also tracking the period itself
logP = pm.Normal("logP", mu=np.log(periods), sd=0.01, shape=1)
rho_star = pm.Normal("rho_star", mu=0.14, sd=0.01, shape=1)
r_star = pm.Normal("r_star", mu=2.7, sd=0.01, shape=1)
period = pm.Deterministic("period", pm.math.exp(logP))
# The Kipping (2013) parameterization for quadratic limb darkening paramters
u = xo.distributions.QuadLimbDark("u", testval=np.array([0.3, 0.2]))
r = pm.Uniform("r", lower=0.01, upper=0.3, shape=1, testval=0.15)
b = xo.distributions.ImpactParameter("b", ror=r, shape=1, testval=0.5)
# Transit jitter & GP parameters
logs2 = pm.Normal("logs2", mu=np.log(np.var(y)), sd=10)
logw0 = pm.Normal("logw0", mu=0, sd=10)
logSw4 = pm.Normal("logSw4", mu=np.log(np.var(y)), sd=10)
# Set up a Keplerian orbit for the planets
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
rho_star=rho_star,
r_star=r_star)
# Compute the model light curve using starry
light_curves = xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit,
r=r,
t=t)
light_curve = pm.math.sum(light_curves, axis=-1) + mean
# Here we track the value of the model light curve for plotting
# purposes
pm.Deterministic("light_curves", light_curves)
S1 = pm.InverseGamma(
"S1", **estimate_inverse_gamma_parameters(0.5**2, 10.0**2))
S2 = pm.InverseGamma(
"S2", **estimate_inverse_gamma_parameters(0.25**2, 1.0**2))
w1 = pm.InverseGamma(
"w1", **estimate_inverse_gamma_parameters(2 * np.pi / 10.0, np.pi))
w2 = pm.InverseGamma(
"w2", **estimate_inverse_gamma_parameters(0.5 * np.pi, 2 * np.pi))
log_Q = pm.Uniform("log_Q", lower=np.log(2), upper=np.log(10))
# Set up the kernel an GP
kernel = terms.SHOTerm(S_tot=S1, w0=w1, Q=1.0 / np.sqrt(2))
kernel += terms.SHOTerm(S_tot=S2, w0=w2, log_Q=log_Q)
gp = GP(kernel, t, yerr**2, mean=mean)
gp.marginal("gp", observed=y)
pm.Deterministic("gp_pred", gp.predict())
# The likelihood function assuming known Gaussian uncertainty
pm.Normal("obs", mu=light_curve, sd=yerr, observed=y)
# Fit for the maximum a posteriori parameters given the simuated
# dataset
map_soln = xo.optimize(start=model.test_point)
return model, map_soln
def fit_shape(self,
run_MCMC=False,
r_start=None,
b_start=None,
verbose=True):
"""
Fit the orbital parameters to the shape of the folded lightcurve data.
Parameters:
run_MCMC -- Boolean to run Monte Carlo Markov Chain methods. This will take longer but yields better results
and gives uncertainties.
r_start -- Starting estimate for the relative radius.
b_start -- Starting estimate for the impact parameter.
"""
if verbose: print("Optimising the shape of the orbital model:")
folded_lc = self.lightcurve.fold(self.p_ref,
self.t0_ref,
ttvs=self.pars['ttvs'])
t = folded_lc.time * self.p_ref
y = folded_lc.flux
sd = folded_lc.flux_err
if r_start is None: r_start = 0.055
if b_start is None: b_start = 0.5
with pm.Model() as model:
mean = pm.Normal("mean", mu=1.0, sd=0.1) # Baseline flux
t0 = pm.Normal("t0", mu=0, sd=0.025)
u = xo.distributions.QuadLimbDark(
"u") # Quadratic limb-darkening parameters
r = pm.Uniform("r", lower=0.01, upper=0.1,
testval=r_start) # radius ratio
b = xo.distributions.ImpactParameter(
"b", ror=r, testval=b_start) # Impact parameter
orbit = xo.orbits.KeplerianOrbit(period=self.p_ref, t0=t0, b=b)
# Compute the model light curve
lc = xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit,
r=r,
t=t)
light_curve = pm.math.sum(lc, axis=-1) + mean
pm.Deterministic(
"light_curve", light_curve
) # track the value of the model light curve for plotting purposes
# The likelihood function
pm.Normal("obs", mu=light_curve, sd=sd, observed=y)
map_soln = xo.optimize(start=model.test_point,
verbose=verbose,
progress_bar=False)
for k in ['mean', 't0', 'u', 'r', 'b']:
self.pars[k] = map_soln[k]
if verbose: print('\t', k, '=', self.pars[k])
if run_MCMC:
np.random.seed(42)
with model:
trace = pm.sample(
tune=3000,
draws=3000,
start=map_soln,
cores=1,
chains=2,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
for k in ['mean', 't0', 'u', 'r', 'b']:
self.pars[k] = np.median(trace[k], axis=0)
self.pars['e_' + k] = self.pars[k] - np.percentile(
trace[k], 16, axis=0)
self.pars['E_' + k] = -self.pars[k] + np.percentile(
trace[k], 84, axis=0)
if verbose:
print(
f"\t{k} = {self.pars[k]} /+ {self.pars['E_'+k]} /- {self.pars['e_'+k]}"
)
# The difference between those two is
print("{:.2f}".format(((t0_obs2 - t0_obs1) * u.day).to(u.minute)))
# as expected.
# As an example, here's a comparison between a transit light curve that includes light travel delay and one that doesnt't:
# +
# Region around transit
t = np.linspace(-0.5, 0.5, 20000)
# No light delay
light_curve1 = (xo.LimbDarkLightCurve([0.5, 0.25]).get_light_curve(orbit=orbit,
r=0.1,
t=t).eval())
# With light delay
light_curve2 = (xo.LimbDarkLightCurve([0.5, 0.25]).get_light_curve(
orbit=orbit, r=0.1, t=t, light_delay=True).eval())
# Plot!
fig = plt.figure(figsize=(8, 3))
plt.plot(t, light_curve1, label="no delay")
plt.plot(t, light_curve2, label="with delay")
plt.xlabel("time [days]")
plt.ylabel("relative flux")
plt.legend(fontsize=10, loc="lower right")
plt.title("Light delay causes transits to occur 8.32 minutes early",
fontsize=14)
lcs = np.zeros((len(t), len(inds) - 1))
for i in range(len(inds) - 1):
lcs[:, i] = np.sum(noisy_lc[:, inds[i]:inds[i + 1]], axis=1)
import exoplanet as xo
orbit = xo.orbits.KeplerianOrbit(period=5.0 * 60 * 60)
u = [0.3, 0.2]
rp = 0.0203
rm = 0.01
t0p = 10
t0m = 5
planet = (xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit,
r=rp,
t=t / (60 * 60) - t0p,
texp=np.mean(np.diff(t)) /
(60 * 60)).eval()).T[0]
moon = (xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit,
r=rm,
t=t / (60 * 60) - t0m,
texp=np.mean(np.diff(t)) /
(60 * 60)).eval()).T[0]
from astropy.table import table
data = lcs + np.mean(lcs, axis=0) * (moon[:, None] + planet[:, None])
data = table.QTable(data, names=["bin{0}".format(i) for i in range(nbins)])
data.add_column(t, name='time')
data.meta['name'] = name
data.meta['S0'] = S0
logP = pm.Normal("logP", mu=np.log(periods), sd=0.1, shape=2)
period = pm.Deterministic("period", pm.math.exp(logP))
# The Kipping (2013) parameterization for quadratic limb darkening paramters
u = xo.distributions.QuadLimbDark("u", testval=us)
r = pm.Uniform(
"r", lower=0.01, upper=0.1, shape=2, testval=rs
)
b = xo.distributions.ImpactParameter(
"b", ror=r, shape=2, testval=bs
)
# Set up a Keplerian orbit for the planets
orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, b=b)
# Compute the model light curve using starry
light_curves = xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=t, texp=texp
)
light_curve = pm.math.sum(light_curves, axis=-1) + mean
# Here we track the value of the model light curve for plotting
# purposes
pm.Deterministic("light_curves", light_curves)
# In this line, we simulate the dataset that we will fit
y = xo.eval_in_model(light_curve)
y += yerr * np.random.randn(len(y))
# The likelihood function assuming known Gaussian uncertainty
pm.Normal("obs", mu=light_curve, sd=yerr, observed=y)
# Fit for the maximum a posteriori parameters given the simuated
# dataset
map_estimate = xo.optimize(start=model.test_point)
# plot the simulated data and the maximum a posteriori model to make sure that
def get_model_transit_quad(paramd,
time_eval,
_tmid,
t_exp=2 / (60 * 24),
includemean=1):
"""
Same as get_model_transit, but for a transit + quadratic trend. The midtime
of the trend must be the same as used in timmy.modelfitter for the a1 and
a2 coefficients to be correctly defined.
"""
import exoplanet as xo
period = paramd['period']
t0 = paramd['t0']
try:
r = paramd['r']
except KeyError:
r = np.exp(paramd['log_r'])
b = paramd['b']
u0 = paramd['u[0]']
u1 = paramd['u[1]']
r_star = paramd['r_star']
logg_star = paramd['logg_star']
try:
mean = paramd['mean']
except KeyError:
mean_key = [k for k in list(paramd.keys()) if 'mean' in k]
assert len(mean_key) == 1
mean_key = mean_key[0]
mean = paramd[mean_key]
a1 = paramd[[k for k in list(paramd.keys()) if '_a1' in k][0]]
a2 = paramd[[k for k in list(paramd.keys()) if '_a2' in k][0]]
# factor * 10**logg / r_star = rho
factor = 5.141596357654149e-05
rho_star = factor * 10**logg_star / r_star
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
rho_star=rho_star)
u = [u0, u1]
mu_transit = xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=time_eval, texp=t_exp).T.flatten()
mu_model = (mean + a1 * (time_eval - _tmid) + a2 * (time_eval - _tmid)**2 +
mu_transit.eval())
if includemean:
mu_trend = (mean + a1 * (time_eval - _tmid) + a2 *
(time_eval - _tmid)**2)
else:
mu_trend = (a1 * (time_eval - _tmid) + a2 * (time_eval - _tmid)**2)
return mu_model, mu_trend
plt.ion()
plt.clf()
plt.plot(t[um], y[um], '.')
t = t[um]
y = y[um]
yerr = np.zeros(len(y)) + 0.001
orbit = xo.orbits.KeplerianOrbit(period=periods,
t0=t0s,
b=0.5,
rho_star=0.14,
r_star=2.7)
u = [0.3, 0.2]
light_curve = (xo.LimbDarkLightCurve(u).get_light_curve(orbit=orbit,
r=0.19,
t=t).eval())
plt.clf()
plt.plot(t, y, '.')
plt.plot(t, light_curve, color='red', lw=2)
plt.xlim([1994.8, 1995.6])
def build_model(mask=None, start=None):
with pm.Model() as model:
# The baseline flux
mean = pm.Normal("mean", mu=0.0, sd=0.00001)
def run_fitting():
times = self.light_curve_data_source.data['Time (BTJD)'].astype(
np.float32)
flux_errors = self.light_curve_data_source.data[
'Normalized PDCSAP flux error']
fluxes = self.light_curve_data_source.data[
'Normalized PDCSAP flux']
relative_times = self.light_curve_data_source.data['Time (days)']
nan_indexes = np.union1d(
np.argwhere(np.isnan(fluxes)),
np.union1d(np.argwhere(np.isnan(times)),
np.argwhere(np.isnan(flux_errors))))
fluxes = np.delete(fluxes, nan_indexes)
flux_errors = np.delete(flux_errors, nan_indexes)
times = np.delete(times, nan_indexes)
relative_times = np.delete(relative_times, nan_indexes)
with pm.Model() as model:
# Stellar parameters
mean = pm.Normal("mean", mu=0.0, sigma=10.0 * 1e-3)
u = xo.distributions.QuadLimbDark("u")
star_params = [mean, u]
# Gaussian process noise model
sigma = pm.InverseGamma("sigma",
alpha=3.0,
beta=2 * np.nanmedian(flux_errors))
log_Sw4 = pm.Normal("log_Sw4", mu=0.0, sigma=10.0)
log_w0 = pm.Normal("log_w0",
mu=np.log(2 * np.pi / 10.0),
sigma=10.0)
kernel = xo.gp.terms.SHOTerm(log_Sw4=log_Sw4,
log_w0=log_w0,
Q=1.0 / 3)
noise_params = [sigma, log_Sw4, log_w0]
# Planet parameters
log_ror = pm.Normal("log_ror",
mu=0.5 * np.log(self_.depth),
sigma=10.0 * 1e-3)
ror = pm.Deterministic("ror", tt.exp(log_ror))
depth = pm.Deterministic('Transit depth (relative flux)',
tt.square(ror))
planet_radius = pm.Deterministic('Planet radius (solar radii)',
ror * self_.star_radius)
# Orbital parameters
log_period = pm.Normal("log_period",
mu=np.log(self_.period),
sigma=1.0)
t0 = pm.Normal('Transit epoch (BTJD)',
mu=self_.transit_epoch,
sigma=1.0)
log_dur = pm.Normal("log_dur", mu=np.log(0.1), sigma=10.0)
b = xo.distributions.ImpactParameter("b", ror=ror)
period = pm.Deterministic('Transit period (days)',
tt.exp(log_period))
dur = pm.Deterministic('Transit duration (days)',
tt.exp(log_dur))
# Set up the orbit
orbit = xo.orbits.KeplerianOrbit(period=period,
duration=dur,
t0=t0,
b=b,
r_star=self.star_radius)
# We're going to track the implied density for reasons that will become clear later
pm.Deterministic("rho_circ", orbit.rho_star)
# Set up the mean transit model
star = xo.LimbDarkLightCurve(u)
def lc_model(t):
return mean + tt.sum(star.get_light_curve(
orbit=orbit, r=ror * self.star_radius, t=t),
axis=-1)
# Finally the GP observation model
gp = xo.gp.GP(kernel,
times, (flux_errors**2) + (sigma**2),
mean=lc_model)
gp.marginal("obs", observed=fluxes)
# Double check that everything looks good - we shouldn't see any NaNs!
print(model.check_test_point())
# Optimize the model
map_soln = model.test_point
map_soln = xo.optimize(map_soln, [sigma])
map_soln = xo.optimize(map_soln, [log_ror, b, log_dur])
map_soln = xo.optimize(map_soln, noise_params)
map_soln = xo.optimize(map_soln, star_params)
map_soln = xo.optimize(map_soln)
with model:
gp_pred, lc_pred = xo.eval_in_model(
[gp.predict(), lc_model(times)], map_soln)
x_fold = (times - map_soln['Transit epoch (BTJD)'] +
0.5 * map_soln['Transit period (days)']
) % map_soln['Transit period (days)'] - 0.5 * map_soln[
'Transit period (days)']
inds = np.argsort(x_fold)
bokeh_document.add_next_tick_callback(
partial(update_initial_fit_figure, fluxes, gp_pred, inds,
lc_pred, map_soln, relative_times, times, x_fold))
self.bokeh_document.add_next_tick_callback(
partial(fit, self, map_soln, model))
def build_model(mask=None):
if mask is None:
mask = np.ones_like(x, dtype=bool)
with pm.Model() as model:
# Stellar parameters
mean = pm.Normal("mean", mu=0.0, sigma=10.0)
u = xo.distributions.QuadLimbDark("u")
# Gaussian process noise model
sigma = pm.InverseGamma("sigma", alpha=3.0, beta=2 * np.median(yerr))
S_tot = pm.InverseGamma("S_tot", alpha=3.0, beta=2 * np.median(yerr))
ell = pm.Lognormal("ell", mu=0.0, sigma=1.0)
Q = 1.0 / 3.0
w0 = 2 * np.pi / ell
S0 = S_tot / (w0 * Q)
kernel = xo.gp.terms.SHOTerm(S0=S0, w0=w0, Q=Q)
# Transit parameters
t0 = pm.Bound(
pm.Normal,
lower=t0_guess - max_duration,
upper=t0_guess + max_duration,
)(
"t0",
mu=t0_guess,
sigma=0.5 * duration_guess,
testval=t0_guess,
shape=num_toi,
)
depth = pm.Lognormal(
"transit_depth",
mu=np.log(depth_guess),
sigma=np.log(1.2),
shape=num_toi,
)
duration = pm.Bound(pm.Lognormal,
lower=min_duration,
upper=max_duration)(
"transit_duration",
mu=np.log(duration_guess),
sigma=np.log(1.2),
shape=num_toi,
testval=min(
max(duration_guess, 2 * min_duration),
0.99 * max_duration),
)
b = xo.distributions.UnitUniform("b", shape=num_toi)
# Dealing with period, treating single transits properly
period_params = []
period_values = []
t_max_values = []
for n in range(num_toi):
if single_transit[n]:
period = pm.Pareto(
f"period_{n}",
m=period_min[n],
alpha=2.0 / 3,
testval=period_guess[n],
)
period_params.append(period)
t_max_values.append(t0[n])
else:
t_max = pm.Bound(
pm.Normal,
lower=t_max_guess[n] - max_duration[n],
upper=t_max_guess[n] + max_duration[n],
)(
f"t_max_{n}",
mu=t_max_guess[n],
sigma=0.5 * duration_guess[n],
testval=t_max_guess[n],
)
period = (t_max - t0[n]) / num_periods[n]
period_params.append(t_max)
t_max_values.append(t_max)
period_values.append(period)
period = pm.Deterministic("period", tt.stack(period_values))
t_max = pm.Deterministic("t_max", tt.stack(t_max_values))
# Compute the radius ratio from the transit depth, impact parameter, and
# limb darkening parameters making the small-planet assumption
u1 = u[0]
u2 = u[1]
mu = tt.sqrt(1 - b**2)
ror = pm.Deterministic(
"ror",
tt.sqrt(1e-3 * depth * (1 - u1 / 3 - u2 / 6) /
(1 - u1 * (1 - mu) - u2 * (1 - mu)**2)),
)
# Set up the orbit
orbit = xo.orbits.KeplerianOrbit(period=period,
duration=duration,
t0=t0,
b=b)
# We're going to track the implied density for reasons that will become clear later
pm.Deterministic("rho_circ", orbit.rho_star)
# Set up the mean transit model
star = xo.LimbDarkLightCurve(u)
lc_model = tess_world.LightCurveModels(mean, star, orbit, ror)
# Finally the GP observation model
gp = xo.gp.GP(kernel, x[mask], yerr[mask]**2 + sigma**2, mean=lc_model)
gp.marginal("obs", observed=y[mask])
# This is a check on the transit depth constraint
D = tt.concatenate(
(
lc_model.light_curves(x[mask]) / depth[None, :],
tt.ones((mask.sum(), 1)),
),
axis=-1,
)
DTD = tt.dot(D.T, gp.apply_inverse(D))
DTy = tt.dot(D.T, gp.apply_inverse(y[mask, None]))
model.w = tt.slinalg.solve(DTD, DTy)[:, 0]
model.sigma_w = tt.sqrt(
tt.diag(tt.slinalg.solve(DTD, tt.eye(num_toi + 1))))
# Double check that everything looks good - we shouldn't see any NaNs!
print(model.check_test_point())
# Optimize the model
map_soln = model.test_point
map_soln = xo.optimize(map_soln, [sigma])
map_soln = xo.optimize(map_soln, [mean, depth, b, duration])
map_soln = xo.optimize(map_soln, [sigma, S_tot, ell])
map_soln = xo.optimize(map_soln, [mean, u])
map_soln = xo.optimize(map_soln, period_params)
map_soln = xo.optimize(map_soln)
# Save some of the key parameters
model.map_soln = map_soln
model.lc_model = lc_model
model.gp = gp
model.mask = mask
model.x = x[mask]
model.y = y[mask]
model.yerr = yerr[mask]
return model
def run_inference(self, prior_d, pklpath, make_threadsafe=True):
# if the model has already been run, pull the result from the
# pickle. otherwise, run it.
if os.path.exists(pklpath):
d = pickle.load(open(pklpath, 'rb'))
self.model = d['model']
self.trace = d['trace']
self.map_estimate = d['map_estimate']
return 1
with pm.Model() as model:
# Fixed data errors.
sigma = self.y_err
# Define priors and PyMC3 random variables to sample over.
# Start with the transit parameters.
mean = pm.Normal("mean",
mu=prior_d['mean'],
sd=1e-2,
testval=prior_d['mean'])
t0 = pm.Normal("t0",
mu=prior_d['t0'],
sd=5e-3,
testval=prior_d['t0'])
period = pm.Normal('period',
mu=prior_d['period'],
sd=5e-3,
testval=prior_d['period'])
u = xo.distributions.QuadLimbDark("u", testval=prior_d['u'])
r = pm.Normal("r",
mu=prior_d['r'],
sd=0.20 * prior_d['r'],
testval=prior_d['r'])
b = xo.distributions.ImpactParameter("b",
ror=r,
testval=prior_d['b'])
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
mstar=self.mstar,
rstar=self.rstar)
mu_transit = pm.Deterministic(
'mu_transit',
xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=self.x_obs,
texp=self.t_exp).T.flatten())
mean_model = mu_transit + mean
if self.modelcomponents == ['transit']:
mu_model = pm.Deterministic('mu_model', mean_model)
likelihood = pm.Normal('obs',
mu=mu_model,
sigma=sigma,
observed=self.y_obs)
if 'gprot' in self.modelcomponents:
# Instantiate the GP parameters.
P_rot = pm.Normal("P_rot",
mu=prior_d['P_rot'],
sigma=1.0,
testval=prior_d['P_rot'])
amp = pm.Uniform('amp', lower=5, upper=40, testval=10)
mix = xo.distributions.UnitUniform("mix")
log_Q0 = pm.Normal("log_Q0",
mu=1.0,
sd=10.0,
testval=prior_d['log_Q0'])
log_deltaQ = pm.Normal("log_deltaQ",
mu=2.0,
sd=10.0,
testval=prior_d['log_deltaQ'])
kernel = terms.RotationTerm(
period=P_rot,
amp=amp,
mix=mix,
log_Q0=log_Q0,
log_deltaQ=log_deltaQ,
)
gp = GP(kernel, self.x_obs, sigma**2, mean=mean_model)
# Condition the GP on the observations and add the marginal likelihood
# to the model. Needed before calling "gp.predict()".
# NOTE: This formally is the definition of the likelihood? It
# would be good to figure out how this works under the hood...
gp.marginal("transit_obs", observed=self.y_obs)
# Compute the mean model prediction for plotting purposes
mu_gprot = pm.Deterministic("mu_gprot", gp.predict())
mu_model = pm.Deterministic("mu_model", mu_gprot + mean_model)
# Optimizing
start = model.test_point
if 'transit' in self.modelcomponents:
map_estimate = xo.optimize(start=start,
vars=[r, b, period, t0])
if 'gprot' in self.modelcomponents:
map_estimate = xo.optimize(
start=map_estimate,
vars=[P_rot, amp, mix, log_Q0, log_deltaQ])
map_estimate = xo.optimize(start=map_estimate)
# map_estimate = pm.find_MAP(model=model)
# Plot the simulated data and the maximum a posteriori model to
# make sure that our initialization looks ok.
self.y_MAP = (map_estimate['mean'] + map_estimate['mu_transit'])
if 'gprot' in self.modelcomponents:
self.y_MAP += map_estimate['mu_gprot']
if make_threadsafe:
pass
else:
# as described in
# https://github.com/matplotlib/matplotlib/issues/15410
# matplotlib is not threadsafe. so do not make plots before
# sampling, because some child processes tries to close a
# cached file, and crashes the sampler.
print(map_estimate)
if self.PLOTDIR is None:
raise NotImplementedError
outpath = os.path.join(self.PLOTDIR,
'test_{}_MAP.png'.format(self.modelid))
plot_MAP_data(self.x_obs, self.y_obs, self.y_MAP, outpath)
# sample from the posterior defined by this model.
trace = pm.sample(
tune=self.N_samples,
draws=self.N_samples,
start=map_estimate,
cores=self.N_cores,
chains=self.N_chains,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
with open(pklpath, 'wb') as buff:
pickle.dump(
{
'model': model,
'trace': trace,
'map_estimate': map_estimate
}, buff)
self.model = model
self.trace = trace
self.map_estimate = map_estimate
def run_inference(self, prior_d, pklpath, make_threadsafe=True):
# if the model has already been run, pull the result from the
# pickle. otherwise, run it.
if os.path.exists(pklpath):
d = pickle.load(open(pklpath, 'rb'))
self.model = d['model']
self.trace = d['trace']
self.map_estimate = d['map_estimate']
return 1
with pm.Model() as model:
# Fixed data errors.
sigma = self.y_err
# Define priors and PyMC3 random variables to sample over.
A_d, B_d, omega_d, phi_d = {}, {}, {}, {}
_A_d, _B_d = {}, {}
for modelcomponent in self.modelcomponents:
if 'transit' in modelcomponent:
BoundedNormal = pm.Bound(pm.Normal, lower=0, upper=3)
m_star = BoundedNormal("m_star",
mu=MSTAR_VANEYKEN,
sd=MSTAR_STDEV)
r_star = BoundedNormal("r_star",
mu=RSTAR_VANEYKEN,
sd=RSTAR_STDEV)
# mean = pm.Normal(
# "mean", mu=prior_d['mean'], sd=0.02, testval=prior_d['mean']
# )
mean = pm.Uniform("mean",
lower=prior_d['mean'] - 1e-2,
upper=prior_d['mean'] + 1e-2,
testval=prior_d['mean'])
t0 = pm.Normal("t0",
mu=prior_d['t0'],
sd=0.002,
testval=prior_d['t0'])
# logP = pm.Normal(
# "logP", mu=np.log(prior_d['period']),
# sd=0.001*np.abs(np.log(prior_d['period'])),
# testval=np.log(prior_d['period'])
# )
# period = pm.Deterministic("period", pm.math.exp(logP))
period = pm.Normal('period',
mu=prior_d['period'],
sd=1e-3,
testval=prior_d['period'])
u = xo.distributions.QuadLimbDark("u",
testval=prior_d['u'])
r = pm.Normal("r",
mu=prior_d['r'],
sd=0.10 * prior_d['r'],
testval=prior_d['r'])
# r = pm.Uniform(
# "r", lower=prior_d['r']-1e-2,
# upper=prior_d['r']+1e-2, testval=prior_d['r']
# )
b = xo.distributions.ImpactParameter("b",
ror=r,
testval=prior_d['b'])
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
b=b,
mstar=m_star,
rstar=r_star)
light_curve = (
mean + xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit, r=r, t=self.x_obs, texp=self.t_exp))
#
# derived quantities
#
# stellar density in cgs
rhostar = pm.Deterministic(
"rhostar",
((m_star / ((4 * np.pi / 3) * r_star**3)) *
(1 * units.Msun / ((1 * units.Rsun)**3)).cgs.value))
# planet radius in jupiter radii
r_planet = pm.Deterministic("r_planet", (r * r_star) *
(1 * units.Rsun /
(1 * units.Rjup)).cgs.value)
#
# eq 30 of winn+2010, ignoring planet density.
#
a_Rs = pm.Deterministic(
"a_Rs", (rhostar * period**2)**(1 / 3) *
(((1 * units.gram /
(1 * units.cm)**3) * (1 * units.day**2) * const.G /
(3 * np.pi))**(1 / 3)).cgs.value)
if 'sincos' in modelcomponent:
if 'Porb' in modelcomponent:
k = 'orb'
elif 'Prot' in modelcomponent:
k = 'rot'
else:
msg = 'expected Porb or Prot for freq specification'
raise NotImplementedError(msg)
omegakey = 'omega{}'.format(k)
if k == 'rot':
omega_d[omegakey] = pm.Normal(
omegakey,
mu=prior_d[omegakey],
sd=0.01 * prior_d[omegakey],
testval=prior_d[omegakey])
P_rot = pm.Deterministic(
'P_rot',
pm.math.dot(1 / omega_d[omegakey], 2 * np.pi))
#omega_d[omegakey] = pm.Uniform(omegakey,
# lower=prior_d[omegakey]-1e-2,
# upper=prior_d[omegakey]+1e-2,
# testval=prior_d[omegakey])
elif k == 'orb':
# For orbital frequency, no need to declare new
# random variable!
omega_d[omegakey] = pm.Deterministic(
omegakey, pm.math.dot(1 / period, 2 * np.pi))
# sin and cosine terms are highly degenerate...
phikey = 'phi{}'.format(k)
if k == 'rot':
phi_d[phikey] = pm.Uniform(
phikey,
lower=prior_d[phikey] - np.pi / 8,
upper=prior_d[phikey] + np.pi / 8,
testval=prior_d[phikey])
#phi_d[phikey] = pm.Uniform(phikey,
# lower=0,
# upper=np.pi,
# testval=prior_d[phikey])
elif k == 'orb':
# For orbital phase, no need to declare new
# random variable!
phi_d[phikey] = pm.Deterministic(
phikey, pm.math.dot(t0 / period, 2 * np.pi))
N_harmonics = int(modelcomponent[0])
for ix in range(N_harmonics):
if LINEAR_AMPLITUDES:
Akey = 'A{}{}'.format(k, ix)
Bkey = 'B{}{}'.format(k, ix)
A_d[Akey] = pm.Uniform(
Akey,
lower=-2 * np.abs(prior_d[Akey]),
upper=2 * np.abs(prior_d[Akey]),
testval=np.abs(prior_d[Akey]))
B_d[Bkey] = pm.Uniform(
Bkey,
lower=-2 * np.abs(prior_d[Bkey]),
upper=2 * np.abs(prior_d[Bkey]),
testval=np.abs(prior_d[Bkey]))
if LOG_AMPLITUDES:
Akey = 'A{}{}'.format(k, ix)
Bkey = 'B{}{}'.format(k, ix)
logAkey = 'logA{}{}'.format(k, ix)
logBkey = 'logB{}{}'.format(k, ix)
if k == 'rot':
mfact = 3
elif k == 'orb':
mfact = 10
_A_d[logAkey] = pm.Uniform(
logAkey,
lower=np.log(prior_d[Akey] / mfact),
upper=np.log(mfact * prior_d[Akey]),
testval=np.log(prior_d[Akey]))
A_d[Akey] = pm.Deterministic(
Akey, pm.math.exp(_A_d[logAkey]))
_B_d[logBkey] = pm.Uniform(
logBkey,
lower=np.log(prior_d[Bkey] / mfact),
upper=np.log(mfact * prior_d[Bkey]),
testval=np.log(prior_d[Bkey]))
B_d[Bkey] = pm.Deterministic(
Bkey, pm.math.exp(_B_d[logBkey]))
harmonic_d = {**A_d, **B_d, **omega_d, **phi_d}
# Build the likelihood
if 'transit' not in self.modelcomponents:
# NOTE: hacky implementation detail: I didn't now how else to
# initialize an "empty" pymc3 random variable, so I assumed
# here that "transit" would be a modelcomponent, and the
# likelihood variable is initialized using it.
msg = 'Expected transit to be a model component.'
raise NotImplementedError(msg)
for modelcomponent in self.modelcomponents:
if 'transit' in modelcomponent:
mu_model = light_curve.flatten()
pm.Deterministic("mu_transit", light_curve.flatten())
if 'sincos' in modelcomponent:
if 'Porb' in modelcomponent:
k = 'orb'
elif 'Prot' in modelcomponent:
k = 'rot'
N_harmonics = int(modelcomponent[0])
for ix in range(N_harmonics):
spnames = [
'A{}{}'.format(k, ix), 'omega{}'.format(k),
'phi{}'.format(k)
]
cpnames = [
'B{}{}'.format(k, ix), 'omega{}'.format(k),
'phi{}'.format(k)
]
sin_params = [harmonic_d[k] for k in spnames]
cos_params = [harmonic_d[k] for k in cpnames]
# harmonic multiplier
mult = ix + 1
sin_params[1] = pm.math.dot(sin_params[1], mult)
cos_params[1] = pm.math.dot(cos_params[1], mult)
s_mod = sin_model(sin_params, self.x_obs)
c_mod = cos_model(cos_params, self.x_obs)
mu_model += s_mod
mu_model += c_mod
# save model components (rot and orb) for plotting
pm.Deterministic("mu_{}sin{}".format(k, ix), s_mod)
pm.Deterministic("mu_{}cos{}".format(k, ix), c_mod)
# track the total model to plot it
pm.Deterministic("mu_model", mu_model)
likelihood = pm.Normal('obs',
mu=mu_model,
sigma=sigma,
observed=self.y_obs)
# Get MAP estimate from model.
map_estimate = pm.find_MAP(model=model)
# Plot the simulated data and the maximum a posteriori model to
# make sure that our initialization looks ok.
self.y_MAP = map_estimate['mu_model'].flatten()
if make_threadsafe:
pass
else:
# as described in
# https://github.com/matplotlib/matplotlib/issues/15410
# matplotlib is not threadsafe. so do not make plots before
# sampling, because some child processes tries to close a
# cached file, and crashes the sampler.
if self.PLOTDIR is None:
raise NotImplementedError
outpath = os.path.join(self.PLOTDIR,
'test_{}_MAP.png'.format(self.modelid))
plot_MAP_data(self.x_obs, self.y_obs, self.y_MAP, outpath)
# sample from the posterior defined by this model.
trace = pm.sample(
tune=self.N_samples,
draws=self.N_samples,
start=map_estimate,
cores=self.N_cores,
chains=self.N_chains,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
with open(pklpath, 'wb') as buff:
pickle.dump(
{
'model': model,
'trace': trace,
'map_estimate': map_estimate
}, buff)
self.model = model
self.trace = trace
self.map_estimate = map_estimate
def getmodel(holds={}, mu={}, sig={}, transform=False, nterms=1):
params = [
'logS0', 'logw', 'alpha', 'logsig', 'mean', 'u', 'logrp', 'logrm',
't0p', 't0m'
]
for p in params:
if p not in holds:
holds[p] = None
if p not in mu:
mu[p] = None
if p not in sig:
sig[p] = None
with pm.Model() as model:
logS0 = pm.Normal("logS0",
mu=mu["logS0"],
sd=sig["logS0"],
observed=holds['logS0'])
logw = pm.Normal("logw",
mu=mu["logw"],
sd=sig["logw"],
observed=holds['logw'])
if np.shape(flux)[0] > 1:
alpha = pm.MvNormal("alpha",
mu=mu["alpha"],
chol=np.diag(sig["alpha"]),
shape=np.shape(flux)[0] - 1,
observed=holds['alpha'])
logsig = pm.MvNormal("logsig",
mu=mu["logsig"],
chol=np.diag(sig["logsig"]),
shape=np.shape(flux)[0],
observed=holds['logsig'])
mean = pm.MvNormal("mean",
mu=mu["mean"],
chol=np.diag(sig["mean"]),
shape=np.shape(flux)[0],
observed=holds['mean'])
u = sgp.distributions.MvUniform("u",
lower=[0, 0],
upper=[1, 1],
testval=[0.5, 0.5],
observed=holds['u'])
if transform:
logrp = pm.Uniform("logrp",
lower=-20.0,
upper=0.0,
testval=mu['logrp'],
observed=holds['logrp'])
logrm = pm.Uniform("logrm",
lower=-20.0,
upper=0.0,
testval=mu['logrm'],
observed=holds['logrm'])
t0p = pm.Uniform("t0p",
lower=t[0],
upper=t[-1],
testval=mu['t0p'],
observed=holds['t0p'])
t0m = pm.Uniform("t0m",
lower=t[0],
upper=t[-1],
testval=mu['t0m'],
observed=holds['t0m'])
else:
logrp = pm.Uniform("logrp",
lower=-20.0,
upper=0.0,
testval=mu['logrp'],
transform=None,
observed=holds['logrp'])
logrm = pm.Uniform("logrm",
lower=-20.0,
upper=0.0,
testval=mu['logrm'],
transform=None,
observed=holds['logrm'])
t0p = pm.Uniform("t0p",
lower=t[0],
upper=t[-1],
testval=mu['t0p'],
transform=None,
observed=holds['t0p'])
t0m = pm.Uniform("t0m",
lower=t[0],
upper=t[-1],
testval=mu['t0m'],
transform=None,
observed=holds['t0m'])
orbit = xo.orbits.KeplerianOrbit(period=5.0 * 60 * 60)
lcp = (xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit,
r=np.exp(logrp),
t=t / (60 * 60) - t0p,
texp=np.mean(np.diff(t)) / (60 * 60)))
lcm = (xo.LimbDarkLightCurve(u).get_light_curve(
orbit=orbit,
r=np.exp(logrm),
t=t / (60 * 60) - t0m,
texp=np.mean(np.diff(t)) / (60 * 60)))
mean = mean[:, None] + lcp.T[0] + lcm.T[0]
term = xo.gp.terms.SHOTerm(log_S0=logS0,
log_w0=logw,
log_Q=-np.log(np.sqrt(2)))
if np.shape(flux)[0] > 1:
a = tt.exp(tt.concatenate([[0.0], alpha]))
kernel = sgp.terms.KronTerm(term, alpha=a)
else:
kernel = term
yerr = tt.exp(2 * logsig)
yerr = yerr[:, None] * tt.ones(len(t))
if np.shape(flux)[0] > 1:
gp = xo.gp.GP(kernel, t, yerr, J=2, mean=sgp.means.KronMean(mean))
else:
gp = xo.gp.GP(kernel, t, yerr[0], J=2, mean=mean)
marg = gp.marginal("gp", observed=obs.T)
return model
