Add folded distribution (#1050)

Author: fehiepsi

docs/source/distributions.rst

@@ -17,6 +17,14 @@ ExpandedDistribution
     :show-inheritance:
     :member-order: bysource
 
+FoldedDistribution
+------------------
+.. autoclass:: numpyro.distributions.distribution.FoldedDistribution
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 ImproperUniform
 ---------------
 .. autoclass:: numpyro.distributions.distribution.ImproperUniform

numpyro/distributions/__init__.py

@@ -64,6 +64,7 @@
     Delta,
     Distribution,
     ExpandedDistribution,
+    FoldedDistribution,
     ImproperUniform,
     Independent,
     MaskedDistribution,
@@ -109,6 +110,7 @@
     "Distribution",
     "Exponential",
     "ExpandedDistribution",
+    "FoldedDistribution",
     "Gamma",
     "GammaPoisson",
     "GaussianRandomWalk",

numpyro/distributions/distribution.py

@@ -35,8 +35,9 @@
 
 from jax import lax, tree_util
 import jax.numpy as jnp
+from jax.scipy.special import logsumexp
 
-from numpyro.distributions.transforms import ComposeTransform, Transform
+from numpyro.distributions.transforms import AbsTransform, ComposeTransform, Transform
 from numpyro.distributions.util import (
     lazy_property,
     promote_shapes,
@@ -1051,6 +1052,38 @@ def tree_flatten(self):
         )
 
 
+class FoldedDistribution(TransformedDistribution):
+    """
+    Equivalent to ``TransformedDistribution(base_dist, AbsTransform())``,
+    but additionally supports :meth:`log_prob` .
+
+    :param Distribution base_dist: A univariate distribution to reflect.
+    """
+
+    support = constraints.positive
+
+    def __init__(self, base_dist, validate_args=None):
+        if base_dist.event_shape:
+            raise ValueError("Only univariate distributions can be folded.")
+        super().__init__(base_dist, AbsTransform(), validate_args=validate_args)
+
+    @validate_sample
+    def log_prob(self, value):
+        dim = max(len(self.batch_shape), jnp.ndim(value))
+        plus_minus = jnp.array([1.0, -1.0]).reshape((2,) + (1,) * dim)
+        return logsumexp(self.base_dist.log_prob(plus_minus * value), axis=0)
+
+    def tree_flatten(self):
+        base_flatten, base_aux = self.base_dist.tree_flatten()
+        return base_flatten, (type(self.base_dist), base_aux)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        base_cls, base_aux = aux_data
+        base_dist = base_cls.tree_unflatten(base_aux, params)
+        return cls(base_dist)
+
+
 class Delta(Distribution):
     arg_constraints = {
         "v": constraints.dependent(is_discrete=False),

test/test_distributions.py

@@ -109,6 +109,19 @@ def __init__(self, rate, *, validate_args=None):
         super().__init__(rate, is_sparse=True, validate_args=validate_args)
 
 
+class FoldedNormal(dist.FoldedDistribution):
+    arg_constraints = {"loc": constraints.real, "scale": constraints.positive}
+
+    def __init__(self, loc, scale, validate_args=None):
+        self.loc = loc
+        self.scale = scale
+        super().__init__(dist.Normal(loc, scale), validate_args=validate_args)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        return dist.FoldedDistribution.tree_unflatten(aux_data, params)
+
+
 _DIST_MAP = {
     dist.BernoulliProbs: lambda probs: osp.bernoulli(p=probs),
     dist.BernoulliLogits: lambda logits: osp.bernoulli(p=_to_probs_bernoulli(logits)),
@@ -189,6 +202,8 @@ def get_sp_dist(jax_dist):
     T(dist.Gumbel, 0.0, 1.0),
     T(dist.Gumbel, 0.5, 2.0),
     T(dist.Gumbel, jnp.array([0.0, 0.5]), jnp.array([1.0, 2.0])),
+    T(FoldedNormal, 2.0, 4.0),
+    T(FoldedNormal, jnp.array([2.0, 50.0]), jnp.array([4.0, 100.0])),
     T(dist.HalfCauchy, 1.0),
     T(dist.HalfCauchy, jnp.array([1.0, 2.0])),
     T(dist.HalfNormal, 1.0),
@@ -1102,6 +1117,8 @@ def fn(*args):
 def test_mean_var(jax_dist, sp_dist, params):
     if jax_dist is _ImproperWrapper:
         pytest.skip("Improper distribution does not has mean/var implemented")
+    if jax_dist is FoldedNormal:
+        pytest.skip("Folded distribution does not has mean/var implemented")
     if jax_dist in (
         _TruncatedNormal,
         dist.LeftTruncatedDistribution,