Refactor Constant

Author: ricardoV94

pymc3/distributions/__init__.py

@@ -63,7 +63,6 @@
     Binomial,
     Categorical,
     Constant,
-    ConstantDist,
     DiscreteUniform,
     DiscreteWeibull,
     Geometric,
@@ -138,7 +137,6 @@
     "Bernoulli",
     "Poisson",
     "NegativeBinomial",
-    "ConstantDist",
     "Constant",
     "ZeroInflatedPoisson",
     "ZeroInflatedNegativeBinomial",

pymc3/distributions/discrete.py

@@ -11,8 +11,6 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import warnings
-
 import aesara.tensor as at
 import numpy as np
 
@@ -51,7 +49,6 @@
     "DiscreteWeibull",
     "Poisson",
     "NegativeBinomial",
-    "ConstantDist",
     "Constant",
     "ZeroInflatedPoisson",
     "ZeroInflatedBinomial",
@@ -1164,6 +1161,23 @@ def logp(value, p):
         )
 
 
+class ConstantRV(RandomVariable):
+    name = "constant"
+    ndim_supp = 0
+    ndims_params = [0]
+    dtype = "floatX"  # Should be treated as a discrete variable!
+    _print_name = ("Constant", "\\operatorname{Constant}")
+
+    @classmethod
+    def rng_fn(cls, rng, c, size=None):
+        if size is None:
+            return c.copy()
+        return np.full(size, c)
+
+
+constant = ConstantRV()
+
+
 class Constant(Discrete):
     r"""
     Constant log-likelihood.
@@ -1174,40 +1188,14 @@ class Constant(Discrete):
         Constant parameter.
     """
 
-    def __init__(self, c, *args, **kwargs):
-        warnings.warn(
-            "Constant has been deprecated. We recommend using a Deterministic object instead.",
-            DeprecationWarning,
-        )
-        super().__init__(*args, **kwargs)
-        self.mean = self.median = self.mode = self.c = c = at.as_tensor_variable(c)
-
-    def random(self, point=None, size=None):
-        r"""
-        Draw random values from Constant 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 = constant
 
-        Returns
-        -------
-        array
-        """
-        # c = draw_values([self.c], point=point, size=size)[0]
-        # dtype = np.array(c).dtype
-        #
-        # def _random(c, dtype=dtype, size=None):
-        #     return np.full(size, fill_value=c, dtype=dtype)
-        #
-        # return generate_samples(_random, c=c, dist_shape=self.shape, size=size).astype(dtype)
+    @classmethod
+    def dist(cls, c, *args, **kwargs):
+        c = at.as_tensor_variable(floatX(c))
+        return super().dist([c], **kwargs)
 
-    def logp(self, value):
+    def logp(value, c):
         r"""
         Calculate log-probability of Constant distribution at specified value.
 
@@ -1221,11 +1209,10 @@ def logp(self, value):
         -------
         TensorVariable
         """
-        c = self.c
-        return bound(0, at.eq(value, c))
-
-
-ConstantDist = Constant
+        return bound(
+            at.zeros_like(value),
+            at.eq(value, c),
+        )
 
 
 class ZeroInflatedPoisson(Discrete):

pymc3/tests/test_distributions.py

@@ -1614,7 +1614,6 @@ def test_bound_poisson(self):
             x = NonZeroPoisson("x", mu=4)
         assert np.isinf(logpt(x, 0).eval())
 
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_constantdist(self):
         self.check_logp(Constant, I, {"c": I}, lambda value, c: np.log(c == value))
 
@@ -2820,13 +2819,17 @@ 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)
+            x = pm.Uniform("x", 0, 2, size=10, transform=None)
         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, size=10)
         assert_almost_equal(m.logp({"x": np.ones(10)}), -np.log(2) * 10)
 
+        with pm.Model(check_bounds=False) as m:
+            x = pm.Constant("x", 1, size=10)
+        assert_almost_equal(m.logp({"x": np.ones(10)}), 0 * 10)
+
 
 @pytest.mark.xfail(reason="DensityDist no longer supported")
 def test_serialize_density_dist():

pymc3/tests/test_distributions_random.py

@@ -40,7 +40,6 @@
 from pymc3.tests.helpers import SeededTest, select_by_precision
 from pymc3.tests.test_distributions import (
     Domain,
-    I,
     Nat,
     PdMatrix,
     PdMatrixChol,
@@ -314,12 +313,6 @@ class TestLogitNormal(BaseTestCases.BaseTestCase):
     params = {"mu": 0.0, "sigma": 1.0}
 
 
-@pytest.mark.xfail(reason="This distribution has not been refactored for v4")
-class TestConstant(BaseTestCases.BaseTestCase):
-    distribution = pm.Constant
-    params = {"c": 3}
-
-
 @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
 class TestZeroInflatedPoisson(BaseTestCases.BaseTestCase):
     distribution = pm.ZeroInflatedPoisson
@@ -918,6 +911,24 @@ def discrete_uniform_rng_fn(self, size, lower, upper, rng):
     ]
 
 
+class TestConstant(BaseTestDistribution):
+    def constant_rng_fn(self, size, c):
+        if size is None:
+            return c
+        return np.full(size, c)
+
+    pymc_dist = pm.Constant
+    pymc_dist_params = {"c": 3}
+    expected_rv_op_params = {"c": 3}
+    reference_dist_params = {"c": 3}
+    reference_dist = lambda self: self.constant_rng_fn
+    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):
@@ -1027,13 +1038,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_constant_dist(self):
-        def ref_rand(size, c):
-            return c * np.ones(size, dtype=int)
-
-        pymc3_random_discrete(pm.Constant, {"c": I}, ref_rand=ref_rand)
-
     @pytest.mark.xfail(reason="This distribution has not been refactored for v4")
     def test_matrix_normal(self):
         def ref_rand(size, mu, rowcov, colcov):