Fixed pymc-devs#3310. Added broadcast_distribution_samples, which helps broadcasting multiple rvs calls with different size and distribution parameter shapes. Added shape guards to other continuous distributions.

Author: lucianopaz

RELEASE-NOTES.md

@@ -36,6 +36,8 @@ This will be the last release to support Python 2.
 - Fixed `Rice` distribution, which inconsistently mixed two parametrizations (#3286).
 - `Rice` distribution now accepts multiple parameters and observations and is usable with NUTS (#3289).
 - `sample_posterior_predictive` no longer calls `draw_values` to initialize the shape of the ppc trace. This called could lead to `ValueError`'s when sampling the ppc from a model with `Flat` or `HalfFlat` prior distributions (Fix issue #3294).
+- Added the `broadcast_distribution_samples` function that helps broadcasting arrays of drawn samples, taking into account the requested `size` and the inferred distribution shape. This sometimes is needed by distributions that call several `rvs` separately within their `random` method, such as the `ZeroInflatedPoisson` (Fix issue #3310).
+- The `Wald`, `Kumaraswamy`, `LogNormal`, `Pareto`, `Cauchy`, `HalfCauchy`, `Weibull` and `ExGaussian` distributions `random` method used a hidden `_random` function that was written with scalars in mind. This could potentially lead to artificial correlations between random draws. Added shape guards and broadcasting of the distribution samples to prevent this (Similar to issue #3310).
 
 
 ### Deprecations

pymc3/distributions/continuous.py

@@ -21,7 +21,8 @@
     alltrue_elemwise, betaln, bound, gammaln, i0e, incomplete_beta, logpow,
     normal_lccdf, normal_lcdf, SplineWrapper, std_cdf, zvalue,
 )
-from .distribution import Continuous, draw_values, generate_samples
+from .distribution import (Continuous, draw_values, generate_samples,
+                           to_tuple, broadcast_distribution_samples)
 
 __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'TruncatedNormal', 'Beta',
            'Kumaraswamy', 'Exponential', 'Laplace', 'StudentT', 'Cauchy',
@@ -937,10 +938,15 @@ def get_mu_lam_phi(self, mu, lam, phi):
                          'mu and lam, mu and phi, or lam and phi.')
 
     def _random(self, mu, lam, alpha, size=None):
-        v = np.random.normal(size=size)**2
+        _size = alpha.shape
+        if size is not None:
+            size = to_tuple(size)
+            if _size[:len(size)] != size:
+                _size = size + _size
+        v = np.random.normal(size=_size)**2
         value = (mu + (mu**2) * v / (2. * lam) - mu / (2. * lam)
                  * np.sqrt(4. * mu * lam * v + (mu * v)**2))
-        z = np.random.uniform(size=size)
+        z = np.random.uniform(size=_size)
         i = np.floor(z - mu / (mu + value)) * 2 + 1
         value = (value**-i) * (mu**(i + 1))
         return value + alpha
@@ -964,6 +970,8 @@ def random(self, point=None, size=None):
         """
         mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha],
                                      point=point, size=size)
+        mu, lam, alpha = broadcast_distribution_samples([mu, lam, alpha],
+                                                        size=size)
         return generate_samples(self._random,
                                 mu, lam, alpha,
                                 dist_shape=self.shape,
@@ -1270,7 +1278,12 @@ def __init__(self, a, b, *args, **kwargs):
         assert_negative_support(b, 'b', 'Kumaraswamy')
 
     def _random(self, a, b, size=None):
-        u = np.random.uniform(size=size)
+        _size = a.shape
+        if size is not None:
+            size = to_tuple(size)
+            if _size[:len(size)] != size:
+                _size = size + _size
+        u = np.random.uniform(size=_size)
         return (1 - (1 - u) ** (1 / b)) ** (1 / a)
 
     def random(self, point=None, size=None):
@@ -1292,6 +1305,7 @@ def random(self, point=None, size=None):
         """
         a, b = draw_values([self.a, self.b],
                            point=point, size=size)
+        a, b = broadcast_distribution_samples([a, b], size=size)
         return generate_samples(self._random, a, b,
                                 dist_shape=self.shape,
                                 size=size)
@@ -1644,7 +1658,12 @@ def __init__(self, mu=0, sd=None, tau=None, *args, **kwargs):
         assert_negative_support(sd, 'sd', 'Lognormal')
 
     def _random(self, mu, tau, size=None):
-        samples = np.random.normal(size=size)
+        _size = tau.shape
+        if size is not None:
+            size = to_tuple(size)
+            if _size[:len(size)] != size:
+                _size = size + _size
+        samples = np.random.normal(size=_size)
         return np.exp(mu + (tau**-0.5) * samples)
 
     def random(self, point=None, size=None):
@@ -1665,6 +1684,7 @@ def random(self, point=None, size=None):
         array
         """
         mu, tau = draw_values([self.mu, self.tau], point=point, size=size)
+        mu, tau = broadcast_distribution_samples([mu, tau], size=size)
         return generate_samples(self._random, mu, tau,
                                 dist_shape=self.shape,
                                 size=size)
@@ -1930,7 +1950,12 @@ def __init__(self, alpha, m, transform='lowerbound', *args, **kwargs):
         super(Pareto, self).__init__(transform=transform, *args, **kwargs)
 
     def _random(self, alpha, m, size=None):
-        u = np.random.uniform(size=size)
+        _size = alpha.shape
+        if size is not None:
+            size = to_tuple(size)
+            if _size[:len(size)] != size:
+                _size = size + _size
+        u = np.random.uniform(size=_size)
         return m * (1. - u)**(-1. / alpha)
 
     def random(self, point=None, size=None):
@@ -1952,6 +1977,7 @@ def random(self, point=None, size=None):
         """
         alpha, m = draw_values([self.alpha, self.m],
                                point=point, size=size)
+        alpha, m = broadcast_distribution_samples([alpha, m], size=size)
         return generate_samples(self._random, alpha, m,
                                 dist_shape=self.shape,
                                 size=size)
@@ -2054,7 +2080,12 @@ def __init__(self, alpha, beta, *args, **kwargs):
         assert_negative_support(beta, 'beta', 'Cauchy')
 
     def _random(self, alpha, beta, size=None):
-        u = np.random.uniform(size=size)
+        _size = alpha.shape
+        if size is not None:
+            size = to_tuple(size)
+            if _size[:len(size)] != size:
+                _size = size + _size
+        u = np.random.uniform(size=_size)
         return alpha + beta * np.tan(np.pi * (u - 0.5))
 
     def random(self, point=None, size=None):
@@ -2076,6 +2107,7 @@ def random(self, point=None, size=None):
         """
         alpha, beta = draw_values([self.alpha, self.beta],
                                   point=point, size=size)
+        alpha, beta = broadcast_distribution_samples([alpha, beta], size=size)
         return generate_samples(self._random, alpha, beta,
                                 dist_shape=self.shape,
                                 size=size)
@@ -2163,7 +2195,12 @@ def __init__(self, beta, *args, **kwargs):
         assert_negative_support(beta, 'beta', 'HalfCauchy')
 
     def _random(self, beta, size=None):
-        u = np.random.uniform(size=size)
+        _size = beta.shape
+        if size is not None:
+            size = to_tuple(size)
+            if _size[:len(size)] != size:
+                _size = size + _size
+        u = np.random.uniform(size=_size)
         return beta * np.abs(np.tan(np.pi * (u - 0.5)))
 
     def random(self, point=None, size=None):
@@ -2637,9 +2674,15 @@ def random(self, point=None, size=None):
         """
         alpha, beta = draw_values([self.alpha, self.beta],
                                   point=point, size=size)
+        alpha, beta = broadcast_distribution_samples([alpha, beta], size=size)
 
         def _random(a, b, size=None):
-            return b * (-np.log(np.random.uniform(size=size)))**(1 / a)
+            _size = a.shape
+            if size is not None:
+                size = to_tuple(size)
+                if _size[:len(size)] != size:
+                    _size = size + _size
+            return b * (-np.log(np.random.uniform(size=_size)))**(1 / a)
 
         return generate_samples(_random, alpha, beta,
                                 dist_shape=self.shape,
@@ -2921,6 +2964,8 @@ def random(self, point=None, size=None):
         """
         mu, sigma, nu = draw_values([self.mu, self.sigma, self.nu],
                                     point=point, size=size)
+        mu, sigma, nu = broadcast_distribution_samples([mu, sigma, nu],
+                                                       size=size)
 
         def _random(mu, sigma, nu, size=None):
             return (np.random.normal(mu, sigma, size=size)

pymc3/distributions/discrete.py

@@ -6,7 +6,8 @@
 
 from pymc3.util import get_variable_name
 from .dist_math import bound, factln, binomln, betaln, logpow, random_choice
-from .distribution import Discrete, draw_values, generate_samples
+from .distribution import (Discrete, draw_values, generate_samples,
+                           broadcast_distribution_samples)
 from pymc3.math import tround, sigmoid, logaddexp, logit, log1pexp
 
 
@@ -345,6 +346,7 @@ def _ppf(self, p):
 
     def _random(self, q, beta, size=None):
         p = np.random.uniform(size=size)
+        p, q, beta = broadcast_distribution_samples([p, q, beta], size=size)
 
         return np.ceil(np.power(np.log(1 - p) / np.log(q), 1. / beta)) - 1
 
@@ -847,7 +849,8 @@ def random(self, point=None, size=None):
         g = generate_samples(stats.poisson.rvs, theta,
                              dist_shape=self.shape,
                              size=size)
-        return g * (np.random.random(np.squeeze(g.shape)) < psi)
+        g, psi = broadcast_distribution_samples([g, psi], size=size)
+        return g * (np.random.random(g.shape) < psi)
 
     def logp(self, value):
         psi = self.psi
@@ -939,7 +942,8 @@ def random(self, point=None, size=None):
         g = generate_samples(stats.binom.rvs, n, p,
                              dist_shape=self.shape,
                              size=size)
-        return g * (np.random.random(np.squeeze(g.shape)) < psi)
+        g, psi = broadcast_distribution_samples([g, psi], size=size)
+        return g * (np.random.random(g.shape) < psi)
 
     def logp(self, value):
         psi = self.psi
@@ -1057,7 +1061,8 @@ def random(self, point=None, size=None):
                              dist_shape=self.shape,
                              size=size)
         g[g == 0] = np.finfo(float).eps  # Just in case
-        return stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
+        g, psi = broadcast_distribution_samples([g, psi], size=size)
+        return stats.poisson.rvs(g) * (np.random.random(g.shape) < psi)
 
     def logp(self, value):
         alpha = self.alpha

pymc3/distributions/distribution.py

@@ -640,3 +640,30 @@ def generate_samples(generator, *args, **kwargs):
     if one_d and samples.shape[-1] == 1:
         samples = samples.reshape(samples.shape[:-1])
     return np.asarray(samples)
+
+
+def broadcast_distribution_samples(samples, size=None):
+    if size is None:
+        return np.broadcast_arrays(*samples)
+    _size = to_tuple(size)
+    try:
+        broadcasted_samples = np.broadcast_arrays(*samples)
+    except ValueError:
+        # Raw samples shapes
+        p_shapes = [p.shape for p in samples]
+        # samples shapes without the size prepend
+        sp_shapes = [s[len(_size):] if _size == s[:len(_size)] else s
+                     for s in p_shapes]
+        broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape
+        broadcasted_samples = []
+        for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
+            if _size == p_shape[:len(_size)]:
+                slicer_head = [slice(None)] * len(_size)
+            else:
+                slicer_head = [np.newaxis] * len(_size)
+            slicer_tail = ([np.newaxis] * (len(broadcast_shape) -
+                                           len(sp_shape)) +
+                           [slice(None)] * len(sp_shape))
+            broadcasted_samples.append(param[tuple(slicer_head + slicer_tail)])
+        broadcasted_samples = np.broadcast_arrays(*broadcasted_samples)
+    return broadcasted_samples

pymc3/tests/test_sampling.py

@@ -467,3 +467,13 @@ def test_shape_edgecase(self):
             x = pm.Normal('x', mu=mu, sd=sd, shape=5)
             prior = pm.sample_prior_predictive(10)
         assert prior['mu'].shape == (10, 5)
+
+    def test_zeroinflatedpoisson(self):
+        with pm.Model():
+            theta = pm.Beta('theta', alpha=1, beta=1)
+            psi = pm.HalfNormal('psi', sd=1)
+            pm.ZeroInflatedPoisson('suppliers', psi=psi, theta=theta, shape=20)
+            gen_data = pm.sample_prior_predictive(samples=5000)
+            assert gen_data['theta'].shape == (5000,)
+            assert gen_data['psi'].shape == (5000,)
+            assert gen_data['suppliers'].shape == (5000, 20)