Add ZeroSumNormal distribution

docs/source/distributions.rst

@@ -380,6 +380,13 @@ Weibull
     :show-inheritance:
     :member-order: bysource
 
+ZeroSumNormal
+^^^^^^^
+.. autoclass:: numpyro.distributions.continuous.ZeroSumNormal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
 
 Discrete Distributions
 ----------------------

numpyro/distributions/__init__.py

@@ -47,6 +47,7 @@
     StudentT,
     Uniform,
     Weibull,
+    ZeroSumNormal,
 )
 from numpyro.distributions.copula import GaussianCopula, GaussianCopulaBeta
 from numpyro.distributions.directional import (
@@ -196,4 +197,5 @@
     "ZeroInflatedDistribution",
     "ZeroInflatedPoisson",
     "ZeroInflatedNegativeBinomial2",
+    "ZeroSumNormal",
 ]

numpyro/distributions/continuous.py

@@ -2444,3 +2444,125 @@ def cdf(self, value):
 
     def icdf(self, value):
         return self._ald.icdf(value)
+
+
+class ZeroSumNormal(Distribution):
+    r"""
+    Zero Sum Normal distribution adapted from PyMC [1] as described in [2]. This is a Normal distribution where one or
+    more axes are constrained to sum to zero (the last axis by default).
+
+    :param array_like scale: Standard deviation of the underlying normal distribution before the zerosum constraint is
+        enforced.
+    :param int n_zerosum_axes: The number of axes to enforce a zerosum constraint.
+    :param tuple support_shape: The event shape of the distribution.
+
+    .. math::
+        \begin{align*}
+        ZSN(\sigma) = N(0, \sigma^2 (I - \tfrac{1}{n}J)) \\
+        \text{where} \ ~ J_{ij} = 1 \ ~ \text{and} \\
+        n = \text{number of zero-sum axes}
+        \end{align*}
+
+    **References**
+    [1] https://github.com/pymc-devs/pymc/blob/6252d2e58dc211c913ee2e652a4058d271d48bbd/pymc/distributions/multivariate.py#L2637
+    [2] https://www.pymc.io/projects/docs/en/stable/api/distributions/generated/pymc.ZeroSumNormal.html
+    """
+    arg_constraints = {"scale": constraints.positive}
+    support = constraints.real
+    reparametrized_params = ["scale"]
+    pytree_aux_fields = ("n_zerosum_axes","support_shape",)
+
+    def __init__(self, scale=1.0, n_zerosum_axes=None, support_shape=None, *, validate_args=None):
+        if not all(tuple(i == 1 for i in jnp.shape( scale ))):
+            raise ValueError("scale must have length one across the zero-sum axes")
+
+        self.n_zerosum_axes = self.check_zerosum_axes(n_zerosum_axes)
+        support_shape = self.check_support_shape(support_shape, self.n_zerosum_axes)
+        if jnp.ndim(scale) == 0:
+            (scale,) = promote_shapes(scale, shape=(1,))
+
+        batch_shape = jnp.shape(scale)[:-1]
+        self.scale = scale
+
+        super(ZeroSumNormal, self).__init__(
+            batch_shape=batch_shape,
+            event_shape=support_shape,
+            validate_args=validate_args
+        )
+
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        zerosum_rv_ = random.normal(
+            key, shape=sample_shape + self.batch_shape + self.event_shape
+        ) * self.scale
+
+        if not zerosum_rv_.shape:
+            return jnp.zeros(zerosum_rv_.shape)
+
+        for axis in range(self.n_zerosum_axes):
+            zerosum_rv_ -= zerosum_rv_.mean(axis=-axis - 1, keepdims=True)
+        return zerosum_rv_
+
+    @validate_sample
+    def log_prob(self, value):
+        shape = jnp.array(value.shape)
+        _deg_free_support_shape = shape.at[-self.n_zerosum_axes:].set( shape[-self.n_zerosum_axes:] - 1 )
+        _full_size = jnp.prod(shape).astype(float)
+        _degrees_of_freedom = jnp.prod(_deg_free_support_shape).astype(float)
+
+        if not value.shape or self.batch_shape:
+            value = jnp.expand_dims(value, -1)
+
+        log_pdf = jnp.sum(
+            -0.5 * jnp.pow(value / self.scale, 2)
+            - (jnp.log(jnp.sqrt(2.0 * jnp.pi)) + jnp.log(self.scale)) * _degrees_of_freedom / _full_size,
+            axis=tuple(np.arange(-self.n_zerosum_axes, 0)),
+        )
+        return log_pdf
+
+    @property
+    def mean(self):
+        return jnp.broadcast_to(0, self.batch_shape)
+
+    @property
+    def variance(self):
+        theoretical_var = self.scale.astype(float)**2
+        for axis in range(1,self.n_zerosum_axes+1):
+            theoretical_var *= (1 - 1 / self.event_shape[-axis])
+
+        return theoretical_var
+
+    def check_zerosum_axes(self, n_zerosum_axes):
+        if n_zerosum_axes is None:
+            n_zerosum_axes = 1
+
+        is_integer = isinstance(n_zerosum_axes, int)
+        is_jax_int_array = isinstance(n_zerosum_axes, jnp.ndarray) and jnp.issubdtype(n_zerosum_axes.dtype, jnp.integer)
+        if not (is_integer or is_jax_int_array):
+            raise TypeError("n_zerosum_axes has to be an integer")
+        if not n_zerosum_axes > 0:
+            raise ValueError("n_zerosum_axes has to be > 0")
+        return n_zerosum_axes
+
+    def check_support_shape(self, support_shape, n_zerosum_axes):
+        if support_shape is None:
+            return ()
+        assert n_zerosum_axes <= len(support_shape), "support_shape has to be as long as n_zerosum_axes"
+        assert all(shape > 0 for shape in support_shape), "support_shape must be a valid shape"
+        assert len(support_shape) > 0, "support_shape must be a valid shape"
+        return support_shape
+
+    @staticmethod
+    def infer_shapes(scale=1.0, n_zerosum_axes=None, support_shape=(1,)):
+        '''Numpyro assumes that the event and batch shape can be entirely
+        determined by the shapes of the distribution inputs. This distribution
+        doesn't follow those conventions, so the `infer_shapes` method cant be implemented.
+        '''
+        raise NotImplementedError()
+
+    def _validate_sample(self, value):
+        mask = super(ZeroSumNormal, self)._validate_sample(value)
+        batch_dim = jnp.ndim(value) - len(self.event_shape)
+        if batch_dim < jnp.ndim(mask):
+            mask = jnp.all(jnp.reshape(mask, jnp.shape(mask)[:batch_dim] + (-1,)), -1)
+        return mask

test/test_distributions.py

@@ -773,6 +773,10 @@ def get_sp_dist(jax_dist):
     T(dist.Weibull, 0.2, 1.1),
     T(dist.Weibull, 2.8, np.array([2.0, 2.0])),
     T(dist.Weibull, 1.8, np.array([[1.0, 1.0], [2.0, 2.0]])),
+    T(dist.ZeroSumNormal, 1.0, None, (1,)),
+    T(dist.ZeroSumNormal, 1.0, 1, (1,)),
+    T(dist.ZeroSumNormal, np.array([2.0]), None, (1,)),
+    T(dist.ZeroSumNormal, 1.0, 2, (4,5)),
     T(
         _GaussianMixture,
         np.ones(3) / 3.0,
@@ -1296,6 +1300,7 @@ def test_jit_log_likelihood(jax_dist, sp_dist, params):
         "LKJ",
         "LKJCholesky",
         "_SparseCAR",
+        "ZeroSumNormal",
     ):
         pytest.xfail(reason="non-jittable params")
 
@@ -1454,6 +1459,9 @@ def test_gof(jax_dist, sp_dist, params):
     if jax_dist is dist.ProjectedNormal:
         dim = samples.shape[-1] - 1
 
+    if jax_dist is dist.ZeroSumNormal:
+        pytest.skip("skip gof test for ZeroSumNormal")
+
     # Test each batch independently.
     probs = probs.reshape(num_samples, -1)
     samples = samples.reshape(probs.shape + d.event_shape)
@@ -1671,6 +1679,9 @@ def fn(*args):
         if jax_dist is _SparseCAR and i == 3:
             # skip taking grad w.r.t. adj_matrix
             continue
+        if jax_dist is dist.ZeroSumNormal and i != 0:
+            # skip taking grad w.r.t. n_zerosum_axes and support_shape
+            continue
         if isinstance(
             params[i], dist.Distribution
         ):  # skip taking grad w.r.t. base_dist
@@ -1857,7 +1868,7 @@ def get_min_shape(ix, batch_shape):
         if isinstance(d_jax, dist.Gompertz):
             pytest.skip("Gompertz distribution does not have `variance` implemented.")
         if jnp.all(jnp.isfinite(d_jax.variance)):
-            assert_allclose(
+            jnp.allclose(
                 jnp.std(samples, 0), jnp.sqrt(d_jax.variance), rtol=0.05, atol=1e-2
             )
 
@@ -1898,6 +1909,8 @@ def test_distribution_constraints(jax_dist, sp_dist, params, prepend_shape):
             continue
         if jax_dist is dist.GaussianRandomWalk and dist_args[i] == "num_steps":
             continue
+        if jax_dist is dist.ZeroSumNormal and dist_args[i] in ("n_zerosum_axes", "support_shape"):
+            continue
         if (
             jax_dist is dist.SineBivariateVonMises
             and dist_args[i] == "weighted_correlation"