def _sample_pymc3(cls, dist, size):
"""Sample from PyMC3."""
import pymc3
pymc3_rv_map = {
'BetaDistribution': lambda dist:
pymc3.Beta('X', alpha=float(dist.alpha), beta=float(dist.beta)),
'CauchyDistribution': lambda dist:
pymc3.Cauchy('X', alpha=float(dist.x0), beta=float(dist.gamma)),
'ChiSquaredDistribution': lambda dist:
pymc3.ChiSquared('X', nu=float(dist.k)),
'ExponentialDistribution': lambda dist:
pymc3.Exponential('X', lam=float(dist.rate)),
'GammaDistribution': lambda dist:
pymc3.Gamma('X', alpha=float(dist.k), beta=1/float(dist.theta)),
'LogNormalDistribution': lambda dist:
pymc3.Lognormal('X', mu=float(dist.mean), sigma=float(dist.std)),
'NormalDistribution': lambda dist:
pymc3.Normal('X', float(dist.mean), float(dist.std)),
'GaussianInverseDistribution': lambda dist:
pymc3.Wald('X', mu=float(dist.mean), lam=float(dist.shape)),
'ParetoDistribution': lambda dist:
pymc3.Pareto('X', alpha=float(dist.alpha), m=float(dist.xm)),
'UniformDistribution': lambda dist:
pymc3.Uniform('X', lower=float(dist.left), upper=float(dist.right))
}
dist_list = pymc3_rv_map.keys()
if dist.__class__.__name__ not in dist_list:
return None
with pymc3.Model():
pymc3_rv_map[dist.__class__.__name__](dist)
return pymc3.sample(size, chains=1, progressbar=False)[:]['X']
def build_model(mask=None):
if mask is None:
mask = np.ones_like(x, dtype=bool)
with pm.Model() as model:
# Stellar parameters
mean = pm.Normal("mean", mu=0.0, sigma=10.0)
u = xo.distributions.QuadLimbDark("u")
# Gaussian process noise model
sigma = pm.InverseGamma("sigma", alpha=3.0, beta=2 * np.median(yerr))
S_tot = pm.InverseGamma("S_tot", alpha=3.0, beta=2 * np.median(yerr))
ell = pm.Lognormal("ell", mu=0.0, sigma=1.0)
Q = 1.0 / 3.0
w0 = 2 * np.pi / ell
S0 = S_tot / (w0 * Q)
kernel = xo.gp.terms.SHOTerm(S0=S0, w0=w0, Q=Q)
# Transit parameters
t0 = pm.Bound(
pm.Normal,
lower=t0_guess - max_duration,
upper=t0_guess + max_duration,
)(
"t0",
mu=t0_guess,
sigma=0.5 * duration_guess,
testval=t0_guess,
shape=num_toi,
)
depth = pm.Lognormal(
"transit_depth",
mu=np.log(depth_guess),
sigma=np.log(1.2),
shape=num_toi,
)
duration = pm.Bound(pm.Lognormal,
lower=min_duration,
upper=max_duration)(
"transit_duration",
mu=np.log(duration_guess),
sigma=np.log(1.2),
shape=num_toi,
testval=min(
max(duration_guess, 2 * min_duration),
0.99 * max_duration),
)
b = xo.distributions.UnitUniform("b", shape=num_toi)
# Dealing with period, treating single transits properly
period_params = []
period_values = []
t_max_values = []
for n in range(num_toi):
if single_transit[n]:
period = pm.Pareto(
f"period_{n}",
m=period_min[n],
alpha=2.0 / 3,
testval=period_guess[n],
)
period_params.append(period)
t_max_values.append(t0[n])
else:
t_max = pm.Bound(
pm.Normal,
lower=t_max_guess[n] - max_duration[n],
upper=t_max_guess[n] + max_duration[n],
)(
f"t_max_{n}",
mu=t_max_guess[n],
sigma=0.5 * duration_guess[n],
testval=t_max_guess[n],
)
period = (t_max - t0[n]) / num_periods[n]
period_params.append(t_max)
t_max_values.append(t_max)
period_values.append(period)
period = pm.Deterministic("period", tt.stack(period_values))
t_max = pm.Deterministic("t_max", tt.stack(t_max_values))
# Compute the radius ratio from the transit depth, impact parameter, and
# limb darkening parameters making the small-planet assumption
u1 = u[0]
u2 = u[1]
mu = tt.sqrt(1 - b**2)
ror = pm.Deterministic(
"ror",
tt.sqrt(1e-3 * depth * (1 - u1 / 3 - u2 / 6) /
(1 - u1 * (1 - mu) - u2 * (1 - mu)**2)),
)
# Set up the orbit
orbit = xo.orbits.KeplerianOrbit(period=period,
duration=duration,
t0=t0,
b=b)
# We're going to track the implied density for reasons that will become clear later
pm.Deterministic("rho_circ", orbit.rho_star)
# Set up the mean transit model
star = xo.LimbDarkLightCurve(u)
lc_model = tess_world.LightCurveModels(mean, star, orbit, ror)
# Finally the GP observation model
gp = xo.gp.GP(kernel, x[mask], yerr[mask]**2 + sigma**2, mean=lc_model)
gp.marginal("obs", observed=y[mask])
# This is a check on the transit depth constraint
D = tt.concatenate(
(
lc_model.light_curves(x[mask]) / depth[None, :],
tt.ones((mask.sum(), 1)),
),
axis=-1,
)
DTD = tt.dot(D.T, gp.apply_inverse(D))
DTy = tt.dot(D.T, gp.apply_inverse(y[mask, None]))
model.w = tt.slinalg.solve(DTD, DTy)[:, 0]
model.sigma_w = tt.sqrt(
tt.diag(tt.slinalg.solve(DTD, tt.eye(num_toi + 1))))
# Double check that everything looks good - we shouldn't see any NaNs!
print(model.check_test_point())
# Optimize the model
map_soln = model.test_point
map_soln = xo.optimize(map_soln, [sigma])
map_soln = xo.optimize(map_soln, [mean, depth, b, duration])
map_soln = xo.optimize(map_soln, [sigma, S_tot, ell])
map_soln = xo.optimize(map_soln, [mean, u])
map_soln = xo.optimize(map_soln, period_params)
map_soln = xo.optimize(map_soln)
# Save some of the key parameters
model.map_soln = map_soln
model.lc_model = lc_model
model.gp = gp
model.mask = mask
model.x = x[mask]
model.y = y[mask]
model.yerr = yerr[mask]
return model
def pymc3_hrchl_fit(data, tune=1000, nsteps=1000, random_seed=0):
asep = data['asep'].values
asep_err = data['asep_err'].values
cr = np.abs(data['cr'].values)
cr_err = data['cr_err'].values
parallax = data['parallax'].values
with pm.Model() as hierarchical_model:
# PRIORS
# -------------
# Separations
# (for normal disribution of log of separations)
width = pm.Gamma('width', mu=1.1, sigma=0.5)
center = pm.Gamma('center', mu=5, sigma=3)
# power_index
power_index = pm.Normal('power_index', mu=1.2, sigma=.1)
# MODELS OF POPULATIONS PHYSICAL PROPERTIES
# --------------
# Gaussian model for separations (in AU)
sep_physical = pm.Lognormal('sep_physical',
mu=center,
sigma=width,
shape=len(asep))
# Mass Ratios - inverted
mass_ratios_inverted = pm.Pareto('mass_ratios_inverted',
alpha=power_index,
m=1,
shape=len(cr))
# MAPPING FROM PHYSICAL TO OBSERVED PROPERTIES
# ---------------
# physical separations to angular separations
sep_angular = pm.Deterministic('sep_angular', sep_physical * parallax)
# inverted mass ratios to inverted contrast ratios
contrast_ratios = pm.Deterministic('contrast_ratios',
imr_to_cr(mass_ratios_inverted))
# LIKELIHOODS, WITH MEASUREMENT ERROR
# -----------------
# separations
sep_observed = pm.Normal('sep_observed',
mu=sep_angular,
sigma=asep_err,
observed=asep)
# contrast ratios
cr_observed = pm.Normal('cr_observed',
mu=contrast_ratios,
sigma=cr_err,
observed=cr)
# ACCOUNTING FOR OBSERVATION LIMITS
# ------------------
# integrate out the points which fall outside of the observation limits from the likelihood
# pm.Potential accounts for truncation and normalizes likelihood
# separation limits
truncated_seps = pm.Potential(
'truncated_seps', -pm.math.logdiffexp(
normal_lcdf(sep_angular, asep_err, sep_max(cr)),
normal_lcdf(sep_angular, asep_err, sep_min(cr))))
# contrast ratio limits
truncated_crs = pm.Potential(
'truncated_crs', -pm.math.logdiffexp(
normal_lcdf(contrast_ratios, cr_err, cr_max(asep)),
normal_lcdf(contrast_ratios, cr_err, cr_min(asep))))
# RUNNING THE FIT
# -------------------
traces = pm.sample(tune=tune,
draws=nsteps,
step=None,
chains=1,
random_seed=random_seed)
# output as dataframe
df = pm.trace_to_dataframe(traces)
return traces, df