From 743137682b2ddb3300e8b843f8ab59828cca3c36 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Louf?= <remi@thetypicalset.com>
Date: Mon, 26 Sep 2022 00:27:23 +0200
Subject: [PATCH] Add the relation between a Half-Cauchy and an Inverse-Gamma
 scale mixture

---
 aemcmc/scale_mixtures.py     | 99 ++++++++++++++++++++++++++++++++++++
 aemcmc/transforms.py         |  4 +-
 tests/test_scale_mixtures.py | 42 +++++++++++++++
 3 files changed, 143 insertions(+), 2 deletions(-)
 create mode 100644 aemcmc/scale_mixtures.py
 create mode 100644 tests/test_scale_mixtures.py

diff --git a/aemcmc/scale_mixtures.py b/aemcmc/scale_mixtures.py
new file mode 100644
index 0000000..c7b7cda
--- /dev/null
+++ b/aemcmc/scale_mixtures.py
@@ -0,0 +1,99 @@
+import aesara.tensor as at
+from etuples import etuple, etuplize
+from kanren import eq, lall, lany, var
+from kanren.core import succeed
+
+
+def halfcauchy_inverse_gamma(in_expr, out_expr):
+    r"""Produce a goal that represents the fact that a half-Cauchy distribution can be
+    expressed as a scale mixture of inverse gamma distributions.
+
+    .. math::
+
+        \begin{equation*}
+            \frac{
+                X^2 | a \sim \operatorname{Gamma^{-1}}\left(1/2, a), \quad
+                a \sim \operatorname{Gamma^{-1}}\left(1/2, 1/A^2\right), \quad
+            }{
+                X \sim \operatorname{C^{+}}\left(0, A)
+            }
+        \end{equation*}
+
+    TODO: This relation is a particular case of a similar relation for the
+    half-t distribution [1]_ which does not have an implementation yet in Aesara.
+    When it becomes available we should replace this relation with the more
+    general one, and implement the relation between the half-t and half-Cauchy
+    distributions.
+
+    Parameters
+    ----------
+    in_expr
+        An expression that represents a half-Cauchy distribution.
+    out_expr
+        An expression that represents the square root of the inverse gamma scale
+        mixture.
+
+    References
+    ----------
+    .. [1]:  Wand, M. P., Ormerod, J. T., Padoan, S. A., & Frühwirth, R. (2011).
+             Mean field variational Bayes for elaborate distributions. Bayesian
+             Analysis, 6(4), 847-900.
+
+    """
+
+    # Half-Cauchy distribution
+    rng_lv, size_lv, type_idx_lv = var(), var(), var()
+    loc_at = at.as_tensor(0)
+    scale_lv = var()
+    X_halfcauchy_et = etuple(
+        etuplize(at.random.halfcauchy), rng_lv, size_lv, type_idx_lv, loc_at, scale_lv
+    )
+
+    # Inverse-Gamma scale mixture
+    rng_inner_lv, size_inner_lv, type_idx_inner_lv = var(), var(), var()
+    rng_outer_lv, size_outer_lv, type_idx_outer_lv = var(), var(), var()
+    a_et = etuple(
+        etuplize(at.random.invgamma),
+        at.as_tensor(0.5),
+        etuple(
+            etuplize(at.true_div),
+            at.as_tensor(1.0),
+            etuple(etuplize(at.pow), scale_lv, at.as_tensor(2)),
+        ),
+        rng=rng_inner_lv,
+        size=size_inner_lv,
+        dtype=type_idx_inner_lv,
+    )
+    X_scale_mixture_square_et = etuple(
+        etuplize(at.random.invgamma),
+        at.as_tensor(0.5),
+        etuple(
+            etuplize(at.true_div),
+            at.as_tensor(1.0),
+            a_et,
+        ),
+        rng=rng_outer_lv,
+        size=size_outer_lv,
+        dtype=type_idx_outer_lv,
+    )
+    X_scale_mixture_et = etuple(etuplize(at.sqrt), X_scale_mixture_square_et)
+
+    return lall(
+        eq(in_expr, X_halfcauchy_et),
+        eq(out_expr, X_scale_mixture_et),
+        eq(rng_inner_lv, rng_lv),
+        eq(type_idx_inner_lv, type_idx_lv),
+        eq(size_inner_lv, size_lv),
+        lany(
+            eq(rng_outer_lv, rng_lv),
+            succeed,
+        ),
+        lany(
+            eq(size_outer_lv, size_lv),
+            succeed,
+        ),
+        lany(
+            eq(type_idx_outer_lv, type_idx_lv),
+            succeed,
+        ),
+    )
diff --git a/aemcmc/transforms.py b/aemcmc/transforms.py
index 6cc9e82..2c79c44 100644
--- a/aemcmc/transforms.py
+++ b/aemcmc/transforms.py
@@ -34,10 +34,10 @@ def location_scale_transform(in_expr, out_expr):
     Parameters
     ----------
     in_expr
-        An expression that represents a random variable whose distribution belongs
+        an expression that represents a random variable whose distribution belongs
         to the location-scale family.
     out_expr
-        An expression for the non-centered representation of this random variable.
+        an expression for the non-centered representation of this random variable.
 
     """
 
diff --git a/tests/test_scale_mixtures.py b/tests/test_scale_mixtures.py
new file mode 100644
index 0000000..f290e31
--- /dev/null
+++ b/tests/test_scale_mixtures.py
@@ -0,0 +1,42 @@
+import aesara.tensor as at
+import pytest
+from aesara.graph.rewriting.unify import eval_if_etuple
+from kanren import run, var
+
+from aemcmc.scale_mixtures import halfcauchy_inverse_gamma
+
+
+def test_halfcauchy_to_inverse_gamma_mixture():
+
+    srng = at.random.RandomStream(0)
+    A = at.scalar("A")
+    X_rv = srng.halfcauchy(0, A)
+
+    q_lv = var()
+    results = run(0, q_lv, halfcauchy_inverse_gamma(X_rv, q_lv))
+
+    found_mixture = False
+    for res in results:
+        try:
+            mixture = eval_if_etuple(res)
+            found_mixture = True
+        except (AttributeError, TypeError):
+            continue
+
+    assert found_mixture is True
+    assert isinstance(mixture.owner.op, type(at.sqrt))
+    assert isinstance(mixture.owner.inputs[0].owner.op, type(at.random.invgamma))
+
+
+@pytest.mark.xfail(
+    reason="Op.__call__ does not dispatch to Op.make_node for some RandomVariable and etuple evaluation returns an error"
+)
+def test_halfcauchy_from_inverse_gamma_mixture():
+
+    srng = at.random.RandomStream(0)
+    A = at.scalar("A")
+    a_rv = srng.invgamma(0.5, 1.0 / A**2)
+    X_rv = at.sqrt(srng.invgamma(0.5, 1.0 / a_rv))
+
+    q_lv = var()
+    run(0, q_lv, halfcauchy_inverse_gamma(q_lv, X_rv))