Skip to content

fix: compute Multinomial pmf/ln_pmf in log space to avoid NaN for large n#381

Open
teddytennant wants to merge 1 commit into
statrs-dev:masterfrom
teddytennant:fix/multinomial-pmf-log-space
Open

fix: compute Multinomial pmf/ln_pmf in log space to avoid NaN for large n#381
teddytennant wants to merge 1 commit into
statrs-dev:masterfrom
teddytennant:fix/multinomial-pmf-log-space

Conversation

@teddytennant

Copy link
Copy Markdown

What

Multinomial::pmf and Multinomial::ln_pmf now compute the distribution entirely in log space, so valid large-n inputs return finite values instead of NaN/+inf.

  • pmf delegates to self.ln_pmf(x).exp() (the length/sum guards are preserved: ln_pmf still panics on a length mismatch and returns NEG_INFINITY when the counts don't sum to n, and NEG_INFINITY.exp() == 0.0).
  • ln_pmf builds the multinomial coefficient directly with factorial::ln_factorial (ln n! - Σ ln x_i!) instead of factorial::multinomial(...).ln().
  • Zero-count categories are skipped in the probability sum so an allowed zero probability no longer produces 0 * ln(0) = NaN.

Why

factorial::multinomial returns the coefficient itself (it exponentiates the stable log value before returning), which overflows f64 to +inf once the log exceeds ~709.78. For n = 2000:

  • pmf computed inf * prod(p_i^x_i), and since the probability product simultaneously underflows to 0.0, the result was inf * 0.0 = NaN.
  • ln_pmf computed inf.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 existing test_pmf cases still pass under the log-space implementation.

$ CARGO_BUILD_JOBS=2 cargo test -p statrs multinomial
test result: ok. 10 passed; 0 failed; 0 ignored; 0 measured; 716 filtered out

cargo clippy -p statrs --all-targets -- -D warnings is clean; cargo fmt --check reports no diffs in the touched file (the only fmt diffs on this tree are pre-existing in geometric.rs).

…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
Copilot AI review requested due to automatic review settings July 8, 2026 17:44

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Multinomial PMF likely converts back from ln to exp too early

2 participants