diff --git a/pymc3/distributions/discrete.py b/pymc3/distributions/discrete.py index 31cf813a5d..807e173bad 100644 --- a/pymc3/distributions/discrete.py +++ b/pymc3/distributions/discrete.py @@ -16,7 +16,14 @@ import aesara.tensor as at import numpy as np -from aesara.tensor.random.basic import bernoulli, binomial, categorical, nbinom, poisson +from aesara.tensor.random.basic import ( + RandomVariable, + bernoulli, + binomial, + categorical, + nbinom, + poisson, +) from scipy import stats from pymc3.aesaraf import floatX, intX, take_along_axis @@ -434,6 +441,22 @@ def _distr_parameters_for_repr(self): return ["p"] +class DiscreteWeibullRV(RandomVariable): + name = "discrete_weibull" + ndim_supp = 0 + ndims_params = [0, 0] + dtype = "int64" + _print_name = ("dWeibull", "\\operatorname{dWeibull}") + + @classmethod + def rng_fn(cls, rng, q, beta, size): + p = rng.uniform(size=size) + return np.ceil(np.power(np.log(1 - p) / np.log(q), 1.0 / beta)) - 1 + + +discrete_weibull = DiscreteWeibullRV() + + class DiscreteWeibull(Discrete): R"""Discrete Weibull log-likelihood @@ -473,51 +496,15 @@ def DiscreteWeibull(q, b, x): Variance :math:`2 \sum_{x = 1}^{\infty} x q^{x^{\beta}} - \mu - \mu^2` ======== ====================== """ + rv_op = discrete_weibull - def __init__(self, q, beta, *args, **kwargs): - super().__init__(*args, defaults=("median",), **kwargs) - - self.q = at.as_tensor_variable(floatX(q)) - self.beta = at.as_tensor_variable(floatX(beta)) - - self.median = self._ppf(0.5) - - def _ppf(self, p): - r""" - The percentile point function (the inverse of the cumulative - distribution function) of the discrete Weibull distribution. - """ - q = self.q - beta = self.beta - - return (at.ceil(at.power(at.log(1 - p) / at.log(q), 1.0 / beta)) - 1).astype("int64") - - def _random(self, q, beta, size=None): - p = np.random.uniform(size=size) - - return np.ceil(np.power(np.log(1 - p) / np.log(q), 1.0 / beta)) - 1 - - def random(self, point=None, size=None): - r""" - Draw random values from DiscreteWeibull 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). - - Returns - ------- - array - """ - # q, beta = draw_values([self.q, self.beta], point=point, size=size) - # return generate_samples(self._random, q, beta, dist_shape=self.shape, size=size) + @classmethod + def dist(cls, q, beta, *args, **kwargs): + q = at.as_tensor_variable(floatX(q)) + beta = at.as_tensor_variable(floatX(beta)) + return super().dist([q, beta], **kwargs) - def logp(self, value): + def logp(value, q, beta): r""" Calculate log-probability of DiscreteWeibull distribution at specified value. @@ -531,8 +518,6 @@ def logp(self, value): ------- TensorVariable """ - q = self.q - beta = self.beta return bound( at.log(at.power(q, at.power(value, beta)) - at.power(q, at.power(value + 1, beta))), 0 <= value, @@ -541,7 +526,7 @@ def logp(self, value): 0 < beta, ) - def logcdf(self, value): + def logcdf(value, q, beta): """ Compute the log of the cumulative distribution function for Discrete Weibull distribution at the specified value. @@ -556,9 +541,6 @@ def logcdf(self, value): ------- TensorVariable """ - q = self.q - beta = self.beta - return bound( at.log1p(-at.power(q, at.power(value + 1, beta))), 0 <= value, diff --git a/pymc3/tests/test_distributions.py b/pymc3/tests/test_distributions.py index dfb215aa53..438bc46b54 100644 --- a/pymc3/tests/test_distributions.py +++ b/pymc3/tests/test_distributions.py @@ -424,10 +424,7 @@ def logpow(v, p): def discrete_weibull_logpmf(value, q, beta): return floatX( - np.log( - np.power(floatX(q), np.power(floatX(value), floatX(beta))) - - np.power(floatX(q), np.power(floatX(value + 1), floatX(beta))) - ) + np.log(np.power(q, np.power(value, beta)) - np.power(q, np.power(value + 1, beta))) ) @@ -1557,7 +1554,6 @@ def test_bernoulli(self): {"p": Unit}, ) - @pytest.mark.xfail(reason="Distribution not refactored yet") def test_discrete_weibull(self): self.check_logp( DiscreteWeibull, diff --git a/pymc3/tests/test_distributions_random.py b/pymc3/tests/test_distributions_random.py index 16b960f1a2..63bddaae89 100644 --- a/pymc3/tests/test_distributions_random.py +++ b/pymc3/tests/test_distributions_random.py @@ -431,7 +431,6 @@ class TestBernoulli(BaseTestCases.BaseTestCase): params = {"p": 0.5} -@pytest.mark.xfail(reason="This distribution has not been refactored for v4") class TestDiscreteWeibull(BaseTestCases.BaseTestCase): distribution = pm.DiscreteWeibull params = {"q": 0.25, "beta": 2.0} @@ -817,7 +816,6 @@ def ref_rand(size, lower, upper): pm.DiscreteUniform, {"lower": -NatSmall, "upper": NatSmall}, ref_rand=ref_rand ) - @pytest.mark.xfail(reason="This distribution has not been refactored for v4") def test_discrete_weibull(self): def ref_rand(size, q, beta): u = np.random.uniform(size=size)