def logp(likelihood):
#p_D_p = tt.log(tt.dot(likelihood, p).prod())
logps = tt.log(p) + likelihood
# ps = tt.exp(logps).sum(axis=1)
# p_D_p = tt.log(ps).sum()
p_D_p = pm.logsumexp(logps, axis=1).sum()
return p_D_p
def lp_mix(xs):
def flip(xs):
# Filp 1st and 2nd features
return tt.stack([xs[:, 1], xs[:, 0]], axis=0).T
return pm.logsumexp(
tt.stack([tt.log(z) + lp_m1(xs),
tt.log(1 - z) + lp_m2(flip(xs))],
axis=0))
def _get_update_log_q_ploidy_sjk_theano_func(self) -> th.compile.function_module.Function:
new_log_q_ploidy_sjk = (self.ploidy_workspace.log_p_ploidy_jk.dimshuffle('x', 0, 1)
+ self.ploidy_workspace.log_ploidy_emission_sjk)
new_log_q_ploidy_sjk -= pm.logsumexp(new_log_q_ploidy_sjk, axis=2)
old_log_q_ploidy_sjk = self.ploidy_workspace.log_q_ploidy_sjk
admixed_new_log_q_ploidy_sjk = commons.safe_logaddexp(
new_log_q_ploidy_sjk + np.log(self.inference_params.caller_external_admixing_rate),
old_log_q_ploidy_sjk + np.log(1.0 - self.inference_params.caller_external_admixing_rate))
update_norm_sj = commons.get_hellinger_distance(admixed_new_log_q_ploidy_sjk, old_log_q_ploidy_sjk)
return th.function(inputs=[],
outputs=[update_norm_sj],
updates=[(self.ploidy_workspace.log_q_ploidy_sjk, admixed_new_log_q_ploidy_sjk)])
def _get_update_log_q_ploidy_sjk_theano_func(self) -> th.compile.function_module.Function:
new_log_q_ploidy_sjk = (self.ploidy_workspace.log_p_ploidy_jk.dimshuffle('x', 0, 1)
+ self.ploidy_workspace.log_ploidy_emission_sjk)
new_log_q_ploidy_sjk -= pm.logsumexp(new_log_q_ploidy_sjk, axis=2)
old_log_q_ploidy_sjk = self.ploidy_workspace.log_q_ploidy_sjk
admixed_new_log_q_ploidy_sjk = commons.safe_logaddexp(
new_log_q_ploidy_sjk + np.log(self.inference_params.caller_external_admixing_rate),
old_log_q_ploidy_sjk + np.log(1.0 - self.inference_params.caller_external_admixing_rate))
update_norm_sj = commons.get_hellinger_distance(admixed_new_log_q_ploidy_sjk, old_log_q_ploidy_sjk)
return th.function(inputs=[],
outputs=[update_norm_sj],
updates=[(self.ploidy_workspace.log_q_ploidy_sjk, admixed_new_log_q_ploidy_sjk)])
def logp(u, v, K2, N2):
num_elements = (K2 + 1) * (K2 + 2) / 2 # Number of values to consider
elems = np.zeros(num_elements)
count = 0
# We define each element via its log for computational reasons
for i in range(K2 + 1):
j_upper = K2 + 1 - i
for j in range(j_upper):
temp =
return pm.logsumexp(elems)
def logp(r):
lp = lw + betabin_lpdf(r, R, a, b).sum(1)
return pmc.logsumexp(lp, 0) #.sum()
def logp(r):
lp = lw + bin_lpdf(r, R, theta).sum(1)
return pmc.logsumexp(lp, 0) #.sum()
def logp_(value):
logps = tt.log(mf) + value
return tt.sum(logsumexp(logps, axis=1))
def duration_to_eccentricity(func, duration, ror, **kwargs):
num_planets = kwargs.pop("num_planets", 1)
orbit_type = kwargs.pop("orbit_type", KeplerianOrbit)
name = kwargs.get("name", "dur_ecc")
inputs = _get_consistent_inputs(
kwargs.get("a", None),
kwargs.get("period", None),
kwargs.get("rho_star", None),
kwargs.get("r_star", None),
kwargs.get("m_star", None),
kwargs.get("rho_star_units", None),
kwargs.get("m_planet", 0.0),
kwargs.get("m_planet_units", None),
)
a, period, rho_star, r_star, m_star, m_planet = inputs
b = kwargs.get("b", 0.0)
s = tt.sin(kwargs["omega"])
umax_inv = tt.switch(tt.lt(s, 0), tt.sqrt(1 - s**2), 1.0)
const = (period * tt.shape_padright(r_star) * tt.sqrt((1 + ror)**2 - b**2))
const /= np.pi * a
u = duration / const
e1 = -s * u**2 / ((s * u)**2 + 1)
e2 = tt.sqrt((s**2 - 1) * u**2 + 1) / ((s * u)**2 + 1)
models = []
logjacs = []
logprobs = []
for args in product(*(zip("np", (-1, 1)) for _ in range(num_planets))):
labels, signs = zip(*args)
# Compute the eccentricity branch
ecc = tt.stack([e1[i] + signs[i] * e2[i] for i in range(num_planets)])
# Work out the Jacobian
valid_ecc = tt.and_(tt.lt(ecc, 1.0), tt.ge(ecc, 0.0))
logjac = tt.switch(
tt.all(valid_ecc),
tt.sum(0.5 * tt.log(1 - ecc**2) + 2 * tt.log(s * ecc + 1) -
tt.log(tt.abs_(s + ecc)) - tt.log(const)),
-np.inf,
)
ecc = tt.switch(valid_ecc, ecc, tt.zeros_like(ecc))
# Create a sub-model to capture this component
with pm.Model(name="dur_ecc_" + "_".join(labels)) as model:
pm.Deterministic("ecc", ecc)
orbit = orbit_type(ecc=ecc, **kwargs)
logprob = tt.sum(func(orbit))
models.append(model)
logjacs.append(logjac)
logprobs.append(logprob)
# Compute the marginalized likelihood
logjacs = tt.stack(logjacs)
logprobs = tt.stack(logprobs)
logprob = tt.switch(
tt.gt(1.0 / u, umax_inv),
tt.sum(pm.logsumexp(logprobs + logjacs)),
-np.inf,
)
pm.Potential(name + "_logp", logprob)
pm.Deterministic(name + "_logjacs", logjacs)
pm.Deterministic(name + "_logprobs", logprobs)
# Loop over models and compute the marginalized values for all the
# parameters and deterministics
norm = tt.sum(pm.logsumexp(logjacs))
logw = tt.switch(
tt.gt(1.0 / u, umax_inv),
logjacs - norm,
-np.inf + tt.zeros_like(logjacs),
)
pm.Deterministic(name + "_logw", logw)
for k in models[0].named_vars.keys():
name = k[len(models[0].name) + 1:]
pm.Deterministic(
name,
sum(
tt.exp(logw[i]) * model.named_vars[model.name + "_" + name]
for i, model in enumerate(models)),
)
def __call__(self, pt):
logps = []
for k in range(self.K):
logp = self.ps[k].logp(pt) + tt.log(self.w[k])
logps.append(logp)
return pm.logsumexp(logps).squeeze()
