Implement icdf for Poisson, Binomial, and NegativeBinomial#8339
Implement icdf for Poisson, Binomial, and NegativeBinomial#8339jacobpstein wants to merge 4 commits into
Conversation
|
Can you split the bernoulli into it's own PR, it's much more trivial to review and approve than the ones requiring search |
1f662e1 to
1fb6f6f
Compare
jacobpstein
changed the title
|
Done! split the Bernoulli |
| whose quantile function has no closed form (e.g. ``Poisson``, ``Binomial``, | ||
| ``NegativeBinomial``). | ||
|
|
||
| The search uses fixed-length ``scan`` loops (rather than a data-dependent |
There was a problem hiding this comment.
Have to think about this, C/Numba are more than happy with data dependent loops.
There was a problem hiding this comment.
Yeahhhh, the real constraint I was designing around is JAX (and static-shape lowering generally), which can't do a data-dependent number of steps.
I went with fixed-length scan so the same icdf graph works across all backends without special-casing, at the cost of always running max_iters iterations even after convergence. One option is to use a data-dependent while loop, accepting that JAX would be unsupported (or fall back) for these icdfs. Something to think about.
| return expr | ||
|
|
||
|
|
||
| def discrete_icdf_via_search(logcdf, value, lower, upper=None, *, max_iters=64): |
There was a problem hiding this comment.
I may be misremembering, but one of the libraries I looked at also allowed some candidate bounds or something to inform the starting point of the search, in case the callers can offer something useful. Not saying we have to, but if there are obvious cases and it's not too much complexity would be good to discuss
There was a problem hiding this comment.
I'm always open to learning some new libraries! Right now upper is already the big lever-- Binomial passes upper=n and skips the expansion phase. For the unbounded cases the obvious hint is the mean (mu for Poisson, n(1-p)/p for NegativeBinomial), so I could add an optional start that seeds the geometric expansion there instead of doubling from 1.
Payoff is modest though since bisection is already logarithmic, so a hint just trims a few iterations off phase 1. Low-complexity to add if you think it's worth it, otherwise happy to leave as a follow-up, and if there's an existing package whose approach we can reference here, even better
|
@jacobpstein implementation looks really clean. I don't have concrete feedback yet, have to marinate a bit on it. Let's some vague questions/thoughts above |
|
Went ahead and added the On the loop question, I dug into this a bit more and unfortunately there's no middle ground. Adding an Fun discovery while stress-testing at large |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #8339 +/- ##
==========================================
+ Coverage 91.73% 91.75% +0.01%
==========================================
Files 128 128
Lines 20697 20736 +39
==========================================
+ Hits 18987 19026 +39
Misses 1710 1710
🚀 New features to boost your workflow:
|
|
Opened pymc-devs/pytensor#2291 for the gammaincc thing. |
We try not to compromise around JAX limitations. In this case we simply haven't implemented while loops for jax which cane be done (just not autodiffed through later) |
|
Having said that I think this is already a great functionality to have and we can fine-tune in subsequent iterations / as demand calls. Btw the implementation I vaguely recalled for these was from the boost library: https://github.com/boostorg/math/blob/8ee12a5355935cbaac5d5338372d0d0e3311b473/include/boost/math/distributions/binomial.hpp#L241-L270 |
| - pip: | ||
| - numdifftools>=0.9.40 | ||
| - mcbackend>=0.4.0 | ||
| # pyarrow 25.0.0 (pulled in transitively by nutpie, installed later via pip |
There was a problem hiding this comment.
rebase from main, we already merged similar pins
|
This is really nice work btw, in case I forgot to say so |
|
Ah I think we need to handle batched parameters a bit more carefully. This crashes because the initial value is broadcasted inside scan: import pymc as pm
from pymc.logprob import icdf
# scalar quantile, vector parameter -> ValueError from pytensor.scan
icdf(pm.Poisson.dist(mu=[2.0, 5.0, 20.0]), 0.5).eval()And a silent failure import numpy as np
import scipy.stats as st
import pymc as pm
from pymc.logprob import icdf
q = np.array([[0.1], [0.5], [0.9]]) # shape (3, 1)
n = np.array([10, 40]) # shape (2,)
p = np.array([0.3, 0.6]) # shape (2,) -> broadcasts to (3, 2)
got = icdf(pm.Binomial.dist(n=n, p=p), q).eval()
ref = st.binom.ppf(q, n, p)
print(got)
print(ref)
assert np.allclose(got, ref) # AssertionError |
ricardoV94
left a comment
ricardoV94
left a comment
There was a problem hiding this comment.
Need to handle the batched cases
The geometric bracket expansion can now be seeded with a caller-provided guess (typically the distribution mean) instead of always doubling from lower + 1. Poisson and NegativeBinomial pass their means; Binomial is unchanged since it already provides upper=n.
icdf graphs are always float-typed because out-of-domain values map to nan. When Truncated uses inverse-CDF sampling of a discrete base distribution, the resulting RV silently became float64, so e.g. Hurdle mixtures rejected it as a continuous component. Cast the icdf draw back to the base RV dtype.
The scan carrying the (lo, hi) search state started at the shape of the quantile value, but logcdf(k) broadcasts k against the distribution parameters. When the parameters had a larger broadcast shape than the value, the state changed shape inside scan, raising a ValueError (scalar value, vector parameters) or silently misbroadcasting (value and parameters broadcasting to a larger common shape). Probe logcdf once to obtain the full broadcast shape and start the search state there. Only the shape of the probe is used, so the probe computation is pruned from the compiled graph.
b30a084 to
93437bc
Compare
|
Good catch, thanks @ricardoV94 ! The scan's search state started at the shape of the quantile, so it changed shape once |
2 participants
Description
Adds
icdf(quantile / inverse-CDF) methods for four discrete distributions:icdf(q) = 0 if q ≤ 1 − p else 1.gammaincinv/betaincinv) don't apply since they invert the CDF in its continuous probability argument rather than over the integer support. This adds a shareddiscrete_icdf_via_searchhelper (indist_math.py) that inverts the monotonelogcdfby bisection, with geometric bracket expansion for the unbounded-support cases (Poisson, NegativeBinomial).The helper uses fixed-length
scanloops (no data-dependentuntil) so the graph remains convertible to the C, Numba, and JAX backends — verified for all three. As a side effect this enablespm.Truncatedfor these distributions, which falls back toicdffor inverse-transform sampling.Tests compare against SciPy's
ppfusing the existingcheck_icdfhelper.A few other distributions in the same family (BetaBinomial, HyperGeometric) could reuse the same helper if you all want! This was very fun to work on. I met Daniel Lee at the Sloan conf back in March and have been using PyMc ever since--though I still enjoy Stan.
Related Issue
Checklist
Type of change