Refactor DiscreteUniform

Author: ricardoV94

pymc3/distributions/discrete.py

@@ -970,6 +970,21 @@ def logcdf(value, good, bad, n):
         )
 
 
+class DiscreteUniformRV(RandomVariable):
+    name = "discrete_uniform"
+    ndim_supp = 0
+    ndims_params = [0, 0]
+    dtype = "int64"
+    _print_name = ("DiscreteUniform", "\\operatorname{DiscreteUniform}")
+
+    @classmethod
+    def rng_fn(cls, rng, lower, upper, size=None):
+        return stats.randint.rvs(lower, upper + 1, size=size, random_state=rng)
+
+
+discrete_uniform = DiscreteUniformRV()
+
+
 class DiscreteUniform(Discrete):
     R"""
     Discrete uniform distribution.
@@ -1010,39 +1025,15 @@ class DiscreteUniform(Discrete):
         Upper limit (upper > lower).
     """
 
-    def __init__(self, lower, upper, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.lower = intX(at.floor(lower))
-        self.upper = intX(at.floor(upper))
-        self.mode = at.maximum(intX(at.floor((upper + lower) / 2.0)), self.lower)
-
-    def _random(self, lower, upper, size=None):
-        # This way seems to be the only to deal with lower and upper
-        # as array-like.
-        samples = stats.randint.rvs(lower, upper + 1, size=size)
-        return samples
-
-    def random(self, point=None, size=None):
-        r"""
-        Draw random values from DiscreteUniform distribution.
-
-        Parameters
-        ----------
-        point: dict, optional
-            Dict of variable values on which random values are to be
-            conditioned (uses default point if not specified).
-        size: int, optional
-            Desired size of random sample (returns one sample if not
-            specified).
+    rv_op = discrete_uniform
 
-        Returns
-        -------
-        array
-        """
-        # lower, upper = draw_values([self.lower, self.upper], point=point, size=size)
-        # return generate_samples(self._random, lower, upper, dist_shape=self.shape, size=size)
+    @classmethod
+    def dist(cls, lower, upper, *args, **kwargs):
+        lower = intX(at.floor(lower))
+        upper = intX(at.floor(upper))
+        return super().dist([lower, upper], **kwargs)
 
-    def logp(self, value):
+    def logp(value, lower, upper):
         r"""
         Calculate log-probability of DiscreteUniform distribution at specified value.
 
@@ -1056,15 +1047,13 @@ def logp(self, value):
         -------
         TensorVariable
         """
-        upper = self.upper
-        lower = self.lower
         return bound(
             at.fill(value, -at.log(upper - lower + 1)),
             lower <= value,
             value <= upper,
         )
 
-    def logcdf(self, value):
+    def logcdf(value, lower, upper):
         """
         Compute the log of the cumulative distribution function for Discrete uniform distribution
         at the specified value.
@@ -1079,8 +1068,6 @@ def logcdf(self, value):
         -------
         TensorVariable
         """
-        upper = self.upper
-        lower = self.lower
 
         return bound(
             at.switch(

pymc3/tests/test_distributions.py

@@ -929,7 +929,6 @@ def test_bound_normal(self):
             x = PositiveNormal("x", mu=0, sigma=1, transform=None)
         assert np.isinf(logpt(x, -1).eval())
 
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_discrete_unif(self):
         self.check_logp(
             DiscreteUniform,
@@ -2817,17 +2816,16 @@ def test_issue_3051(self, dims, dist_cls, kwargs):
         assert isinstance(actual_a, np.ndarray)
         assert actual_a.shape == (X.shape[0],)
 
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_issue_4499(self):
         # Test for bug in Uniform and DiscreteUniform logp when setting check_bounds = False
         # https://github.com/pymc-devs/pymc3/issues/4499
         with pm.Model(check_bounds=False) as m:
             x = pm.Uniform("x", 0, 2, shape=10, transform=None)
-        assert_almost_equal(m.logp_array(np.ones(10)), -np.log(2) * 10)
+        assert_almost_equal(m.logp({"x": np.ones(10)}), -np.log(2) * 10)
 
         with pm.Model(check_bounds=False) as m:
-            x = pm.DiscreteUniform("x", 0, 1, shape=10)
-        assert_almost_equal(m.logp_array(np.ones(10)), -np.log(2) * 10)
+            x = pm.DiscreteUniform("x", 0, 1, size=10)
+        assert_almost_equal(m.logp({"x": np.ones(10)}), -np.log(2) * 10)
 
 
 @pytest.mark.xfail(reason="DensityDist no longer supported")

pymc3/tests/test_distributions_random.py

@@ -42,7 +42,6 @@
     Domain,
     I,
     Nat,
-    NatSmall,
     PdMatrix,
     PdMatrixChol,
     R,
@@ -339,12 +338,6 @@ class TestZeroInflatedBinomial(BaseTestCases.BaseTestCase):
     params = {"n": 10, "p": 0.6, "psi": 0.3}
 
 
-@pytest.mark.xfail(reason="This distribution has not been refactored for v4")
-class TestDiscreteUniform(BaseTestCases.BaseTestCase):
-    distribution = pm.DiscreteUniform
-    params = {"lower": 0.0, "upper": 10.0}
-
-
 @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
 class TestMoyal(BaseTestCases.BaseTestCase):
     distribution = pm.Moyal
@@ -907,6 +900,24 @@ class TestBetaBinomial(BaseTestDistribution):
     ]
 
 
+class TestDiscreteUniform(BaseTestDistribution):
+    def discrete_uniform_rng_fn(self, size, lower, upper, rng):
+        return st.randint.rvs(lower, upper + 1, size=size, random_state=rng)
+
+    pymc_dist = pm.DiscreteUniform
+    pymc_dist_params = {"lower": -1, "upper": 9}
+    expected_rv_op_params = {"lower": -1, "upper": 9}
+    reference_dist_params = {"lower": -1, "upper": 9}
+    reference_dist = lambda self: functools.partial(
+        self.discrete_uniform_rng_fn, rng=self.get_random_state()
+    )
+    tests_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_pymc_draws_match_reference",
+        "check_rv_size",
+    ]
+
+
 class TestScalarParameterSamples(SeededTest):
     @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
     def test_bounded(self):
@@ -1016,15 +1027,6 @@ def test_half_flat(self):
             with pytest.raises(ValueError):
                 f.random(1)
 
-    @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
-    def test_discrete_uniform(self):
-        def ref_rand(size, lower, upper):
-            return st.randint.rvs(lower, upper + 1, size=size)
-
-        pymc3_random_discrete(
-            pm.DiscreteUniform, {"lower": -NatSmall, "upper": NatSmall}, ref_rand=ref_rand
-        )
-
     @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
     def test_constant_dist(self):
         def ref_rand(size, c):