Add size everywhere (#2984)

* Add size arg everywhere

* remove print statement

* No size in transforms

* Remove test

Author: ColCarroll

pymc3/distributions/bound.py

@@ -68,17 +68,17 @@ def random(self, point=None, size=None):
         if self.lower is None and self.upper is None:
             return self._wrapped.random(point=point, size=size)
         elif self.lower is not None and self.upper is not None:
-            lower, upper = draw_values([self.lower, self.upper], point=point)
+            lower, upper = draw_values([self.lower, self.upper], point=point, size=size)
             return generate_samples(self._random, lower, upper, point,
                                     dist_shape=self.shape,
                                     size=size)
         elif self.lower is not None:
-            lower = draw_values([self.lower], point=point)
+            lower = draw_values([self.lower], point=point, size=size)
             return generate_samples(self._random, lower, np.inf, point,
                                     dist_shape=self.shape,
                                     size=size)
         else:
-            upper = draw_values([self.upper], point=point)
+            upper = draw_values([self.upper], point=point, size=size)
             return generate_samples(self._random, -np.inf, upper, point,
                                     dist_shape=self.shape,
                                     size=size)

pymc3/distributions/continuous.py

@@ -305,7 +305,7 @@ def __init__(self, mu=0, sd=None, tau=None, **kwargs):
 
     def random(self, point=None, size=None):
         mu, tau, _ = draw_values([self.mu, self.tau, self.sd],
-                                 point=point)
+                                 point=point, size=size)
         return generate_samples(stats.norm.rvs, loc=mu, scale=tau**-0.5,
                                 dist_shape=self.shape,
                                 size=size)
@@ -406,7 +406,7 @@ def __init__(self, sd=None, tau=None, *args, **kwargs):
         assert_negative_support(sd, 'sd', 'HalfNormal')
 
     def random(self, point=None, size=None):
-        sd = draw_values([self.sd], point=point)[0]
+        sd = draw_values([self.sd], point=point, size=size)[0]
         return generate_samples(stats.halfnorm.rvs, loc=0., scale=sd,
                                 dist_shape=self.shape,
                                 size=size)
@@ -547,7 +547,7 @@ def _random(self, mu, lam, alpha, size=None):
 
     def random(self, point=None, size=None):
         mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha],
-                                     point=point)
+                                     point=point, size=size)
         return generate_samples(self._random,
                                 mu, lam, alpha,
                                 dist_shape=self.shape,
@@ -672,7 +672,7 @@ def get_alpha_beta(self, alpha=None, beta=None, mu=None, sd=None):
 
     def random(self, point=None, size=None):
         alpha, beta = draw_values([self.alpha, self.beta],
-                                  point=point)
+                                  point=point, size=size)
         return generate_samples(stats.beta.rvs, alpha, beta,
                                 dist_shape=self.shape,
                                 size=size)
@@ -2239,7 +2239,7 @@ def __init__(self, nu=None, sd=None, *args, **kwargs):
 
     def random(self, point=None, size=None, repeat=None):
         nu, sd = draw_values([self.nu, self.sd],
-                             point=point)
+                             point=point, size=size)
         return generate_samples(stats.rice.rvs, b=nu, scale=sd, loc=0,
                                 dist_shape=self.shape, size=size)
 

pymc3/distributions/discrete.py

@@ -166,7 +166,7 @@ def _random(self, alpha, beta, n, size=None):
 
     def random(self, point=None, size=None):
         alpha, beta, n = \
-            draw_values([self.alpha, self.beta, self.n], point=point)
+            draw_values([self.alpha, self.beta, self.n], point=point, size=size)
         return generate_samples(self._random, alpha=alpha, beta=beta, n=n,
                                 dist_shape=self.shape,
                                 size=size)
@@ -247,7 +247,7 @@ def __init__(self, p=None, logit_p=None, *args, **kwargs):
         self.mode = tt.cast(tround(self.p), 'int8')
 
     def random(self, point=None, size=None):
-        p = draw_values([self.p], point=point)[0]
+        p = draw_values([self.p], point=point, size=size)[0]
         return generate_samples(stats.bernoulli.rvs, p,
                                 dist_shape=self.shape,
                                 size=size)
@@ -343,7 +343,7 @@ def _random(self, q, beta, size=None):
         return np.ceil(np.power(np.log(1 - p) / np.log(q), 1. / beta)) - 1
 
     def random(self, point=None, size=None):
-        q, beta = draw_values([self.q, self.beta], point=point)
+        q, beta = draw_values([self.q, self.beta], point=point, size=size)
 
         return generate_samples(self._random, q, beta,
                                 dist_shape=self.shape,
@@ -410,7 +410,7 @@ def __init__(self, mu, *args, **kwargs):
         self.mode = tt.floor(mu).astype('int32')
 
     def random(self, point=None, size=None):
-        mu = draw_values([self.mu], point=point)[0]
+        mu = draw_values([self.mu], point=point, size=size)[0]
         return generate_samples(stats.poisson.rvs, mu,
                                 dist_shape=self.shape,
                                 size=size)
@@ -490,7 +490,7 @@ def __init__(self, mu, alpha, *args, **kwargs):
         self.mode = tt.floor(mu).astype('int32')
 
     def random(self, point=None, size=None):
-        mu, alpha = draw_values([self.mu, self.alpha], point=point)
+        mu, alpha = draw_values([self.mu, self.alpha], point=point, size=size)
         g = generate_samples(stats.gamma.rvs, alpha, scale=mu / alpha,
                              dist_shape=self.shape,
                              size=size)
@@ -564,7 +564,7 @@ def __init__(self, p, *args, **kwargs):
         self.mode = 1
 
     def random(self, point=None, size=None):
-        p = draw_values([self.p], point=point)[0]
+        p = draw_values([self.p], point=point, size=size)[0]
         return generate_samples(np.random.geometric, p,
                                 dist_shape=self.shape,
                                 size=size)
@@ -636,7 +636,7 @@ def _random(self, lower, upper, size=None):
         return samples
 
     def random(self, point=None, size=None):
-        lower, upper = draw_values([self.lower, self.upper], point=point)
+        lower, upper = draw_values([self.lower, self.upper], point=point, size=size)
         return generate_samples(self._random,
                                 lower, upper,
                                 dist_shape=self.shape,
@@ -714,7 +714,7 @@ def random_choice(k, *args, **kwargs):
             else:
                 return np.random.choice(k, *args, **kwargs)
 
-        p, k = draw_values([self.p, self.k], point=point)
+        p, k = draw_values([self.p, self.k], point=point, size=size)
         return generate_samples(partial(random_choice, np.arange(k)),
                                 p=p,
                                 broadcast_shape=p.shape[:-1] or (1,),
@@ -764,7 +764,7 @@ def __init__(self, c, *args, **kwargs):
         self.mean = self.median = self.mode = self.c = c = tt.as_tensor_variable(c)
 
     def random(self, point=None, size=None):
-        c = draw_values([self.c], point=point)[0]
+        c = draw_values([self.c], point=point, size=size)[0]
         dtype = np.array(c).dtype
 
         def _random(c, dtype=dtype, size=None):
@@ -845,7 +845,7 @@ def __init__(self, psi, theta, *args, **kwargs):
         self.mode = self.pois.mode
 
     def random(self, point=None, size=None):
-        theta, psi = draw_values([self.theta, self.psi], point=point)
+        theta, psi = draw_values([self.theta, self.psi], point=point, size=size)
         g = generate_samples(stats.poisson.rvs, theta,
                              dist_shape=self.shape,
                              size=size)
@@ -938,7 +938,7 @@ def __init__(self, psi, n, p, *args, **kwargs):
         self.mode = self.bin.mode
 
     def random(self, point=None, size=None):
-        n, p, psi = draw_values([self.n, self.p, self.psi], point=point)
+        n, p, psi = draw_values([self.n, self.p, self.psi], point=point, size=size)
         g = generate_samples(stats.binom.rvs, n, p,
                              dist_shape=self.shape,
                              size=size)
@@ -1056,7 +1056,7 @@ def __init__(self, psi, mu, alpha, *args, **kwargs):
 
     def random(self, point=None, size=None):
         mu, alpha, psi = draw_values(
-            [self.mu, self.alpha, self.psi], point=point)
+            [self.mu, self.alpha, self.psi], point=point, size=size)
         g = generate_samples(stats.gamma.rvs, alpha, scale=mu / alpha,
                              dist_shape=self.shape,
                              size=size)

pymc3/distributions/distribution.py

@@ -177,8 +177,8 @@ def __init__(self, shape=(), dtype=None, defaults=('median', 'mean', 'mode'),
 
 class DensityDist(Distribution):
     """Distribution based on a given log density function.
-        
-        A distribution with the passed log density function is created. 
+
+        A distribution with the passed log density function is created.
         Requires a custom random function passed as kwarg `random` to
         enable sampling.
 
@@ -200,7 +200,7 @@ def __init__(self, logp, shape=(), dtype=None, testval=0, random=None, *args, **
             shape, dtype, testval, *args, **kwargs)
         self.logp = logp
         self.rand = random
-    
+
     def random(self, *args, **kwargs):
         if self.rand is not None:
             return self.rand(*args, **kwargs)
@@ -345,9 +345,7 @@ def _draw_value(param, point=None, givens=None, size=None):
     size : int, optional
         Number of samples
     """
-    if isinstance(param, numbers.Number):
-        return param
-    elif isinstance(param, np.ndarray):
+    if isinstance(param, (numbers.Number, np.ndarray)):
         return param
     elif isinstance(param, tt.TensorConstant):
         return param.value
@@ -357,7 +355,7 @@ def _draw_value(param, point=None, givens=None, size=None):
         if point and hasattr(param, 'model') and param.name in point:
             return point[param.name]
         elif hasattr(param, 'random') and param.random is not None:
-            return param.random(point=point, size=None)
+            return param.random(point=point, size=size)
         else:
             if givens:
                 variables, values = list(zip(*givens))

pymc3/distributions/multivariate.py

@@ -234,7 +234,7 @@ def random(self, point=None, size=None):
                 size = [size]
 
         if self._cov_type == 'cov':
-            mu, cov = draw_values([self.mu, self.cov], point=point)
+            mu, cov = draw_values([self.mu, self.cov], point=point, size=size)
             if mu.shape[-1] != cov.shape[-1]:
                 raise ValueError("Shapes for mu and cov don't match")
 
@@ -246,15 +246,15 @@ def random(self, point=None, size=None):
                 return np.nan * np.zeros(size)
             return dist.rvs(size)
         elif self._cov_type == 'chol':
-            mu, chol = draw_values([self.mu, self.chol_cov], point=point)
+            mu, chol = draw_values([self.mu, self.chol_cov], point=point, size=size)
             if mu.shape[-1] != chol[0].shape[-1]:
                 raise ValueError("Shapes for mu and chol don't match")
 
             size.append(mu.shape[-1])
             standard_normal = np.random.standard_normal(size)
             return mu + np.dot(standard_normal, chol.T)
         else:
-            mu, tau = draw_values([self.mu, self.tau], point=point)
+            mu, tau = draw_values([self.mu, self.tau], point=point, size=size)
             if mu.shape[-1] != tau[0].shape[-1]:
                 raise ValueError("Shapes for mu and tau don't match")
 
@@ -338,15 +338,15 @@ def __init__(self, nu, Sigma=None, mu=None, cov=None, tau=None, chol=None,
         self.mean = self.median = self.mode = self.mu = self.mu
 
     def random(self, point=None, size=None):
-        nu, mu = draw_values([self.nu, self.mu], point=point)
+        nu, mu = draw_values([self.nu, self.mu], point=point, size=size)
         if self._cov_type == 'cov':
-            cov, = draw_values([self.cov], point=point)
+            cov, = draw_values([self.cov], point=point, size=size)
             dist = MvNormal.dist(mu=np.zeros_like(mu), cov=cov)
         elif self._cov_type == 'tau':
-            tau, = draw_values([self.tau], point=point)
+            tau, = draw_values([self.tau], point=point, size=size)
             dist = MvNormal.dist(mu=np.zeros_like(mu), tau=tau)
         else:
-            chol, = draw_values([self.chol_cov], point=point)
+            chol, = draw_values([self.chol_cov], point=point, size=size)
             dist = MvNormal.dist(mu=np.zeros_like(mu), chol=chol)
 
         samples = dist.random(point, size)
@@ -422,7 +422,7 @@ def __init__(self, a, transform=transforms.stick_breaking,
                               np.nan)
 
     def random(self, point=None, size=None):
-        a = draw_values([self.a], point=point)[0]
+        a = draw_values([self.a], point=point, size=size)[0]
 
         def _random(a, size=None):
             return stats.dirichlet.rvs(a, None if size == a.shape else size)
@@ -545,7 +545,7 @@ def _random(self, n, p, size=None):
         return randnum.astype(original_dtype)
 
     def random(self, point=None, size=None):
-        n, p = draw_values([self.n, self.p], point=point)
+        n, p = draw_values([self.n, self.p], point=point, size=size)
         samples = generate_samples(self._random, n, p,
                                    dist_shape=self.shape,
                                    size=size)
@@ -679,7 +679,7 @@ def __init__(self, nu, V, *args, **kwargs):
                               np.nan)
 
     def random(self, point=None, size=None):
-        nu, V = draw_values([self.nu, self.V], point=point)
+        nu, V = draw_values([self.nu, self.V], point=point, size=size)
         size= 1 if size is None else size
         return generate_samples(stats.wishart.rvs, np.asscalar(nu), V,
                                     broadcast_shape=(size,))
@@ -1071,7 +1071,7 @@ def _random(self, n, eta, size=None):
         return C.transpose((1, 2, 0))[np.triu_indices(n, k=1)].T
 
     def random(self, point=None, size=None):
-        n, eta = draw_values([self.n, self.eta], point=point)
+        n, eta = draw_values([self.n, self.eta], point=point, size=size)
         size= 1 if size is None else size
         samples = generate_samples(self._random, n, eta,
                                    broadcast_shape=(size,))
@@ -1264,8 +1264,8 @@ def random(self, point=None, size=None):
 
         mu, colchol, rowchol = draw_values(
                                 [self.mu, self.colchol_cov, self.rowchol_cov],
-                                point=point
-                                )
+                                point=point,
+                                size=size)
         standard_normal = np.random.standard_normal(size)
         return mu + np.matmul(rowchol, np.matmul(standard_normal, colchol.T))
 

pymc3/tests/test_distributions_random.py

@@ -97,20 +97,6 @@ def test_random_sample_returns_nd_array(self):
         assert isinstance(mu, np.ndarray)
         assert isinstance(tau, np.ndarray)
 
-    def test_random_sample_returns_correctly(self):
-        # Based on what we discovered in #GH2909
-        with pm.Model():
-            a = pm.Uniform('a', lower=0, upper=1, shape=10)
-            b = pm.Binomial('b', n=1, p=a, shape=10)
-        array_of_uniform = a.random(size=10000).mean(axis=0)
-        array_of_binomial = b.random(size=10000).mean(axis=0)
-        npt.assert_allclose(array_of_uniform, [0.49886929, 0.49949713, 0.49946077, 0.49922606, 0.49927498, 0.50003914,
-                                               0.49980687, 0.50180495, 0.500905, 0.50035121], rtol=1e-2, atol=0)
-        npt.assert_allclose(array_of_binomial, [0.7232, 0.131 , 0.9457, 0.8279, 0.2911, 0.8686, 0.57  , 0.9184,
-                                                0.8177, 0.1625], rtol=1e-2, atol=0)
-        assert isinstance(array_of_binomial, np.ndarray)
-        assert isinstance(array_of_uniform, np.ndarray)
-
 
 class BaseTestCases(object):
     class BaseTestCase(SeededTest):
@@ -741,17 +727,17 @@ def test_lkj(self):
         for n in [2, 10, 50]:
             #pylint: disable=cell-var-from-loop
             shape = n*(n-1)//2
-            
+
             def ref_rand(size, eta):
                 beta = eta - 1 + n/2
                 return (st.beta.rvs(size=(size, shape), a=beta, b=beta)-.5)*2
 
             class TestedLKJCorr (pm.LKJCorr):
-                
+
                 def __init__(self, **kwargs):
                     kwargs.pop('shape', None)
                     super(TestedLKJCorr, self).__init__(
-                            n=n, 
+                            n=n,
                             **kwargs
                     )
 
@@ -767,12 +753,12 @@ def ref_rand(size, w, mu, sd):
 
         pymc3_random(pm.NormalMixture, {'w': Simplex(2),
                      'mu': Domain([[.05, 2.5], [-5., 1.]], edges=(None, None)),
-                     'sd': Domain([[1, 1], [1.5, 2.]], edges=(None, None))}, 
+                     'sd': Domain([[1, 1], [1.5, 2.]], edges=(None, None))},
                      size=1000,
                      ref_rand=ref_rand)
         pymc3_random(pm.NormalMixture, {'w': Simplex(3),
                      'mu': Domain([[-5., 1., 2.5]], edges=(None, None)),
-                     'sd': Domain([[1.5, 2., 3.]], edges=(None, None))}, 
+                     'sd': Domain([[1.5, 2., 3.]], edges=(None, None))},
                      size=1000,
                      ref_rand=ref_rand)