Refactor OrderedLogistic

Author: ricardoV94

pymc3/distributions/discrete.py

@@ -1717,11 +1717,14 @@ class OrderedLogistic(Categorical):
 
     """
 
-    def __init__(self, eta, cutpoints, *args, **kwargs):
-        self.eta = at.as_tensor_variable(floatX(eta))
-        self.cutpoints = at.as_tensor_variable(cutpoints)
+    rv_op = categorical
 
-        pa = sigmoid(self.cutpoints - at.shape_padright(self.eta))
+    @classmethod
+    def dist(cls, eta, cutpoints, *args, **kwargs):
+        eta = at.as_tensor_variable(floatX(eta))
+        cutpoints = at.as_tensor_variable(cutpoints)
+
+        pa = sigmoid(cutpoints - at.shape_padright(eta))
         p_cum = at.concatenate(
             [
                 at.zeros_like(at.shape_padright(pa[..., 0])),
@@ -1732,7 +1735,7 @@ def __init__(self, eta, cutpoints, *args, **kwargs):
         )
         p = p_cum[..., 1:] - p_cum[..., :-1]
 
-        super().__init__(p=p, *args, **kwargs)
+        return super().dist(p, **kwargs)
 
 
 class OrderedProbit(Categorical):

pymc3/tests/test_distributions.py

@@ -101,6 +101,7 @@
     continuous,
     logcdf,
     logpt,
+    logpt_sum,
 )
 from pymc3.math import kronecker, logsumexp
 from pymc3.model import Deterministic, Model, Point
@@ -2297,7 +2298,6 @@ def test_categorical(self, n):
         )
 
     @pytest.mark.parametrize("n", [2, 3, 4])
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_orderedlogistic(self, n):
         self.check_logp(
             OrderedLogistic,
@@ -2738,27 +2738,34 @@ def test_discrete_trafo():
         err.match("Transformations for discrete distributions")
 
 
+# TODO: Is this test working as expected / still relevant?
 @pytest.mark.parametrize("shape", [tuple(), (1,), (3, 1), (3, 2)], ids=str)
-@pytest.mark.xfail(reason="Distribution not refactored yet")
 def test_orderedlogistic_dimensions(shape):
     # Test for issue #3535
     loge = np.log10(np.exp(1))
     size = 7
     p = np.ones(shape + (10,)) / 10
     cutpoints = np.tile(logit(np.linspace(0, 1, 11)[1:-1]), shape + (1,))
-    obs = np.random.randint(0, 1, size=(size,) + shape)
+    obs = np.random.randint(0, 2, size=(size,) + shape)
     with Model():
         ol = OrderedLogistic(
-            "ol", eta=np.zeros(shape), cutpoints=cutpoints, size=shape, observed=obs
-        )
-        c = Categorical("c", p=p, size=shape, observed=obs)
-    ologp = logpt(ol, 1).eval() * loge
-    clogp = logpt(c, 1) * loge
+            "ol",
+            eta=np.zeros(shape),
+            cutpoints=cutpoints,
+            observed=obs,
+        )
+        c = Categorical(
+            "c",
+            p=p,
+            observed=obs,
+        )
+    ologp = logpt_sum(ol, np.ones_like(obs)).eval() * loge
+    clogp = logpt_sum(c, np.ones_like(obs)).eval() * loge
     expected = -np.prod((size,) + shape)
 
-    assert c.distribution.p.ndim == (len(shape) + 1)
+    assert c.owner.inputs[3].ndim == (len(shape) + 1)
     assert np.allclose(clogp, expected)
-    assert ol.distribution.p.ndim == (len(shape) + 1)
+    assert ol.owner.inputs[3].ndim == (len(shape) + 1)
     assert np.allclose(ologp, expected)
 
 

pymc3/tests/test_distributions_random.py

@@ -929,6 +929,16 @@ def constant_rng_fn(self, size, c):
     ]
 
 
+class TestOrderedLogistic(BaseTestDistribution):
+    pymc_dist = pm.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])}
+    tests_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_rv_size",
+    ]
+
+
 class TestScalarParameterSamples(SeededTest):
     @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
     def test_bounded(self):