import aesara.tensor as at
from aesara.graph.rewriting.basic import in2out, node_rewriter
from aesara.graph.rewriting.db import LocalGroupDB
from aesara.graph.rewriting.unify import eval_if_etuple
from aesara.tensor.random.basic import BinomialRV, NegBinomialRV, PoissonRV
from etuples import etuple, etuplize
from kanren import eq, lall, run
from unification import var
from aemcmc.rewriting import sampler_finder_db
[docs]def gamma_poisson_conjugateo(observed_val, observed_rv_expr, posterior_expr):
r"""Produce a goal that represents the application of Bayes theorem
for a beta prior with a binomial observation model.
.. math::
\frac{
Y \sim \operatorname{Poisson}\left(\lambda\right), \quad
\lambda \sim \operatorname{Gamma}\left(\alpha, \beta\right)
}{
\left(\lambda|Y=y\right) \sim \operatorname{Gamma}\left(\alpha+y, \beta+1\right)
}
Parameters
----------
observed_val
The observed value.
observed_rv_expr
An expression that represents the observed variable.
posterior_exp
An expression that represents the posterior distribution of the latent
variable.
"""
# Gamma-poisson observation model
alpha_lv, beta_lv = var(), var()
z_rng_lv = var()
z_size_lv = var()
z_type_idx_lv = var()
z_et = etuple(
etuplize(at.random.gamma), z_rng_lv, z_size_lv, z_type_idx_lv, alpha_lv, beta_lv
)
Y_et = etuple(etuplize(at.random.poisson), var(), var(), var(), z_et)
# Posterior distribution for p
new_alpha_et = etuple(etuplize(at.add), alpha_lv, observed_val)
new_beta_et = etuple(etuplize(at.add), beta_lv, 1)
z_posterior_et = etuple(
etuplize(at.random.gamma),
new_alpha_et,
new_beta_et,
rng=z_rng_lv,
size=z_size_lv,
dtype=z_type_idx_lv,
)
return lall(
eq(observed_rv_expr, Y_et),
eq(posterior_expr, z_posterior_et),
)
@node_rewriter([PoissonRV])
def local_gamma_poisson_posterior(fgraph, node):
sampler_mappings = getattr(fgraph, "sampler_mappings", None)
rv_var = node.outputs[1]
key = ("local_gamma_poisson_posterior", rv_var)
if sampler_mappings is None or key in sampler_mappings.rvs_seen:
return None # pragma: no cover
q = var()
rv_et = etuplize(rv_var)
res = run(None, q, gamma_poisson_conjugateo(rv_var, rv_et, q))
res = next(res, None)
if res is None:
return None # pragma: no cover
gamma_rv = rv_et[-1].evaled_obj
gamma_posterior = eval_if_etuple(res)
sampler_mappings.rvs_to_samplers.setdefault(gamma_rv, []).append(
("local_gamma_poisson_posterior", gamma_posterior, None)
)
sampler_mappings.rvs_seen.add(key)
return rv_var.owner.outputs
[docs]def beta_binomial_conjugateo(observed_val, observed_rv_expr, posterior_expr):
r"""Produce a goal that represents the application of Bayes theorem
for a beta prior with a binomial observation model.
.. math::
\frac{
Y \sim \operatorname{Binom}\left(N, p\right), \quad
p \sim \operatorname{Beta}\left(\alpha, \beta\right)
}{
\left(p|Y=y\right) \sim \operatorname{Beta}\left(\alpha+y, \beta+N-y\right)
}
Parameters
----------
observed_val
The observed value.
observed_rv_expr
An expression that represents the observed variable.
posterior_exp
An expression that represents the posterior distribution of the latent
variable.
"""
# Beta-binomial observation model
alpha_lv, beta_lv = var(), var()
p_rng_lv = var()
p_size_lv = var()
p_type_idx_lv = var()
p_et = etuple(
etuplize(at.random.beta), p_rng_lv, p_size_lv, p_type_idx_lv, alpha_lv, beta_lv
)
n_lv = var()
Y_et = etuple(etuplize(at.random.binomial), var(), var(), var(), n_lv, p_et)
# Posterior distribution for p
new_alpha_et = etuple(etuplize(at.add), alpha_lv, observed_val)
new_beta_et = etuple(
etuplize(at.sub), etuple(etuplize(at.add), beta_lv, n_lv), observed_val
)
p_posterior_et = etuple(
etuplize(at.random.beta),
new_alpha_et,
new_beta_et,
rng=p_rng_lv,
size=p_size_lv,
dtype=p_type_idx_lv,
)
return lall(
eq(observed_rv_expr, Y_et),
eq(posterior_expr, p_posterior_et),
)
@node_rewriter([BinomialRV])
def local_beta_binomial_posterior(fgraph, node):
sampler_mappings = getattr(fgraph, "sampler_mappings", None)
rv_var = node.outputs[1]
key = ("local_beta_binomial_posterior", rv_var)
if sampler_mappings is None or key in sampler_mappings.rvs_seen:
return None # pragma: no cover
q = var()
rv_et = etuplize(rv_var)
res = run(None, q, beta_binomial_conjugateo(rv_var, rv_et, q))
res = next(res, None)
if res is None:
return None # pragma: no cover
beta_rv = rv_et[-1].evaled_obj
beta_posterior = eval_if_etuple(res)
sampler_mappings.rvs_to_samplers.setdefault(beta_rv, []).append(
("local_beta_binomial_posterior", beta_posterior, None)
)
sampler_mappings.rvs_seen.add(key)
return rv_var.owner.outputs
def beta_negative_binomial_conjugateo(observed_val, observed_rv_expr, posterior_expr):
r"""Produce a goal that represents the application of Bayes theorem
for a beta prior with a negative binomial observation model.
.. math::
\frac{
Y \sim \operatorname{NB}\left(k, p\right), \quad
p \sim \operatorname{Beta}\left(\alpha, \beta\right)
}{
\left(p \mid Y=y\right) \sim \operatorname{Beta}\left(\alpha + \sum^{n}_{i=1} y_i, \beta + k n\right)
}
where :math:`k` is the number of successes before experiment ended.
Parameters
----------
observed_val
The observed value.
observed_rv_expr
An expression that represents the observed variable.
posterior_exp
An expression that represents the posterior distribution of the latent
variable.
"""
# beta-negative_binomial observation model
alpha_lv, beta_lv = var(), var()
p_rng_lv = var()
p_size_lv = var()
p_type_idx_lv = var()
p_et = etuple(
etuplize(at.random.beta), p_rng_lv, p_size_lv, p_type_idx_lv, alpha_lv, beta_lv
)
n_lv = var() # success
Y_et = etuple(
etuplize(at.random.negative_binomial), var(), var(), var(), n_lv, p_et
)
new_alpha_et = etuple(etuplize(at.add), alpha_lv, observed_val)
new_beta_et = etuple(etuplize(at.add), beta_lv, n_lv)
p_posterior_et = etuple(
etuplize(at.random.beta),
new_alpha_et,
new_beta_et,
rng=p_rng_lv,
size=p_size_lv,
dtype=p_type_idx_lv,
)
return lall(
eq(observed_rv_expr, Y_et),
eq(posterior_expr, p_posterior_et),
)
@node_rewriter([NegBinomialRV])
def local_beta_negative_binomial_posterior(fgraph, node):
sampler_mappings = getattr(fgraph, "sampler_mappings", None)
rv_var = node.outputs[1]
key = ("local_beta_negative_binomial_posterior", rv_var)
if sampler_mappings is None or key in sampler_mappings.rvs_seen:
return None # pragma: no cover
q = var()
rv_et = etuplize(rv_var)
res = run(None, q, beta_negative_binomial_conjugateo(rv_var, rv_et, q))
res = next(res, None)
if res is None:
return None # pragma: no cover
beta_rv = rv_et[-1].evaled_obj
beta_posterior = eval_if_etuple(res)
sampler_mappings.rvs_to_samplers.setdefault(beta_rv, []).append(
("local_beta_negative_binomial_posterior", beta_posterior, None)
)
sampler_mappings.rvs_seen.add(key)
return rv_var.owner.outputs
conjugates_db = LocalGroupDB(apply_all_rewrites=True)
conjugates_db.name = "conjugates_db"
conjugates_db.register("beta_binomial", local_beta_binomial_posterior, "basic")
conjugates_db.register("gamma_poisson", local_gamma_poisson_posterior, "basic")
conjugates_db.register(
"negative_binomial", local_beta_negative_binomial_posterior, "basic"
)
sampler_finder_db.register(
"conjugates", in2out(conjugates_db.query("+basic"), name="gibbs"), "basic"
)