Fix for shape issue in prior sampling for bounded distributions (#3451)

* Fix for shape issue in prior sampling for bounded distributions

* Fixed unrelated test failure in test_normal_scalar

* Fixed test_normal_scalar failure

Author: fonnesbeck

pymc3/distributions/bound.py

@@ -5,11 +5,16 @@
 import theano
 
 from pymc3.distributions.distribution import (
-    Distribution, Discrete, Continuous, draw_values, generate_samples)
+    Distribution,
+    Discrete,
+    Continuous,
+    draw_values,
+    generate_samples,
+)
 from pymc3.distributions import transforms
 from pymc3.distributions.dist_math import bound
 
-__all__ = ['Bound']
+__all__ = ["Bound"]
 
 
 class _Bounded(Distribution):
@@ -23,15 +28,16 @@ def __init__(self, distribution, lower, upper, default, *args, **kwargs):
             for name in defaults:
                 setattr(self, name, getattr(self._wrapped, name))
         else:
-            defaults = ('_default',)
+            defaults = ("_default",)
             self._default = default
 
         super().__init__(
             shape=self._wrapped.shape,
             dtype=self._wrapped.dtype,
             testval=self._wrapped.testval,
             defaults=defaults,
-            transform=self._wrapped.transform)
+            transform=self._wrapped.transform,
+        )
 
     def logp(self, value):
         logp = self._wrapped.logp(value)
@@ -49,49 +55,52 @@ def _random(self, lower, upper, point=None, size=None):
         lower = np.asarray(lower)
         upper = np.asarray(upper)
         if lower.size > 1 or upper.size > 1:
-            raise ValueError('Drawing samples from distributions with '
-                             'array-valued bounds is not supported.')
-        samples = np.zeros(size, dtype=self.dtype).flatten()
-        i, n = 0, len(samples)
-        while i < len(samples):
-            sample = np.atleast_1d(self._wrapped.random(point=point, size=n))
+            raise ValueError(
+                "Drawing samples from distributions with "
+                "array-valued bounds is not supported."
+            )
+        total_size = np.prod(size)
+        samples = []
+        s = 0
+        while s < total_size:
+            sample = np.atleast_1d(
+                self._wrapped.random(point=point, size=total_size)
+            ).flatten()
 
             select = sample[np.logical_and(sample >= lower, sample <= upper)]
-            samples[i:(i + len(select))] = select[:]
-            i += len(select)
-            n -= len(select)
+            samples.append(select)
+            s += len(select)
         if size is not None:
-            return np.reshape(samples, size)
+            return np.reshape(np.concatenate(samples)[:total_size], size)
         else:
-            return samples
+            return samples[0]
 
     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, size=size)
-            return generate_samples(self._random, lower, upper, point,
-                                    dist_shape=self.shape,
-                                    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, size=size)
-            return generate_samples(self._random, lower, np.inf, point,
-                                    dist_shape=self.shape,
-                                    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, size=size)
-            return generate_samples(self._random, -np.inf, upper, point,
-                                    dist_shape=self.shape,
-                                    size=size)
+            return generate_samples(
+                self._random, -np.inf, upper, point, dist_shape=self.shape, size=size
+            )
 
 
 class _DiscreteBounded(_Bounded, Discrete):
-    def __init__(self, distribution, lower, upper,
-                 transform='infer', *args, **kwargs):
-        if transform == 'infer':
+    def __init__(self, distribution, lower, upper, transform="infer", *args, **kwargs):
+        if transform == "infer":
             transform = None
         if transform is not None:
-            raise ValueError('Can not transform discrete variable.')
+            raise ValueError("Can not transform discrete variable.")
 
         if lower is None and upper is None:
             default = None
@@ -103,12 +112,12 @@ def __init__(self, distribution, lower, upper,
             default = lower + 1
 
         super().__init__(
-            distribution, lower, upper,
-            default, *args, transform=transform, **kwargs)
+            distribution, lower, upper, default, *args, transform=transform, **kwargs
+        )
 
 
 class _ContinuousBounded(_Bounded, Continuous):
-    R"""
+    r"""
     An upper, lower or upper+lower bounded distribution
 
     Parameters
@@ -125,16 +134,15 @@ class _ContinuousBounded(_Bounded, Continuous):
         See pymc3.distributions.transforms for more information.
     """
 
-    def __init__(self, distribution, lower, upper,
-                 transform='infer', *args, **kwargs):
-        dtype = kwargs.get('dtype', theano.config.floatX)
+    def __init__(self, distribution, lower, upper, transform="infer", *args, **kwargs):
+        dtype = kwargs.get("dtype", theano.config.floatX)
 
         if lower is not None:
             lower = tt.as_tensor_variable(lower).astype(dtype)
         if upper is not None:
             upper = tt.as_tensor_variable(upper).astype(dtype)
 
-        if transform == 'infer':
+        if transform == "infer":
             if lower is None and upper is None:
                 transform = None
                 default = None
@@ -151,12 +159,12 @@ def __init__(self, distribution, lower, upper,
             default = None
 
         super().__init__(
-            distribution, lower, upper,
-            default, *args, transform=transform, **kwargs)
+            distribution, lower, upper, default, *args, transform=transform, **kwargs
+        )
 
 
 class Bound:
-    R"""
+    r"""
     Create a Bound variable object that can be applied to create
     a new upper, lower, or upper and lower bounded distribution.
 
@@ -207,31 +215,49 @@ def __init__(self, distribution, lower=None, upper=None):
         self.upper = upper
 
     def __call__(self, name, *args, **kwargs):
-        if 'observed' in kwargs:
-            raise ValueError('Observed Bound distributions are not supported. '
-                             'If you want to model truncated data '
-                             'you can use a pm.Potential in combination '
-                             'with the cumulative probability function. See '
-                             'pymc3/examples/censored_data.py for an example.')
-
-        transform = kwargs.pop('transform', 'infer')
+        if "observed" in kwargs:
+            raise ValueError(
+                "Observed Bound distributions are not supported. "
+                "If you want to model truncated data "
+                "you can use a pm.Potential in combination "
+                "with the cumulative probability function. See "
+                "pymc3/examples/censored_data.py for an example."
+            )
+
+        transform = kwargs.pop("transform", "infer")
         if issubclass(self.distribution, Continuous):
-            return _ContinuousBounded(name, self.distribution, self.lower,
-                                    self.upper, transform, *args, **kwargs)
+            return _ContinuousBounded(
+                name,
+                self.distribution,
+                self.lower,
+                self.upper,
+                transform,
+                *args,
+                **kwargs
+            )
         elif issubclass(self.distribution, Discrete):
-            return _DiscreteBounded(name, self.distribution, self.lower,
-                                    self.upper, transform, *args, **kwargs)
+            return _DiscreteBounded(
+                name,
+                self.distribution,
+                self.lower,
+                self.upper,
+                transform,
+                *args,
+                **kwargs
+            )
         else:
-            raise ValueError(
-                'Distribution is neither continuous nor discrete.')
+            raise ValueError("Distribution is neither continuous nor discrete.")
 
     def dist(self, *args, **kwargs):
         if issubclass(self.distribution, Continuous):
             return _ContinuousBounded.dist(
-                self.distribution, self.lower, self.upper, *args, **kwargs)
+                self.distribution, self.lower, self.upper, *args, **kwargs
+            )
 
         elif issubclass(self.distribution, Discrete):
             return _DiscreteBounded.dist(
-                self.distribution, self.lower, self.upper, *args, **kwargs)
+                self.distribution, self.lower, self.upper, *args, **kwargs
+            )
         else:
-            raise ValueError('Distribution is neither continuous nor discrete.')
+            raise ValueError("Distribution is neither continuous nor discrete.")
+