Add goal for beta-binomial observation model

Author: rlouf

aemcmc/conjugates.py

@@ -0,0 +1,59 @@
+import aesara.tensor as at
+from etuples import etuple, etuplize
+from kanren.facts import Relation, fact
+from unification import var
+
+conjugate = Relation("conjugate")
+
+
+def _create_beta_binomial_goal():
+    r"""Produce a goal that represents the application of Bayes theorem
+    for a beta prior with a binomial observation model.
+
+    .. math::
+
+        \begin{align*}
+            p &\sim \operatorname{Beta}\left(\alpha, \beta\right)\\
+            y &\sim \operatorname{Binomial}\left(n, p\right)
+        \end{align*}
+
+    If we observe :math:`y=Y`, then :math:`p` follows a beta distribution:
+
+    .. math::
+
+        p \sim \operatorname{Beta}\left(\alpha + Y, \beta + n - Y\right)
+
+    """
+
+    # Beta-binomial observation model
+    alpha_lv, beta_lv = var(), var()
+    p_srng_lv = var()
+    p_size_lv = var()
+    p_type_idx_lv = var()
+    p_lv = etuple(
+        etuplize(at.random.beta), p_srng_lv, p_size_lv, p_type_idx_lv, alpha_lv, beta_lv
+    )
+    n_lv = var()
+    Y_lv = etuple(etuplize(at.random.binomial), var(), var(), var(), n_lv, p_lv)
+
+    y_vv = var()  # observation
+
+    # Posterior distribution for p
+    new_alpha_lv = etuple(etuplize(at.add), alpha_lv, y_vv)
+    new_beta_lv = etuple(
+        etuplize(at.add), beta_lv, n_lv, etuple(etuplize(at.neg), y_vv)
+    )
+    p_posterior_lv = etuple(
+        etuplize(at.random.beta),
+        p_srng_lv,
+        p_size_lv,
+        p_type_idx_lv,
+        new_alpha_lv,
+        new_beta_lv,
+    )
+
+    return (Y_lv, y_vv, p_posterior_lv)
+
+
+model, observation, posterior = _create_beta_binomial_goal()
+fact(conjugate, (model, observation), posterior)

tests/test_conjugates.py

@@ -0,0 +1,28 @@
+import aesara.tensor as at
+from aesara.tensor.random import RandomStream
+from kanren import run
+from kanren.graph import walko
+from unification import var
+
+from aemcmc.conjugates import conjugate
+
+
+def test_beta_binomial_conjugate():
+    """Test that we can produce the closed-form posterior for
+    the binomial observation model with a beta prior.
+
+    """
+
+    srng = RandomStream(0)
+
+    alpha = at.scalar("alpha")
+    beta = at.scalar("alpha")
+    p = srng.beta(alpha, beta, name="p")
+
+    n = at.iscalar("n")
+    Y_rv = srng.binomial(n, p)
+    y_vv = Y_rv.clone()
+
+    q_lv = var()
+    (posterior_expr,) = run(1, q_lv, walko(conjugate, (Y_rv, y_vv), q_lv))
+    posterior_expr.evaled_obj