Move math out of distribution random methods

RELEASE-NOTES.md

@@ -5,6 +5,9 @@
 ### New features
 - Distinguish between `Data` and `Deterministic` variables when graphing models with graphviz. PR [#3491](https://github.com/pymc-defs/pymc3/pulls/3491).
 
+### Maintenance
+- Moved math operations out of `Rice`, `TruncatedNormal`, `Triangular` and `ZeroInflatedNegativeBinomial` `random` methods. Math operations on values returned by `draw_values` might not broadcast well, and all the `size` aware broadcasting is left to `generate_samples`. Fixes [#3481](https://github.com/pymc-devs/pymc3/issues/3481) and [#3508](https://github.com/pymc-devs/pymc3/issues/3508)
+
 ## PyMC3 3.7 (May 29 2019)
 
 ### New features

pymc3/distributions/continuous.py

@@ -664,16 +664,34 @@ def random(self, point=None, size=None):
         -------
         array
         """
-        mu_v, std_v, a_v, b_v = draw_values(
-            [self.mu, self.sigma, self.lower, self.upper], point=point, size=size)
-        return generate_samples(stats.truncnorm.rvs,
-                                a=(a_v - mu_v)/std_v,
-                                b=(b_v - mu_v) / std_v,
-                                loc=mu_v,
-                                scale=std_v,
-                                dist_shape=self.shape,
-                                size=size,
-                                )
+        mu, sigma, lower, upper = draw_values(
+            [self.mu, self.sigma, self.lower, self.upper],
+            point=point,
+            size=size
+        )
+        return generate_samples(
+            self._random,
+            mu=mu,
+            sigma=sigma,
+            lower=lower,
+            upper=upper,
+            dist_shape=self.shape,
+            size=size,
+        )
+
+    def _random(self, mu, sigma, lower, upper, size):
+        """ Wrapper around stats.truncnorm.rvs that converts TruncatedNormal's
+        parametrization to scipy.truncnorm. All parameter arrays should have
+        been broadcasted properly by generate_samples at this point and size is
+        the scipy.rvs representation.
+        """
+        return stats.truncnorm.rvs(
+            a=(lower - mu) / sigma,
+            b=(upper - mu) / sigma,
+            loc=mu,
+            scale=sigma,
+            size=size,
+        )
 
     def logp(self, value):
         """
@@ -3582,11 +3600,23 @@ def random(self, point=None, size=None):
         """
         c, lower, upper = draw_values([self.c, self.lower, self.upper],
                                       point=point, size=size)
-        scale = upper - lower
-        c_ = (c - lower) / scale
-        return generate_samples(stats.triang.rvs, c=c_, loc=lower, scale=scale,
+        return generate_samples(self._random, c=c, lower=lower, upper=upper,
                                 size=size, dist_shape=self.shape)
 
+    def _random(self, c, lower, upper, size):
+        """ Wrapper around stats.triang.rvs that converts Triangular's
+        parametrization to scipy.triang. All parameter arrays should have
+        been broadcasted properly by generate_samples at this point and size is
+        the scipy.rvs representation.
+        """
+        scale = upper - lower
+        return stats.triang.rvs(
+            c=(c - lower) / scale,
+            loc=lower,
+            scale=scale,
+            size=size,
+        )
+
     def logp(self, value):
         """
         Calculate log-probability of Triangular distribution at specified value.
@@ -3875,9 +3905,21 @@ def random(self, point=None, size=None):
         """
         nu, sigma = draw_values([self.nu, self.sigma],
                              point=point, size=size)
-        return generate_samples(stats.rice.rvs, b=nu / sigma, scale=sigma, loc=0,
+        return generate_samples(self._random, nu=nu, sigma=sigma,
                                 dist_shape=self.shape, size=size)
 
+    def _random(self, nu, sigma, size):
+        """ Wrapper around stats.rice.rvs that converts Rice's
+        parametrization to scipy.rice. All parameter arrays should have
+        been broadcasted properly by generate_samples at this point and size is
+        the scipy.rvs representation.
+        """
+        return stats.rice.rvs(
+            b=nu / sigma,
+            scale=sigma,
+            size=size,
+        )
+
     def logp(self, value):
         """
         Calculate log-probability of Rice distribution at specified value.

pymc3/distributions/discrete.py

@@ -1427,13 +1427,29 @@ def random(self, point=None, size=None):
         """
         mu, alpha, psi = draw_values(
             [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)
+        g = generate_samples(
+            self._random,
+            mu=mu,
+            alpha=alpha,
+            dist_shape=self.shape,
+            size=size
+        )
         g[g == 0] = np.finfo(float).eps  # Just in case
         g, psi = broadcast_distribution_samples([g, psi], size=size)
         return stats.poisson.rvs(g) * (np.random.random(g.shape) < psi)
 
+    def _random(self, mu, alpha, size):
+        """ Wrapper around stats.gamma.rvs that converts NegativeBinomial's
+        parametrization to scipy.gamma. All parameter arrays should have
+        been broadcasted properly by generate_samples at this point and size is
+        the scipy.rvs representation.
+        """
+        return stats.gamma.rvs(
+            a=alpha,
+            scale=mu / alpha,
+            size=size,
+        )
+
     def logp(self, value):
         """
         Calculate log-probability of ZeroInflatedNegativeBinomial distribution at specified value.

pymc3/tests/test_distributions_random.py

@@ -944,3 +944,163 @@ def test_density_dist_without_random_not_sampleable():
     samples = 500
     with pytest.raises(ValueError):
         pm.sample_posterior_predictive(trace, samples=samples, model=model, size=100)
+
+
+class TestNestedRandom(SeededTest):
+    def build_model(self, distribution, shape, nested_rvs_info):
+        with pm.Model() as model:
+            nested_rvs = {}
+            for rv_name, info in nested_rvs_info.items():
+                try:
+                    value, nested_shape = info
+                    loc = 0.
+                except ValueError:
+                    value, nested_shape, loc = info
+                if value is None:
+                    nested_rvs[rv_name] = pm.Uniform(
+                        rv_name,
+                        0 + loc,
+                        1 + loc,
+                        shape=nested_shape,
+                    )
+                else:
+                    nested_rvs[rv_name] = value * np.ones(nested_shape)
+            rv = distribution(
+                "target",
+                shape=shape,
+                **nested_rvs,
+            )
+        return model, rv, nested_rvs
+
+    def sample_prior(
+        self,
+        distribution,
+        shape,
+        nested_rvs_info,
+        prior_samples
+    ):
+        model, rv, nested_rvs = self.build_model(
+            distribution,
+            shape,
+            nested_rvs_info,
+        )
+        with model:
+            return pm.sample_prior_predictive(prior_samples)
+
+    @pytest.mark.parametrize(
+        ["prior_samples", "shape", "psi", "mu", "alpha"],
+        [
+            [10, (3,), (0.5, tuple()), (None, tuple()), (None, (3,))],
+            [10, (3,), (0.5, (3,)), (None, tuple()), (None, (3,))],
+            [10, (3,), (0.5, tuple()), (None, (3,)), (None, tuple())],
+            [10, (3,), (0.5, (3,)), (None, (3,)), (None, tuple())],
+            [10, (4, 3,), (0.5, (3,)), (None, (3,)), (None, (3,))],
+            [10, (4, 3,), (0.5, (3,)), (None, (3,)), (None, (4, 3))],
+        ],
+        ids=str,
+    )
+    def test_ZeroInflatedNegativeBinomial(
+        self,
+        prior_samples,
+        shape,
+        psi,
+        mu,
+        alpha,
+    ):
+        prior = self.sample_prior(
+            distribution=pm.ZeroInflatedNegativeBinomial,
+            shape=shape,
+            nested_rvs_info=dict(psi=psi, mu=mu, alpha=alpha),
+            prior_samples=prior_samples,
+        )
+        assert prior["target"].shape == (prior_samples,) + shape
+
+    @pytest.mark.parametrize(
+        ["prior_samples", "shape", "nu", "sigma"],
+        [
+            [10, (3,), (None, tuple()), (None, (3,))],
+            [10, (3,), (None, tuple()), (None, (3,))],
+            [10, (3,), (None, (3,)), (None, tuple())],
+            [10, (3,), (None, (3,)), (None, tuple())],
+            [10, (4, 3,), (None, (3,)), (None, (3,))],
+            [10, (4, 3,), (None, (3,)), (None, (4, 3))],
+        ],
+        ids=str,
+    )
+    def test_Rice(
+        self,
+        prior_samples,
+        shape,
+        nu,
+        sigma,
+    ):
+        prior = self.sample_prior(
+            distribution=pm.Rice,
+            shape=shape,
+            nested_rvs_info=dict(nu=nu, sigma=sigma),
+            prior_samples=prior_samples,
+        )
+        assert prior["target"].shape == (prior_samples,) + shape
+
+    @pytest.mark.parametrize(
+        ["prior_samples", "shape", "mu", "sigma", "lower", "upper"],
+        [
+            [10, (3,), (None, tuple()), (1., tuple()), (None, tuple(), -1), (None, (3,))],
+            [10, (3,), (None, tuple()), (1., tuple()), (None, tuple(), -1), (None, (3,))],
+            [10, (3,), (None, tuple()), (1., tuple()), (None, (3,), -1), (None, tuple())],
+            [10, (3,), (None, tuple()), (1., tuple()), (None, (3,), -1), (None, tuple())],
+            [10, (4, 3,), (None, (3,)), (1., tuple()), (None, (3,), -1), (None, (3,))],
+            [10, (4, 3,), (None, (3,)), (1., tuple()), (None, (3,), -1), (None, (4, 3))],
+            [10, (3,), (0., tuple()), (None, tuple()), (None, tuple(), -1), (None, (3,))],
+            [10, (3,), (0., tuple()), (None, tuple()), (None, tuple(), -1), (None, (3,))],
+            [10, (3,), (0., tuple()), (None, tuple()), (None, (3,), -1), (None, tuple())],
+            [10, (3,), (0., tuple()), (None, tuple()), (None, (3,), -1), (None, tuple())],
+            [10, (4, 3,), (0., tuple()), (None, (3,)), (None, (3,), -1), (None, (3,))],
+            [10, (4, 3,), (0., tuple()), (None, (3,)), (None, (3,), -1), (None, (4, 3))],
+        ],
+        ids=str,
+    )
+    def test_TruncatedNormal(
+        self,
+        prior_samples,
+        shape,
+        mu,
+        sigma,
+        lower,
+        upper,
+    ):
+        prior = self.sample_prior(
+            distribution=pm.TruncatedNormal,
+            shape=shape,
+            nested_rvs_info=dict(mu=mu, sigma=sigma, lower=lower, upper=upper),
+            prior_samples=prior_samples,
+        )
+        assert prior["target"].shape == (prior_samples,) + shape
+
+
+    @pytest.mark.parametrize(
+        ["prior_samples", "shape", "c", "lower", "upper"],
+        [
+            [10, (3,), (None, tuple()), (-1., (3,)), (2, tuple())],
+            [10, (3,), (None, tuple()), (-1., tuple()), (None, tuple(), 1)],
+            [10, (3,), (None, (3,)), (-1., tuple()), (None, tuple(), 1)],
+            [10, (4, 3,), (None, (3,)), (-1., tuple()), (None, (3,), 1)],
+            [10, (4, 3,), (None, (3,)), (None, tuple(), -1), (None, (3,), 1)],
+        ],
+        ids=str,
+    )
+    def test_Triangular(
+        self,
+        prior_samples,
+        shape,
+        c,
+        lower,
+        upper,
+    ):
+        prior = self.sample_prior(
+            distribution=pm.Triangular,
+            shape=shape,
+            nested_rvs_info=dict(c=c, lower=lower, upper=upper),
+            prior_samples=prior_samples,
+        )
+        assert prior["target"].shape == (prior_samples,) + shape