Proposal: Dist shape refactor

pymc3/distributions/discrete.py

@@ -1,20 +1,26 @@
 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
-
-__all__ = ['Binomial',  'BetaBinomial',  'Bernoulli',  'Poisson',
+from .distribution import (
+    Univariate,
+    Discrete,
+    draw_values,
+    generate_samples,
+    has_const_inputs)
+
+__all__ = ['Binomial', 'BetaBinomial', 'Bernoulli', 'Poisson',
            'NegativeBinomial', 'ConstantDist', 'Constant', 'ZeroInflatedPoisson',
            'ZeroInflatedNegativeBinomial', 'DiscreteUniform', 'Geometric',
            'Categorical']
 
 
-class Binomial(Discrete):
+class Binomial(Univariate, Discrete):
     R"""
     Binomial log-likelihood.
 
@@ -39,10 +45,14 @@ class Binomial(Discrete):
     """
 
     def __init__(self, n, p, *args, **kwargs):
-        super(Binomial, self).__init__(*args, **kwargs)
-        self.n = n
-        self.p = p
-        self.mode = tt.cast(tt.round(n * p), self.dtype)
+        self.n = tt.as_tensor_variable(n)
+        self.p = tt.as_tensor_variable(p)
+
+        self.mode = tt.cast(tt.round(n * p), 'int32')
+
+        dist_params = (self.n, self.p)
+
+        super(Binomial, self).__init__(dist_params, *args, **kwargs)
 
     def random(self, point=None, size=None, repeat=None):
         n, p = draw_values([self.n, self.p], point=point)
@@ -60,7 +70,7 @@ def logp(self, value):
             0 <= p, p <= 1)
 
 
-class BetaBinomial(Discrete):
+class BetaBinomial(Univariate, Discrete):
     R"""
     Beta-binomial log-likelihood.
 
@@ -90,11 +100,16 @@ class BetaBinomial(Discrete):
     """
 
     def __init__(self, alpha, beta, n, *args, **kwargs):
-        super(BetaBinomial, self).__init__(*args, **kwargs)
-        self.alpha = alpha
-        self.beta = beta
-        self.n = n
-        self.mode = tt.cast(tt.round(alpha / (alpha + beta)), 'int8')
+        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(self.alpha / (self.alpha + self.beta)),
+                            'int8')
+
+        dist_params = (self.alpha, self.beta, self.n)
+
+        super(BetaBinomial, self).__init__(dist_params, *args, **kwargs)
 
     def _random(self, alpha, beta, n, size=None):
         size = size or 1
@@ -125,7 +140,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
@@ -146,10 +161,14 @@ class Bernoulli(Discrete):
     """
 
     def __init__(self, p, *args, **kwargs):
-        super(Bernoulli, self).__init__(*args, **kwargs)
-        self.p = p
+        self.p = tt.as_tensor_variable(p)
+
+        dist_params = (self.p,)
+
         self.mode = tt.cast(tt.round(p), 'int8')
 
+        super(Bernoulli, self).__init__(dist_params, *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 +183,7 @@ def logp(self, value):
             p >= 0, p <= 1)
 
 
-class Poisson(Discrete):
+class Poisson(Univariate, Discrete):
     R"""
     Poisson log-likelihood.
 
@@ -192,10 +211,14 @@ class Poisson(Discrete):
     """
 
     def __init__(self, mu, *args, **kwargs):
-        super(Poisson, self).__init__(*args, **kwargs)
-        self.mu = mu
+        self.mu = tt.as_tensor_variable(mu)
+
         self.mode = tt.floor(mu).astype('int32')
 
+        dist_params = (self.mu,)
+
+        super(Poisson, self).__init__(dist_params, *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 +235,7 @@ def logp(self, value):
                          0, log_prob)
 
 
-class NegativeBinomial(Discrete):
+class NegativeBinomial(Univariate, Discrete):
     R"""
     Negative binomial log-likelihood.
 
@@ -239,11 +262,16 @@ class NegativeBinomial(Discrete):
     """
 
     def __init__(self, mu, alpha, *args, **kwargs):
-        super(NegativeBinomial, self).__init__(*args, **kwargs)
-        self.mu = mu
-        self.alpha = alpha
+        self.mu = tt.as_tensor_variable(mu)
+        self.alpha = tt.as_tensor_variable(alpha)
+
         self.mode = tt.floor(mu).astype('int32')
 
+        dist_params = (self.mu, self.alpha)
+
+        super(NegativeBinomial, self).__init__(dist_params,
+                                               *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 +294,7 @@ def logp(self, value):
                          negbinom)
 
 
-class Geometric(Discrete):
+class Geometric(Univariate, Discrete):
     R"""
     Geometric log-likelihood.
 
@@ -288,9 +316,14 @@ class Geometric(Discrete):
     """
 
     def __init__(self, p, *args, **kwargs):
-        super(Geometric, self).__init__(*args, **kwargs)
-        self.p = p
-        self.mode = 1
+        self.p = tt.as_tensor_variable(p)
+
+        self.mode = tt.as_tensor_variable(1)
+
+        dist_params = (self.p,)
+
+        super(Geometric, self).__init__(
+              dist_params, *args, **kwargs)
 
     def random(self, point=None, size=None, repeat=None):
         p = draw_values([self.p], point=point)
@@ -304,7 +337,7 @@ def logp(self, value):
                      0 <= p, p <= 1, value >= 1)
 
 
-class DiscreteUniform(Discrete):
+class DiscreteUniform(Univariate, Discrete):
     R"""
     Discrete uniform distribution.
 
@@ -325,12 +358,16 @@ class DiscreteUniform(Discrete):
     """
 
     def __init__(self, lower, upper, *args, **kwargs):
-        super(DiscreteUniform, self).__init__(*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)
 
+        dist_params = (self.lower, self.upper)
+
+        super(DiscreteUniform, self).__init__(dist_params, *args, **kwargs)
+
     def _random(self, lower, upper, size=None):
         # This way seems to be the only to deal with lower and upper
         # as array-like.
@@ -352,7 +389,7 @@ def logp(self, value):
                      lower <= value, value <= upper)
 
 
-class Categorical(Discrete):
+class Categorical(Univariate, Discrete):
     R"""
     Categorical log-likelihood.
 
@@ -372,15 +409,23 @@ class Categorical(Discrete):
     """
 
     def __init__(self, p, *args, **kwargs):
-        super(Categorical, self).__init__(*args, **kwargs)
-        try:
-            self.k = tt.shape(p)[-1].tag.test_value
-        except AttributeError:
-            self.k = tt.shape(p)[-1]
+
+        self.k = tt.shape(p)[-1]
+
+        if has_const_inputs(self.k):
+            # FIXME: This feels like a hack.  Seems like it would be better to
+            # evaluate somewhere else (e.g. exactly where a value is needed).
+            self.k = self.k.eval()
+
         self.p = tt.as_tensor_variable(p)
         self.p = (p.T / tt.sum(p, -1)).T
+
         self.mode = tt.argmax(p)
 
+        dist_params = (self.p,)
+
+        super(Categorical, self).__init__(dist_params, *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 +458,7 @@ def logp(self, value):
                      sumto1)
 
 
-class Constant(Discrete):
+class Constant(Univariate, Discrete):
     """
     Constant log-likelihood.
 
@@ -424,8 +469,10 @@ class Constant(Discrete):
     """
 
     def __init__(self, c, *args, **kwargs):
-        super(Constant, self).__init__(*args, **kwargs)
-        self.mean = self.median = self.mode = self.c = c
+        self.mean = self.median = self.mode = self.c = tt.as_tensor_variable(c)
+        dist_params = (self.c,)
+
+        super(Constant, self).__init__(dist_params, *args, **kwargs)
 
     def random(self, point=None, size=None, repeat=None):
         c = draw_values([self.c], point=point)
@@ -441,13 +488,14 @@ 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)
+                  DeprecationWarning)
     return Constant(*args, **kwargs)
 
 
-class ZeroInflatedPoisson(Discrete):
+class ZeroInflatedPoisson(Univariate, Discrete):
     R"""
     Zero-inflated Poisson log-likelihood.
 
@@ -478,12 +526,15 @@ class ZeroInflatedPoisson(Discrete):
     """
 
     def __init__(self, theta, psi, *args, **kwargs):
-        super(ZeroInflatedPoisson, self).__init__(*args, **kwargs)
-        self.theta = theta
-        self.psi = psi
+        self.theta = tt.as_tensor_variable(theta)
+        self.psi = tt.as_tensor_variable(psi)
         self.pois = Poisson.dist(theta)
         self.mode = self.pois.mode
 
+        dist_params = (self.theta, self.psi)
+
+        super(ZeroInflatedPoisson, self).__init__(dist_params, *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 +548,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.
 
@@ -529,13 +580,18 @@ class ZeroInflatedNegativeBinomial(Discrete):
     """
 
     def __init__(self, mu, alpha, psi, *args, **kwargs):
-        super(ZeroInflatedNegativeBinomial, self).__init__(*args, **kwargs)
-        self.mu = mu
-        self.alpha = alpha
-        self.psi = psi
+        self.mu = tt.as_tensor_variable(mu)
+        self.alpha = tt.as_tensor_variable(alpha)
+        self.psi = tt.as_tensor_variable(psi)
+
         self.nb = NegativeBinomial.dist(mu, alpha)
         self.mode = self.nb.mode
 
+        dist_params = (self.mu, self.alpha, self.psi)
+
+        super(ZeroInflatedNegativeBinomial, self).__init__(
+            dist_params, *args, **kwargs)
+
     def random(self, point=None, size=None, repeat=None):
         mu, alpha, psi = draw_values(
             [self.mu, self.alpha, self.psi], point=point)