Make all univariate distributions use bound_elemwise. Addd tests.

Author: twiecki

pymc3/distributions/continuous.py

@@ -243,7 +243,7 @@ def logp(self, value):
         sd = self.sd
         tau = self.tau
         mu = self.mu
-        return bound((-tau * (value - mu)**2 + tt.log(tau / np.pi / 2.)) / 2.,
+        return bound_elemwise((-tau * (value - mu)**2 + tt.log(tau / np.pi / 2.)) / 2.,
                      sd > 0)
 
 
@@ -289,7 +289,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         tau = self.tau
         sd = self.sd
-        return bound(-0.5 * tau * value**2 + 0.5 * tt.log(tau * 2. / np.pi),
+        return bound_elemwise(-0.5 * tau * value**2 + 0.5 * tt.log(tau * 2. / np.pi),
                      value >= 0,
                      tau > 0, sd > 0)
 
@@ -402,7 +402,7 @@ def logp(self, value):
         lam = self.lam
         alpha = self.alpha
         # value *must* be iid. Otherwise this is wrong.
-        return bound(logpow(lam / (2. * np.pi), 0.5)
+        return bound_elemwise(logpow(lam / (2. * np.pi), 0.5)
                      - logpow(value - alpha, 1.5)
                      - (0.5 * lam / (value - alpha)
                         * ((value - alpha - mu) / mu)**2),
@@ -492,7 +492,7 @@ def logp(self, value):
         alpha = self.alpha
         beta = self.beta
 
-        return bound(logpow(value, alpha - 1) + logpow(1 - value, beta - 1)
+        return bound_elemwise(logpow(value, alpha - 1) + logpow(1 - value, beta - 1)
                      - betaln(alpha, beta),
                      value >= 0, value <= 1,
                      alpha > 0, beta > 0)
@@ -537,7 +537,7 @@ def random(self, point=None, size=None, repeat=None):
 
     def logp(self, value):
         lam = self.lam
-        return bound(tt.log(lam) - lam * value, value > 0, lam > 0)
+        return bound_elemwise(tt.log(lam) - lam * value, value > 0, lam > 0)
 
 
 class Laplace(Continuous):
@@ -641,7 +641,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         mu = self.mu
         tau = self.tau
-        return bound(-0.5 * tau * (tt.log(value) - mu)**2
+        return bound_elemwise(-0.5 * tau * (tt.log(value) - mu)**2
                      + 0.5 * tt.log(tau / (2. * np.pi))
                      - tt.log(value),
                      tau > 0)
@@ -702,7 +702,7 @@ def logp(self, value):
         lam = self.lam
         sd = self.sd
 
-        return bound(gammaln((nu + 1.0) / 2.0)
+        return bound_elemwise(gammaln((nu + 1.0) / 2.0)
                      + .5 * tt.log(lam / (nu * np.pi))
                      - gammaln(nu / 2.0)
                      - (nu + 1.0) / 2.0 * tt.log1p(lam * (value - mu)**2 / nu),
@@ -765,7 +765,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         alpha = self.alpha
         m = self.m
-        return bound(tt.log(alpha) + logpow(m, alpha)
+        return bound_elemwise(tt.log(alpha) + logpow(m, alpha)
                      - logpow(value, alpha + 1),
                      value >= m, alpha > 0, m > 0)
 
@@ -817,7 +817,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         alpha = self.alpha
         beta = self.beta
-        return bound(- tt.log(np.pi) - tt.log(beta)
+        return bound_elemwise(- tt.log(np.pi) - tt.log(beta)
                      - tt.log1p(((value - alpha) / beta)**2),
                      beta > 0)
 
@@ -863,7 +863,7 @@ def random(self, point=None, size=None, repeat=None):
 
     def logp(self, value):
         beta = self.beta
-        return bound(tt.log(2) - tt.log(np.pi) - tt.log(beta)
+        return bound_elemwise(tt.log(2) - tt.log(np.pi) - tt.log(beta)
                      - tt.log1p((value / beta)**2),
                      value >= 0, beta > 0)
 
@@ -943,7 +943,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         alpha = self.alpha
         beta = self.beta
-        return bound(
+        return bound_elemwise(
             -gammaln(alpha) + logpow(
                 beta, alpha) - beta * value + logpow(value, alpha - 1),
 
@@ -1007,7 +1007,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         alpha = self.alpha
         beta = self.beta
-        return bound(logpow(beta, alpha) - gammaln(alpha) - beta / value
+        return bound_elemwise(logpow(beta, alpha) - gammaln(alpha) - beta / value
                      + logpow(value, -alpha - 1),
                      value > 0, alpha > 0, beta > 0)
 
@@ -1088,7 +1088,7 @@ def _random(a, b, size=None):
     def logp(self, value):
         alpha = self.alpha
         beta = self.beta
-        return bound(tt.log(alpha) - tt.log(beta)
+        return bound_elemwise(tt.log(alpha) - tt.log(beta)
                      + (alpha - 1) * tt.log(value / beta)
                      - (value / beta)**alpha,
                      value >= 0, alpha > 0, beta > 0)
@@ -1131,12 +1131,12 @@ def __init__(self, distribution, lower, upper, transform='infer', *args, **kwarg
                     self.testval = 0.5 * (upper + lower)
 
             if not np.isinf(lower) and np.isinf(upper):
-                self.transform = transforms.lowerbound(lower)
+                self.transform = transforms.lowerbound_elemwise(lower)
                 if default <= lower:
                     self.testval = lower + 1
 
             if np.isinf(lower) and not np.isinf(upper):
-                self.transform = transforms.upperbound(upper)
+                self.transform = transforms.upperbound_elemwise(upper)
                 if default >= upper:
                     self.testval = upper - 1
 
@@ -1161,7 +1161,7 @@ def random(self, point=None, size=None, repeat=None):
                                 size=size)
 
     def logp(self, value):
-        return bound(self.dist.logp(value),
+        return bound_elemwise(self.dist.logp(value),
                      value >= self.lower, value <= self.upper)
 
 
@@ -1286,7 +1286,7 @@ def logp(self, value):
                        + logpow(std_cdf((value - mu) / sigma - sigma / nu), 1.),
                        - tt.log(sigma * tt.sqrt(2 * np.pi))
                        - 0.5 * ((value - mu) / sigma)**2)
-        return bound(lp, sigma > 0., nu > 0.)
+        return bound_elemwise(lp, sigma > 0., nu > 0.)
 
 
 class VonMises(Continuous):
@@ -1335,7 +1335,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         mu = self.mu
         kappa = self.kappa
-        return bound(kappa * tt.cos(mu - value) - tt.log(2 * np.pi * i0(kappa)), value >= -np.pi, value <= np.pi, kappa >= 0)
+        return bound_elemwise(kappa * tt.cos(mu - value) - tt.log(2 * np.pi * i0(kappa)), value >= -np.pi, value <= np.pi, kappa >= 0)
 
 
 class SkewNormal(Continuous):
@@ -1401,7 +1401,7 @@ def logp(self, value):
         sd = self.sd
         mu = self.mu
         alpha = self.alpha
-        return bound(
+        return bound_elemwise(
             tt.log(1 +
             tt.erf(((value - mu) * tt.sqrt(tau) * alpha) / tt.sqrt(2)))
             + (-tau * (value - mu)**2

pymc3/distributions/discrete.py

@@ -54,7 +54,7 @@ def logp(self, value):
         n = self.n
         p = self.p
 
-        return bound(
+        return bound_elemwise(
             binomln(n, value) + logpow(p, value) + logpow(1 - p, n - value),
             0 <= value, value <= n,
             0 <= p, p <= 1)
@@ -118,7 +118,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         alpha = self.alpha
         beta = self.beta
-        return bound(binomln(self.n, value)
+        return bound_elemwise(binomln(self.n, value)
                      + betaln(value + alpha, self.n - value + beta)
                      - betaln(alpha, beta),
                      value >= 0, value <= self.n,
@@ -158,7 +158,7 @@ def random(self, point=None, size=None, repeat=None):
 
     def logp(self, value):
         p = self.p
-        return bound(
+        return bound_elemwise(
             tt.switch(value, tt.log(p), tt.log(1 - p)),
             value >= 0, value <= 1,
             p >= 0, p <= 1)
@@ -204,7 +204,7 @@ def random(self, point=None, size=None, repeat=None):
 
     def logp(self, value):
         mu = self.mu
-        log_prob = bound(
+        log_prob = bound_elemwise(
             logpow(mu, value) - factln(value) - mu,
             mu >= 0, value >= 0)
         # Return zero when mu and value are both zero
@@ -255,7 +255,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         mu = self.mu
         alpha = self.alpha
-        negbinom = bound(binomln(value + alpha - 1, value)
+        negbinom = bound_elemwise(binomln(value + alpha - 1, value)
                          + logpow(mu / (mu + alpha), value)
                          + logpow(alpha / (mu + alpha), alpha),
                          value >= 0, mu > 0, alpha > 0)
@@ -300,7 +300,7 @@ def random(self, point=None, size=None, repeat=None):
 
     def logp(self, value):
         p = self.p
-        return bound(tt.log(p) + logpow(1 - p, value - 1),
+        return bound_elemwise(tt.log(p) + logpow(1 - p, value - 1),
                      0 <= p, p <= 1, value >= 1)
 
 
@@ -348,7 +348,7 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         upper = self.upper
         lower = self.lower
-        return bound_elemwise(-tt.log(upper - lower + 1) * tt.ones_like(value),
+        return bound_elemwise(-tt.log(upper - lower + 1),
                               lower <= value, value <= upper)
 
 
@@ -408,7 +408,7 @@ def logp(self, value):
             a = tt.log(p[tt.arange(p.shape[0]), value])
         else:
             a = tt.log(p[value])
-        return bound(a,
+        return bound_elemwise(a,
                      value >= 0, value <= (k - 1),
                      sumto1)
 
@@ -439,7 +439,7 @@ def _random(c, dtype=dtype, size=None):
 
     def logp(self, value):
         c = self.c
-        return bound(0, tt.eq(value, c))
+        return bound_elemwise(0, tt.eq(value, c))
 
 def ConstantDist(*args, **kwargs):
     warnings.warn("ConstantDist has been deprecated. In future, use Constant instead.",

pymc3/tests/test_dist_math.py

@@ -1,36 +1,81 @@
 import numpy as np
 import theano.tensor as tt
+import pymc3 as pm
 
-from ..distributions.dist_math import alltrue
-
-
-def test_alltrue():
-    assert alltrue([]).eval()
-    assert alltrue([True]).eval()
-    assert alltrue([tt.ones(10)]).eval()
-    assert alltrue([tt.ones(10),
-                    5 * tt.ones(101)]).eval()
-    assert alltrue([np.ones(10),
-                    5 * tt.ones(101)]).eval()
-    assert alltrue([np.ones(10),
-                    True,
-                    5 * tt.ones(101)]).eval()
-    assert alltrue([np.array([1, 2, 3]),
-                    True,
-                    5 * tt.ones(101)]).eval()
-
-    assert not alltrue([False]).eval()
-    assert not alltrue([tt.zeros(10)]).eval()
-    assert not alltrue([True,
-                        False]).eval()
-    assert not alltrue([np.array([0, -1]),
-                        tt.ones(60)]).eval()
-    assert not alltrue([np.ones(10),
-                        False,
-                        5 * tt.ones(101)]).eval()
-
-
-def test_alltrue_shape():
-    vals = [True, tt.ones(10), tt.zeros(5)]
-
-    assert alltrue(vals).eval().shape == ()
+from ..distributions import Discrete
+from ..distributions.dist_math import bound_elemwise, bound, factln
+
+
+def test_bound_elemwise():
+    logp = tt.ones((10, 10))
+    cond = tt.ones((10, 10))
+    assert np.all(bound_elemwise(logp, cond).eval() == logp.eval())
+
+    logp = tt.ones((10, 10))
+    cond = tt.zeros((10, 10))
+    assert np.all(bound_elemwise(logp, cond).eval() == (-np.inf * logp).eval())
+
+    logp = tt.ones((10, 10))
+    cond = True
+    assert np.all(bound_elemwise(logp, cond).eval() == logp.eval())
+
+    logp = tt.ones(3)
+    cond = np.array([1, 0, 1])
+    assert not np.all(bound_elemwise(logp, cond).eval() == 1)
+    assert np.prod(bound_elemwise(logp, cond).eval()) == -np.inf
+
+    logp = tt.ones((2, 3))
+    cond = np.array([[1, 1, 1], [1, 0, 1]])
+    assert not np.all(bound_elemwise(logp, cond).eval() == 1)
+    assert np.prod(bound_elemwise(logp, cond).eval()) == -np.inf
+
+
+class MultinomialA(Discrete):
+    def __init__(self, n, p, *args, **kwargs):
+        super(MultinomialA, self).__init__(*args, **kwargs)
+
+        self.n = n
+        self.p = p
+
+    def logp(self, value):
+        n = self.n
+        p = self.p
+
+        return bound(factln(n) - factln(value).sum() + (value * tt.log(p)).sum(),
+                     value >= 0,
+                     0 <= p, p <= 1,
+                     tt.isclose(p.sum(), 1))
+
+
+class MultinomialB(Discrete):
+    def __init__(self, n, p, *args, **kwargs):
+        super(MultinomialB, self).__init__(*args, **kwargs)
+
+        self.n = n
+        self.p = p
+
+    def logp(self, value):
+        n = self.n
+        p = self.p
+
+        return bound(factln(n) - factln(value).sum() + (value * tt.log(p)).sum(),
+                     tt.all(value >= 0),
+                     tt.all(0 <= p), tt.all(p <= 1),
+                     tt.isclose(p.sum(), 1))
+
+
+def test_multinomial_bound():
+
+    x = np.array([1, 5])
+    n = x.sum()
+
+    with pm.Model() as modelA:
+        p_a = pm.Dirichlet('p', np.ones(2))
+        x_obs_a = MultinomialA('x', n, p_a, observed=x)
+
+    with pm.Model() as modelB:
+        p_b = pm.Dirichlet('p', np.ones(2))
+        x_obs_b = MultinomialB('x', n, p_b, observed=x)
+
+    assert np.isclose(modelA.logp({'p_stickbreaking_': [0]}),
+                      modelB.logp({'p_stickbreaking_': [0]}))