fix: compute Multinomial pmf/ln_pmf in log space to avoid NaN for large n#381
Open
teddytennant wants to merge 1 commit into
Open
fix: compute Multinomial pmf/ln_pmf in log space to avoid NaN for large n#381teddytennant wants to merge 1 commit into
teddytennant wants to merge 1 commit into
Conversation
…ge n For large `n`, `factorial::multinomial` overflows f64 to +inf because it exponentiates the coefficient's logarithm before returning. In `pmf` this inf was multiplied by an underflowing (0.0) probability product, yielding NaN; `ln_pmf` computed `inf.ln()` = +inf. Both are wrong for valid inputs such as n=2000. Keep the whole computation in log space: build the multinomial coefficient directly from `ln_factorial`, have `pmf` delegate to `ln_pmf(x).exp()`, and skip zero-count categories so an allowed zero probability no longer produces `0 * ln(0) = NaN`. Adds a regression test for n=2000. Fixes statrs-dev#370
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Multinomial::pmfandMultinomial::ln_pmfnow compute the distribution entirely in log space, so valid large-ninputs return finite values instead ofNaN/+inf.pmfdelegates toself.ln_pmf(x).exp()(the length/sum guards are preserved:ln_pmfstill panics on a length mismatch and returnsNEG_INFINITYwhen the counts don't sum ton, andNEG_INFINITY.exp() == 0.0).ln_pmfbuilds the multinomial coefficient directly withfactorial::ln_factorial(ln n! - Σ ln x_i!) instead offactorial::multinomial(...).ln().0 * ln(0) = NaN.Why
factorial::multinomialreturns the coefficient itself (it exponentiates the stable log value before returning), which overflowsf64to+infonce the log exceeds ~709.78. Forn = 2000:pmfcomputedinf * prod(p_i^x_i), and since the probability product simultaneously underflows to0.0, the result wasinf * 0.0 = NaN.ln_pmfcomputedinf.ln() = +inf.Both are wrong for perfectly valid inputs; SciPy returns finite values for the same parameters. Fixes #370.
Testing
Added
test_pmf_large_n_no_overflow(n=2000, p=[0.5, 0.5], x=[1000, 1000]), which asserts the pmf is finite and matches the lgamma-computed reference values (pmf ≈ 0.0178390111, ln_pmf ≈ -4.0263675824). Verified it fails on the previous implementation (assert!(p.is_finite())panics) and passes after the fix. The existingtest_pmfcases still pass under the log-space implementation.cargo clippy -p statrs --all-targets -- -D warningsis clean;cargo fmt --checkreports no diffs in the touched file (the only fmt diffs on this tree are pre-existing ingeometric.rs).