Proposal: Dist shape refactor

pymc3/distributions/discrete.py

@@ -1,20 +1,21 @@
 from functools import partial
+import warnings
 
 import numpy as np
 import theano
 import theano.tensor as tt
 from scipy import stats
 
 from .dist_math import bound, factln, binomln, betaln, logpow
-from .distribution import Discrete, draw_values, generate_samples
+from .distribution import Univariate, Discrete, draw_values, generate_samples
 
 __all__ = ['Binomial',  'BetaBinomial',  'Bernoulli',  'Poisson',
            'NegativeBinomial', 'ConstantDist', 'Constant', 'ZeroInflatedPoisson',
            'ZeroInflatedNegativeBinomial', 'DiscreteUniform', 'Geometric',
            'Categorical']
 
 
-class Binomial(Discrete):
+class Binomial(Univariate, Discrete):
     R"""
     Binomial log-likelihood.
 
@@ -37,13 +38,16 @@ class Binomial(Discrete):
     p : float
         Probability of success in each trial (0 < p < 1).
     """
-
-    def __init__(self, n, p, *args, **kwargs):
-        super(Binomial, self).__init__(*args, **kwargs)
-        self.n = n
-        self.p = p
+    def __init__(self, n, p, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.n = tt.as_tensor_variable(n)
+        self.p = tt.as_tensor_variable(p)
         self.mode = tt.cast(tt.round(n * p), self.dtype)
 
+        self.dist_params = (self.n, self.p)
+
+        super(Binomial, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
+
     def random(self, point=None, size=None, repeat=None):
         n, p = draw_values([self.n, self.p], point=point)
         return generate_samples(stats.binom.rvs, n=n, p=p,
@@ -60,7 +64,7 @@ def logp(self, value):
             0 <= p, p <= 1)
 
 
-class BetaBinomial(Discrete):
+class BetaBinomial(Univariate, Discrete):
     R"""
     Beta-binomial log-likelihood.
 
@@ -89,13 +93,17 @@ class BetaBinomial(Discrete):
         beta > 0.
     """
 
-    def __init__(self, alpha, beta, n, *args, **kwargs):
-        super(BetaBinomial, self).__init__(*args, **kwargs)
-        self.alpha = alpha
-        self.beta = beta
-        self.n = n
+    def __init__(self, alpha, beta, n, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.alpha = tt.as_tensor_variable(alpha)
+        self.beta = tt.as_tensor_variable(beta)
+        self.n = tt.as_tensor_variable(n)
         self.mode = tt.cast(tt.round(alpha / (alpha + beta)), 'int8')
 
+        self.dist_params = (self.alpha, self.beta, self.n)
+
+        super(BetaBinomial, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
+
     def _random(self, alpha, beta, n, size=None):
         size = size or 1
         p = np.atleast_1d(stats.beta.rvs(a=alpha, b=beta, size=np.prod(size)))
@@ -125,7 +133,7 @@ def logp(self, value):
                      alpha > 0, beta > 0)
 
 
-class Bernoulli(Discrete):
+class Bernoulli(Univariate, Discrete):
     R"""Bernoulli log-likelihood
 
     The Bernoulli distribution describes the probability of successes
@@ -144,12 +152,15 @@ class Bernoulli(Discrete):
     p : float
         Probability of success (0 < p < 1).
     """
-
-    def __init__(self, p, *args, **kwargs):
-        super(Bernoulli, self).__init__(*args, **kwargs)
-        self.p = p
+    def __init__(self, p, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.p = tt.as_tensor_variable(p)
         self.mode = tt.cast(tt.round(p), 'int8')
 
+        self.dist_params = (self.p,)
+        # FIXME: bcast, etc.
+        super(Bernoulli, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
+
     def random(self, point=None, size=None, repeat=None):
         p = draw_values([self.p], point=point)
         return generate_samples(stats.bernoulli.rvs, p,
@@ -164,7 +175,7 @@ def logp(self, value):
             p >= 0, p <= 1)
 
 
-class Poisson(Discrete):
+class Poisson(Univariate, Discrete):
     R"""
     Poisson log-likelihood.
 
@@ -190,12 +201,15 @@ class Poisson(Discrete):
     The Poisson distribution can be derived as a limiting case of the
     binomial distribution.
     """
-
-    def __init__(self, mu, *args, **kwargs):
-        super(Poisson, self).__init__(*args, **kwargs)
-        self.mu = mu
+    def __init__(self, mu, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.mu = tt.as_tensor_variable(mu)
         self.mode = tt.floor(mu).astype('int32')
 
+        self.dist_params = (self.mu,)
+
+        super(Poisson, self).__init__(self.dist_params,
+                                      ndim, size, dtype, *args, **kwargs)
+
     def random(self, point=None, size=None, repeat=None):
         mu = draw_values([self.mu], point=point)
         return generate_samples(stats.poisson.rvs, mu,
@@ -212,7 +226,7 @@ def logp(self, value):
                          0, log_prob)
 
 
-class NegativeBinomial(Discrete):
+class NegativeBinomial(Univariate, Discrete):
     R"""
     Negative binomial log-likelihood.
 
@@ -237,13 +251,16 @@ class NegativeBinomial(Discrete):
     alpha : float
         Gamma distribution parameter (alpha > 0).
     """
-
-    def __init__(self, mu, alpha, *args, **kwargs):
-        super(NegativeBinomial, self).__init__(*args, **kwargs)
-        self.mu = mu
-        self.alpha = alpha
+    def __init__(self, mu, alpha, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.mu = tt.as_tensor_variable(mu)
+        self.alpha = tt.as_tensor_variable(alpha)
         self.mode = tt.floor(mu).astype('int32')
 
+        self.dist_params = (self.mu, self.alpha)
+
+        super(NegativeBinomial, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
+
     def random(self, point=None, size=None, repeat=None):
         mu, alpha = draw_values([self.mu, self.alpha], point=point)
         g = generate_samples(stats.gamma.rvs, alpha, scale=mu / alpha,
@@ -266,7 +283,7 @@ def logp(self, value):
                          negbinom)
 
 
-class Geometric(Discrete):
+class Geometric(Univariate, Discrete):
     R"""
     Geometric log-likelihood.
 
@@ -286,11 +303,13 @@ class Geometric(Discrete):
     p : float
         Probability of success on an individual trial (0 < p <= 1).
     """
+    def __init__(self, p, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.p = tt.as_tensor_variable(p)
+        self.mode = tt.as_tensor_variable(1)
+        self.dist_params = (self.p,)
 
-    def __init__(self, p, *args, **kwargs):
-        super(Geometric, self).__init__(*args, **kwargs)
-        self.p = p
-        self.mode = 1
+        super(Geometric, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
 
     def random(self, point=None, size=None, repeat=None):
         p = draw_values([self.p], point=point)
@@ -304,7 +323,7 @@ def logp(self, value):
                      0 <= p, p <= 1, value >= 1)
 
 
-class DiscreteUniform(Discrete):
+class DiscreteUniform(Univariate, Discrete):
     R"""
     Discrete uniform distribution.
 
@@ -324,12 +343,15 @@ class DiscreteUniform(Discrete):
         Upper limit (upper > lower).
     """
 
-    def __init__(self, lower, upper, *args, **kwargs):
-        super(DiscreteUniform, self).__init__(*args, **kwargs)
+    def __init__(self, lower, upper, ndim=None, size=None, dtype=None, *args, **kwargs):
         self.lower = tt.floor(lower).astype('int32')
         self.upper = tt.floor(upper).astype('int32')
         self.mode = tt.maximum(
             tt.floor((upper - lower) / 2.).astype('int32'), self.lower)
+        self.dist_params = (self.lower, self.upper)
+
+        super(DiscreteUniform, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
 
     def _random(self, lower, upper, size=None):
         # This way seems to be the only to deal with lower and upper
@@ -352,7 +374,7 @@ def logp(self, value):
                      lower <= value, value <= upper)
 
 
-class Categorical(Discrete):
+class Categorical(Univariate, Discrete):
     R"""
     Categorical log-likelihood.
 
@@ -370,9 +392,8 @@ class Categorical(Discrete):
         p > 0 and the elements of p must sum to 1. They will be automatically
         rescaled otherwise.
     """
+    def __init__(self, p, ndim=None, size=None, dtype=None, *args, **kwargs):
 
-    def __init__(self, p, *args, **kwargs):
-        super(Categorical, self).__init__(*args, **kwargs)
         try:
             self.k = tt.shape(p)[-1].tag.test_value
         except AttributeError:
@@ -381,6 +402,14 @@ def __init__(self, p, *args, **kwargs):
         self.p = (p.T / tt.sum(p, -1)).T
         self.mode = tt.argmax(p)
 
+        self.shape_support = tt.shape(self.p)[-1]
+
+        self.dist_params = (self.p,)
+        # FIXME: The stated support is univariate?
+
+        super(Categorical, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
+
     def random(self, point=None, size=None, repeat=None):
         def random_choice(k, *args, **kwargs):
             if len(kwargs['p'].shape) > 1:
@@ -413,7 +442,7 @@ def logp(self, value):
                      sumto1)
 
 
-class Constant(Discrete):
+class Constant(Univariate, Discrete):
     """
     Constant log-likelihood.
 
@@ -423,9 +452,11 @@ class Constant(Discrete):
         Constant parameter.
     """
 
-    def __init__(self, c, *args, **kwargs):
-        super(Constant, self).__init__(*args, **kwargs)
-        self.mean = self.median = self.mode = self.c = c
+    def __init__(self, c, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.mean = self.median = self.mode = self.c = tt.as_tensor_variable(c)
+        self.dist_params = (self.c,)
+        super(Constant, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
 
     def random(self, point=None, size=None, repeat=None):
         c = draw_values([self.c], point=point)
@@ -441,13 +472,13 @@ def logp(self, value):
         c = self.c
         return bound(0, tt.eq(value, c))
 
+
 def ConstantDist(*args, **kwargs):
     warnings.warn("ConstantDist has been deprecated. In future, use Constant instead.",
                 DeprecationWarning)
     return Constant(*args, **kwargs)
 
-
-class ZeroInflatedPoisson(Discrete):
+class ZeroInflatedPoisson(Univariate, Discrete):
     R"""
     Zero-inflated Poisson log-likelihood.
 
@@ -476,14 +507,17 @@ class ZeroInflatedPoisson(Discrete):
         Expected proportion of Poisson variates (0 < psi < 1)
 
     """
-
-    def __init__(self, theta, psi, *args, **kwargs):
-        super(ZeroInflatedPoisson, self).__init__(*args, **kwargs)
-        self.theta = theta
-        self.psi = psi
+    def __init__(self, theta, psi, ndim=None, size=None, dtype=None, *args, **kwargs):
+        self.theta = tt.as_tensor_variable(theta)
+        self.psi = tt.as_tensor_variable(psi)
         self.pois = Poisson.dist(theta)
         self.mode = self.pois.mode
 
+        self.dist_params = (self.theta, self.psi)
+
+        super(ZeroInflatedPoisson, self).__init__(
+            self.dist_params, ndim, size, dtype, *args, **kwargs)
+
     def random(self, point=None, size=None, repeat=None):
         theta, psi = draw_values([self.theta, self.psi], point=point)
         g = generate_samples(stats.poisson.rvs, theta,
@@ -497,7 +531,7 @@ def logp(self, value):
                          tt.log((1. - self.psi) + self.psi * tt.exp(-self.theta)))
 
 
-class ZeroInflatedNegativeBinomial(Discrete):
+class ZeroInflatedNegativeBinomial(Univariate, Discrete):
     R"""
     Zero-Inflated Negative binomial log-likelihood.
 
@@ -528,13 +562,14 @@ class ZeroInflatedNegativeBinomial(Discrete):
         Expected proportion of NegativeBinomial variates (0 < psi < 1)
     """
 
-    def __init__(self, mu, alpha, psi, *args, **kwargs):
-        super(ZeroInflatedNegativeBinomial, self).__init__(*args, **kwargs)
+    def __init__(self, mu, alpha, psi, ndim=None, size=None, dtype=None, *args, **kwargs):
         self.mu = mu
         self.alpha = alpha
         self.psi = psi
         self.nb = NegativeBinomial.dist(mu, alpha)
         self.mode = self.nb.mode
+        self.dist_params = (mu, alpha, psi)
+        super(ZeroInflatedNegativeBinomial, self).__init__(self.dist_params, ndim, size, dtype, *args, **kwargs)
 
     def random(self, point=None, size=None, repeat=None):
         mu, alpha, psi = draw_values(