OrderedLogistic and OrderedProbit no longer subclass Distribution

This fixes a bug when auto-imputation takes place, in which case the logic to retrieve `p` wasn't valid. This will also show `p` as an input of the underlying Categorical in the graphviz.

Author: ricardoV94

pymc/distributions/discrete.py

@@ -1197,30 +1197,9 @@ class _OrderedLogistic(Categorical):
     See docs for the OrderedLogistic wrapper class for more details on how to use it in models.
     """
 
-    rv_op = categorical
-
-    @classmethod
-    def dist(cls, eta, cutpoints, *args, **kwargs):
-        eta = pt.as_tensor_variable(eta)
-        cutpoints = pt.as_tensor_variable(cutpoints)
-
-        pa = sigmoid(cutpoints - pt.shape_padright(eta))
-        p_cum = pt.concatenate(
-            [
-                pt.zeros_like(pt.shape_padright(pa[..., 0])),
-                pa,
-                pt.ones_like(pt.shape_padright(pa[..., 0])),
-            ],
-            axis=-1,
-        )
-        p = p_cum[..., 1:] - p_cum[..., :-1]
-
-        return super().dist(p, *args, **kwargs)
-
 
 class OrderedLogistic:
-    R"""
-    Wrapper class for Ordered Logistic distributions.
+    R"""Ordered Logistic distribution.
 
     Useful for regression on ordinal data values whose values range
     from 1 to K as a function of some predictor, :math:`\eta`. The
@@ -1287,50 +1266,39 @@ class OrderedLogistic:
         plt.hist(posterior["cutpoints"][1], 80, alpha=0.2, color='k');
     """
 
-    def __new__(cls, name, *args, compute_p=True, **kwargs):
-        out_rv = _OrderedLogistic(name, *args, **kwargs)
+    def __new__(cls, name, eta, cutpoints, compute_p=True, **kwargs):
+        p = cls.compute_p(eta, cutpoints)
         if compute_p:
-            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[3], dims=kwargs.get("dims"))
+            p = pm.Deterministic(f"{name}_probs", p, dims=kwargs.get("dims"))
+        out_rv = Categorical(name, p=p, **kwargs)
         return out_rv
 
     @classmethod
-    def dist(cls, *args, **kwargs):
-        return _OrderedLogistic.dist(*args, **kwargs)
-
-
-class _OrderedProbit(Categorical):
-    r"""
-    Underlying class for ordered probit distributions.
-    See docs for the OrderedProbit wrapper class for more details on how to use it in models.
-    """
-
-    rv_op = categorical
+    def dist(cls, eta, cutpoints, **kwargs):
+        p = cls.compute_p(eta, cutpoints)
+        return Categorical.dist(p=p, **kwargs)
 
     @classmethod
-    def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs):
+    def compute_p(cls, eta, cutpoints):
         eta = pt.as_tensor_variable(eta)
         cutpoints = pt.as_tensor_variable(cutpoints)
 
-        probits = pt.shape_padright(eta) - cutpoints
-        _log_p = pt.concatenate(
+        pa = sigmoid(cutpoints - pt.shape_padright(eta))
+        p_cum = pt.concatenate(
             [
-                pt.shape_padright(normal_lccdf(0, sigma, probits[..., 0])),
-                log_diff_normal_cdf(
-                    0, pt.shape_padright(sigma), probits[..., :-1], probits[..., 1:]
-                ),
-                pt.shape_padright(normal_lcdf(0, sigma, probits[..., -1])),
+                pt.zeros_like(pt.shape_padright(pa[..., 0])),
+                pa,
+                pt.ones_like(pt.shape_padright(pa[..., 0])),
             ],
             axis=-1,
         )
-        _log_p = pt.as_tensor_variable(_log_p)
-        p = pt.exp(_log_p)
-
-        return super().dist(p, *args, **kwargs)
+        p = p_cum[..., 1:] - p_cum[..., :-1]
+        return p
 
 
 class OrderedProbit:
     R"""
-    Wrapper class for Ordered Probit distributions.
+    Ordered Probit distributions.
 
     Useful for regression on ordinal data values whose values range
     from 1 to K as a function of some predictor, :math:`\eta`. The
@@ -1402,12 +1370,33 @@ class OrderedProbit:
         plt.hist(posterior["cutpoints"][1], 80, alpha=0.2, color='k');
     """
 
-    def __new__(cls, name, *args, compute_p=True, **kwargs):
-        out_rv = _OrderedProbit(name, *args, **kwargs)
+    def __new__(cls, name, eta, cutpoints, sigma=1, compute_p=True, **kwargs):
+        p = cls.compute_p(eta, cutpoints, sigma)
         if compute_p:
-            pm.Deterministic(f"{name}_probs", out_rv.owner.inputs[3], dims=kwargs.get("dims"))
+            p = pm.Deterministic(f"{name}_probs", p, dims=kwargs.get("dims"))
+        out_rv = Categorical(name, p=p, **kwargs)
         return out_rv
 
     @classmethod
-    def dist(cls, *args, **kwargs):
-        return _OrderedProbit.dist(*args, **kwargs)
+    def dist(cls, eta, cutpoints, sigma=1, **kwargs):
+        p = cls.compute_p(eta, cutpoints, sigma)
+        return Categorical.dist(p=p, **kwargs)
+
+    @classmethod
+    def compute_p(cls, eta, cutpoints, sigma):
+        eta = pt.as_tensor_variable(eta)
+        cutpoints = pt.as_tensor_variable(cutpoints)
+
+        probits = pt.shape_padright(eta) - cutpoints
+        log_p = pt.concatenate(
+            [
+                pt.shape_padright(normal_lccdf(0, sigma, probits[..., 0])),
+                log_diff_normal_cdf(
+                    0, pt.shape_padright(sigma), probits[..., :-1], probits[..., 1:]
+                ),
+                pt.shape_padright(normal_lcdf(0, sigma, probits[..., -1])),
+            ],
+            axis=-1,
+        )
+        p = pt.exp(log_p)
+        return p

tests/distributions/test_discrete.py

@@ -29,7 +29,7 @@
 
 import pymc as pm
 
-from pymc.distributions.discrete import _OrderedLogistic, _OrderedProbit
+from pymc.distributions.discrete import OrderedLogistic, OrderedProbit
 from pymc.logprob.basic import icdf, logcdf, logp
 from pymc.logprob.utils import ParameterValueError
 from pymc.pytensorf import floatX
@@ -481,26 +481,6 @@ def test_orderedlogistic_dimensions(shape):
     assert np.allclose(ologp, expected)
 
 
-def test_ordered_logistic_probs():
-    with pm.Model() as m:
-        pm.OrderedLogistic("ol_p", cutpoints=np.array([-2, 0, 2]), eta=0)
-        pm.OrderedLogistic("ol_no_p", cutpoints=np.array([-2, 0, 2]), eta=0, compute_p=False)
-    assert len(m.deterministics) == 1
-
-    x = pm.OrderedLogistic.dist(cutpoints=np.array([-2, 0, 2]), eta=0)
-    assert isinstance(x, TensorVariable)
-
-
-def test_ordered_probit_probs():
-    with pm.Model() as m:
-        pm.OrderedProbit("op_p", cutpoints=np.array([-2, 0, 2]), eta=0, sigma=1)
-        pm.OrderedProbit("op_no_p", cutpoints=np.array([-2, 0, 2]), eta=0, sigma=1, compute_p=False)
-    assert len(m.deterministics) == 1
-
-    x = pm.OrderedProbit.dist(cutpoints=np.array([-2, 0, 2]), eta=0, sigma=1)
-    assert isinstance(x, TensorVariable)
-
-
 class TestMoments:
     @pytest.mark.parametrize(
         "p, size, expected",
@@ -857,14 +837,12 @@ def test_implied_degenerate_shape(self):
         assert x.eval().shape == (1,)
 
 
-class TestOrderedLogistic(BaseTestDistributionRandom):
-    pymc_dist = _OrderedLogistic
-    pymc_dist_params = {"eta": 0, "cutpoints": np.array([-2, 0, 2])}
-    expected_rv_op_params = {"p": np.array([0.11920292, 0.38079708, 0.38079708, 0.11920292])}
-    checks_to_run = [
-        "check_pymc_params_match_rv_op",
-        "check_rv_size",
-    ]
+class TestOrderedLogistic:
+    def test_expected_categorical(self):
+        categorical = OrderedLogistic.dist(eta=0, cutpoints=np.array([-2, 0, 2]))
+        p = categorical.owner.inputs[3].eval()
+        expected_p = np.array([0.11920292, 0.38079708, 0.38079708, 0.11920292])
+        np.testing.assert_allclose(p, expected_p)
 
     @pytest.mark.parametrize(
         "eta, cutpoints, expected",
@@ -881,22 +859,34 @@ def test_shape_inputs(self, eta, cutpoints, expected):
         """
         This test checks when providing different shapes for `eta` parameters.
         """
-        categorical = _OrderedLogistic.dist(
+        categorical = OrderedLogistic.dist(
             eta=eta,
             cutpoints=cutpoints,
         )
-        p = categorical.owner.inputs[3].eval()
-        assert p.shape == expected
+        p_shape = tuple(categorical.owner.inputs[-1].shape.eval())
+        assert p_shape == expected
 
+    def test_compute_p(self):
+        with pm.Model() as m:
+            pm.OrderedLogistic("ol_p", cutpoints=np.array([-2, 0, 2]), eta=0)
+            pm.OrderedLogistic("ol_no_p", cutpoints=np.array([-2, 0, 2]), eta=0, compute_p=False)
+        assert len(m.deterministics) == 1
 
-class TestOrderedProbit(BaseTestDistributionRandom):
-    pymc_dist = _OrderedProbit
-    pymc_dist_params = {"eta": 0, "cutpoints": np.array([-2, 0, 2])}
-    expected_rv_op_params = {"p": np.array([0.02275013, 0.47724987, 0.47724987, 0.02275013])}
-    checks_to_run = [
-        "check_pymc_params_match_rv_op",
-        "check_rv_size",
-    ]
+        x = pm.OrderedLogistic.dist(cutpoints=np.array([-2, 0, 2]), eta=0)
+        assert isinstance(x, TensorVariable)
+
+        # Test it works with auto-imputation
+        with pm.Model() as m:
+            pm.OrderedLogistic("ol", cutpoints=np.array([-2, 0, 2]), eta=0, observed=[0, np.nan, 1])
+        assert len(m.deterministics) == 2  # One from the auto-imputation, the other from compute_p
+
+
+class TestOrderedProbit:
+    def test_expected_categorical(self):
+        categorical = OrderedProbit.dist(eta=0, cutpoints=np.array([-2, 0, 2]))
+        p = categorical.owner.inputs[3].eval()
+        expected_p = np.array([0.02275013, 0.47724987, 0.47724987, 0.02275013])
+        np.testing.assert_allclose(p, expected_p)
 
     @pytest.mark.parametrize(
         "eta, cutpoints, sigma, expected",
@@ -914,10 +904,28 @@ def test_shape_inputs(self, eta, cutpoints, sigma, expected):
         """
         This test checks when providing different shapes for `eta` and `sigma` parameters.
         """
-        categorical = _OrderedProbit.dist(
+        categorical = OrderedProbit.dist(
             eta=eta,
             cutpoints=cutpoints,
             sigma=sigma,
         )
-        p = categorical.owner.inputs[3].eval()
-        assert p.shape == expected
+        p_shape = tuple(categorical.owner.inputs[-1].shape.eval())
+        assert p_shape == expected
+
+    def test_compute_p(self):
+        with pm.Model() as m:
+            pm.OrderedProbit("op_p", cutpoints=np.array([-2, 0, 2]), eta=0, sigma=1)
+            pm.OrderedProbit(
+                "op_no_p", cutpoints=np.array([-2, 0, 2]), eta=0, sigma=1, compute_p=False
+            )
+        assert len(m.deterministics) == 1
+
+        x = pm.OrderedProbit.dist(cutpoints=np.array([-2, 0, 2]), eta=0, sigma=1)
+        assert isinstance(x, TensorVariable)
+
+        # Test it works with auto-imputation
+        with pm.Model() as m:
+            pm.OrderedProbit(
+                "op", cutpoints=np.array([-2, 0, 2]), eta=0, sigma=1, observed=[0, np.nan, 1]
+            )
+        assert len(m.deterministics) == 2  # One from the auto-imputation, the other from compute_p