From c711b2b67f5beae27088d12074aea6a091d9c512 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     | 82 ++++++++++++++++++++++++++++++++++++
 aemcmc/transforms.py         |  4 +-
 tests/test_scale_mixtures.py | 39 +++++++++++++++++
 3 files changed, 123 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..c324fc0
--- /dev/null
+++ b/aemcmc/scale_mixtures.py
@@ -0,0 +1,82 @@
+import aesara.tensor as at
+from etuples import etuple, etuplize
+from kanren import eq, lall, var
+
+
+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()
+    X_square_scale_mixture_et = etuple(
+        etuplize(at.random.invgamma),
+        at.as_tensor(0.5),
+        etuple(
+            etuplize(at.true_div),
+            at.as_tensor(1),
+            etuple(
+                etuplize(at.random.invgamma),
+                at.as_tensor(0.5),
+                etuple(
+                    etuplize(at.true_div),
+                    at.as_tensor(1),
+                    etuple(etuplize(at.pow), scale_lv, at.as_tensor(2)),
+                ),
+                rng=rng_inner_lv,
+                size=size_inner_lv,
+                dtype=type_idx_inner_lv,
+            ),
+        ),
+        rng=rng_outer_lv,
+        size=size_outer_lv,
+        dtype=type_idx_outer_lv,
+    )
+    X_scale_mixture_et = etuple(etuplize(at.sqrt), X_square_scale_mixture_et)
+
+    return lall(
+        eq(in_expr, X_halfcauchy_et),
+        eq(out_expr, X_scale_mixture_et),
+    )
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..603d7a2
--- /dev/null
+++ b/tests/test_scale_mixtures.py
@@ -0,0 +1,39 @@
+import aesara.tensor as at
+import pytest
+from aesara.graph.fg import FunctionGraph
+from aesara.graph.kanren import KanrenRelationSub
+
+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)
+
+    fgraph = FunctionGraph(outputs=[X_rv], clone=False)
+    res = KanrenRelationSub(halfcauchy_inverse_gamma).transform(
+        fgraph, fgraph.outputs[0].owner
+    )[0]
+
+    assert isinstance(res.owner.op, type(at.sqrt))
+    assert isinstance(res.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 / A**2)
+    X_rv = at.sqrt(srng.invgamma(0.5, 1 / a_rv))
+
+    fgraph = FunctionGraph(outputs=[X_rv], clone=False)
+    res = KanrenRelationSub(lambda x, y: halfcauchy_inverse_gamma(y, x)).transform(
+        fgraph, fgraph.outputs[0].owner
+    )[0]
+
+    assert isinstance(res.owner, type(at.random.halfcauchy))