def make_data(t, log_S0, log_w0, log_Q, logsig,
t0=5.0, r=0.1, d=1.0, tin=0.1, a=2.0):
transit = utils.theano_transit(t,
t0,
r,
d,
tin)
kernel = xo.gp.terms.SHOTerm(
log_S0 = log_S0,
log_w0 = log_w0,
log_Q = log_Q
)
J = 4
q = np.array([1, a])
Q = q[:, None]*q[None, :]
diag = np.exp(2*logsig)*np.ones((2, len(t)))
kernel = xo.gp.terms.KroneckerTerm(kernel, Q)
gp = GP(kernel, t, diag, J=J)
n = gp.dot_l(np.random.randn(2*len(t), 1)).eval()
transit = tt.reshape(tt.tile(transit, (2, 1)).T,
(1, 2*transit.shape[0])).T
n += transit.eval()
return n, np.sum([n[i::2].T[0] for i in range(len(q))], axis=0)/2
def run_mcmc_1d(t, data, logS0_init,
logw0_init, logQ_init,
logsig_init, t0_init,
r_init, d_init, tin_init):
with pm.Model() as model:
#logsig = pm.Uniform("logsig", lower=-20.0, upper=0.0, testval=logsig_init)
# The parameters of the SHOTerm kernel
#logS0 = pm.Uniform("logS0", lower=-50.0, upper=0.0, testval=logS0_init)
#logQ = pm.Uniform("logQ", lower=-50.0, upper=20.0, testval=logQ_init)
#logw0 = pm.Uniform("logw0", lower=-50.0, upper=20.0, testval=logw0_init)
# The parameters for the transit mean function
t0 = pm.Uniform("t0", lower=t[0], upper=t[-1], testval=t0_init)
r = pm.Uniform("r", lower=0.0, upper=1.0, testval=r_init)
d = pm.Uniform("d", lower=0.0, upper=10.0, testval=d_init)
tin = pm.Uniform("tin", lower=0.0, upper=10.0, testval=tin_init)
# Deterministics
# mean = pm.Deterministic("mean", utils.transit(t, t0, r, d, tin))
transit = utils.theano_transit(t, t0, r, d, tin)
# Set up the Gaussian Process model
kernel = xo.gp.terms.SHOTerm(
log_S0 = logS0_init,
log_w0 = logw0_init,
log_Q=logQ_init
)
diag = np.exp(2*logsig_init)*tt.ones((1, len(t)))
gp = GP(kernel, t, diag, J=2)
# Compute the Gaussian Process likelihood and add it into the
# the PyMC3 model as a "potential"
pm.Potential("loglike", gp.log_likelihood(data - transit))
# Compute the mean model prediction for plotting purposes
#pm.Deterministic("mu", gp.predict())
map_soln = xo.optimize(start=model.test_point, verbose=False)
with model:
map_soln = xo.optimize(start=model.test_point)
with model:
trace = pm.sample(
tune=500,
draws=500,
start=map_soln,
cores=2,
chains=2,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
return trace
def gp_fit(t, y, yerr, t_grid, integrated=False, exp_time=60.):
# optimize kernel hyperparameters and return fit + predictions
with pm.Model() as model:
logS0 = pm.Normal("logS0", mu=0.4, sd=5.0, testval=np.log(np.var(y)))
logw0 = pm.Normal("logw0", mu=-3.9, sd=0.1)
logQ = pm.Normal("logQ", mu=3.5, sd=5.0)
# Set up the kernel and GP
kernel = terms.SHOTerm(log_S0=logS0, log_w0=logw0, log_Q=logQ)
if integrated:
kernel_int = terms.IntegratedTerm(kernel, exp_time)
gp = GP(kernel_int, t, yerr**2)
else:
gp = GP(kernel, t, yerr**2)
# Add a custom "potential" (log probability function) with the GP likelihood
pm.Potential("gp", gp.log_likelihood(y))
with model:
map_soln = xo.optimize(start=model.test_point)
mu, var = xo.eval_in_model(gp.predict(t_grid, return_var=True),
map_soln)
sd = np.sqrt(var)
y_pred = xo.eval_in_model(gp.predict(t), map_soln)
return map_soln, mu, sd, y_pred
def gp_predict(t,
y,
yerr,
t_grid,
logS0=0.4,
logw0=-3.9,
logQ=3.5,
integrated=False,
exp_time=60.):
# take kernel hyperparameters as fixed inputs, train + predict
with pm.Model() as model:
kernel = terms.SHOTerm(log_S0=logS0, log_w0=logw0, log_Q=logQ)
if integrated:
kernel_int = terms.IntegratedTerm(kernel, exp_time)
gp = GP(kernel_int, t, yerr**2)
else:
gp = GP(kernel, t, yerr**2)
gp.condition(y)
mu, var = xo.eval_in_model(gp.predict(t_grid, return_var=True))
sd = np.sqrt(var)
y_pred = xo.eval_in_model(gp.predict(t))
return y_pred, mu, sd
def multi_gp_predict(t, y, yerr, t_grid, integrated=False, exp_time=60.):
# this code is GARBAGE. but in principle does gp_predict() for a full comb of modes.
a_max = 0.55 # amplitude of central mode in m/s
nu_max = 3.1e-3 # peak frequency in Hz
c_env = 0.331e-3 # envelope width in Hz
delta_nu = 0.00013 # Hz
gamma = 1. / (2 * 24. * 60. * 60.) # s^-1 ; 2-day damping timescale
freq_grid = np.arange(nu_max - 0.001, nu_max + 0.001,
delta_nu) # magic numbers
amp_grid = a_max**2 * np.exp(-(freq_grid - nu_max)**2 /
(2. * c_env**2)) # amplitudes in m/s
driving_amp_grid = np.sqrt(amp_grid * gamma * dt)
log_S0_grid = [
np.log(d**2 / (dt * o)) for o, d in zip(omega_grid, driving_amp_grid)
]
with pm.Model() as model:
kernel = None
for o, lS in zip(omega_grid, log_S0_grid):
if kernel is None:
kernel = terms.SHOTerm(log_S0=lS,
log_w0=np.log(o),
log_Q=np.log(o / gamma))
else:
kernel += terms.SHOTerm(log_S0=lS,
log_w0=np.log(o),
log_Q=np.log(o / gamma))
if integrated:
kernel_int = terms.IntegratedTerm(kernel, exp_time)
gp = GP(kernel_int, t, yerr**2)
else:
gp = GP(kernel, t, yerr**2)
gp.condition(y)
mu, var = xo.eval_in_model(gp.predict(t_grid, return_var=True))
sd = np.sqrt(var)
y_pred = xo.eval_in_model(gp.predict(t))
return y_pred, mu, sd
mean = pm.Normal("mean", mu=0.0, sigma=1.0)
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)
# Condition the GP on the observations and add the marginal likelihood
# to the model
gp.marginal("gp", observed=y)
with model:
map_soln = xo.optimize(start=model.test_point)
with model:
mu, var = xo.eval_in_model(
gp.predict(true_t, return_var=True, predict_mean=True), map_soln)
# Plot the prediction and the 1-sigma uncertainty
plt.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0, label="data")
plt.plot(true_t, true_y, "k", lw=1.5, alpha=0.3, label="truth")
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 load_models(t, data1, data2, logS0_init,
logw0_init, logQ_init,
logsig_init, t0_init,
r_init, d_init, tin_init, a):
with pm.Model() as model1d:
logsig = pm.Uniform("logsig", lower=-20.0, upper=0.0, testval=logsig_init)
# The parameters of the SHOTerm kernel
#logS0 = pm.Uniform("logS0", lower=-50.0, upper=0.0, testval=logS0_init)
#logQ = pm.Uniform("logQ", lower=-50.0, upper=20.0, testval=logQ_init)
#logw0 = pm.Uniform("logw0", lower=-50.0, upper=20.0, testval=logw0_init)
# The parameters for the transit mean function
t0 = pm.Uniform("t0", lower=t[0], upper=t[-1], testval=t0_init)
r = pm.Uniform("r", lower=0.0, upper=1.0, testval=r_init)
d = pm.Uniform("d", lower=0.0, upper=10.0, testval=d_init)
tin = pm.Uniform("tin", lower=0.0, upper=10.0, testval=tin_init)
# Deterministics
# mean = pm.Deterministic("mean", utils.transit(t, t0, r, d, tin))
transit = utils.theano_transit(t, t0, r, d, tin)
# Set up the Gaussian Process model
kernel = xo.gp.terms.SHOTerm(
log_S0 = logS0_init,
log_w0 = logw0_init,
log_Q=logQ_init
)
diag = np.exp(2*logsig_init)*tt.ones(len(t))
gp = GP(kernel, t, diag, J=2)
# Compute the Gaussian Process likelihood and add it into the
# the PyMC3 model as a "potential"
pm.Potential("loglike", gp.log_likelihood(data1 - transit))
# Compute the mean model prediction for plotting purposes
#pm.Deterministic("mu", gp.predict())
#map_soln = xo.optimize(start=model1d.test_point, verbose=False)
with pm.Model() as model2d:
#logsig = pm.Uniform("logsig", lower=-20.0, upper=0.0, testval=logsig_init)
# The parameters of the SHOTerm kernel
#logS0 = pm.Uniform("logS0", lower=-50.0, upper=0.0, testval=logS0_init)
#logQ = pm.Uniform("logQ", lower=-50.0, upper=20.0, testval=logQ_init)
#logw0 = pm.Uniform("logw0", lower=-50.0, upper=20.0, testval=logw0_init)
a = pm.Uniform("a", lower=1.0, upper=10.0, testval=2.0)
# The parameters for the transit mean function
t0 = pm.Uniform("t0", lower=t[0], upper=t[-1], testval=t0_init)
r = pm.Uniform("r", lower=0.0, upper=1.0, testval=r_init)
d = pm.Uniform("d", lower=0.0, upper=10.0, testval=d_init)
tin = pm.Uniform("tin", lower=0.0, upper=10.0, testval=tin_init)
# Deterministics
# mean = pm.Deterministic("mean", utils.transit(t, t0, r, d, tin))
transit = utils.theano_transit(t, t0, r, d, tin)
transit = tt.reshape(tt.tile(transit, (2, 1)).T, (1, 2*transit.shape[0])).T
# Set up the Gaussian Process model
kernel = xo.gp.terms.SHOTerm(
log_S0 = logS0_init,
log_w0 = logw0_init,
log_Q=logQ_init
)
q = tt.stack(1, a)
Q = q[:, None]*q[None, :]
kernel = xo.gp.terms.KroneckerTerm(kernel, Q)
diag = np.exp(2*logsig_init)*tt.ones((2, len(t)))
gp = GP(kernel, t, diag, J=4)
# Compute the Gaussian Process likelihood and add it into the
# the PyMC3 model as a "potential"
pm.Potential("loglike", gp.log_likelihood((data2 - transit).T))
# Compute the mean model prediction for plotting purposes
#pm.Deterministic("mu", gp.predict())
#map_soln = xo.optimize(start=model2d.test_point, verbose=False)
return model1d, model2d
def run_gp_single(Sgv, wgv, S1v, w1v, Q1v, opt=opt):
if (opt == 1):
print('Running Gp Single Optimiziation', 'Sgv', Sgv, 'wgv', wgv, 'S1v',
S1v, 'w1v', w1v, 'Q1v', Q1v)
with pm.Model() as model:
logs2 = pm.Normal("logs2",
mu=2 * np.log(np.mean(yerr)),
sigma=100.0,
testval=100)
logSg = pm.Normal("logSg", mu=Sgv, sigma=100.0, testval=Sgv)
logwg = pm.Normal("logwg", mu=wgv, sigma=100.0, testval=wgv)
logS1 = pm.Normal("logS1", mu=S1v, sigma=100.0, testval=S1v)
logw1 = pm.Normal("logw1", mu=w1v, sigma=100.0, testval=w1v)
logQ1 = pm.Normal("logQ1", mu=Q1v, sigma=100.0, testval=Q1v)
# Set up the kernel an GP
bg_kernel = terms.SHOTerm(log_S0=logSg,
log_w0=logwg,
Q=1.0 / np.sqrt(2))
star_kernel1 = terms.SHOTerm(log_S0=logS1, log_w0=logw1, log_Q=logQ1)
kernel = star_kernel1 + bg_kernel
gp = GP(kernel, t, yerr**2 + pm.math.exp(logs2))
gp_star1 = GP(star_kernel1, t, yerr**2 + pm.math.exp(logs2))
gp_bg = GP(bg_kernel, t, yerr**2 + pm.math.exp(logs2))
# Condition the GP on the observations and add the marginal likelihood
# to the model
gp.marginal("gp", observed=y)
with model:
val = gp.kernel.psd(omega)
psd_init = xo.eval_in_model(val)
bg_val = gp_bg.kernel.psd(omega)
star_val_1 = gp_star1.kernel.psd(omega)
bg_psd_init = xo.eval_in_model(bg_val)
star_1_psd_init = xo.eval_in_model(star_val_1)
# print('done_init_plot')
map_soln = model.test_point
if (opt == 1):
map_soln = xo.optimize(start=map_soln, vars=[logSg])
#ask about this, do i need to scale when I show this?
map_soln = xo.optimize(start=map_soln, vars=[logwg])
map_soln = xo.optimize(start=map_soln, vars=[logw1])
map_soln = xo.optimize(start=map_soln, vars=[logS1])
map_soln = xo.optimize(start=map_soln)
print(map_soln.values())
mu, var = xo.eval_in_model(gp.predict(t, return_var=True),
map_soln)
plt.figure()
plt.errorbar(t, y, yerr=yerr, fmt=".k", capsize=0, label="data")
sd = np.sqrt(var)
art = plt.fill_between(t, mu + sd, mu - sd, color="C1", alpha=0.3)
art.set_edgecolor("none")
plt.plot(t, mu, color="C1", label="prediction")
plt.legend(fontsize=12)
plt.xlabel("t")
plt.ylabel("y")
plt.xlim(0, 10)
_ = plt.ylim(-2.5, 2.5)
psd_final = xo.eval_in_model(gp.kernel.psd(omega), map_soln)
bg_psd_fin = xo.eval_in_model(bg_val, map_soln)
star_1_psd_fin = xo.eval_in_model(star_val_1, map_soln)
return psd_init, star_1_psd_init, bg_psd_init, psd_final, star_1_psd_fin, bg_psd_fin, map_soln
def run_gp_binary(Sg, wg, S1, w1, Q1, S2, w2, Q2, opt=opt):
with pm.Model() as model:
logs2 = pm.Normal("logs2",
mu=2 * np.log(np.mean(yerr)),
sigma=100.0,
testval=-100)
mean = pm.Normal("mean", mu=np.mean(y), sigma=1.0)
logSg = pm.Normal("logSg", mu=0.0, sigma=15.0, testval=Sg)
logwg = pm.Normal("logwg",
mu=0.0,
sigma=15.0,
testval=wg - np.log(1e6))
logS1 = pm.Normal("logS1", mu=0.0, sigma=15.0, testval=S1)
logw1 = pm.Normal("logw1",
mu=0.0,
sigma=15.0,
testval=w1 - np.log(1e6))
logQ1 = pm.Normal("logQ1", mu=0.0, sigma=15.0, testval=Q1)
logS2 = pm.Normal("logS2", mu=0.0, sigma=15.0, testval=S2)
logw2 = pm.Normal("logw2",
mu=0.0,
sigma=15.0,
testval=w2 - np.log(1e6))
logQ2 = pm.Normal("logQ2", mu=0.0, sigma=15.0, testval=Q2)
# Set up the kernel an GP
bg_kernel = terms.SHOTerm(log_S0=logSg,
log_w0=logwg,
Q=1.0 / np.sqrt(2))
star_kernel1 = terms.SHOTerm(log_S0=logS1, log_w0=logw1, log_Q=logQ1)
star_kernel2 = terms.SHOTerm(log_S0=logS2, log_w0=logw2, log_Q=logQ2)
kernel = star_kernel1 + star_kernel2 + bg_kernel
gp = GP(kernel, t, yerr**2 + pm.math.exp(logs2), mean=mean)
gp_star1 = GP(star_kernel1, t, yerr**2 + pm.math.exp(logs2), mean=mean)
gp_bg = GP(bg_kernel, t, yerr**2 + pm.math.exp(logs2), mean=mean)
gp_star2 = GP(star_kernel2, t, yerr**2 + pm.math.exp(logs2), mean=mean)
# Condition the GP on the observations and add the marginal likelihood
# to the model
gp.marginal("gp", observed=y)
with model:
val = gp.kernel.psd(omega)
psd_init = xo.eval_in_model(val)
bg_val = gp_bg.kernel.psd(omega)
star_val_1 = gp_star1.kernel.psd(omega)
star_val_2 = gp_star2.kernel.psd(omega)
bg_psd_init = xo.eval_in_model(bg_val)
star_1_psd_init = xo.eval_in_model(star_val_1)
star_2_psd_init = xo.eval_in_model(star_val_2)
# print('done_init_plot')
map_soln = model.test_point
if (opt == 1):
print('running opt')
map_soln = xo.optimize(start=map_soln, vars=[logSg])
#ask about this, do i need to scale when I show this?
map_soln = xo.optimize(start=map_soln, vars=[logwg])
#map_soln = xo.optimize(start=map_soln, vars=[logS1,logw1])
#map_soln = xo.optimize(start=map_soln, vars=[logS2,logw2])
psd_final = xo.eval_in_model(gp.kernel.psd(omega), map_soln)
bg_psd_fin = xo.eval_in_model(bg_val, map_soln)
star_1_psd_fin = xo.eval_in_model(star_val_1, map_soln)
star_2_psd_fin = xo.eval_in_model(star_val_2, map_soln)
return psd_init, star_1_psd_init, star_2_psd_init, bg_psd_init, psd_final, star_1_psd_fin, star_2_psd_fin, bg_psd_fin, map_soln
import theano.tensor as tt
from exoplanet.gp import terms, GP
with pm.Model() as model:
mean = pm.Normal("mean", mu=0.0, sigma=1.0)
logS1 = pm.Normal("logS1", mu=0.0, sigma=15.0, testval=np.log(np.var(y)))
logw1 = pm.Normal("logw1", mu=0.0, sigma=15.0, testval=np.log(3.0))
logS2 = pm.Normal("logS2", mu=0.0, sigma=15.0, testval=np.log(np.var(y)))
logw2 = pm.Normal("logw2", mu=0.0, sigma=15.0, testval=np.log(3.0))
logQ = pm.Normal("logQ", mu=0.0, sigma=15.0, testval=0)
# Set up the kernel an GP
kernel = terms.SHOTerm(log_S0=logS1, log_w0=logw1, Q=1.0 / np.sqrt(2))
kernel += terms.SHOTerm(log_S0=logS2, log_w0=logw2, log_Q=logQ)
gp = GP(kernel, t, yerr**2, mean=mean)
# Condition the GP on the observations and add the marginal likelihood
# to the model
gp.marginal("gp", observed=y)
# %% [markdown]
# A few comments here:
#
# 1. The `term` interface in *exoplanet* only accepts keyword arguments with names given by the `parameter_names` property of the term. But it will also interpret keyword arguments with the name prefaced by `log_` to be the log of the parameter. For example, in this case, we used `log_S0` as the parameter for each term, but `S0=tt.exp(log_S0)` would have been equivalent. This is useful because many of the parameters are required to be positive so fitting the log of those parameters is often best.
# 2. The third argument to the :class:`exoplanet.gp.GP` constructor should be the *variance* to add along the diagonal, not the standard deviation as in the original [celerite implementation](https://celerite.readthedocs.io).
# 3. Finally, the :class:`exoplanet.gp.GP` constructor takes an optional argument `J` which specifies the width of the problem if it is known at compile time. Just to be confusing, this is actually two times the `J` from [the celerite paper](https://arxiv.org/abs/1703.09710). There are various technical reasons why this is difficult to work out in general and this code will always work if you don't provide a value for `J`, but you can get much better performance (especially for small `J`) if you know what it will be for your problem. In general, most terms cost `J=2` with the exception of a :class:`exoplanet.gp.terms.RealTerm` (which costs `J=1`) and a :class:`exoplanet.gp.terms.RotationTerm` (which costs `J=4`).
#
# To start, let's fit for the maximum a posteriori (MAP) parameters and look the the predictions that those make.
# %%
def model_offsets(self):
"""
Define the GP offset model.
"""
stdev = self.stdev
# stdev = 2.0
nsteps = len(self.steps)
with pm.Model() as model:
# Parameters
logsigma = pm.Normal("logsigma", mu=0.0, sd=15.0)
logrho = pm.Normal("logrho", mu=0.0, sd=5.0)
# Define step variables
step1 = pm.Normal("step1", mu=self.steps[0], sd=stdev)
steps = step1
if nsteps > 1:
step2 = pm.Normal("step2", mu=self.steps[1], sd=stdev)
steps = [step1, step2]
if nsteps > 2:
step3 = pm.Normal("step3", mu=self.steps[2], sd=stdev)
steps = [step1, step2, step3]
if nsteps > 3:
step4 = pm.Normal("step4", mu=self.steps[3], sd=stdev)
steps = [step1, step2, step3, step4]
if nsteps > 4:
step5 = pm.Normal("step5", mu=self.steps[4], sd=stdev)
steps = [step1, step2, step3, step4, step5]
if nsteps > 5:
step6 = pm.Normal("step6", mu=self.steps[5], sd=stdev)
steps = [step1, step2, step3, step4, step5, step6]
if nsteps > 6:
step7 = pm.Normal("step7", mu=self.steps[6], sd=stdev)
steps = [step1, step2, step3, step4, step5, step6, step7]
if nsteps > 7:
step8 = pm.Normal("step8", mu=self.steps[7], sd=stdev)
steps = [
step1, step2, step3, step4, step5, step6, step7, step8
]
if nsteps > 8:
step9 = pm.Normal("step9", mu=self.steps[8], sd=stdev)
steps = [
step1, step2, step3, step4, step5, step6, step7, step8,
step9
]
if nsteps > 9:
step10 = pm.Normal("step10", mu=self.steps[9], sd=stdev)
steps = [
step1, step2, step3, step4, step5, step6, step7, step8,
step9, step10
]
if nsteps > 10:
step11 = pm.Normal("step11", mu=self.steps[10], sd=stdev)
steps = [
step1, step2, step3, step4, step5, step6, step7, step8,
step9, step10, step11
]
if nsteps > 11:
step12 = pm.Normal("step12", mu=self.steps[11], sd=stdev)
steps = [
step1, step2, step3, step4, step5, step6, step7, step8,
step9, step10, step11, step12
]
if nsteps > 12:
step13 = pm.Normal("step13", mu=self.steps[12], sd=stdev)
steps = [
step1, step2, step3, step4, step5, step6, step7, step8,
step9, step10, step11, step12, step13
]
if nsteps > 13:
step14 = pm.Normal("step14", mu=self.steps[13], sd=stdev)
steps = [
step1, step2, step3, step4, step5, step6, step7, step8,
step9, step10, step11, step12, step13, step14
]
# The step model
mu = step_model(self.t, self.gap_times, steps)
# The likelihood function assuming known Gaussian uncertainty
# pm.Normal("obs", mu=mu, sd=self.yerr, observed=self.y)
# Set up the kernel an GP
kernel = terms.Matern32Term(log_sigma=logsigma, log_rho=logrho)
gp = GP(kernel, self.t, self.yerr**2)
# Add a custom "potential" (log probability function) with the GP
# likelihood
pm.Potential("gp", gp.log_likelihood(self.y - mu))
self.gp = gp
self.model = model
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