Implement icdf for Poisson, Binomial, and NegativeBinomial

pymc/distributions/discrete.py

@@ -40,6 +40,7 @@
     check_icdf_parameters,
     check_icdf_value,
     check_parameters,
+    discrete_icdf_via_search,
     factln,
     log_diff_normal_cdf,
     logpow,
@@ -174,6 +175,23 @@ def logcdf(value, n, p):
             msg="n >= 0, 0 <= p <= 1",
         )
 
+    def icdf(value, n, p):
+        dist = Binomial.dist(n=n, p=p)
+        res = discrete_icdf_via_search(
+            lambda k: logcdf(dist, k),
+            value,
+            lower=0,
+            upper=n,
+        )
+        res = check_icdf_value(res, value)
+        return check_icdf_parameters(
+            res,
+            n >= 0,
+            0 <= p,
+            p <= 1,
+            msg="n >= 0, 0 <= p <= 1",
+        )
+
 
 class BetaBinomial(Discrete):
     R"""
@@ -614,6 +632,21 @@ def logcdf(value, mu):
             msg="mu >= 0",
         )
 
+    def icdf(value, mu):
+        dist = Poisson.dist(mu=mu)
+        res = discrete_icdf_via_search(
+            lambda k: logcdf(dist, k),
+            value,
+            lower=0,
+            start=mu,
+        )
+        res = check_icdf_value(res, value)
+        return check_icdf_parameters(
+            res,
+            mu >= 0,
+            msg="mu >= 0",
+        )
+
 
 class NegativeBinomial(Discrete):
     R"""
@@ -762,6 +795,23 @@ def logcdf(value, n, p):
             msg="n > 0, 0 <= p <= 1",
         )
 
+    def icdf(value, n, p):
+        dist = NegativeBinomial.dist(n=n, p=p)
+        res = discrete_icdf_via_search(
+            lambda k: logcdf(dist, k),
+            value,
+            lower=0,
+            start=n * (1 - p) / p,
+        )
+        res = check_icdf_value(res, value)
+        return check_icdf_parameters(
+            res,
+            n > 0,
+            0 <= p,
+            p <= 1,
+            msg="n > 0, 0 <= p <= 1",
+        )
+
 
 class Geometric(Discrete):
     R"""

pymc/distributions/dist_math.py

@@ -89,6 +89,95 @@ def check_icdf_value(expr: Variable, value: Variable) -> Variable:
     return expr
 
 
+def discrete_icdf_via_search(logcdf, value, lower, upper=None, *, start=None, max_iters=64):
+    """Numerically invert a discrete CDF by bisecting the (monotone) ``logcdf``.
+
+    Returns the smallest integer ``k`` in the support such that
+    ``cdf(k) >= value``, i.e. the quantile function evaluated at ``value``.
+    This is the discrete counterpart of inverting a continuous CDF with a
+    special function (e.g. ``betaincinv``), and is meant for distributions
+    whose quantile function has no closed form (e.g. ``Poisson``, ``Binomial``,
+    ``NegativeBinomial``).
+
+    The search uses fixed-length ``scan`` loops rather than a data-dependent
+    ``until`` condition. Each step is idempotent once an element has converged,
+    so running the full ``max_iters`` iterations does not change the result.
+
+    Parameters
+    ----------
+    logcdf : callable
+        Function mapping an integer-valued tensor ``k`` to ``logcdf(k)`` of the
+        target distribution.
+    value : tensor_like
+        Cumulative probabilities at which to evaluate the quantile function.
+        Assumed to lie in the open interval ``(0, 1)``; edge handling for
+        ``value`` is the responsibility of :func:`check_icdf_value`.
+    lower : tensor_like
+        Smallest value in the support of the distribution.
+    upper : tensor_like, optional
+        Largest value in the support. If ``None`` (unbounded support), a finite
+        upper bracket is first found by geometric expansion.
+    start : tensor_like, optional
+        Initial guess for the upper bracket when ``upper`` is ``None``,
+        typically the distribution mean. The geometric expansion starts from
+        ``max(floor(start), lower + 1)`` instead of ``lower + 1``, so that
+        elements whose quantile is near or below the guess converge in fewer
+        effective iterations. The guess does not need to be accurate for the
+        result to be exact. Ignored when ``upper`` is given.
+    max_iters : int
+        Number of search iterations per phase. This needs to exceed ``log2`` of
+        the largest reachable quantile; the default of 64 covers the full range
+        of integers exactly representable as float64.
+    """
+    log_value = pt.log(value)
+    lower = pt.as_tensor_variable(lower).astype("float64")
+
+    # The scan below carries (lo, hi) as recurrent state, whose shape must stay
+    # fixed across iterations. Since logcdf(k) broadcasts k against the
+    # distribution parameters, the state must start at the full broadcast shape
+    # of value and the parameters. Probing logcdf gives that shape; only the
+    # shape is used, so the probe computation itself is pruned from the graph.
+    result_shape = logcdf(pt.broadcast_arrays(lower, log_value)[0]).shape
+    log_value = pt.broadcast_to(log_value, result_shape)
+    lower = pt.broadcast_to(lower, result_shape)
+
+    if upper is None:
+        hi_init = lower + 1.0
+        if start is not None:
+            start = pt.broadcast_to(pt.as_tensor_variable(start).astype("float64"), result_shape)
+            hi_init = pt.maximum(pt.floor(start), hi_init)
+
+        # Phase 1: geometrically expand an upper bracket until cdf(hi) >= value.
+        def expand(lo, hi):
+            insufficient = logcdf(hi) < log_value
+            return pt.switch(insufficient, hi + 1.0, lo), pt.switch(insufficient, hi * 2.0, hi)
+
+        lo, hi = pytensor.scan(
+            expand,
+            outputs_info=[lower, hi_init],
+            n_steps=max_iters,
+            return_updates=False,
+        )
+        lo, hi = lo[-1], hi[-1]
+    else:
+        lo = lower
+        hi = pt.broadcast_to(pt.as_tensor_variable(upper).astype("float64"), result_shape)
+
+    # Phase 2: bisect [lo, hi] keeping the invariant cdf(hi) >= value > cdf(lo - 1).
+    def bisect(lo, hi):
+        mid = pt.floor((lo + hi) / 2.0)
+        covers = logcdf(mid) >= log_value
+        return pt.switch(covers, lo, mid + 1.0), pt.switch(covers, mid, hi)
+
+    lo, _ = pytensor.scan(
+        bisect,
+        outputs_info=[lo, hi],
+        n_steps=max_iters,
+        return_updates=False,
+    )
+    return lo[-1]
+
+
 def logpow(x, m):
     """Calculate log(x**m) since m*log(x) will fail when m, x = 0."""
     # return m * log(x)

pymc/distributions/truncated.py

@@ -124,7 +124,9 @@ def rv_op(cls, dist, lower, upper, max_n_steps, *, size=None):
                 size=rv_.shape,
                 return_next_rng=True,
             )
-            truncated_rv_ = icdf(rv_, uniform_, warn_rvs=False)
+            # icdf graphs are always float-typed (they use nan for invalid values),
+            # so restore the base RV dtype for discrete distributions
+            truncated_rv_ = icdf(rv_, uniform_, warn_rvs=False).astype(rv_.type.dtype)
             return TruncatedRV(
                 base_rv_op=dist.owner.op,
                 inputs=graph_inputs_,

tests/distributions/test_discrete.py

@@ -228,6 +228,17 @@ def scipy_mu_alpha_logcdf(value, mu, alpha):
             Nat,
             {"mu": Rplus, "alpha": Rplus},
         )
+        check_icdf(
+            pm.NegativeBinomial,
+            {"p": Unit, "n": Rplus},
+            lambda q, p, n: st.nbinom.ppf(q, n, p),
+        )
+        # Large mean, where the icdf search relies on the mean-seeded bracket
+        q = np.array([0.001, 0.5, 0.999])
+        np.testing.assert_array_equal(
+            icdf(pm.NegativeBinomial.dist(n=10, p=1e-3), q).eval(),
+            st.nbinom.ppf(q, 10, 1e-3),
+        )
 
     @pytest.mark.parametrize(
         "mu, p, alpha, n, expected",
@@ -267,6 +278,11 @@ def test_binomial(self):
             Nat,
             {"n": NatSmall, "p": Unit},
         )
+        check_icdf(
+            pm.Binomial,
+            {"n": NatSmall, "p": Unit},
+            lambda q, n, p: st.binom.ppf(q, n, p),
+        )
 
     def test_beta_binomial(self):
         check_logp(
@@ -369,6 +385,40 @@ def test_poisson(self):
             Nat,
             {"mu": Rplus},
         )
+        check_icdf(
+            pm.Poisson,
+            {"mu": Rplus},
+            st.poisson.ppf,
+        )
+        # Large mu, where the icdf search relies on the mean-seeded bracket.
+        # mu is kept at 1e4 because the C implementation of gammaincc (used by
+        # logcdf inside the search) loses precision for larger arguments.
+        q = np.array([0.001, 0.5, 0.999])
+        np.testing.assert_array_equal(
+            icdf(pm.Poisson.dist(mu=1e4), q).eval(),
+            st.poisson.ppf(q, 1e4),
+        )
+
+    def test_discrete_icdf_batched_params(self):
+        # The search state must be broadcast against the distribution
+        # parameters, not just the quantile, or scan fails / misbroadcasts
+        mu = np.array([2.0, 5.0, 20.0])
+        np.testing.assert_array_equal(
+            icdf(pm.Poisson.dist(mu=mu), 0.5).eval(),
+            st.poisson.ppf(0.5, mu),
+        )
+
+        q = np.array([[0.1], [0.5], [0.9]])
+        n = np.array([10, 40])
+        p = np.array([0.3, 0.6])  # broadcasts with q to shape (3, 2)
+        np.testing.assert_array_equal(
+            icdf(pm.Binomial.dist(n=n, p=p), q).eval(),
+            st.binom.ppf(q, n, p),
+        )
+        np.testing.assert_array_equal(
+            icdf(pm.NegativeBinomial.dist(n=n, p=p), q).eval(),
+            st.nbinom.ppf(q, n, p),
+        )
 
     @pytest.mark.parametrize("n", [2, 3, 4])
     def test_categorical(self, n):