perf: sample Binomial via BINV/BTPE instead of n Bernoulli trials - #409
perf: sample Binomial via BINV/BTPE instead of n Bernoulli trials#409agene0001 wants to merge 6 commits into
Conversation
`Binomial::sample` ran one Bernoulli trial per trial, so a single draw was O(n). The reported case in statrs-dev#183 - 5 categories, n = 100_000, 100 samples via `Multinomial` - took 22.4 ms; numpy was measured at ~400x faster. Replaces it with Kachitvichyanukul & Schmeiser (1988): BINV inversion for `n * min(p, 1-p) < 10`, BTPE rejection from a triangle/parallelogram/two exponential tails envelope otherwise, plus exact shortcuts for p in {0, 1} and a Poisson-limit path when `1 - p` rounds to 1. Binomial(0.4, 1_000).sample() 3.6 us -> 56 ns Binomial(0.4, 100_000).sample() 361 us -> 39 ns Binomial(0.4, 1e9).sample() ~3.6 s -> 39 ns Multinomial(5 cats, n=1e5) x100 22.4 ms -> 22.8 us (981x) `Multinomial` benefits automatically, since it draws conditioned binomials per category - which is what the issue asked for. The implementation is adapted from `rand_distr` (MIT/Apache-2.0, attributed in-source). That matters for one detail: the final acceptance test uses the Stirling-series signs from GSL, which differ from the published paper and were verified by one of the algorithm's original designers. Correctness is checked by a chi-square goodness-of-fit test against the exact pmf over five parameter sets covering every code path (BINV, BTPE, and both flipped variants), with expected-count pooling and a 6-sigma threshold on fixed seeds - so it is deterministic and a failure means a real defect. Plus moment tests for n = 1e9 and the Poisson-limit path, and degenerate-parameter tests. Closes statrs-dev#183
The test sizes its histogram and its pooled cells from n and p at run time, so it needs `Vec`, which does not exist under `--no-default-features`: the crate is `no_std` without `alloc`. It was gated on `rand` alone, so `--no-default-features --features rand` failed to compile the test binary. Gated on `std` as well rather than rewritten around fixed-size arrays, since a goodness-of-fit test with dynamic binning genuinely wants allocation. The sampler still has no-std coverage from `test_sample_extreme_parameters_moments`, which uses no collections. CI did not catch this because the test job runs only default features, and the feature-powerset job passes `--no-dev-deps`, so test code is never compiled without std. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #409 +/- ##
==========================================
+ Coverage 94.77% 94.93% +0.15%
==========================================
Files 61 62 +1
Lines 13402 13916 +514
==========================================
+ Hits 12702 13211 +509
- Misses 700 705 +5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| if q == 1.0 { | ||
| // p is below ~2^-54 (and not flipped, since p <= 0.5): the | ||
| // distribution is indistinguishable from Poisson(n * p) to O(p) | ||
| return super::poisson::sample_unchecked(rng, n as f64 * p) as u64; |
|
made some comments on the implementation. On a broader note, I'm not sure how I feel about AI driven development for this project. Do we need to develop faster? I'm not going to be reviewing much faster, since the speed isn't my goal. If you want to keep one or two of your PRs not as a draft, that would be appreciated. Even broader, there's mentions of GSL here - while my understanding is that it's complex and developing around licensing from the training corpus, I don't really know how to regard if the model is referring to it. Can you post the chat here? |
Both points from @YeungOnion's review on statrs-dev#409. Reflect the Poisson-limit result. That early return skipped the `flipped` reflection every other path applies. It turns out not to be a live bug -- entering the branch needs `p <= 2^-54`, while reflecting can only produce `p >= 2^-53`, the largest f64 below one being `1 - 2^-53` -- but the margin is a single bit, and the comment that stood there gave the wrong reason for why it was safe. Applying the reflection makes the branch correct on its own terms instead of resting on a spacing argument that a later change to the guard could quietly invalidate. The result is also clamped to `n`, since a Poisson variate has no upper bound where a binomial one does. Two tests come with it: one pinning the one-bit reachability margin so that any future change making the branch reachable fails loudly, and one sampling at `p = 1 - 2^-53` -- the closest reachable point on the flipped side, one bit clear of the guard -- which counts failures rather than successes and so exercises the reflection directly. Build the goodness-of-fit test on `stats_tests::chisquare` rather than the hand-rolled statistic and 6-sigma threshold, which gives a real p-value and reuses the crate's own tested implementation. Binning now covers the whole support instead of a 6-sigma window, which is also what makes the observed and expected totals agree, as `chisquare` requires; the final cell is pinned to the exact remainder because summing the pmf across cells otherwise drifts past its `relative_eq` check. The test was mutation-checked rather than assumed meaningful: biasing BINV's ratio by 0.5% takes it to chi-square 2174 over 15 cells at p = 0. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
@YeungOnion Sorry about that. I was using it for a project of mine and though I might spend some time updating it. I didn't mean to dump so many pr's at once. The moved most of the pr's to draft so its not as much as a burden and for the gsl question, I was able to get a response from the session to hopefully clear up any misconceptions.
So the GSL mentions in my code are paraphrases of comments in MIT-licensed code, describing a mathematical fact — which signs are correct. Anyone can download the |
|
@YeungOnion that kind of function shouldnt be like different backend/trait or sth? i prefer share decission with users about decisisoin about sampling etc |
| /// The margin is a single bit, so pin it -- a later change to either the | ||
| /// guard or the reflection that makes the branch reachable should fail here | ||
| /// rather than silently return an unreflected variate. | ||
| #[test] |
There was a problem hiding this comment.
I think this one tests float implementation, perhaps we can expose the logic within a private function? See other comment.
| let p = if flipped { 1.0 - p } else { p }; | ||
| let q = 1.0 - p; | ||
|
|
||
| if q == 1.0 { |
There was a problem hiding this comment.
maybe something like this would enable testing in that short circuit -> poisson distribution
if let Some(s) = short_circuit_function(p) { return p };
|
Overall, good PR, and would like to bring it in, moving out of draft |
|
@agene0001 appreciate it! For the licensing, we'll assume it's okay, we've at least some record of investigation. @day01 We could allow the user selection of algorithm, but these are much more developed algorithms than what we have currently, which is an unoptimized coin flip counter. But exposing the option to select algorithm regardless of As far as defaults go, it's good if the need is expressed by "I want to sample the Binomial distribution" instead of something more granular, and the docs mention it as well. |
Follow-up to @YeungOnion's review on statrs-dev#409. `Plan::new(n, p)` now chooses the algorithm and reports whether the result needs reflecting, with no RNG involved, and `sample_unchecked` just dispatches on it. Two things fall out. The reflection now happens in exactly one place, at the end of `sample_unchecked`, so a path cannot be added or changed without inheriting it. That is structurally the bug caught earlier in this review, where the Poisson short circuit returned early and skipped it. And the branch selection is testable without going through the sampler. The two tests that previously probed f64 bit patterns to reach the Poisson limit are replaced by tests over `Plan`: one covering every path, including both reflected ones and the BINV/BTPE threshold, and one asserting that no `p` at all reaches the Poisson limit in its reflected form -- which is the property that made the earlier bug latent rather than live, now checked by walking the 64 f64 values below one rather than by asserting a spacing identity. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Addresses @day01's request to share the sampling decision with users and @YeungOnion's "exposing the option to select algorithm regardless of n*p could be good for some". `Binomial::sampler(algorithm)` returns a `BinomialSampler` bound to a `BinomialAlgorithm` -- `Automatic`, `Inversion` (BINV) or `Rejection` (BTPE). It implements `rand`'s `Distribution`, so it composes with `sample_iter` and the rest. Sampling a `Binomial` directly is untouched and still chooses for you, per the note about the default being expressed as "I want to sample the Binomial distribution"; its docs now point at the explicit route. One thing did not survive contact with the algorithm, and it is worth being precise about: BTPE cannot be forced regardless of n*p. Its triangle region has radius `floor(2.195 sqrt(npq) - 4.6 q) + 0.5`, which goes negative once the mean is small, inverting the region boundaries. The sampler still terminates and still returns values in range, so nothing looks wrong -- but the distribution is not binomial. Measured against the exact pmf at n = 100, p = 0.04, chi-square is 1.1e7 on 14 degrees of freedom; at p = 0.044 it is 11.6 on 15, which is where `2.195 sqrt(npq) - 4.6 q` crosses zero. The failure is silent, so `sampler` validates that expression and returns `BinomialAlgorithmError::RejectionMeanTooSmall` rather than handing back a sampler that quietly lies. BINV has no such restriction and can be forced anywhere, including where `Automatic` would pick BTPE. Degenerate parameters and the Poisson limit run neither algorithm, so they ignore the request instead of failing. `sample_unchecked` and `BinomialSampler` now share `sample_plan`, so both apply the p -> 1 - p reflection through the same code. Tests: the chi-square goodness-of-fit is run for every valid (parameters, algorithm) pair, including inversion forced at n = 500, p = 0.85 where the default would pick BTPE; the validity predicate is checked against BTPE's own radius expression over 6 values of n by 99 of p, together with the invariant that `Automatic` never selects BTPE where it is undefined. Making the predicate unconditionally true fails both. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Added the algorithm selection, atop this PR as suggested.
let dist = Binomial::new(0.3, 1_000).unwrap();
let sampler = dist.sampler(BinomialAlgorithm::Rejection).unwrap();
let draws: Vec<u64> = sampler.sample_iter(&mut rng).take(5).collect();Sampling a One caveat on "regardless of
|
| n | p | n*p | chi-square | df |
|---|---|---|---|---|
| 100 | 0.04 | 4.0 | 11,042,868 | 14 |
| 100 | 0.044 | 4.4 | 11.6 | 15 |
| 100 | 0.05 | 5.0 | 18.3 | 16 |
| 100 | 0.10 | 10.0 | 26.0 | 24 |
The cliff is exactly where 2.195*sqrt(n*p*q) - 4.6*q crosses zero, at n*p ≈ 4.39*q.
Since the failure is silent, sampler validates that expression and returns BinomialAlgorithmError::RejectionMeanTooSmall rather than handing back a sampler that quietly lies. BINV has no such restriction and can be forced anywhere, including where Automatic would pick BTPE. Degenerate parameters and the Poisson limit run neither algorithm, so they ignore the request instead of failing.
Tests
The chi-square goodness-of-fit now runs for every valid (parameters, algorithm) pair, including inversion forced at n = 500, p = 0.85 where the default would pick BTPE. The validity predicate is checked against BTPE's own radius expression across 6 values of n by 99 of p, along with the invariant that Automatic never selects BTPE where it's undefined. Making the predicate unconditionally true fails both tests.
sample_unchecked and BinomialSampler share a sample_plan helper, so both apply the p -> 1 - p reflection through the same code.
| @@ -115,14 +115,437 @@ impl core::fmt::Display for Binomial { | |||
| #[cfg(feature = "rand")] | |||
| #[cfg_attr(docsrs, doc(cfg(feature = "rand")))] | |||
| impl ::rand::distr::Distribution<u64> for Binomial { | |||
| /// Samples in `O(1)` expected time (independent of `n`) via the BINV | |||
There was a problem hiding this comment.
lot of code in one file. maybe some module? or sth? currently 1k is hard to read for human, and llm too tbh.
https://arxiv.org/pdf/2601.03731#:~:text=Our%20results%20unveil%20a%20critical%20Aggregation%20Deficit,loss%20of%20consistency%20beyond%20100%20execution%20steps.
https://cbarkinozer.medium.com/agentic-code-generation-papers-part-2-23d6482da032
pls split it to a few smaller
There was a problem hiding this comment.
Seems a good call, moving Binomial algorithms and their tests into module would be a good grouping
@day01 and @YeungOnion both asked for this; the grouping is the one @YeungOnion described. `binomial.rs` was 1058 lines. It is now: binomial/mod.rs 556 the distribution: constructors, moments, pmf/cdf, and their tests binomial/sampling.rs 824 BINV and BTPE, path selection, the algorithm selection API, and their tests Pure code movement -- no behaviour change. The 27 tests are the same 27, split between the two modules; `--list` before and after gives an identical set of names. Two fixups carried along. `Binomial::sampler` and the `Distribution` impls moved with the sampling code, since they are its entry points, and `sample_unchecked` stays crate-internal as before. And the doc comment describing the chi-square goodness-of-fit test had ended up above `test_forced_algorithms_sample_correctly` when that test was inserted in the previous commit; it is back on the test it describes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
I like it! Thanks @day01 and @agene0001 |
3 participants
Binomial::sampleran one Bernoulli trial per trial, so a single draw was O(n). The case reported in #183 (5 categories, n = 100 000, sampled viaMultinomial) took 22.4 ms per 100 draws here; numpy was reported ~400x faster.This replaces it with Kachitvichyanukul & Schmeiser (1988): BINV inversion for
n·min(p,1−p) < 10, BTPE rejection otherwise, exact shortcuts for p ∈ {0, 1}, and a Poisson-limit path when1 − prounds to 1.0.Binomial(0.4, 1_000).sample()Binomial(0.4, 100_000).sample()Binomial(0.4, 1e9).sample()Multinomial(5 cats, n=1e5)×100Multinomialbenefits automatically since it draws conditioned binomials per category.The implementation is adapted from
rand_distr(MIT/Apache-2.0, attributed in-source). One detail worth review attention: the final acceptance test uses the Stirling-series signs from GSL, which differ from the published paper and were verified by one of the algorithm's original designers (see the in-source comment).Correctness: chi-square goodness-of-fit against the exact pmf over five parameter sets covering every code path (BINV, BTPE, both flipped variants), with expected-count pooling and a 6σ threshold on fixed seeds — deterministic, so a failure means a real defect, not chance. Plus moment tests for n = 1e9 and the Poisson-limit path, and degenerate-parameter tests. Relates to the sampling-verification wish list in #288.
Closes #183
🤖 Generated with Claude Code