Fixed mixture random method (and also multinomial when using a Cholesky covariance matrix).

Author: lucianopaz

pymc3/distributions/distribution.py

@@ -262,27 +262,6 @@ def __init__(self):
         self.drawn_vars = dict()
 
 
-class _DrawValuesContextDetacher(_DrawValuesContext,
-                                 metaclass=InitContextMeta):
-    """
-    Context manager that starts a new drawn variables context copying the
-    parent's context drawn_vars dict. The following changes do not affect the
-    parent contexts but do affect the subsequent calls. This can be used to
-    iterate the same random method many times to get different results, while
-    respecting the drawn variables from previous contexts.
-    """
-    def __new__(cls, *args, **kwargs):
-        return super().__new__(cls)
-
-    def __init__(self):
-        self.drawn_vars = self.drawn_vars.copy()
-
-    def update_parent(self):
-        parent = self.parent
-        if parent is not None:
-            parent.drawn_vars.update(self.drawn_vars)
-
-
 def is_fast_drawable(var):
     return isinstance(var, (numbers.Number,
                             np.ndarray,
@@ -660,20 +639,30 @@ def generate_samples(generator, *args, **kwargs):
             samples = generator(size=broadcast_shape, *args, **kwargs)
         elif dist_shape == broadcast_shape:
             samples = generator(size=size_tup + dist_shape, *args, **kwargs)
-        elif len(dist_shape) == 0 and size_tup and broadcast_shape[:len(size_tup)] == size_tup:
-            # Input's dist_shape is scalar, but it has size repetitions.
-            # So now the size matches but we have to manually broadcast to
-            # the right dist_shape
-            samples = [generator(*args, **kwargs)]
-            if samples[0].shape == broadcast_shape:
-                samples = samples[0]
+        elif len(dist_shape) == 0 and size_tup and broadcast_shape:
+            # There is no dist_shape (scalar distribution) but the parameters
+            # broadcast shape and size_tup determine the size to provide to
+            # the generator
+            if broadcast_shape[:len(size_tup)] == size_tup:
+                # Input's dist_shape is scalar, but it has size repetitions.
+                # So now the size matches but we have to manually broadcast to
+                # the right dist_shape
+                samples = [generator(*args, **kwargs)]
+                if samples[0].shape == broadcast_shape:
+                    samples = samples[0]
+                else:
+                    suffix = broadcast_shape[len(size_tup):] + dist_shape
+                    samples.extend([generator(*args, **kwargs).
+                                    reshape(broadcast_shape)[..., np.newaxis]
+                                    for _ in range(np.prod(suffix,
+                                                           dtype=int) - 1)])
+                    samples = np.hstack(samples).reshape(size_tup + suffix)
             else:
-                suffix = broadcast_shape[len(size_tup):] + dist_shape
-                samples.extend([generator(*args, **kwargs).
-                                reshape(broadcast_shape)[..., np.newaxis]
-                                for _ in range(np.prod(suffix,
-                                                       dtype=int) - 1)])
-                samples = np.hstack(samples).reshape(size_tup + suffix)
+                # The parameter shape is given, but we have to concatenate it
+                # with the size tuple
+                samples = generator(size=size_tup + broadcast_shape,
+                                    *args,
+                                    **kwargs)
         else:
             samples = None
     # Args have been broadcast correctly, can just ask for the right shape out

pymc3/distributions/mixture.py

@@ -9,8 +9,7 @@
 from .distribution import (Discrete, Distribution, draw_values,
                            generate_samples, _DrawValuesContext,
                            _DrawValuesContextBlocker, to_tuple,
-                           broadcast_distribution_samples,
-                           _DrawValuesContextDetacher)
+                           broadcast_distribution_samples)
 from .continuous import get_tau_sigma, Normal
 from ..theanof import _conversion_map
 
@@ -192,7 +191,11 @@ def generator(*args, **kwargs):
                     # differently from scipy.*.rvs, and generate_samples
                     # follows the latter usage pattern. For this reason we
                     # decorate (horribly hack) the size kwarg of
-                    # comp_dist.random
+                    # comp_dist.random. We also have to disable pylint W0640
+                    # because comp_dist is changed at each iteration of the
+                    # for loop, and this generator function must be defined
+                    # for each comp_dist.
+                    # pylint: disable=W0640
                     if len(args) > 2:
                         args[1] = size
                     else:

pymc3/distributions/multivariate.py

@@ -267,7 +267,7 @@ def random(self, point=None, size=None):
             else:
                 std_norm_shape = mu.shape
             standard_normal = np.random.standard_normal(std_norm_shape)
-            return mu + np.tensordot(standard_normal, chol, axes=[[-1], [-1]])
+            return mu + np.einsum('...ij,...j->...i', chol, standard_normal)
         else:
             mu, tau = draw_values([self.mu, self.tau], point=point, size=size)
             if mu.shape[-1] != tau[0].shape[-1]:

pymc3/tests/test_sampling.py

@@ -223,8 +223,9 @@ def test_normal_scalar(self):
             ppc = pm.sample_posterior_predictive(trace, samples=1000, vars=[a])
             assert 'a' in ppc
             assert ppc['a'].shape == (1000,)
-        _, pval = stats.kstest(ppc['a'],
-                               stats.norm(loc=0, scale=np.sqrt(2)).cdf)
+        # mu's standard deviation may have changed thanks to a's observed
+        _, pval = stats.kstest(ppc['a'] - trace['mu'],
+                               stats.norm(loc=0, scale=1).cdf)
         assert pval > 0.001
 
         with model: