pyro-ppl / numpyro / 34692431167

93%

push

github

Add TruncatedGamma distributions (#2257)

* Add TruncatedGamma distributions

Add LeftTruncatedGamma, RightTruncatedGamma and TwoSidedTruncatedGamma,
dispatched by TruncatedGamma(concentration, rate, low=, high=), following the
structure outlined in #969. Supersedes the inactive #1187.

The two-sided normalizer switches to the upper-tail form when the interval lies
above the median: the lower-tail difference F(high) - F(low) is exactly zero in
single precision for Gamma(2, 1) on [30, 40]. Adds gammainccinv to
distributions/util.py, wrapping tfp.math.igammacinv.

* Address review: per-order tail switch, bisection icdf, guards

- `_partial_moment` re-decides the tail switch at `concentration + order`.
  Gamma(alpha+k) has a larger median, so an interval above the base median can
  sit deep in the lower tail at the shifted order, where the upper-tail form is
  the one that cancels. Reusing the order-zero branch collapsed the second
  moment to zero and produced a negative variance.
- `icdf` bisects the cdf instead of calling an inverse incomplete gamma, which
  saturates below a tail probability of ~3e-8 and returned draws outside the
  support once the retained mass got small. A Newton step from the bracketed
  root supplies the implicit derivative, so sampling stays differentiable. This
  removes the need for the `gammainccinv` wrapper, so `util.py` is untouched.
- `mean` and `variance` now return nan for an underflowed normalizer, matching
  `cdf`; documented why that differs from `log_prob` returning -inf.
- Documented the float32 cancellation in `E[X^2] - E[X]^2` for narrow intervals
  far from the origin.
- Fixed two line-wrapped reST cross-references, gave the base-distribution
  assert a message pointing at TruncatedGamma, dropped a test assertion that
  measured scipy's accuracy rather than ours, and restored the file's
  left/right/two-sided ordering in batch_util.

* Conform to the Array return annotations from #2206

Narrow the new classes' r... (continued)

247 of 247 new or added lines in 2 files covered. (100.0%)

15144 of 16348 relevant lines covered (92.64%)

1.69 hits per line