Implement partial derivative wrt to shape parameter for `gamma_inc` as a convergent series

[SamuelBrand1]

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

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 (_grad_p_a_series), ported from Stan. Core loop:

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.

[seabbs]

This would be great to have implemented. For my part I have no prior experience writing rules and most/all the work was LLM driven as a port from the stan implementation. Implementing this would allow AD backends to work correctly with Distributions.cdf(Gamma) on top of our needs for the incomplete gamma directly.

Very happy to support or make a PR.