JuliaMath / JuliaMath/SpecialFunctions.jl
Implement partial derivative wrt to shape parameter for `gamma_inc` as a convergent series
- Dominant language
- Julia
- Stars
- 381
- Forks
- 113
- Avg merge
- 1d 23h
- Merged PRs (30d)
- 1
Description
### Summary
The shape-parameter (`a`) partial of `gamma_inc` is currently marked `@not_implemented`, tracked in #317. That issue looks at using hypergeometric functions (for mathematical context cf https://en.wikipedia.org/wiki/Incomplete_gamma_function#Partial_derivatives ), and I've personally tried out using the Meijer-G function from https://github.com/TSGut/MeijerG.jl , which got the correct partial derivate but was falling back on a slow contour integral approach.
For the regularized lower incomplete gamma `P(a, x)` we have an absolutely-convergent series representation. We can differentiate that termwise into a series that needs only `digamma` and `loggamma` which are in `SpecialFunctions.jl`. This is the approach Stan uses in `grad_reg_inc_gamma` on the small-`x` branch.
We have a working implementation of this at https://github.com/epiforecasts/BVDOutbreakSize/blob/main/src/gamma_cdf.jl and @seabbs has implemented this into https://github.com/EpiAware/CensoredDistributions.jl . Rather than just having a custom chain rule following specific projects around I think it would be better to jimmy this into something that fits for `SpecialFunctions.jl`
### Current behavior
```julia
ChainRulesCore.@scalar_rule(
gamma_inc(a, x, IND),
@setup(z = exp(-x) * x^(a - 1) / gamma(a)),
(
ChainRulesCore.@not_implemented(INCOMPLETE_GAMMA_INFO), # ∂/∂a ← this one
z,
ChainRulesCore.NoTangent(),
),
(
ChainRulesCore.@not_implemented(INCOMPLETE_GAMMA_INFO), # ∂/∂a ← and this one
-z,
ChainRulesCore.NoTangent(),
),
)
```
The `x` partials are the gamma density (and its negative for `Q`); only the `a` partials are missing.
### The approach
For real params the regularized lower incomplete gamma has the series representation (Temme, 1994 eq 2.1)
$$
P(a, x) = x^{a} e^{-x} \sum_{n=0}^{\infty} \frac{x^{n}}{\Gamma(a + n + 1)}.
$$
Because this converges we can do the term by term partials, apply some standard tricks and:
$$
\partial_a P(a, x) = \log(x) P(a, x) - x^{a} e^{-x} \sum_{n=0}^{\infty} \frac{\psi(a + n + 1) x^{n}}{\Gamma(a + n + 1)}.
$$
For the upper/complement, $\partial_a Q(a, z) = -\partial_a P(a, z)$.
Implementation-wise we just resolve the series above, with as standard stopping rule at a chosen series convergence. We be a bit careful not to recalculate factors where not necessary so using the recurrence $\psi(a+n+1) = \psi(a+n) + \tfrac{1}{a+n}$, each term costs a division, an add and two multiplies — no per-iteration `gamma`/`digamma` calls.
Whilst this is pretty basic (we haven't done any branching of approach as per Temme for the primal calculation), it does mean no new dependencies. Only `digamma` and `loggamma` are needed, both already here. I think this addresses the main concern in #317 avoiding a hypergeometric dependency.
For large $x$ the same branch logic Stan uses applies: compute via the complement or a continued-fraction route like in Temme so a complete implementation might want the branch selection, not just the series. This is the one thing worth agreeing on before a PR. But this simple implementation has worked well for us.
### Reference implementation
We use this in [epiforecasts/BVDOutbreakSize](https://github.com/epiforecasts/BVDOutbreakSize/blob/main/src/gamma_cdf.jl) (`_grad_p_a_series`), ported from Stan. Core loop:
```julia
function _grad_p_a_series(a, z; rtol = 1e-14, maxiter = 10_000)
z <= zero(z) && return zero(a) * zero(z)
log_term0 = a * log(z) - z - loggamma(a + 1)
term = exp(log_term0)
ψ = digamma(a + 1)
P = term
S = term * ψ
for n in 1:maxiter
term *= z / (a + n)
ψ += 1 / (a + n)
P += term
S += term * ψ
abs(term * ψ) <= rtol * abs(S) && abs(term) <= rtol * abs(P) && break
end
return log(z) * P - S
end
```
### Proposal
Replace the `@not_implemented(INCOMPLETE_GAMMA_INFO)` entries for the `a` partial with a `_grad_reg_inc_gamma_a(a, x)` helper implementing the series (small-`z`) and complement (large-`z`) branches, returning `+grad` for `P` and `-grad` for `Q`. Happy to PR this with tests against finite differences and against Stan's reference values if maintainers are on board, and to coordinate with #317.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.