Added ZeroSumNormal

Co-authored-by: Adrian Seyboldt <aseyboldt@users.noreply.github.com>

Author: kc611

pymc3/distributions/__init__.py

@@ -48,6 +48,7 @@
     VonMises,
     Wald,
     Weibull,
+    ZeroSumNormal,
 )
 from pymc3.distributions.discrete import (
     Bernoulli,
@@ -123,6 +124,7 @@
     "HalfStudentT",
     "ChiSquared",
     "HalfNormal",
+    "ZeroSumNormal",
     "Wald",
     "Pareto",
     "InverseGamma",

pymc3/distributions/continuous.py

@@ -65,6 +65,7 @@
     "Lognormal",
     "ChiSquared",
     "HalfNormal",
+    "ZeroSumNormal",
     "Wald",
     "Pareto",
     "InverseGamma",
@@ -924,6 +925,67 @@ def logcdf(self, value):
         )
 
 
+class ZeroSumNormal(Continuous):
+    def __init__(self, sigma=1, zerosum_dims=None, zerosum_axes=None, **kwargs):
+        shape = kwargs.get("shape", ())
+        dims = kwargs.get("dims", None)
+        if isinstance(shape, int):
+            shape = (shape,)
+
+        if isinstance(dims, str):
+            dims = (dims,)
+
+        self.mu = self.median = self.mode = tt.zeros(shape)
+        self.sigma = tt.as_tensor_variable(sigma)
+
+        if zerosum_dims is None and zerosum_axes is None:
+            if shape:
+                zerosum_axes = (-1,)
+            else:
+                zerosum_axes = ()
+
+        if isinstance(zerosum_axes, int):
+            zerosum_axes = (zerosum_axes,)
+
+        if isinstance(zerosum_dims, str):
+            zerosum_dims = (zerosum_dims,)
+
+        if zerosum_axes is not None and zerosum_dims is not None:
+            raise ValueError("Only one of zerosum_axes and zerosum_dims can be specified.")
+
+        if zerosum_dims is not None:
+            if dims is None:
+                raise ValueError("zerosum_dims can only be used with the dims kwargs.")
+            zerosum_axes = []
+            for dim in zerosum_dims:
+                zerosum_axes.append(dims.index(dim))
+        self.zerosum_axes = [a if a >= 0 else len(shape) + a for a in zerosum_axes]
+
+        if "transform" not in kwargs or kwargs["transform"] == None:
+            kwargs["transform"] = transforms.ZeroSumTransform(zerosum_axes)
+
+        super().__init__(**kwargs)
+
+    def logp(self, value):
+        return Normal.dist(sigma=self.sigma).logp(value)
+
+    def _random(self, scale, size):
+        samples = stats.norm.rvs(loc=0, scale=scale, size=size)
+        for axis in self.zerosum_axes:
+            samples -= np.mean(samples, axis=axis, keepdims=True)
+        return samples
+
+    def random(self, point=None, size=None):
+        (sigma,) = draw_values([self.sigma], point=point, size=size)
+        return generate_samples(self._random, scale=sigma, dist_shape=self.shape, size=size)
+
+    def _distr_parameters_for_repr(self):
+        return ["sigma"]
+
+    def logcdf(self, value):
+        raise NotImplementedError()
+
+
 class Wald(PositiveContinuous):
     r"""
     Wald log-likelihood.

pymc3/distributions/distribution.py

@@ -98,7 +98,7 @@ def __new__(cls, name, *args, **kwargs):
             raise TypeError("observed needs to be data but got: {}".format(type(data)))
         total_size = kwargs.pop("total_size", None)
 
-        dims = kwargs.pop("dims", None)
+        dims = kwargs["dims"] if "dims" in kwargs else None
         has_shape = "shape" in kwargs
         shape = kwargs.pop("shape", None)
         if dims is not None:

pymc3/distributions/transforms.py

@@ -14,6 +14,8 @@
 
 import warnings
 
+from typing import List
+
 import numpy as np
 import theano.tensor as tt
 
@@ -565,3 +567,78 @@ def jacobian_det(self, y):
             else:
                 det += det_
         return det
+
+
+def _extend_axis(array, axis):
+    n = array.shape[axis] + 1
+    sum_vals = array.sum(axis, keepdims=True)
+    norm = sum_vals / (np.sqrt(n) + n)
+    fill_val = norm - sum_vals / np.sqrt(n)
+
+    out = tt.concatenate([array, fill_val], axis=axis)
+    return out - norm
+
+
+def _extend_axis_rev(array, axis):
+    if axis < 0:
+        axis = axis % array.ndim
+    assert axis >= 0 and axis < array.ndim
+
+    n = array.shape[axis]
+    last = tt.take(array, [-1], axis=axis)
+
+    sum_vals = -last * np.sqrt(n)
+    norm = sum_vals / (np.sqrt(n) + n)
+    slice_before = (slice(None, None),) * axis
+    return array[slice_before + (slice(None, -1),)] + norm
+
+
+def _extend_axis_val(array, axis):
+    n = array.shape[axis] + 1
+    sum_vals = array.sum(axis, keepdims=True)
+    norm = sum_vals / (np.sqrt(n) + n)
+    fill_val = norm - sum_vals / np.sqrt(n)
+
+    out = np.concatenate([array, fill_val], axis=axis)
+    return out - norm
+
+
+def _extend_axis_rev_val(array, axis):
+    n = array.shape[axis]
+    last = np.take(array, [-1], axis=axis)
+
+    sum_vals = -last * np.sqrt(n)
+    norm = sum_vals / (np.sqrt(n) + n)
+    slice_before = (slice(None, None),) * len(array.shape[:axis])
+    return array[slice_before + (slice(None, -1),)] + norm
+
+
+class ZeroSumTransform(Transform):
+    name = "zerosum"
+
+    _zerosum_axes: List[int]
+
+    def __init__(self, zerosum_axes):
+        self._zerosum_axes = zerosum_axes
+
+    def forward(self, x):
+        for axis in self._zerosum_axes:
+            x = _extend_axis_rev(x, axis=axis)
+        return floatX(x)
+
+    def forward_val(self, x, point):
+        for axis in self._zerosum_axes:
+            x = _extend_axis_rev_val(x, axis=axis)
+        return x
+
+    def backward(self, z):
+        z = tt.as_tensor_variable(z)
+        for axis in self._zerosum_axes:
+            z = _extend_axis(z, axis=axis)
+        return floatX(z)
+
+    def jacobian_det(self, x):
+        return tt.constant(0.0)
+
+
+zerosum = ZeroSumTransform

pymc3/tests/test_distributions.py

@@ -93,8 +93,10 @@
     ZeroInflatedBinomial,
     ZeroInflatedNegativeBinomial,
     ZeroInflatedPoisson,
+    ZeroSumNormal,
     continuous,
 )
+from pymc3.distributions.transforms import zerosum
 from pymc3.math import kronecker, logsumexp
 from pymc3.model import Deterministic, Model, Point
 from pymc3.tests.helpers import select_by_precision
@@ -556,6 +558,7 @@ def check_logp(
         n_samples=100,
         extra_args=None,
         scipy_args=None,
+        transform=None,
     ):
         """
         Generic test for PyMC3 logp methods
@@ -599,13 +602,18 @@ def check_logp(
 
         def logp_reference(args):
             args.update(scipy_args)
+            if transform:
+                args["value"] = args.pop(f"value_{transform.name}__")
             return scipy_logp(**args)
 
         model = build_model(pymc3_dist, domain, paramdomains, extra_args)
         logp = model.fastlogp
 
         domains = paramdomains.copy()
-        domains["value"] = domain
+        if transform:
+            domains[f"value_{transform.name}__"] = domain
+        else:
+            domains["value"] = domain
         for pt in product(domains, n_samples=n_samples):
             pt = Point(pt, model=model)
             assert_almost_equal(
@@ -932,6 +940,22 @@ def test_half_normal(self):
             lambda value, sigma: sp.halfnorm.logcdf(value, scale=sigma),
         )
 
+    def test_zerosum_normal(self):
+        zerosum_axes = [-1]
+
+        def ref_fn(value, sigma):
+            mu = 0
+            return sp.norm.logpdf(value, mu, sigma)
+
+        self.check_logp(
+            ZeroSumNormal,
+            R,
+            {"sigma": Rplus},
+            ref_fn,
+            decimal=select_by_precision(float64=6, float32=1),
+            transform=zerosum(zerosum_axes),
+        )
+
     def test_chi_squared(self):
         self.check_logp(
             ChiSquared,

pymc3/tests/test_distributions_random.py

@@ -343,6 +343,11 @@ class TestHalfNormal(BaseTestCases.BaseTestCase):
     params = {"tau": 1.0}
 
 
+class TestZeroSumNormal(BaseTestCases.BaseTestCase):
+    distribution = pm.ZeroSumNormal
+    params = {"sigma": 1.0}
+
+
 class TestUniform(BaseTestCases.BaseTestCase):
     distribution = pm.Uniform
     params = {"lower": 0.0, "upper": 1.0}
@@ -622,6 +627,21 @@ def ref_rand(size, tau):
 
         pymc3_random(pm.HalfNormal, {"tau": Rplus}, ref_rand=ref_rand)
 
+    def test_zerosum_normal(self):
+        def ref_rand(size, sigma):
+            shape = sigma.shape
+            zerosum_axes = (-1,) if shape else ()
+            zerosum_axes = [a if a >= 0 else len(shape) + a for a in zerosum_axes]
+            n = shape[-1] if shape else 1
+            samples = st.multivariate_normal.rvs(
+                cov=sigma ** 2 * (np.eye(n) - np.ones(n) / n), size=n
+            )
+            for axis in zerosum_axes:
+                samples -= np.mean(samples, axis=axis, keepdims=True)
+            return samples
+
+        pymc3_random(pm.ZeroSumNormal, {"sigma": PdMatrix(3)}, ref_rand=ref_rand)
+
     def test_wald(self):
         # Cannot do anything too exciting as scipy wald is a
         # location-scale model of the *standard* wald with mu=1 and lam=1

pymc3/tests/test_transforms.py

@@ -242,6 +242,14 @@ def test_chain():
     close_to_logical(np.diff(vals) >= 0, True, tol)
 
 
+def test_zerosum():
+    zerosum_axes = [0]
+    zerosum_transf = tr.ZeroSumTransform(zerosum_axes)
+
+    vals = get_values(zerosum_transf, Vector(R, 5), tt.dvector, np.random.random(5))
+    close_to_logical(np.mean(vals) >= 0, True, tol)
+
+
 class TestElementWiseLogp(SeededTest):
     def build_model(self, distfam, params, shape, transform, testval=None):
         if testval is not None: