def run_sampling(pm_model,
output_dir,
ncores=1,
nchains=2,
max_attempts=2,
filename="trace"):
# Log file output
logging.basicConfig(
filename=output_dir + "/sampling.log",
filemode="w",
format="%(name)s - %(levelname)s - %(message)s",
)
# Sample the model
divperc = 20
with pm_model:
# Run initial chain
try:
trace = pm.sample(
tune=1000,
draws=4000,
cores=ncores,
chains=nchains,
step=xo.get_dense_nuts_step(),
)
except pm.exceptions.SamplingError:
logging.error("Sampling failed, model misspecified")
return None
# Check for divergences, restart sampling if necessary
divergent = trace["diverging"]
divperc = divergent.nonzero()[0].size / len(trace) * 100
n_attempts = 1
while divperc > 15.0 and n_attempts <= max_attempts:
# Run sampling
trace = pm.sample(
tune=2000,
draws=n_attempts * 10000,
cores=ncores,
chains=nchains,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
n_attempts += 1
if divperc > 15:
logging.warning(f"{divperc} of samples are diverging.")
df = pm.trace_to_dataframe(trace, include_transformed=True)
df.to_csv(output_dir + f"/{filename}.csv")
return None
else:
df = pm.trace_to_dataframe(trace, include_transformed=True)
df.to_csv(output_dir + f"/{filename}.csv")
return trace
def fit_model(samples, samples_logp, output_dir):
# Convert samples to theano.tensor
samples_tensor = T.as_tensor_variable(samples) # for performance reasons
samples_logp_tensor = T.as_tensor_variable(samples_logp)
model = HierarchicalModel(samples_tensor, samples_logp_tensor)
with model:
# Print initial logps
print(model.test_point)
# Run sampling
trace = pm.sample(tune=100,
draws=1000,
cores=4,
step=xo.get_dense_nuts_step())
# Save the samples to disk
samples_hyper = np.stack((
trace["lam_ln_A0"],
trace["mu_ln_delta_t0"],
trace["sig_ln_delta_t0"],
trace["mu_ln_tE"],
trace["sig_ln_tE"],
trace["alpha_f"],
trace["beta_f"],
)).T
np.save(output_dir + "samples_hyper.npy", samples_hyper)
def run_pymc3_model(pos, pos_err, proper, proper_err, mean, cov):
M = get_tangent_basis(pos[0] * 2 * np.pi / 360, pos[1] * 2 * np.pi / 360)
# mean, cov = get_prior()
with pm.Model() as model:
vxyzD = pm.MvNormal("vxyzD", mu=mean, cov=cov, shape=4)
vxyz = pm.Deterministic("vxyz", vxyzD[:3])
log_D = pm.Deterministic("log_D", vxyzD[3])
D = pm.Deterministic("D", tt.exp(log_D))
xyz = pm.Deterministic("xyz", tt_eqtogal(pos[0], pos[1], D)[:, 0])
pm_from_v, rv_from_v = tt_get_icrs_from_galactocentric(
xyz, vxyz, pos[0], pos[1], D, M)
pm.Normal("proper_motion",
mu=pm_from_v,
sigma=np.array(proper_err),
observed=np.array(proper))
pm.Normal("parallax", mu=1. / D, sigma=pos_err[2], observed=pos[2])
map_soln = xo.optimize()
trace = pm.sample(tune=1500,
draws=1000,
start=map_soln,
step=xo.get_dense_nuts_step(target_accept=0.9))
return trace
def predict(self, model_type=None):
"""
Predict the period of stellar variability via Gaussian Process fitting
Returns all samples of paramters after mcmc fitting
"""
if model_type == None:
if self.lctype == "rotation":
model, map_soln = self.rotation_model()
elif self.lctype == "granulation":
model, map_soln = self.granulation_model()
elif self.lctype == "hybrid":
model, map_soln = self.hybrid_model()
else:
if model_type == "rotation":
model, map_soln = self.rotation_model()
elif model_type == "granulation":
model, map_soln = self.granulation_model()
elif model_type == "hybrid":
model, map_soln = self.hybrid_model()
np.random.seed(42)
with model:
trace = pm.sample(tune=2000,
draws=2000,
start=map_soln,
step=xo.get_dense_nuts_step(target_accept=0.99),
progressbar=True)
self._trace = trace
return trace
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 sample_from_model(model,
map_soln,
tune=500,
draws=200,
chains=5,
cores=None,
step=None):
"""
Sample from the transit light curve model.
Parameters
----------
model : `~pymc3.model`
A model object.
map_soln : dict
A dictionary with the maximum a posteriori estimates of the variables.
tune : int, optional
The number of iterations to tune.
draws : int, optional
The number of samples to draw.
chains : int, optional
The number of chains to sample.
cores : int, optional
The number of cores to run in parallel.
step : function, optional
A step function.
Returns
-------
trace : `~pymc3.backends.base.MultiTrace`
A ``MultiTrace`` object that contains the samples.
"""
# Use 1 CPU thread per chain, unless specified otherwise
if cores is None:
cores = min(chains, mp.cpu_count())
# Ignore FutureWarnings
warnings.simplefilter('ignore', FutureWarning)
with model:
if step is None:
step = xo.get_dense_nuts_step(target_accept=0.95)
trace = pm.sample(tune=tune,
draws=draws,
start=map_soln,
chains=chains,
cores=cores,
step=step)
# Reset warnings
warnings.resetwarnings()
return trace
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 fit(self, map_soln, model):
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']['Transit epoch (BTJD)'],
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']['Transit epoch (BTJD)'] = epoch
self.bokeh_document.add_next_tick_callback(
partial(self.update_parameters_table, trace_summary))
with pd.option_context('display.max_columns', None,
'display.max_rows', None):
print(trace_summary)
print(f'Star radius: {self.star_radius}')
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 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
ax[3].set_xlim(t_fine[0], t_fine[-1])
_ = ax[0].set_title("map orbit")
# %% [markdown]
# Now let's sample the posterior.
# %%
np.random.seed(1234)
with model:
trace = pm.sample(
tune=5000,
draws=4000,
start=map_soln,
cores=2,
chains=2,
step=xo.get_dense_nuts_step(target_accept=0.9, adaptation_window=201),
)
# %% [markdown]
# First we can check the convergence for some of the key parameters.
# %%
pm.summary(
trace, varnames=["P", "tperi", "a_ang", "omega", "Omega", "incl", "ecc"]
)
# %% [markdown]
# That looks pretty good.
# Now here's a corner plot showing the covariances between parameters.
# %%
plt.xlabel("t")
plt.ylabel("y")
plt.xlim(0, 10)
_ = plt.ylim(-2.5, 2.5)
plt.savefig('../results/test_results/gp_model/prediction_uncert.png',
dpi=200)
plt.close('all')
with model:
trace = pm.sample(
tune=2000,
draws=2000,
start=map_soln,
cores=2,
chains=2,
step=xo.get_dense_nuts_step(target_accept=0.9),
)
print(pm.summary(trace))
# truth = np.concatenate(
# xo.eval_in_model([period, r], model.test_point, model=model)
# )
fig = corner.corner(
samples,
# truths=truth,
labels=["S1", "S2", "w1", "w2", "log_Q"],
)
fig.savefig('../results/test_results/gp_model/test_gp_corner.png')
plt.close('all')
def worker(task):
(i1, i2), data, model_kw, basename = task
g = GaiaData(data)
cache_filename = os.path.abspath(f'../cache/tmp-{basename}_{i1}-{i2}.fits')
if os.path.exists(cache_filename):
print(f"({pid}) cache filename exists for index range: "
f"{cache_filename}")
return cache_filename
print(f"({pid}) setting up model")
helper = ComovingHelper(g)
niter = 0
while niter < 10:
try:
model = helper.get_model(**model_kw)
break
except OSError:
print(f"{pid} failed to compile - trying again in 2sec...")
time.sleep(5)
niter += 1
continue
else:
print(f"{pid} never successfully compiled. aborting")
import socket
print(socket.gethostname(), socket.getfqdn(),
os.path.exists("/cm/shared/sw/pkg/devel/gcc/7.4.0/bin/g++"))
return ''
print(f"({pid}) done init model - running {len(g)} stars")
probs = np.full(helper.N, np.nan)
for n in range(helper.N):
with model:
pm.set_data({
'y': helper.ys[n],
'Cinv': helper.Cinvs[n],
'M': helper.Ms[n]
})
test_pt = {
'vxyz': helper.test_vxyz[n],
'r': helper.test_r[n],
'w': np.array([0.5, 0.5])
}
try:
print("starting optimize")
res = xo.optimize(start=test_pt,
progress_bar=False,
verbose=False)
print("done optimize - starting sample")
trace = pm.sample(
start=res,
tune=2000,
draws=1000,
cores=1,
chains=1,
step=xo.get_dense_nuts_step(target_accept=0.95),
progressbar=False)
except Exception as e:
print(str(e))
continue
# print("done sample - computing prob")
ll_fg = trace.get_values(model.group_logp)
ll_bg = trace.get_values(model.field_logp)
post_prob = np.exp(ll_fg - np.logaddexp(ll_fg, ll_bg))
probs[n] = post_prob.sum() / len(post_prob)
# write probs to cache filename
tbl = at.Table()
tbl['source_id'] = g.source_id
tbl['prob'] = probs
tbl.write(cache_filename)
return cache_filename
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]}"
)
def main(c, prior, metadata_row, overwrite=False):
mcmc_cache_path = os.path.join(c.run_path, 'mcmc')
os.makedirs(mcmc_cache_path, exist_ok=True)
apogee_id = metadata_row['APOGEE_ID']
this_cache_path = os.path.join(mcmc_cache_path, apogee_id)
if os.path.exists(this_cache_path) and not overwrite:
logger.info(f"{apogee_id} already done!")
# Assume it's already done
return
# Set up The Joker:
joker = tj.TheJoker(prior)
# Load the data:
logger.debug(f"{apogee_id}: Loading all data")
allstar, allvisit = c.load_alldata()
allstar = allstar[np.isin(allstar['APOGEE_ID'].astype(str), apogee_id)]
allvisit = allvisit[np.isin(allvisit['APOGEE_ID'].astype(str),
allstar['APOGEE_ID'].astype(str))]
visits = allvisit[allvisit['APOGEE_ID'] == apogee_id]
data = get_rvdata(visits)
t0 = time.time()
# Read MAP sample:
MAP_sample = extract_MAP_sample(metadata_row)
logger.log(1, f"{apogee_id}: MAP sample loaded")
# Run MCMC:
with joker.prior.model as model:
logger.log(1, f"{apogee_id}: Setting up MCMC...")
mcmc_init = joker.setup_mcmc(data, MAP_sample)
logger.log(1, f"{apogee_id}: ...setup complete")
if 'ln_prior' not in model.named_vars:
ln_prior_var = None
for k in joker.prior._nonlinear_equiv_units:
var = model.named_vars[k]
try:
if ln_prior_var is None:
ln_prior_var = var.distribution.logp(var)
else:
ln_prior_var = ln_prior_var + var.distribution.logp(
var)
except Exception as e:
logger.warning("Cannot auto-compute log-prior value for "
f"parameter {var}.")
print(e)
continue
pm.Deterministic('ln_prior', ln_prior_var)
logger.log(1, f"{apogee_id}: setting up ln_prior in pymc3 model")
if 'logp' not in model.named_vars:
pm.Deterministic('logp', model.logpt)
logger.log(1, f"{apogee_id}: setting up logp in pymc3 model")
logger.debug(f"{apogee_id}: Starting MCMC sampling")
trace = pm.sample(start=mcmc_init,
chains=4,
cores=1,
step=xo.get_dense_nuts_step(target_accept=0.95),
tune=c.tune,
draws=c.draws)
pm.save_trace(trace, directory=this_cache_path, overwrite=True)
logger.debug(
"{apogee_id}: Finished MCMC sampling ({time:.2f} seconds)".format(
apogee_id=apogee_id, time=time.time() - t0))
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}')
cov = np.dot(L, L.T)
# %% [markdown]
# And then we can sample this using PyMC3 and :func:`exoplanet.get_dense_nuts_step`:
# %%
import pymc3 as pm
import exoplanet as xo
with pm.Model() as model:
pm.MvNormal("x", mu=np.zeros(ndim), chol=L, shape=(ndim, ))
trace = pm.sample(tune=2000,
draws=2000,
chains=2,
cores=2,
step=xo.get_dense_nuts_step())
# %% [markdown]
# This is a little more verbose than the standard use of PyMC3, but the performance is several orders of magnitude better than you would get without the mass matrix tuning.
# As you can see from the `pymc3.summary`, the autocorrelation time of this chain is about 1 as we would expect for a simple problem like this.
# %%
pm.summary(trace)
# %% [markdown]
# ## Evaluating model components for specific samples
#
# I find that when I'm debugging a PyMC3 model, I often want to inspect the value of some part of the model for a given set of parameters.
# As far as I can tell, there isn't a simple way to do this in PyMC3, so *exoplanet* comes with a hack for doing this: :func:`exoplanet.eval_in_model`.
# This function handles the mapping between named PyMC3 variables and the input required by the Theano function that can evaluate the requested variable or tensor.
#
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 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_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
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
