Move stickbreaking tests

These tests were introduced between pre-existing Dirichlet/Multinomial/DirichletMultionmial tests that belong to conceptually related distributions

Author: ricardoV94

pymc/tests/test_distributions.py

@@ -2122,40 +2122,6 @@ def test_dirichlet_invalid(self):
         valid_dist = Dirichlet.dist(a=[1, 1, 1])
         assert np.all(np.isfinite(pm.logp(valid_dist, value).eval()) == np.array([True, False]))
 
-    @pytest.mark.parametrize(
-        "value,alpha,K,logp",
-        [
-            (np.array([5, 4, 3, 2, 1]) / 15, 0.5, 4, 1.5126301307277439),
-            (np.tile(1, 13) / 13, 2, 12, 13.980045245672827),
-            (np.array([0.001] * 10 + [0.99]), 0.1, 10, -22.971662448814723),
-            (np.append(0.5 ** np.arange(1, 20), 0.5**20), 5, 19, 94.20462772778092),
-            (
-                (np.array([[7, 5, 3, 2], [19, 17, 13, 11]]) / np.array([[17], [60]])),
-                2.5,
-                3,
-                np.array([1.29317672, 1.50126157]),
-            ),
-        ],
-    )
-    def test_stickbreakingweights_logp(self, value, alpha, K, logp):
-        with Model() as model:
-            sbw = StickBreakingWeights("sbw", alpha=alpha, K=K, transform=None)
-        pt = {"sbw": value}
-        assert_almost_equal(
-            pm.logp(sbw, value).eval(),
-            logp,
-            decimal=select_by_precision(float64=6, float32=2),
-            err_msg=str(pt),
-        )
-
-    def test_stickbreakingweights_invalid(self):
-        sbw = pm.StickBreakingWeights.dist(3.0, 3)
-        sbw_wrong_K = pm.StickBreakingWeights.dist(3.0, 7)
-        assert pm.logp(sbw, np.array([0.4, 0.3, 0.2, 0.15])).eval() == -np.inf
-        assert pm.logp(sbw, np.array([1.1, 0.3, 0.2, 0.1])).eval() == -np.inf
-        assert pm.logp(sbw, np.array([0.4, 0.3, 0.2, -0.1])).eval() == -np.inf
-        assert pm.logp(sbw_wrong_K, np.array([0.4, 0.3, 0.2, 0.1])).eval() == -np.inf
-
     @pytest.mark.parametrize(
         "a",
         [
@@ -2318,6 +2284,40 @@ def test_dirichlet_multinomial_vectorized(self, n, a, size):
             err_msg=f"vals={vals}",
         )
 
+    @pytest.mark.parametrize(
+        "value,alpha,K,logp",
+        [
+            (np.array([5, 4, 3, 2, 1]) / 15, 0.5, 4, 1.5126301307277439),
+            (np.tile(1, 13) / 13, 2, 12, 13.980045245672827),
+            (np.array([0.001] * 10 + [0.99]), 0.1, 10, -22.971662448814723),
+            (np.append(0.5 ** np.arange(1, 20), 0.5**20), 5, 19, 94.20462772778092),
+            (
+                (np.array([[7, 5, 3, 2], [19, 17, 13, 11]]) / np.array([[17], [60]])),
+                2.5,
+                3,
+                np.array([1.29317672, 1.50126157]),
+            ),
+        ],
+    )
+    def test_stickbreakingweights_logp(self, value, alpha, K, logp):
+        with Model() as model:
+            sbw = StickBreakingWeights("sbw", alpha=alpha, K=K, transform=None)
+        pt = {"sbw": value}
+        assert_almost_equal(
+            pm.logp(sbw, value).eval(),
+            logp,
+            decimal=select_by_precision(float64=6, float32=2),
+            err_msg=str(pt),
+        )
+
+    def test_stickbreakingweights_invalid(self):
+        sbw = pm.StickBreakingWeights.dist(3.0, 3)
+        sbw_wrong_K = pm.StickBreakingWeights.dist(3.0, 7)
+        assert pm.logp(sbw, np.array([0.4, 0.3, 0.2, 0.15])).eval() == -np.inf
+        assert pm.logp(sbw, np.array([1.1, 0.3, 0.2, 0.1])).eval() == -np.inf
+        assert pm.logp(sbw, np.array([0.4, 0.3, 0.2, -0.1])).eval() == -np.inf
+        assert pm.logp(sbw_wrong_K, np.array([0.4, 0.3, 0.2, 0.1])).eval() == -np.inf
+
     @aesara.config.change_flags(compute_test_value="raise")
     def test_categorical_bounds(self):
         with Model():

pymc/tests/test_distributions_random.py

@@ -1309,41 +1309,6 @@ class TestDirichlet(BaseTestDistributionRandom):
     ]
 
 
-class TestStickBreakingWeights(BaseTestDistributionRandom):
-    pymc_dist = pm.StickBreakingWeights
-    pymc_dist_params = {"alpha": 2.0, "K": 19}
-    expected_rv_op_params = {"alpha": 2.0, "K": 19}
-    sizes_to_check = [None, 17, (5,), (11, 5), (3, 13, 5)]
-    sizes_expected = [
-        (20,),
-        (17, 20),
-        (
-            5,
-            20,
-        ),
-        (11, 5, 20),
-        (3, 13, 5, 20),
-    ]
-    checks_to_run = [
-        "check_pymc_params_match_rv_op",
-        "check_rv_size",
-        "check_basic_properties",
-    ]
-
-    def check_basic_properties(self):
-        default_rng = aesara.shared(np.random.default_rng(1234))
-        draws = pm.StickBreakingWeights.dist(
-            alpha=3.5,
-            K=19,
-            size=(2, 3, 5),
-            rng=default_rng,
-        ).eval()
-
-        assert np.allclose(draws.sum(-1), 1)
-        assert np.all(draws >= 0)
-        assert np.all(draws <= 1)
-
-
 class TestMultinomial(BaseTestDistributionRandom):
     pymc_dist = pm.Multinomial
     pymc_dist_params = {"n": 85, "p": np.array([0.28, 0.62, 0.10])}
@@ -1400,6 +1365,41 @@ class TestDirichletMultinomial_1D_n_2D_a(BaseTestDistributionRandom):
     checks_to_run = ["check_rv_size"]
 
 
+class TestStickBreakingWeights(BaseTestDistributionRandom):
+    pymc_dist = pm.StickBreakingWeights
+    pymc_dist_params = {"alpha": 2.0, "K": 19}
+    expected_rv_op_params = {"alpha": 2.0, "K": 19}
+    sizes_to_check = [None, 17, (5,), (11, 5), (3, 13, 5)]
+    sizes_expected = [
+        (20,),
+        (17, 20),
+        (
+            5,
+            20,
+        ),
+        (11, 5, 20),
+        (3, 13, 5, 20),
+    ]
+    checks_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_rv_size",
+        "check_basic_properties",
+    ]
+
+    def check_basic_properties(self):
+        default_rng = aesara.shared(np.random.default_rng(1234))
+        draws = pm.StickBreakingWeights.dist(
+            alpha=3.5,
+            K=19,
+            size=(2, 3, 5),
+            rng=default_rng,
+        ).eval()
+
+        assert np.allclose(draws.sum(-1), 1)
+        assert np.all(draws >= 0)
+        assert np.all(draws <= 1)
+
+
 class TestCategorical(BaseTestDistributionRandom):
     pymc_dist = pm.Categorical
     pymc_dist_params = {"p": np.array([0.28, 0.62, 0.10])}