From 592488030f2102afc13ac984624b95f338a74ad1 Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Fri, 21 Dec 2018 08:49:55 +0100 Subject: [PATCH 1/4] Fixed #3310. Added broadcast_distribution_samples, which helps broadcasting multiple rvs calls with different size and distribution parameter shapes. Added shape guards to other continuous distributions. --- RELEASE-NOTES.md | 2 + pymc3/distributions/continuous.py | 63 ++++++++++++++++++++++++----- pymc3/distributions/discrete.py | 13 ++++-- pymc3/distributions/distribution.py | 27 +++++++++++++ pymc3/tests/test_sampling.py | 10 +++++ 5 files changed, 102 insertions(+), 13 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 8a6fdbf358..7df4a49981 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -36,6 +36,8 @@ This will be the last release to support Python 2. - Fixed `Rice` distribution, which inconsistently mixed two parametrizations (#3286). - `Rice` distribution now accepts multiple parameters and observations and is usable with NUTS (#3289). - `sample_posterior_predictive` no longer calls `draw_values` to initialize the shape of the ppc trace. This called could lead to `ValueError`'s when sampling the ppc from a model with `Flat` or `HalfFlat` prior distributions (Fix issue #3294). +- Added the `broadcast_distribution_samples` function that helps broadcasting arrays of drawn samples, taking into account the requested `size` and the inferred distribution shape. This sometimes is needed by distributions that call several `rvs` separately within their `random` method, such as the `ZeroInflatedPoisson` (Fix issue #3310). +- The `Wald`, `Kumaraswamy`, `LogNormal`, `Pareto`, `Cauchy`, `HalfCauchy`, `Weibull` and `ExGaussian` distributions `random` method used a hidden `_random` function that was written with scalars in mind. This could potentially lead to artificial correlations between random draws. Added shape guards and broadcasting of the distribution samples to prevent this (Similar to issue #3310). ### Deprecations diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py index 91c52e38c4..91f6acd5a3 100644 --- a/pymc3/distributions/continuous.py +++ b/pymc3/distributions/continuous.py @@ -21,7 +21,8 @@ alltrue_elemwise, betaln, bound, gammaln, i0e, incomplete_beta, logpow, normal_lccdf, normal_lcdf, SplineWrapper, std_cdf, zvalue, ) -from .distribution import Continuous, draw_values, generate_samples +from .distribution import (Continuous, draw_values, generate_samples, + to_tuple, broadcast_distribution_samples) __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'TruncatedNormal', 'Beta', 'Kumaraswamy', 'Exponential', 'Laplace', 'StudentT', 'Cauchy', @@ -937,10 +938,15 @@ def get_mu_lam_phi(self, mu, lam, phi): 'mu and lam, mu and phi, or lam and phi.') def _random(self, mu, lam, alpha, size=None): - v = np.random.normal(size=size)**2 + _size = alpha.shape + if size is not None: + size = to_tuple(size) + if _size[:len(size)] != size: + _size = size + _size + v = np.random.normal(size=_size)**2 value = (mu + (mu**2) * v / (2. * lam) - mu / (2. * lam) * np.sqrt(4. * mu * lam * v + (mu * v)**2)) - z = np.random.uniform(size=size) + z = np.random.uniform(size=_size) i = np.floor(z - mu / (mu + value)) * 2 + 1 value = (value**-i) * (mu**(i + 1)) return value + alpha @@ -964,6 +970,8 @@ def random(self, point=None, size=None): """ mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha], point=point, size=size) + mu, lam, alpha = broadcast_distribution_samples([mu, lam, alpha], + size=size) return generate_samples(self._random, mu, lam, alpha, dist_shape=self.shape, @@ -1270,7 +1278,12 @@ def __init__(self, a, b, *args, **kwargs): assert_negative_support(b, 'b', 'Kumaraswamy') def _random(self, a, b, size=None): - u = np.random.uniform(size=size) + _size = a.shape + if size is not None: + size = to_tuple(size) + if _size[:len(size)] != size: + _size = size + _size + u = np.random.uniform(size=_size) return (1 - (1 - u) ** (1 / b)) ** (1 / a) def random(self, point=None, size=None): @@ -1292,6 +1305,7 @@ def random(self, point=None, size=None): """ a, b = draw_values([self.a, self.b], point=point, size=size) + a, b = broadcast_distribution_samples([a, b], size=size) return generate_samples(self._random, a, b, dist_shape=self.shape, size=size) @@ -1644,7 +1658,12 @@ def __init__(self, mu=0, sd=None, tau=None, *args, **kwargs): assert_negative_support(sd, 'sd', 'Lognormal') def _random(self, mu, tau, size=None): - samples = np.random.normal(size=size) + _size = tau.shape + if size is not None: + size = to_tuple(size) + if _size[:len(size)] != size: + _size = size + _size + samples = np.random.normal(size=_size) return np.exp(mu + (tau**-0.5) * samples) def random(self, point=None, size=None): @@ -1665,6 +1684,7 @@ def random(self, point=None, size=None): array """ mu, tau = draw_values([self.mu, self.tau], point=point, size=size) + mu, tau = broadcast_distribution_samples([mu, tau], size=size) return generate_samples(self._random, mu, tau, dist_shape=self.shape, size=size) @@ -1930,7 +1950,12 @@ def __init__(self, alpha, m, transform='lowerbound', *args, **kwargs): super(Pareto, self).__init__(transform=transform, *args, **kwargs) def _random(self, alpha, m, size=None): - u = np.random.uniform(size=size) + _size = alpha.shape + if size is not None: + size = to_tuple(size) + if _size[:len(size)] != size: + _size = size + _size + u = np.random.uniform(size=_size) return m * (1. - u)**(-1. / alpha) def random(self, point=None, size=None): @@ -1952,6 +1977,7 @@ def random(self, point=None, size=None): """ alpha, m = draw_values([self.alpha, self.m], point=point, size=size) + alpha, m = broadcast_distribution_samples([alpha, m], size=size) return generate_samples(self._random, alpha, m, dist_shape=self.shape, size=size) @@ -2054,7 +2080,12 @@ def __init__(self, alpha, beta, *args, **kwargs): assert_negative_support(beta, 'beta', 'Cauchy') def _random(self, alpha, beta, size=None): - u = np.random.uniform(size=size) + _size = alpha.shape + if size is not None: + size = to_tuple(size) + if _size[:len(size)] != size: + _size = size + _size + u = np.random.uniform(size=_size) return alpha + beta * np.tan(np.pi * (u - 0.5)) def random(self, point=None, size=None): @@ -2076,6 +2107,7 @@ def random(self, point=None, size=None): """ alpha, beta = draw_values([self.alpha, self.beta], point=point, size=size) + alpha, beta = broadcast_distribution_samples([alpha, beta], size=size) return generate_samples(self._random, alpha, beta, dist_shape=self.shape, size=size) @@ -2163,7 +2195,12 @@ def __init__(self, beta, *args, **kwargs): assert_negative_support(beta, 'beta', 'HalfCauchy') def _random(self, beta, size=None): - u = np.random.uniform(size=size) + _size = beta.shape + if size is not None: + size = to_tuple(size) + if _size[:len(size)] != size: + _size = size + _size + u = np.random.uniform(size=_size) return beta * np.abs(np.tan(np.pi * (u - 0.5))) def random(self, point=None, size=None): @@ -2637,9 +2674,15 @@ def random(self, point=None, size=None): """ alpha, beta = draw_values([self.alpha, self.beta], point=point, size=size) + alpha, beta = broadcast_distribution_samples([alpha, beta], size=size) def _random(a, b, size=None): - return b * (-np.log(np.random.uniform(size=size)))**(1 / a) + _size = a.shape + if size is not None: + size = to_tuple(size) + if _size[:len(size)] != size: + _size = size + _size + return b * (-np.log(np.random.uniform(size=_size)))**(1 / a) return generate_samples(_random, alpha, beta, dist_shape=self.shape, @@ -2921,6 +2964,8 @@ def random(self, point=None, size=None): """ mu, sigma, nu = draw_values([self.mu, self.sigma, self.nu], point=point, size=size) + mu, sigma, nu = broadcast_distribution_samples([mu, sigma, nu], + size=size) def _random(mu, sigma, nu, size=None): return (np.random.normal(mu, sigma, size=size) diff --git a/pymc3/distributions/discrete.py b/pymc3/distributions/discrete.py index 86db441ae1..a959656cc1 100644 --- a/pymc3/distributions/discrete.py +++ b/pymc3/distributions/discrete.py @@ -6,7 +6,8 @@ from pymc3.util import get_variable_name from .dist_math import bound, factln, binomln, betaln, logpow, random_choice -from .distribution import Discrete, draw_values, generate_samples +from .distribution import (Discrete, draw_values, generate_samples, + broadcast_distribution_samples) from pymc3.math import tround, sigmoid, logaddexp, logit, log1pexp @@ -345,6 +346,7 @@ def _ppf(self, p): def _random(self, q, beta, size=None): p = np.random.uniform(size=size) + p, q, beta = broadcast_distribution_samples([p, q, beta], size=size) return np.ceil(np.power(np.log(1 - p) / np.log(q), 1. / beta)) - 1 @@ -847,7 +849,8 @@ def random(self, point=None, size=None): g = generate_samples(stats.poisson.rvs, theta, dist_shape=self.shape, size=size) - return g * (np.random.random(np.squeeze(g.shape)) < psi) + g, psi = broadcast_distribution_samples([g, psi], size=size) + return g * (np.random.random(g.shape) < psi) def logp(self, value): psi = self.psi @@ -939,7 +942,8 @@ def random(self, point=None, size=None): g = generate_samples(stats.binom.rvs, n, p, dist_shape=self.shape, size=size) - return g * (np.random.random(np.squeeze(g.shape)) < psi) + g, psi = broadcast_distribution_samples([g, psi], size=size) + return g * (np.random.random(g.shape) < psi) def logp(self, value): psi = self.psi @@ -1057,7 +1061,8 @@ def random(self, point=None, size=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) + g, psi = broadcast_distribution_samples([g, psi], size=size) + return stats.poisson.rvs(g) * (np.random.random(g.shape) < psi) def logp(self, value): alpha = self.alpha diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index 0dcb6bdb48..22ab6656fe 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -640,3 +640,30 @@ def generate_samples(generator, *args, **kwargs): if one_d and samples.shape[-1] == 1: samples = samples.reshape(samples.shape[:-1]) return np.asarray(samples) + + +def broadcast_distribution_samples(samples, size=None): + if size is None: + return np.broadcast_arrays(*samples) + _size = to_tuple(size) + try: + broadcasted_samples = np.broadcast_arrays(*samples) + except ValueError: + # Raw samples shapes + p_shapes = [p.shape for p in samples] + # samples shapes without the size prepend + sp_shapes = [s[len(_size):] if _size == s[:len(_size)] else s + for s in p_shapes] + broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape + broadcasted_samples = [] + for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes): + if _size == p_shape[:len(_size)]: + slicer_head = [slice(None)] * len(_size) + else: + slicer_head = [np.newaxis] * len(_size) + slicer_tail = ([np.newaxis] * (len(broadcast_shape) - + len(sp_shape)) + + [slice(None)] * len(sp_shape)) + broadcasted_samples.append(param[tuple(slicer_head + slicer_tail)]) + broadcasted_samples = np.broadcast_arrays(*broadcasted_samples) + return broadcasted_samples diff --git a/pymc3/tests/test_sampling.py b/pymc3/tests/test_sampling.py index 4b43bedad1..b626e1abaa 100644 --- a/pymc3/tests/test_sampling.py +++ b/pymc3/tests/test_sampling.py @@ -467,3 +467,13 @@ def test_shape_edgecase(self): x = pm.Normal('x', mu=mu, sd=sd, shape=5) prior = pm.sample_prior_predictive(10) assert prior['mu'].shape == (10, 5) + + def test_zeroinflatedpoisson(self): + with pm.Model(): + theta = pm.Beta('theta', alpha=1, beta=1) + psi = pm.HalfNormal('psi', sd=1) + pm.ZeroInflatedPoisson('suppliers', psi=psi, theta=theta, shape=20) + gen_data = pm.sample_prior_predictive(samples=5000) + assert gen_data['theta'].shape == (5000,) + assert gen_data['psi'].shape == (5000,) + assert gen_data['suppliers'].shape == (5000, 20) From 77c2a6c042934813807a01138e7df3ac7cc6872c Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Fri, 21 Dec 2018 10:01:02 +0100 Subject: [PATCH 2/4] Fixed broken continuous distributions. Did not notice that _random got a parameter shape aware size input thanks to generate_samples. --- pymc3/distributions/continuous.py | 54 +++++++------------------------ 1 file changed, 11 insertions(+), 43 deletions(-) diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py index 91f6acd5a3..c08ea07f1f 100644 --- a/pymc3/distributions/continuous.py +++ b/pymc3/distributions/continuous.py @@ -938,15 +938,10 @@ def get_mu_lam_phi(self, mu, lam, phi): 'mu and lam, mu and phi, or lam and phi.') def _random(self, mu, lam, alpha, size=None): - _size = alpha.shape - if size is not None: - size = to_tuple(size) - if _size[:len(size)] != size: - _size = size + _size - v = np.random.normal(size=_size)**2 + v = np.random.normal(size=size)**2 value = (mu + (mu**2) * v / (2. * lam) - mu / (2. * lam) * np.sqrt(4. * mu * lam * v + (mu * v)**2)) - z = np.random.uniform(size=_size) + z = np.random.uniform(size=size) i = np.floor(z - mu / (mu + value)) * 2 + 1 value = (value**-i) * (mu**(i + 1)) return value + alpha @@ -970,8 +965,11 @@ def random(self, point=None, size=None): """ mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha], point=point, size=size) + print(mu.shape, lam.shape, alpha.shape) mu, lam, alpha = broadcast_distribution_samples([mu, lam, alpha], size=size) + print(mu.shape, lam.shape, alpha.shape) + print(self.shape, size) return generate_samples(self._random, mu, lam, alpha, dist_shape=self.shape, @@ -1278,12 +1276,7 @@ def __init__(self, a, b, *args, **kwargs): assert_negative_support(b, 'b', 'Kumaraswamy') def _random(self, a, b, size=None): - _size = a.shape - if size is not None: - size = to_tuple(size) - if _size[:len(size)] != size: - _size = size + _size - u = np.random.uniform(size=_size) + u = np.random.uniform(size=size) return (1 - (1 - u) ** (1 / b)) ** (1 / a) def random(self, point=None, size=None): @@ -1658,12 +1651,7 @@ def __init__(self, mu=0, sd=None, tau=None, *args, **kwargs): assert_negative_support(sd, 'sd', 'Lognormal') def _random(self, mu, tau, size=None): - _size = tau.shape - if size is not None: - size = to_tuple(size) - if _size[:len(size)] != size: - _size = size + _size - samples = np.random.normal(size=_size) + samples = np.random.normal(size=size) return np.exp(mu + (tau**-0.5) * samples) def random(self, point=None, size=None): @@ -1950,12 +1938,7 @@ def __init__(self, alpha, m, transform='lowerbound', *args, **kwargs): super(Pareto, self).__init__(transform=transform, *args, **kwargs) def _random(self, alpha, m, size=None): - _size = alpha.shape - if size is not None: - size = to_tuple(size) - if _size[:len(size)] != size: - _size = size + _size - u = np.random.uniform(size=_size) + u = np.random.uniform(size=size) return m * (1. - u)**(-1. / alpha) def random(self, point=None, size=None): @@ -2080,12 +2063,7 @@ def __init__(self, alpha, beta, *args, **kwargs): assert_negative_support(beta, 'beta', 'Cauchy') def _random(self, alpha, beta, size=None): - _size = alpha.shape - if size is not None: - size = to_tuple(size) - if _size[:len(size)] != size: - _size = size + _size - u = np.random.uniform(size=_size) + u = np.random.uniform(size=size) return alpha + beta * np.tan(np.pi * (u - 0.5)) def random(self, point=None, size=None): @@ -2195,12 +2173,7 @@ def __init__(self, beta, *args, **kwargs): assert_negative_support(beta, 'beta', 'HalfCauchy') def _random(self, beta, size=None): - _size = beta.shape - if size is not None: - size = to_tuple(size) - if _size[:len(size)] != size: - _size = size + _size - u = np.random.uniform(size=_size) + u = np.random.uniform(size=size) return beta * np.abs(np.tan(np.pi * (u - 0.5))) def random(self, point=None, size=None): @@ -2677,12 +2650,7 @@ def random(self, point=None, size=None): alpha, beta = broadcast_distribution_samples([alpha, beta], size=size) def _random(a, b, size=None): - _size = a.shape - if size is not None: - size = to_tuple(size) - if _size[:len(size)] != size: - _size = size + _size - return b * (-np.log(np.random.uniform(size=_size)))**(1 / a) + return b * (-np.log(np.random.uniform(size=size)))**(1 / a) return generate_samples(_random, alpha, beta, dist_shape=self.shape, From 7d2cd9fa901785ff640a6342bcf9ee731831dee3 Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Fri, 21 Dec 2018 10:28:35 +0100 Subject: [PATCH 3/4] Fixed lint error --- pymc3/distributions/continuous.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py index c08ea07f1f..9e214a4bff 100644 --- a/pymc3/distributions/continuous.py +++ b/pymc3/distributions/continuous.py @@ -22,7 +22,7 @@ normal_lccdf, normal_lcdf, SplineWrapper, std_cdf, zvalue, ) from .distribution import (Continuous, draw_values, generate_samples, - to_tuple, broadcast_distribution_samples) + broadcast_distribution_samples) __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'TruncatedNormal', 'Beta', 'Kumaraswamy', 'Exponential', 'Laplace', 'StudentT', 'Cauchy', From e6e56f83a7814a6a9825cb6499d54b31f8d56574 Mon Sep 17 00:00:00 2001 From: lucianopaz Date: Sat, 22 Dec 2018 21:02:14 +0100 Subject: [PATCH 4/4] Addressed comments --- RELEASE-NOTES.md | 5 +++-- pymc3/distributions/continuous.py | 3 --- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 7e011b9b74..593a965044 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -6,6 +6,9 @@ ### Maintenance +- Added the `broadcast_distribution_samples` function that helps broadcasting arrays of drawn samples, taking into account the requested `size` and the inferred distribution shape. This sometimes is needed by distributions that call several `rvs` separately within their `random` method, such as the `ZeroInflatedPoisson` (Fix issue #3310). +- The `Wald`, `Kumaraswamy`, `LogNormal`, `Pareto`, `Cauchy`, `HalfCauchy`, `Weibull` and `ExGaussian` distributions `random` method used a hidden `_random` function that was written with scalars in mind. This could potentially lead to artificial correlations between random draws. Added shape guards and broadcasting of the distribution samples to prevent this (Similar to issue #3310). + ### Deprecations ## PyMC3 3.6 (Dec 21 2018) @@ -44,8 +47,6 @@ This will be the last release to support Python 2. - Fixed `Rice` distribution, which inconsistently mixed two parametrizations (#3286). - `Rice` distribution now accepts multiple parameters and observations and is usable with NUTS (#3289). - `sample_posterior_predictive` no longer calls `draw_values` to initialize the shape of the ppc trace. This called could lead to `ValueError`'s when sampling the ppc from a model with `Flat` or `HalfFlat` prior distributions (Fix issue #3294). -- Added the `broadcast_distribution_samples` function that helps broadcasting arrays of drawn samples, taking into account the requested `size` and the inferred distribution shape. This sometimes is needed by distributions that call several `rvs` separately within their `random` method, such as the `ZeroInflatedPoisson` (Fix issue #3310). -- The `Wald`, `Kumaraswamy`, `LogNormal`, `Pareto`, `Cauchy`, `HalfCauchy`, `Weibull` and `ExGaussian` distributions `random` method used a hidden `_random` function that was written with scalars in mind. This could potentially lead to artificial correlations between random draws. Added shape guards and broadcasting of the distribution samples to prevent this (Similar to issue #3310). ### Deprecations diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py index 2ee1c01632..c5184849f4 100644 --- a/pymc3/distributions/continuous.py +++ b/pymc3/distributions/continuous.py @@ -958,11 +958,8 @@ def random(self, point=None, size=None): """ mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha], point=point, size=size) - print(mu.shape, lam.shape, alpha.shape) mu, lam, alpha = broadcast_distribution_samples([mu, lam, alpha], size=size) - print(mu.shape, lam.shape, alpha.shape) - print(self.shape, size) return generate_samples(self._random, mu, lam, alpha, dist_shape=self.shape,