Skip to content

Commit

Permalink
Add the relation between a Half-Cauchy and an Inverse-Gamma scale mix…
Browse files Browse the repository at this point in the history
…ture
  • Loading branch information
rlouf committed Sep 26, 2022
1 parent 5b5ae2f commit aeeda45
Show file tree
Hide file tree
Showing 3 changed files with 143 additions and 2 deletions.
99 changes: 99 additions & 0 deletions aemcmc/scale_mixtures.py
Original file line number Diff line number Diff line change
@@ -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,
),
)
4 changes: 2 additions & 2 deletions aemcmc/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""

Expand Down
42 changes: 42 additions & 0 deletions tests/test_scale_mixtures.py
Original file line number Diff line number Diff line change
@@ -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))

0 comments on commit aeeda45

Please sign in to comment.