def build_model(mask):
with pm.Model() as model:
# Systemic parameters
mean_lc = pm.Normal("mean_lc", mu=0.0, sd=5.0)
mean_rv = pm.Normal("mean_rv", mu=0.0, sd=50.0)
u1 = xo.QuadLimbDark("u1")
u2 = xo.QuadLimbDark("u2")
# Parameters describing the primary
M1 = pm.Lognormal("M1", mu=0.0, sigma=10.0, testval=0.696)
R1 = pm.Lognormal("R1", mu=0.0, sigma=10.0, testval=0.687)
# Secondary ratios
k = pm.Lognormal("k", mu=0.0, sigma=10.0,
testval=0.9403) # radius ratio
q = pm.Lognormal("q", mu=0.0, sigma=10.0, testval=0.9815) # mass ratio
s = pm.Lognormal("s", mu=np.log(0.5),
sigma=10.0) # surface brightness ratio
# Prior on flux ratio
pm.Normal(
"flux_prior",
mu=lit_flux_ratio[0],
sigma=lit_flux_ratio[1],
observed=k**2 * s,
)
# Parameters describing the orbit
b = xo.ImpactParameter("b", ror=k, testval=1.5)
period = pm.Lognormal("period", mu=np.log(lit_period), sigma=1.0)
t0 = pm.Normal("t0", mu=lit_t0, sigma=1.0)
# Parameters describing the eccentricity: ecs = [e * cos(w), e * sin(w)]
ecs = xo.UnitDisk("ecs", testval=np.array([1e-5, 0.0]))
ecc = pm.Deterministic("ecc", tt.sqrt(tt.sum(ecs**2)))
omega = pm.Deterministic("omega", tt.arctan2(ecs[1], ecs[0]))
# Build the orbit
R2 = pm.Deterministic("R2", k * R1)
M2 = pm.Deterministic("M2", q * M1)
orbit = xo.orbits.KeplerianOrbit(
period=period,
t0=t0,
ecc=ecc,
omega=omega,
b=b,
r_star=R1,
m_star=M1,
m_planet=M2,
)
# Track some other orbital elements
pm.Deterministic("incl", orbit.incl)
pm.Deterministic("a", orbit.a)
# Noise model for the light curve
sigma_lc = pm.InverseGamma("sigma_lc",
testval=1.0,
**xo.estimate_inverse_gamma_parameters(
0.1, 2.0))
S_tot_lc = pm.InverseGamma("S_tot_lc",
testval=2.5,
**xo.estimate_inverse_gamma_parameters(
1.0, 5.0))
ell_lc = pm.InverseGamma("ell_lc",
testval=2.0,
**xo.estimate_inverse_gamma_parameters(
1.0, 5.0))
kernel_lc = xo.gp.terms.SHOTerm(S_tot=S_tot_lc,
w0=2 * np.pi / ell_lc,
Q=1.0 / 3)
# Noise model for the radial velocities
sigma_rv1 = pm.InverseGamma("sigma_rv1",
testval=1.0,
**xo.estimate_inverse_gamma_parameters(
0.5, 5.0))
sigma_rv2 = pm.InverseGamma("sigma_rv2",
testval=1.0,
**xo.estimate_inverse_gamma_parameters(
0.5, 5.0))
S_tot_rv = pm.InverseGamma("S_tot_rv",
testval=2.5,
**xo.estimate_inverse_gamma_parameters(
1.0, 5.0))
ell_rv = pm.InverseGamma("ell_rv",
testval=2.0,
**xo.estimate_inverse_gamma_parameters(
1.0, 5.0))
kernel_rv = xo.gp.terms.SHOTerm(S_tot=S_tot_rv,
w0=2 * np.pi / ell_rv,
Q=1.0 / 3)
# Set up the light curve model
lc = xo.SecondaryEclipseLightCurve(u1, u2, s)
def model_lc(t):
return (
mean_lc + 1e3 *
lc.get_light_curve(orbit=orbit, r=R2, t=t, texp=texp)[:, 0])
# Condition the light curve model on the data
gp_lc = xo.gp.GP(kernel_lc,
x[mask],
tt.zeros(mask.sum())**2 + sigma_lc**2,
mean=model_lc)
gp_lc.marginal("obs_lc", observed=y[mask])
# Set up the radial velocity model
def model_rv1(t):
return mean_rv + 1e-3 * orbit.get_radial_velocity(t)
def model_rv2(t):
return mean_rv - 1e-3 * orbit.get_radial_velocity(t) / q
# Condition the radial velocity model on the data
gp_rv1 = xo.gp.GP(kernel_rv,
x_rv,
tt.zeros(len(x_rv))**2 + sigma_rv1**2,
mean=model_rv1)
gp_rv1.marginal("obs_rv1", observed=y1_rv)
gp_rv2 = xo.gp.GP(kernel_rv,
x_rv,
tt.zeros(len(x_rv))**2 + sigma_rv2**2,
mean=model_rv2)
gp_rv2.marginal("obs_rv2", observed=y2_rv)
# Optimize the logp
map_soln = model.test_point
# First the RV parameters
map_soln = xo.optimize(map_soln, [mean_rv, q])
map_soln = xo.optimize(
map_soln, [mean_rv, sigma_rv1, sigma_rv2, S_tot_rv, ell_rv])
# Then the LC parameters
map_soln = xo.optimize(map_soln, [mean_lc, R1, k, s, b])
map_soln = xo.optimize(map_soln, [mean_lc, R1, k, s, b, u1, u2])
map_soln = xo.optimize(map_soln, [mean_lc, sigma_lc, S_tot_lc, ell_lc])
map_soln = xo.optimize(map_soln, [t0, period])
# Then all the parameters together
map_soln = xo.optimize(map_soln)
model.gp_lc = gp_lc
model.model_lc = model_lc
model.gp_rv1 = gp_rv1
model.model_rv1 = model_rv1
model.gp_rv2 = gp_rv2
model.model_rv2 = model_rv2
model.x = x[mask]
model.y = y[mask]
return model, map_soln
logK = pm.Uniform(
"logK",
lower=0,
upper=np.log(200),
testval=np.log(0.5 * (np.max(rv) - np.min(rv))),
)
logP = pm.Uniform("logP",
lower=0,
upper=np.log(10),
testval=np.log(lit_period))
phi = pm.Uniform("phi", lower=0, upper=2 * np.pi, testval=0.1)
# Parameterize the eccentricity using:
# h = sqrt(e) * sin(w)
# k = sqrt(e) * cos(w)
hk = xo.UnitDisk("hk", testval=np.array([0.01, 0.01]))
e = pm.Deterministic("e", hk[0]**2 + hk[1]**2)
w = pm.Deterministic("w", tt.arctan2(hk[1], hk[0]))
rv0 = pm.Normal("rv0", mu=0.0, sd=10.0, testval=0.0)
rvtrend = pm.Normal("rvtrend", mu=0.0, sd=10.0, testval=0.0)
# Deterministic transformations
n = 2 * np.pi * tt.exp(-logP)
P = pm.Deterministic("P", tt.exp(logP))
K = pm.Deterministic("K", tt.exp(logK))
cosw = tt.cos(w)
sinw = tt.sin(w)
t0 = (phi + w) / n
# The RV model
def run_rv_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.
# Stellar parameters. (Following tess.world notebooks).
logg_star = pm.Normal("logg_star",
mu=prior_d['logg_star'][0],
sd=prior_d['logg_star'][1])
r_star = pm.Bound(pm.Normal, lower=0.0)("r_star",
mu=prior_d['r_star'][0],
sd=prior_d['r_star'][1])
rho_star = pm.Deterministic("rho_star",
factor * 10**logg_star / r_star)
# RV parameters.
# Chen & Kipping predicted M: 49.631 Mearth, based on Rp of 8Re. It
# could be bigger, e.g., 94m/s if 1 Mjup.
# Predicted K: 14.26 m/s
#K = pm.Lognormal("K", mu=np.log(prior_d['K'][0]),
# sigma=prior_d['K'][1])
log_K = pm.Uniform('log_K',
lower=prior_d['log_K'][0],
upper=prior_d['log_K'][1])
K = pm.Deterministic('K', tt.exp(log_K))
period = pm.Normal("period",
mu=prior_d['period'][0],
sigma=prior_d['period'][1])
ecs = xo.UnitDisk("ecs", testval=np.array([0.7, -0.3]))
ecc = pm.Deterministic("ecc", tt.sum(ecs**2))
omega = pm.Deterministic("omega", tt.arctan2(ecs[1], ecs[0]))
phase = xo.UnitUniform("phase")
# use time of transit, rather than time of periastron. we do, after
# all, know it.
t0 = pm.Normal("t0",
mu=prior_d['t0'][0],
sd=prior_d['t0'][1],
testval=prior_d['t0'][0])
orbit = xo.orbits.KeplerianOrbit(period=period,
t0=t0,
rho_star=rho_star,
ecc=ecc,
omega=omega)
#FIXME edit these
# noise model parameters: FIXME what are these?
S_tot = pm.Lognormal("S_tot",
mu=np.log(prior_d['S_tot'][0]),
sigma=prior_d['S_tot'][1])
ell = pm.Lognormal("ell",
mu=np.log(prior_d['ell'][0]),
sigma=prior_d['ell'][1])
# per instrument parameters
means = pm.Normal(
"means",
mu=np.array([
np.median(self.y_obs[self.telvec == u])
for u in self.uniqueinstrs
]),
sigma=500,
shape=self.num_inst,
)
# different instruments have different intrinsic jitters. assign
# those based on the reported error bars. (NOTE: might inflate or
# overwrite these, for say, CHIRON)
sigmas = pm.HalfNormal("sigmas",
sigma=np.array([
np.median(self.y_err[self.telvec == u])
for u in self.uniqueinstrs
]),
shape=self.num_inst)
# Compute the RV offset and jitter for each data point depending on
# its instrument
mean = tt.zeros(len(self.x_obs))
diag = tt.zeros(len(self.x_obs))
for i, u in enumerate(self.uniqueinstrs):
mean += means[i] * (self.telvec == u)
diag += (self.y_err**2 + sigmas[i]**2) * (self.telvec == u)
pm.Deterministic("mean", mean)
pm.Deterministic("diag", diag)
# NOTE: local function definition is jank
def rv_model(x):
return orbit.get_radial_velocity(x, K=K)
kernel = xo.gp.terms.SHOTerm(S_tot=S_tot,
w0=2 * np.pi / ell,
Q=1.0 / 3)
# NOTE temp
gp = xo.gp.GP(kernel, self.x_obs, diag, mean=rv_model)
# gp = xo.gp.GP(kernel, self.x_obs, diag,
# mean=orbit.get_radial_velocity(self.x_obs, K=K))
# the actual "conditioning" step, i.e. the likelihood definition
gp.marginal("obs", observed=self.y_obs - mean)
pm.Deterministic("gp_pred", gp.predict())
map_estimate = model.test_point
map_estimate = xo.optimize(map_estimate, [means])
map_estimate = xo.optimize(map_estimate, [means, phase])
map_estimate = xo.optimize(map_estimate, [means, phase, log_K])
map_estimate = xo.optimize(map_estimate,
[means, t0, log_K, period, ecs])
map_estimate = xo.optimize(map_estimate, [sigmas, S_tot, ell])
map_estimate = xo.optimize(map_estimate)
#
# Derived parameters
#
#TODO
# # planet radius in jupiter radii
# r_planet = pm.Deterministic(
# "r_planet", (r*r_star)*( 1*units.Rsun/(1*units.Rjup) ).cgs.value
# )
# Plot the simulated data and the maximum a posteriori model to
# make sure that our initialization looks ok.
# i.e., "detrended". the "rv data" are y_obs - mean. The "trend" model
# is a GP. FIXME: AFAIK, it doesn't do much as-implemented.
self.y_MAP = (self.y_obs - map_estimate["mean"] -
map_estimate["gp_pred"])
t_pred = np.linspace(self.x_obs.min() - 10,
self.x_obs.max() + 10, 10000)
with model:
# NOTE temp
y_pred_MAP = xo.eval_in_model(rv_model(t_pred), map_estimate)
# # NOTE temp
# y_pred_MAP = xo.eval_in_model(
# orbit.get_radial_velocity(t_pred, K=K), map_estimate
# )
self.x_pred = t_pred
self.y_pred_MAP = y_pred_MAP
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_rv(self.x_obs, self.y_obs, self.y_MAP, self.y_err,
self.telcolors, self.x_pred, self.y_pred_MAP,
map_estimate, outpath)
with model:
# 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
# Prior on flux ratio
pm.Normal(
"flux_prior",
mu=lit_flux_ratio[0],
sigma=lit_flux_ratio[1],
observed=k**2 * s,
)
# Parameters describing the orbit
b = xo.ImpactParameter("b", ror=k, testval=1.5)
period = pm.Lognormal("period", mu=np.log(lit_period), sigma=1.0)
t0 = pm.Normal("t0", mu=lit_t0, sigma=1.0)
# Parameters describing the eccentricity: ecs = [e * cos(w), e * sin(w)]
ecs = xo.UnitDisk("ecs", testval=np.array([1e-5, 0.0]))
ecc = pm.Deterministic("ecc", tt.sqrt(tt.sum(ecs**2)))
omega = pm.Deterministic("omega", tt.arctan2(ecs[1], ecs[0]))
# Build the orbit
R2 = pm.Deterministic("R2", k * R1)
M2 = pm.Deterministic("M2", q * M1)
orbit = xo.orbits.KeplerianOrbit(
period=period,
t0=t0,
ecc=ecc,
omega=omega,
b=b,
r_star=R1,
m_star=M1,
m_planet=M2,