Move ZSN tests to test_multivariate.py

Author: AlexAndorra

pymc/tests/distributions/test_continuous.py

@@ -24,15 +24,13 @@
 
 from aeppl.logprob import ParameterValueError
 from aesara.compile.mode import Mode
-from numpy import AxisError
 
 import pymc as pm
 
 from pymc.aesaraf import floatX
 from pymc.distributions import logcdf, logp
 from pymc.distributions.continuous import get_tau_sigma, interpolated
 from pymc.distributions.dist_math import clipped_beta_rvs
-from pymc.distributions.shape_utils import change_dist_size
 from pymc.tests.distributions.util import (
     BaseTestDistributionRandom,
     Circ,
@@ -965,19 +963,6 @@ def test_normal_moment(self, mu, sigma, size, expected):
             pm.Normal("x", mu=mu, sigma=sigma, size=size)
         assert_moment_is_expected(model, expected)
 
-    @pytest.mark.parametrize(
-        "shape, zerosum_axes, expected",
-        [
-            ((2, 5), None, np.zeros((2, 5))),
-            ((2, 5, 6), None, np.zeros((2, 5, 6))),
-            ((2, 5, 6), (0, 1), np.zeros((2, 5, 6))),
-        ],
-    )
-    def test_zerosum_normal_moment(self, shape, zerosum_axes, expected):
-        with pm.Model() as model:
-            pm.ZeroSumNormal("x", shape=shape, zerosum_axes=zerosum_axes)
-        assert_moment_is_expected(model, expected)
-
     @pytest.mark.parametrize(
         "sigma, size, expected",
         [
@@ -1817,100 +1802,6 @@ class TestTruncatedNormalUpperArray(BaseTestDistributionRandom):
     ]
 
 
-COORDS = {
-    "regions": ["a", "b", "c"],
-    "answers": ["yes", "no", "whatever", "don't understand question"],
-}
-
-
-class TestZeroSumNormal:
-    @pytest.mark.parametrize(
-        "dims, zerosum_axes, shape",
-        [
-            (("regions", "answers"), "answers", None),
-            (("regions", "answers"), ("regions", "answers"), None),
-            (("regions", "answers"), 0, None),
-            (("regions", "answers"), -1, None),
-            (("regions", "answers"), (0, 1), None),
-            (None, -2, (len(COORDS["regions"]), len(COORDS["answers"]))),
-        ],
-    )
-    def test_zsn_dims_shape(self, dims, zerosum_axes, shape):
-        with pm.Model(coords=COORDS) as m:
-            v = pm.ZeroSumNormal("v", dims=dims, shape=shape, zerosum_axes=zerosum_axes)
-            s = pm.sample(10, chains=1, tune=100)
-
-        # to test forward graph
-        random_samples = pm.draw(
-            v,
-            draws=10,
-        )
-
-        assert s.posterior.v.shape == (1, 10, len(COORDS["regions"]), len(COORDS["answers"]))
-
-        if not isinstance(zerosum_axes, (list, tuple)):
-            zerosum_axes = [zerosum_axes]
-
-        if isinstance(zerosum_axes[0], str):
-            for ax in zerosum_axes:
-                for samples in [
-                    s.posterior.v.mean(dim=ax),
-                    random_samples.mean(axis=dims.index(ax) + 1),
-                ]:
-                    assert np.isclose(
-                        samples, 0
-                    ).all(), f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
-
-            nonzero_axes = list(set(dims).difference(zerosum_axes))
-            if nonzero_axes:
-                for ax in nonzero_axes:
-                    for samples in [
-                        s.posterior.v.mean(dim=ax),
-                        random_samples.mean(axis=dims.index(ax) + 1),
-                    ]:
-                        assert not np.isclose(
-                            samples, 0
-                        ).all(), f"{ax} is not a zerosum_axis, but is nonetheless summing to 0 across all samples."
-
-        else:
-            for ax in zerosum_axes:
-                if ax < 0:
-                    assert np.isclose(
-                        s.posterior.v.mean(axis=ax), 0
-                    ).all(), f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
-                else:
-                    ax = ax + 2  # because 'chain' and 'draw' are added as new axes after sampling
-                    assert np.isclose(
-                        s.posterior.v.mean(axis=ax), 0
-                    ).all(), f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
-
-    @pytest.mark.parametrize(
-        "dims, zerosum_axes",
-        [
-            (("regions", "answers"), 2),
-            (("regions", "answers"), (0, -2)),
-        ],
-    )
-    def test_zsn_fail_axis(self, dims, zerosum_axes):
-        if isinstance(zerosum_axes, (list, tuple)):
-            with pytest.raises(ValueError, match="repeated axis"):
-                with pm.Model(coords=COORDS) as m:
-                    _ = pm.ZeroSumNormal("v", dims=dims, zerosum_axes=zerosum_axes)
-        else:
-            with pytest.raises(AxisError, match="out of bounds"):
-                with pm.Model(coords=COORDS) as m:
-                    _ = pm.ZeroSumNormal("v", dims=dims, zerosum_axes=zerosum_axes)
-
-    def test_zsn_change_dist_size(self):
-        base_dist = pm.ZeroSumNormal.dist(shape=(4, 9))
-
-        new_dist = change_dist_size(base_dist, new_size=(5, 3), expand=False)
-        assert new_dist.eval().shape == (5, 3)
-
-        new_dist = change_dist_size(base_dist, new_size=(5, 3), expand=True)
-        assert new_dist.eval().shape == (5, 3, 4, 9)
-
-
 class TestWald(BaseTestDistributionRandom):
     pymc_dist = pm.Wald
     mu, lam, alpha = 1.0, 1.0, 0.0

pymc/tests/distributions/test_multivariate.py

@@ -28,6 +28,7 @@
 from aeppl.logprob import ParameterValueError
 from aesara.tensor import TensorVariable
 from aesara.tensor.random.utils import broadcast_params
+from numpy import AxisError
 
 import pymc as pm
 
@@ -754,7 +755,12 @@ def test_car_logp(self, sparse, size):
 
         # d x d adjacency matrix for a square (d=4) of rook-adjacent sites
         W = np.array(
-            [[0.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0]]
+            [
+                [0.0, 1.0, 1.0, 0.0],
+                [1.0, 0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0, 1.0],
+                [0.0, 1.0, 1.0, 0.0],
+            ]
         )
 
         tau = 2
@@ -1007,6 +1013,19 @@ def test_mv_normal_moment(self, mu, cov, size, expected):
         # MvNormal logp is only implemented for up to 2D variables
         assert_moment_is_expected(model, expected, check_finite_logp=x.ndim < 3)
 
+    @pytest.mark.parametrize(
+        "shape, zerosum_axes, expected",
+        [
+            ((2, 5), None, np.zeros((2, 5))),
+            ((2, 5, 6), None, np.zeros((2, 5, 6))),
+            ((2, 5, 6), (0, 1), np.zeros((2, 5, 6))),
+        ],
+    )
+    def test_zerosum_normal_moment(self, shape, zerosum_axes, expected):
+        with pm.Model() as model:
+            pm.ZeroSumNormal("x", shape=shape, zerosum_axes=zerosum_axes)
+        assert_moment_is_expected(model, expected)
+
     @pytest.mark.parametrize(
         "mu, size, expected",
         [
@@ -1026,7 +1045,12 @@ def test_mv_normal_moment(self, mu, cov, size, expected):
     )
     def test_car_moment(self, mu, size, expected):
         W = np.array(
-            [[0.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0]]
+            [
+                [0.0, 1.0, 1.0, 0.0],
+                [1.0, 0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0, 1.0],
+                [0.0, 1.0, 1.0, 0.0],
+            ]
         )
         tau = 2
         alpha = 0.5
@@ -1367,6 +1391,100 @@ def test_issue_3706(self):
         assert prior_pred["X"].shape == (1, N, 2)
 
 
+COORDS = {
+    "regions": ["a", "b", "c"],
+    "answers": ["yes", "no", "whatever", "don't understand question"],
+}
+
+
+class TestZeroSumNormal:
+    @pytest.mark.parametrize(
+        "dims, zerosum_axes, shape",
+        [
+            (("regions", "answers"), "answers", None),
+            (("regions", "answers"), ("regions", "answers"), None),
+            (("regions", "answers"), 0, None),
+            (("regions", "answers"), -1, None),
+            (("regions", "answers"), (0, 1), None),
+            (None, -2, (len(COORDS["regions"]), len(COORDS["answers"]))),
+        ],
+    )
+    def test_zsn_dims_shape(self, dims, zerosum_axes, shape):
+        with pm.Model(coords=COORDS) as m:
+            v = pm.ZeroSumNormal("v", dims=dims, shape=shape, zerosum_axes=zerosum_axes)
+            s = pm.sample(10, chains=1, tune=100)
+
+        # to test forward graph
+        random_samples = pm.draw(
+            v,
+            draws=10,
+        )
+
+        assert s.posterior.v.shape == (1, 10, len(COORDS["regions"]), len(COORDS["answers"]))
+
+        if not isinstance(zerosum_axes, (list, tuple)):
+            zerosum_axes = [zerosum_axes]
+
+        if isinstance(zerosum_axes[0], str):
+            for ax in zerosum_axes:
+                for samples in [
+                    s.posterior.v.mean(dim=ax),
+                    random_samples.mean(axis=dims.index(ax) + 1),
+                ]:
+                    assert np.isclose(
+                        samples, 0
+                    ).all(), f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
+
+            nonzero_axes = list(set(dims).difference(zerosum_axes))
+            if nonzero_axes:
+                for ax in nonzero_axes:
+                    for samples in [
+                        s.posterior.v.mean(dim=ax),
+                        random_samples.mean(axis=dims.index(ax) + 1),
+                    ]:
+                        assert not np.isclose(
+                            samples, 0
+                        ).all(), f"{ax} is not a zerosum_axis, but is nonetheless summing to 0 across all samples."
+
+        else:
+            for ax in zerosum_axes:
+                if ax < 0:
+                    assert np.isclose(
+                        s.posterior.v.mean(axis=ax), 0
+                    ).all(), f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
+                else:
+                    ax = ax + 2  # because 'chain' and 'draw' are added as new axes after sampling
+                    assert np.isclose(
+                        s.posterior.v.mean(axis=ax), 0
+                    ).all(), f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
+
+    @pytest.mark.parametrize(
+        "dims, zerosum_axes",
+        [
+            (("regions", "answers"), 2),
+            (("regions", "answers"), (0, -2)),
+        ],
+    )
+    def test_zsn_fail_axis(self, dims, zerosum_axes):
+        if isinstance(zerosum_axes, (list, tuple)):
+            with pytest.raises(ValueError, match="repeated axis"):
+                with pm.Model(coords=COORDS) as m:
+                    _ = pm.ZeroSumNormal("v", dims=dims, zerosum_axes=zerosum_axes)
+        else:
+            with pytest.raises(AxisError, match="out of bounds"):
+                with pm.Model(coords=COORDS) as m:
+                    _ = pm.ZeroSumNormal("v", dims=dims, zerosum_axes=zerosum_axes)
+
+    def test_zsn_change_dist_size(self):
+        base_dist = pm.ZeroSumNormal.dist(shape=(4, 9))
+
+        new_dist = change_dist_size(base_dist, new_size=(5, 3), expand=False)
+        assert new_dist.eval().shape == (5, 3)
+
+        new_dist = change_dist_size(base_dist, new_size=(5, 3), expand=True)
+        assert new_dist.eval().shape == (5, 3, 4, 9)
+
+
 class TestMvStudentTCov(BaseTestDistributionRandom):
     def mvstudentt_rng_fn(self, size, nu, mu, cov, rng):
         mv_samples = rng.multivariate_normal(np.zeros_like(mu), cov, size=size)