Add sample_generative function

RELEASE-NOTES.md

@@ -18,6 +18,8 @@
   avoids pickleing issues on UNIX, and allows us to show a progress bar
   for all chains. If parallel sampling is interrupted, we now return
   partial results.
+- Add `sample_prior_predictive` which allows for efficient sampling from
+  the unconditioned model.
 
 ### Fixes
 

pymc3/distributions/bound.py

@@ -54,7 +54,8 @@ def _random(self, lower, upper, point=None, size=None):
         samples = np.zeros(size, dtype=self.dtype).flatten()
         i, n = 0, len(samples)
         while i < len(samples):
-            sample = self._wrapped.random(point=point, size=n)
+            sample = np.atleast_1d(self._wrapped.random(point=point, size=n))
+
             select = sample[np.logical_and(sample >= lower, sample <= upper)]
             samples[i:(i + len(select))] = select[:]
             i += len(select)

pymc3/distributions/discrete.py

@@ -1,4 +1,3 @@
-from functools import partial
 import numpy as np
 import theano
 import theano.tensor as tt
@@ -7,7 +6,7 @@
 
 from pymc3.util import get_variable_name
 from .dist_math import bound, factln, binomln, betaln, logpow
-from .distribution import Discrete, draw_values, generate_samples, reshape_sampled
+from .distribution import Discrete, draw_values, generate_samples
 from pymc3.math import tround, sigmoid, logaddexp, logit, log1pexp
 
 
@@ -154,13 +153,20 @@ def __init__(self, alpha, beta, n, *args, **kwargs):
 
     def _random(self, alpha, beta, n, size=None):
         size = size or 1
-        p = np.atleast_1d(stats.beta.rvs(a=alpha, b=beta, size=np.prod(size)))
+        p = stats.beta.rvs(a=alpha, b=beta, size=size).flatten()
         # Sometimes scipy.beta returns nan. Ugh.
         while np.any(np.isnan(p)):
             i = np.isnan(p)
             p[i] = stats.beta.rvs(a=alpha, b=beta, size=np.sum(i))
         # Sigh...
-        _n, _p, _size = np.atleast_1d(n).flatten(), p.flatten(), np.prod(size)
+        _n, _p, _size = np.atleast_1d(n).flatten(), p.flatten(), p.shape[0]
+
+        quotient, remainder = divmod(_p.shape[0], _n.shape[0])
+        if remainder != 0:
+            raise TypeError('n has a bad size! Was cast to {}, must evenly divide {}'.format(
+                _n.shape[0], _p.shape[0]))
+        if quotient != 1:
+            _n = np.tile(_n, quotient)
         samples = np.reshape(stats.binom.rvs(n=_n, p=_p, size=_size), size)
         return samples
 
@@ -186,7 +192,7 @@ def _repr_latex_(self, name=None, dist=None):
         alpha = dist.alpha
         beta = dist.beta
         name = r'\text{%s}' % name
-        return r'${} \sim \text{{NegativeBinomial}}(\mathit{{alpha}}={},~\mathit{{beta}}={})$'.format(name,
+        return r'${} \sim \text{{BetaBinomial}}(\mathit{{alpha}}={},~\mathit{{beta}}={})$'.format(name,
                                                 get_variable_name(alpha),
                                                 get_variable_name(beta))
 
@@ -495,7 +501,7 @@ def random(self, point=None, size=None):
                              dist_shape=self.shape,
                              size=size)
         g[g == 0] = np.finfo(float).eps  # Just in case
-        return reshape_sampled(stats.poisson.rvs(g), size, self.shape)
+        return np.asarray(stats.poisson.rvs(g)).reshape(g.shape)
 
     def logp(self, value):
         mu = self.mu
@@ -700,22 +706,23 @@ def __init__(self, p, *args, **kwargs):
             self.k = tt.shape(p)[-1].tag.test_value
         except AttributeError:
             self.k = tt.shape(p)[-1]
-        self.p = p = tt.as_tensor_variable(p)
+        p = tt.as_tensor_variable(p)
         self.p = (p.T / tt.sum(p, -1)).T
         self.mode = tt.argmax(p)
 
-    def random(self, point=None, size=None):
-        def random_choice(k, *args, **kwargs):
-            if len(kwargs['p'].shape) > 1:
-                return np.asarray(
-                    [np.random.choice(k, p=p)
-                     for p in kwargs['p']]
-                )
-            else:
-                return np.random.choice(k, *args, **kwargs)
+    def _random(self, k, p, size=None):
+        if len(p.shape) > 1:
+            return np.asarray(
+                [np.random.choice(k, p=pp, size=size)
+                    for pp in p]
+            )
+        else:
+            return np.asarray(np.random.choice(k, p=p, size=size))
 
+    def random(self, point=None, size=None):
         p, k = draw_values([self.p, self.k], point=point, size=size)
-        return generate_samples(partial(random_choice, np.arange(k)),
+        return generate_samples(self._random,
+                                k=k,
                                 p=p,
                                 broadcast_shape=p.shape[:-1] or (1,),
                                 dist_shape=self.shape,
@@ -849,8 +856,7 @@ def random(self, point=None, size=None):
         g = generate_samples(stats.poisson.rvs, theta,
                              dist_shape=self.shape,
                              size=size)
-        sampled = g * (np.random.random(np.squeeze(g.shape)) < psi)
-        return reshape_sampled(sampled, size, self.shape)
+        return g * (np.random.random(np.squeeze(g.shape)) < psi)
 
     def logp(self, value):
         psi = self.psi
@@ -942,8 +948,7 @@ def random(self, point=None, size=None):
         g = generate_samples(stats.binom.rvs, n, p,
                              dist_shape=self.shape,
                              size=size)
-        sampled = g * (np.random.random(np.squeeze(g.shape)) < psi)
-        return reshape_sampled(sampled, size, self.shape)
+        return g * (np.random.random(np.squeeze(g.shape)) < psi)
 
     def logp(self, value):
         psi = self.psi
@@ -1061,8 +1066,7 @@ def random(self, point=None, size=None):
                              dist_shape=self.shape,
                              size=size)
         g[g == 0] = np.finfo(float).eps  # Just in case
-        sampled = stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
-        return reshape_sampled(sampled, size, self.shape)
+        return stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
 
     def logp(self, value):
         alpha = self.alpha

pymc3/distributions/distribution.py

@@ -1,4 +1,6 @@
+import collections
 import numbers
+
 import numpy as np
 import theano.tensor as tt
 from theano import function
@@ -254,7 +256,7 @@ def draw_values(params, point=None, size=None):
 
     # Init givens and the stack of nodes to try to `_draw_value` from
     givens = {}
-    stored = set([])  # Some nodes
+    stored = set()  # Some nodes
     stack = list(leaf_nodes.values())  # A queue would be more appropriate
     while stack:
         next_ = stack.pop(0)
@@ -279,13 +281,14 @@ def draw_values(params, point=None, size=None):
             # The named node's children givens values must also be taken
             # into account.
             children = named_nodes_children[next_]
-            temp_givens = [givens[k] for k in givens.keys() if k in children]
+            temp_givens = [givens[k] for k in givens if k in children]
             try:
                 # This may fail for autotransformed RVs, which don't
                 # have the random method
                 givens[next_.name] = (next_, _draw_value(next_,
                                                          point=point,
-                                                         givens=temp_givens, size=size))
+                                                         givens=temp_givens,
+                                                         size=size))
                 stored.add(next_.name)
             except theano.gof.fg.MissingInputError:
                 # The node failed, so we must add the node's parents to
@@ -295,10 +298,31 @@ def draw_values(params, point=None, size=None):
                               if node is not None and
                               node.name not in stored and
                               node not in params])
-    values = []
-    for param in params:
-        values.append(_draw_value(param, point=point, givens=givens.values(), size=size))
-    return values
+
+    # the below makes sure the graph is evaluated in order
+    # test_distributions_random::TestDrawValues::test_draw_order fails without it
+    params = dict(enumerate(params))  # some nodes are not hashable
+    evaluated = {}
+    to_eval = set()
+    missing_inputs = set(params)
+    while to_eval or missing_inputs:
+        if to_eval == missing_inputs:
+            raise ValueError('Cannot resolve inputs for {}'.format([str(params[j]) for j in to_eval]))
+        to_eval = set(missing_inputs)
+        missing_inputs = set()
+        for param_idx in to_eval:
+            param = params[param_idx]
+            if hasattr(param, 'name') and param.name in givens:
+                evaluated[param_idx] = givens[param.name][1]
+            else:
+                try:  # might evaluate in a bad order,
+                    evaluated[param_idx] = _draw_value(param, point=point, givens=givens.values(), size=size)
+                    if isinstance(param, collections.Hashable) and named_nodes_parents.get(param):
+                        givens[param.name] = (param, evaluated[param_idx])
+                except theano.gof.fg.MissingInputError:
+                    missing_inputs.add(param_idx)
+
+    return [evaluated[j] for j in params] # set the order back
 
 
 @memoize
@@ -356,43 +380,26 @@ def _draw_value(param, point=None, givens=None, size=None):
             return point[param.name]
         elif hasattr(param, 'random') and param.random is not None:
             return param.random(point=point, size=size)
+        elif (hasattr(param, 'distribution') and
+                hasattr(param.distribution, 'random') and
+                param.distribution.random is not None):
+            return param.distribution.random(point=point, size=size)
         else:
             if givens:
                 variables, values = list(zip(*givens))
             else:
                 variables = values = []
             func = _compile_theano_function(param, variables)
-            return func(*values)
+            if size and values and not all(var.dshape == val.shape for var, val in zip(variables, values)):
+                return np.array([func(*v) for v in zip(*values)])
+            else:
+                return func(*values)
     else:
         raise ValueError('Unexpected type in draw_value: %s' % type(param))
 
 
-def broadcast_shapes(*args):
-    """Return the shape resulting from broadcasting multiple shapes.
-    Represents numpy's broadcasting rules.
-
-    Parameters
-    ----------
-    *args : array-like of int
-        Tuples or arrays or lists representing the shapes of arrays to be broadcast.
-
-    Returns
-    -------
-    Resulting shape or None if broadcasting is not possible.
-    """
-    x = list(np.atleast_1d(args[0])) if args else ()
-    for arg in args[1:]:
-        y = list(np.atleast_1d(arg))
-        if len(x) < len(y):
-            x, y = y, x
-        x[-len(y):] = [j if i == 1 else i if j == 1 else i if i == j else 0
-                       for i, j in zip(x[-len(y):], y)]
-        if not all(x):
-            return None
-    return tuple(x)
-
-
-def infer_shape(shape):
+def to_tuple(shape):
+    """Convert ints, arrays, and Nones to tuples"""
     try:
         shape = tuple(shape or ())
     except TypeError:  # If size is an int
@@ -401,27 +408,14 @@ def infer_shape(shape):
         shape = tuple(shape)
     return shape
 
-
-def reshape_sampled(sampled, size, dist_shape):
-    dist_shape = infer_shape(dist_shape)
-    repeat_shape = infer_shape(size)
-
-    if np.size(sampled) == 1 or repeat_shape or dist_shape:
-        return np.reshape(sampled, repeat_shape + dist_shape)
-    else:
-        return sampled
-
-
-def replicate_samples(generator, size, repeats, *args, **kwargs):
-    n = int(np.prod(repeats))
-    if n == 1:
-        samples = generator(size=size, *args, **kwargs)
-    else:
-        samples = np.array([generator(size=size, *args, **kwargs)
-                            for _ in range(n)])
-        samples = np.reshape(samples, tuple(repeats) + tuple(size))
-    return samples
-
+def _is_one_d(dist_shape):
+    if hasattr(dist_shape, 'dshape') and dist_shape.dshape in ((), (0,), (1,)):
+        return True
+    elif hasattr(dist_shape, 'shape') and dist_shape.shape in ((), (0,), (1,)):
+        return True
+    elif dist_shape == ():
+        return True
+    return False
 
 def generate_samples(generator, *args, **kwargs):
     """Generate samples from the distribution of a random variable.
@@ -453,42 +447,60 @@ def generate_samples(generator, *args, **kwargs):
     Any remaining *args and **kwargs are passed on to the generator function.
     """
     dist_shape = kwargs.pop('dist_shape', ())
+    one_d = _is_one_d(dist_shape)
     size = kwargs.pop('size', None)
     broadcast_shape = kwargs.pop('broadcast_shape', None)
-    params = args + tuple(kwargs.values())
-
-    if broadcast_shape is None:
-        broadcast_shape = broadcast_shapes(*[np.atleast_1d(p).shape for p in params
-                                             if not isinstance(p, tuple)])
-    if broadcast_shape == ():
-        broadcast_shape = (1,)
+    if size is None:
+        size = 1
 
     args = tuple(p[0] if isinstance(p, tuple) else p for p in args)
+
     for key in kwargs:
         p = kwargs[key]
         kwargs[key] = p[0] if isinstance(p, tuple) else p
 
-    if np.all(dist_shape[-len(broadcast_shape):] == broadcast_shape):
-        prefix_shape = tuple(dist_shape[:-len(broadcast_shape)])
-    else:
-        prefix_shape = tuple(dist_shape)
-
-    repeat_shape = infer_shape(size)
-
-    if broadcast_shape == (1,) and prefix_shape == ():
-        if size is not None:
-            samples = generator(size=size, *args, **kwargs)
+    if broadcast_shape is None:
+        inputs = args + tuple(kwargs.values())
+        broadcast_shape = np.broadcast(*inputs).shape  # size of generator(size=1)
+
+    dist_shape = to_tuple(dist_shape)
+    broadcast_shape = to_tuple(broadcast_shape)
+    size_tup = to_tuple(size)
+
+    # All inputs are scalars, end up size (size_tup, dist_shape)
+    if broadcast_shape in {(), (0,), (1,)}:
+        samples = generator(size=size_tup + dist_shape, *args, **kwargs)
+    # Inputs already have the right shape. Just get the right size.
+    elif broadcast_shape[-len(dist_shape):] == dist_shape or len(dist_shape) == 0:
+        if size == 1 or (broadcast_shape == size_tup + dist_shape):
+            samples = generator(size=broadcast_shape, *args, **kwargs)
+        elif dist_shape == broadcast_shape:
+            samples = generator(size=size_tup + dist_shape, *args, **kwargs)
         else:
-            samples = generator(size=1, *args, **kwargs)
+            samples = None
+    # Args have been broadcast correctly, can just ask for the right shape out
+    elif dist_shape[-len(broadcast_shape):] == broadcast_shape:
+        samples = generator(size=size_tup + dist_shape, *args, **kwargs)
+    # Inputs have the right size, have to manually broadcast to the right dist_shape
+    elif broadcast_shape[:len(size_tup)] == size_tup:
+        suffix = broadcast_shape[len(size_tup):] + dist_shape
+        samples = [generator(*args, **kwargs).reshape(size_tup + (1,)) for _ in range(np.prod(suffix, dtype=int))]
+        samples = np.hstack(samples).reshape(size_tup + suffix)
     else:
-        if size is not None:
-            samples = replicate_samples(generator,
-                                        broadcast_shape,
-                                        repeat_shape + prefix_shape,
-                                        *args, **kwargs)
-        else:
-            samples = replicate_samples(generator,
-                                        broadcast_shape,
-                                        prefix_shape,
-                                        *args, **kwargs)
-    return reshape_sampled(samples, size, dist_shape)
+        samples = None
+
+    if samples is None:
+        raise TypeError('''Attempted to generate values with incompatible shapes:
+            size: {size}
+            dist_shape: {dist_shape}
+            broadcast_shape: {broadcast_shape}
+        '''.format(size=size, dist_shape=dist_shape, broadcast_shape=broadcast_shape))
+
+    # reshape samples here
+    if samples.shape[0] == 1 and size == 1:
+        if len(samples.shape) > len(dist_shape) and samples.shape[-len(dist_shape):] == dist_shape:
+            samples = samples.reshape(samples.shape[1:])
+
+    if one_d and samples.shape[-1] == 1:
+        samples = samples.reshape(samples.shape[:-1])
+    return np.asarray(samples)