Skip to content

Commit 0456f39

Browse files
NegativeBinomial distribution sample generator wrapper (#3866)
* Added wrapper for gamma sampling in negative binomial distribution. * Fixed typo in posterior sampling warnings. * Drop nuts init method from pm.sample (#3863) * Dropped nuts init method * Dropped nuts init method from tests * Refined doc string and added release note * Added wrapper for gamma sampling in negative binomial distribution. * Fixed typo in posterior sampling warnings. * Added test for negative binomial input shape. * Updating release notes with changes to negative binomial sampling. Co-authored-by: Alexandre ANDORRA <[email protected]>
1 parent bf8552e commit 0456f39

File tree

4 files changed

+41
-3
lines changed

4 files changed

+41
-3
lines changed

RELEASE-NOTES.md

+1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
- In named models, `pm.Data` objects now get model-relative names (see [#3843](https://github.com/pymc-devs/pymc3/pull/3843)).
2020
- `pm.sample` now takes 1000 draws and 1000 tuning samples by default, instead of 500 previously (see [#3855](https://github.com/pymc-devs/pymc3/pull/3855)).
2121
- Dropped the outdated 'nuts' initialization method for `pm.sample` (see [#3863](https://github.com/pymc-devs/pymc3/pull/3863)).
22+
- Moved argument division out of `NegativeBinomial` `random` method. Fixes [#3864](https://github.com/pymc-devs/pymc3/issues/3864) in the style of [#3509](https://github.com/pymc-devs/pymc3/pull/3509).
2223

2324
## PyMC3 3.8 (November 29 2019)
2425

pymc3/distributions/discrete.py

+13-1
Original file line numberDiff line numberDiff line change
@@ -673,12 +673,24 @@ def random(self, point=None, size=None):
673673
array
674674
"""
675675
mu, alpha = draw_values([self.mu, self.alpha], point=point, size=size)
676-
g = generate_samples(stats.gamma.rvs, alpha, scale=mu / alpha,
676+
g = generate_samples(self._random, mu=mu, alpha=alpha,
677677
dist_shape=self.shape,
678678
size=size)
679679
g[g == 0] = np.finfo(float).eps # Just in case
680680
return np.asarray(stats.poisson.rvs(g)).reshape(g.shape)
681681

682+
def _random(self, mu, alpha, size):
683+
""" Wrapper around stats.gamma.rvs that converts NegativeBinomial's
684+
parametrization to scipy.gamma. All parameter arrays should have
685+
been broadcasted properly by generate_samples at this point and size is
686+
the scipy.rvs representation.
687+
"""
688+
return stats.gamma.rvs(
689+
a=alpha,
690+
scale=mu / alpha,
691+
size=size,
692+
)
693+
682694
def logp(self, value):
683695
"""
684696
Calculate log-probability of NegativeBinomial distribution at specified value.

pymc3/sampling.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -1567,9 +1567,9 @@ def sample_posterior_predictive(
15671567
nchain = 1
15681568

15691569
if keep_size and samples is not None:
1570-
raise IncorrectArgumentsError("Should not specify both keep_size and samples argukments")
1570+
raise IncorrectArgumentsError("Should not specify both keep_size and samples arguments")
15711571
if keep_size and size is not None:
1572-
raise IncorrectArgumentsError("Should not specify both keep_size and size argukments")
1572+
raise IncorrectArgumentsError("Should not specify both keep_size and size arguments")
15731573

15741574
if samples is None:
15751575
if isinstance(trace, MultiTrace):

pymc3/tests/test_distributions_random.py

+25
Original file line numberDiff line numberDiff line change
@@ -1176,6 +1176,31 @@ def sample_prior(
11761176
with model:
11771177
return pm.sample_prior_predictive(prior_samples)
11781178

1179+
@pytest.mark.parametrize(
1180+
["prior_samples", "shape", "mu", "alpha"],
1181+
[
1182+
[10, (3,), (None, tuple()), (None, (3,))],
1183+
[10, (3,), (None, (3,)), (None, tuple())],
1184+
[10, (4, 3,), (None, (3,)), (None, (3,))],
1185+
[10, (4, 3,), (None, (3,)), (None, (4, 3))],
1186+
],
1187+
ids=str,
1188+
)
1189+
def test_NegativeBinomial(
1190+
self,
1191+
prior_samples,
1192+
shape,
1193+
mu,
1194+
alpha,
1195+
):
1196+
prior = self.sample_prior(
1197+
distribution=pm.NegativeBinomial,
1198+
shape=shape,
1199+
nested_rvs_info=dict(mu=mu, alpha=alpha),
1200+
prior_samples=prior_samples,
1201+
)
1202+
assert prior["target"].shape == (prior_samples,) + shape
1203+
11791204
@pytest.mark.parametrize(
11801205
["prior_samples", "shape", "psi", "mu", "alpha"],
11811206
[

0 commit comments

Comments
 (0)