tests passed

Author: ferrine

pymc3/distributions/discrete.py

@@ -6,7 +6,7 @@
 from scipy import stats
 
 from .dist_math import bound, factln, binomln, betaln, logpow
-from .distribution import Discrete, draw_values, generate_samples
+from .distribution import Discrete, draw_values, generate_samples, reshape_sampled
 
 __all__ = ['Binomial',  'BetaBinomial',  'Bernoulli',  'Poisson',
            'NegativeBinomial', 'ConstantDist', 'Constant', 'ZeroInflatedPoisson',
@@ -250,7 +250,7 @@ def random(self, point=None, size=None, repeat=None):
                              dist_shape=self.shape,
                              size=size)
         g[g == 0] = np.finfo(float).eps  # Just in case
-        return stats.poisson.rvs(g)
+        return reshape_sampled(stats.poisson.rvs(g), size, self.shape)
 
     def logp(self, value):
         mu = self.mu
@@ -441,9 +441,11 @@ def logp(self, value):
         c = self.c
         return bound(0, tt.eq(value, c))
 
+
 def ConstantDist(*args, **kwargs):
+    import warnings
     warnings.warn("ConstantDist has been deprecated. In future, use Constant instead.",
-                DeprecationWarning)
+                  DeprecationWarning)
     return Constant(*args, **kwargs)
 
 
@@ -489,7 +491,8 @@ def random(self, point=None, size=None, repeat=None):
         g = generate_samples(stats.poisson.rvs, theta,
                              dist_shape=self.shape,
                              size=size)
-        return g * (np.random.random(np.squeeze(g.shape)) < psi)
+        sampled = g * (np.random.random(np.squeeze(g.shape)) < psi)
+        return reshape_sampled(sampled, size, self.shape)
 
     def logp(self, value):
         return tt.switch(value > 0,
@@ -543,7 +546,8 @@ def random(self, point=None, size=None, repeat=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)
+        sampled = stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
+        return reshape_sampled(sampled, size, self.shape)
 
     def logp(self, value):
         return tt.switch(value > 0,

pymc3/distributions/distribution.py

@@ -10,9 +10,11 @@
 __all__ = ['DensityDist', 'Distribution', 'Continuous',
            'Discrete', 'NoDistribution', 'TensorType', 'draw_values']
 
+
 class _Unpickling(object):
     pass
 
+
 class Distribution(object):
     """Statistical distribution"""
     def __new__(cls, name, *args, **kwargs):
@@ -129,12 +131,10 @@ def __init__(self, logp, shape=(), dtype='float64', testval=0, *args, **kwargs):
 
 
 class MultivariateContinuous(Continuous):
-
     pass
 
 
 class MultivariateDiscrete(Discrete):
-
     pass
 
 
@@ -265,6 +265,22 @@ def broadcast_shapes(*args):
     return tuple(x)
 
 
+def infer_shape(shape):
+    try:
+        shape = tuple(shape or ())
+    except TypeError:  # If size is an int
+        shape = tuple((shape,))
+    except ValueError:  # If size is np.array
+        shape = tuple(shape)
+    return shape
+
+
+def reshape_sampled(sampled, size, dist_shape):
+    dist_shape = infer_shape(dist_shape)
+    repeat_shape = infer_shape(size)
+    return np.reshape(sampled, repeat_shape + dist_shape)
+
+
 def replicate_samples(generator, size, repeats, *args, **kwargs):
     n = int(np.prod(repeats))
     if n == 1:
@@ -326,10 +342,7 @@ def generate_samples(generator, *args, **kwargs):
     else:
         prefix_shape = tuple(dist_shape)
 
-    try:
-        repeat_shape = tuple(size or ())
-    except TypeError:  # If size is an int
-        repeat_shape = tuple((size,))
+    repeat_shape = infer_shape(size)
 
     if broadcast_shape == (1,) and prefix_shape == ():
         if size is not None:
@@ -342,13 +355,9 @@ def generate_samples(generator, *args, **kwargs):
                                         broadcast_shape,
                                         repeat_shape + prefix_shape,
                                         *args, **kwargs)
-            if broadcast_shape == (1,) and not prefix_shape == ():
-                samples = np.reshape(samples, repeat_shape + prefix_shape)
         else:
             samples = replicate_samples(generator,
                                         broadcast_shape,
                                         prefix_shape,
                                         *args, **kwargs)
-            if broadcast_shape == (1,):
-                samples = np.reshape(samples, prefix_shape)
-    return samples
+    return reshape_sampled(samples, size, dist_shape)

pymc3/tests/test_distributions_random.py

@@ -178,8 +178,10 @@ def test_different_shapes_and_sample_sizes(self):
                         except TypeError:
                             s = [size]
                     s.extend(shape)
-                    expected.append(tuple(s))
-                    actual.append(self.sample_random_variable(rv, size).shape)
+                    e = tuple(s)
+                    a = self.sample_random_variable(rv, size).shape
+                    expected.append(e)
+                    actual.append(a)
             self.assertSequenceEqual(expected, actual)
 
 
@@ -332,8 +334,8 @@ class TestCategorical(BaseTestCases.BaseTestCase):
     distribution = pm.Categorical
     params = {'p': np.ones(BaseTestCases.BaseTestCase.shape)}
 
-    def get_random_variable(self, shape, with_vector_params=False):  # don't transform categories
-        return super(TestCategorical, self).get_random_variable(shape, with_vector_params=False)
+    def get_random_variable(self, shape, with_vector_params=False, **kwargs):  # don't transform categories
+        return super(TestCategorical, self).get_random_variable(shape, with_vector_params=False, **kwargs)
 
 
 @attr('scalar_parameter_samples')