Accurate derivatives of gamma_inc and gamma_inc_dx

Project.toml

@@ -1,6 +1,6 @@
 name = "ForwardDiff"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
-version = "1.5.0"
+version = "1.5.1"
 
 [deps]
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"

src/dual.jl

@@ -877,12 +877,52 @@ function SpecialFunctions.logabsgamma(d::Dual{T,<:Real}) where {T}
 end
 
 # Derivatives wrt to first parameter and precision setting are not supported
-function SpecialFunctions.gamma_inc(a::Real, d::Dual{T,<:Real}, ind::Integer) where {T}
-    x = value(d)
+function SpecialFunctions.gamma_inc(a::Real, d::Dual{T}, ind::Integer) where {T}
+    x = d.value
     p, q = SpecialFunctions.gamma_inc(a, x, ind)
-    ∂p = exp(-x) * x^(a - 1) / SpecialFunctions.gamma(a) * partials(d)
+    ∂p = _gamma_inc_dx(a, x) * d.partials
     return (Dual{T}(p, ∂p), Dual{T}(q, -∂p))
 end
+SpecialFunctions.gamma_inc(a::Dual, x::Real, ind::Integer) = _gamma_inc_constant_a(a, x, ind)
+SpecialFunctions.gamma_inc(a::Dual, x::Dual, ind::Integer) = _gamma_inc_constant_a(a, x, ind)
+function _gamma_inc_constant_a(a::Dual, x::Real, ind::Integer)
+    iszero(a.partials) || throw(ArgumentError("derivatives with respect to the first argument `a` are not supported"))
+    return SpecialFunctions.gamma_inc(a.value, x, ind)
+end
+
+# TODO: Remove fallback once SpecialFunctions >= 2.10 is required
+if isdefined(SpecialFunctions, :gamma_inc_dx)
+    _gamma_inc_dx(a, x) = SpecialFunctions.gamma_inc_dx(a, x)
+    # The first argument `x` is the shape `a`, the second argument `y` is `x`
+    @define_binary_dual_op(
+        SpecialFunctions.gamma_inc_dx,
+        begin
+            va, vx = value(Txy, x), value(Txy, y)
+            f = SpecialFunctions.gamma_inc_dx(va, vx)
+            ∂f∂a = SpecialFunctions._gamma_inc_dx_da(va, vx, f)
+            ∂f∂x = SpecialFunctions._gamma_inc_dx_dx(va, vx, f)
+            return dual_definition_retval(Val{Txy}(), f, ∂f∂a, partials(Txy, x), ∂f∂x, partials(Txy, y))
+        end,
+        begin
+            va = value(Tx, x)
+            f = SpecialFunctions.gamma_inc_dx(va, y)
+            ∂f∂a = SpecialFunctions._gamma_inc_dx_da(va, y, f)
+            return dual_definition_retval(Val{Tx}(), f, ∂f∂a, partials(Tx, x))
+        end,
+        begin
+            vx = value(Ty, y)
+            f = SpecialFunctions.gamma_inc_dx(x, vx)
+            ∂f∂x = SpecialFunctions._gamma_inc_dx_dx(x, vx, f)
+            return dual_definition_retval(Val{Ty}(), f, ∂f∂x, partials(Ty, y))
+        end,
+    )
+else
+    function _gamma_inc_dx(a, x)
+        f = exp(-x) * x^(a - 1) / SpecialFunctions.gamma(a)
+        # avoid overflow for large `a`
+        return isfinite(f) || iszero(x) ? f : exp((a - 1) * log(x) - x - SpecialFunctions.loggamma(a))
+    end
+end
 
 ###################
 # Pretty Printing #

test/DualTest.jl

@@ -10,6 +10,7 @@ using NaNMath, SpecialFunctions, LogExpFunctions
 using DiffRules
 
 import Calculus
+import DiffResults
 import LinearAlgebra
 
 struct TestTag end
@@ -730,6 +731,107 @@ end
     @test ForwardDiff.derivative(x -> sum(1 .+ x .* (0:0.1:1)), 1) == 5.5
 end
 
+@testset "gamma_inc" begin
+    # Accurate derivatives for large `a` require SpecialFunctions >= 2.10
+    has_gamma_inc_dx = isdefined(SpecialFunctions, :gamma_inc_dx)
+    # ∂/∂x gamma_inc(a, x)[1] = x^(a - 1) e^(-x) / Γ(a)
+    ∂p(a, x) = exp((big(a) - 1) * log(big(x)) - big(x)) / gamma(big(a))
+
+    for V in (Float32, Float64), (a, x) in ((2.5, 1.0), (150, 150), (200, 180), (1e6, 1e6)), ind in 0:2
+        dp = ForwardDiff.derivative(x -> gamma_inc(V(a), x, ind)[1], V(x))
+        dq = ForwardDiff.derivative(x -> gamma_inc(V(a), x, ind)[2], V(x))
+        @test dp isa V
+        @test dp ≈ ∂p(a, x) rtol = 4 * eps(V) broken = !has_gamma_inc_dx && a > 100
+        @test isequal(dq, -dp)
+    end
+
+    d2p(a, x) = ForwardDiff.derivative(x -> ForwardDiff.derivative(x -> gamma_inc(a, x, 0)[1], x), x)
+    d3p(a, x) = ForwardDiff.derivative(x -> d2p(a, x), x)
+    for (a, x) in ((150.0, 150.0), (1e6, 1e6), (1e6, 999_000.0), (3.0, 1e-160))
+        @test @inferred(d2p(a, x)) ≈ ∂p(a, x) * (big(a) - 1 - big(x)) / big(x) rtol = 8 * eps() broken = !has_gamma_inc_dx && a > 100
+    end
+    @test ForwardDiff.derivative(x -> gamma_inc(1.0, x, 0)[1], 0.0) == 1
+    @test d2p(2.0, 0.0) == 1
+    @test d3p(3.0, 0.0) == 1
+
+    # Infinite derivative at zero
+    for f in (
+            x -> gamma_inc(0.5, zero(x), 0)[1],
+            x -> gamma_inc(0.5, zero(x), 0)[2],
+            y -> ForwardDiff.derivative(x -> gamma_inc(0.5, zero(x), 0)[1], y),
+        )
+        df = ForwardDiff.derivative(f, 1.0)
+        @test ForwardDiff.NANSAFE_MODE_ENABLED ? iszero(df) : isnan(df)
+    end
+
+    # Duals with zero partials are constants
+    a = Dual{TestTag}(2.0, 0.0)
+    x = Dual{TestTag}(1.5, 1.0)
+    @test gamma_inc(a, 1.5, 0) === gamma_inc(2.0, 1.5, 0)
+    @test gamma_inc(a, x, 0) === gamma_inc(2.0, x, 0)
+    @test gamma_inc(Dual{OuterTestTag}(a, zero(a)), x, 0) === gamma_inc(2.0, x, 0)
+    msg = "ArgumentError: derivatives with respect to the first argument `a` are not supported"
+    @test_throws msg ForwardDiff.derivative(a -> gamma_inc(a, 1.5, 0)[1], 2.0)
+    @test_throws msg ForwardDiff.derivative(a -> ForwardDiff.derivative(x -> gamma_inc(a, x, 0)[1], 1.5), 2.0)
+end
+
+@testset "gamma_inc_dx" begin
+    if isdefined(SpecialFunctions, :gamma_inc_dx)
+        # The x-derivative of `gamma_inc` is `gamma_inc_dx`
+        for (a, x) in ((2.5, 1.0), (150.0, 150.0))
+            @test ForwardDiff.derivative(x -> gamma_inc(a, x, 0)[1], x) === SpecialFunctions.gamma_inc_dx(a, x)
+        end
+
+        # f = x^(a - 1) e^(-x) / Γ(a), ∂f/∂a = f * (log(x) - ψ(a)), ∂f/∂x = f * (a - 1 - x) / x
+        function ∇f(a, x)
+            a, x = big(a), big(x)
+            l = (a - 1) * log(x) - x
+            e = exp(l)
+            # At the poles of ψ(a), f = 0 and ∂f/∂a = (-1)^a Γ(1 - a) x^(a - 1) e^(-x)
+            isinteger(a) && a <= 0 && return [cospi(a) * e * gamma(1 - a), zero(a)]
+            # 1 / Γ(a) = sin(πa) Γ(1 - a) / π
+            f = a > 0 ? exp(l - loggamma(a)) : e * sinpi(a) * gamma(1 - a) / big(π)
+            return [f * (log(x) - digamma(a)), f * (a - 1 - x) / x]
+        end
+        function ∇²f(a, x)
+            setprecision(BigFloat, 512) do
+                h = big(1e-40)
+                return hcat((∇f(a + h, x) - ∇f(a - h, x)) / 2h, (∇f(a, x + h) - ∇f(a, x - h)) / 2h)
+            end
+        end
+
+        for V in (Float32, Float64), (a, x) in ((2.5, 1.0), (0.3, 2.0), (-2.0, 1.5), (150.0, 150.0), (150.0, 149.5), (1e6, 1e6), (1e10, 1e10))
+            ∂a, ∂x = setprecision(() -> V.(∇f(V(a), V(x))), BigFloat, 512)
+            dfda = ForwardDiff.derivative(a -> SpecialFunctions.gamma_inc_dx(a, V(x)), V(a))
+            dfdx = ForwardDiff.derivative(x -> SpecialFunctions.gamma_inc_dx(V(a), x), V(x))
+            @test dfda isa V
+            @test dfda ≈ ∂a rtol = 8 * eps(V)
+            @test dfdx isa V
+            @test dfdx ≈ ∂x rtol = 8 * eps(V)
+            @test ForwardDiff.gradient(v -> SpecialFunctions.gamma_inc_dx(v[1], v[2]), [V(a), V(x)]) == [dfda, dfdx]
+        end
+
+        # `a` and `x` with different tags
+        d2fdadx(a, x) = ForwardDiff.derivative(a -> ForwardDiff.derivative(x -> SpecialFunctions.gamma_inc_dx(a, x), x), a)
+        d2fdxda(a, x) = ForwardDiff.derivative(x -> ForwardDiff.derivative(a -> SpecialFunctions.gamma_inc_dx(a, x), a), x)
+        for (a, x) in ((2.5, 1.0), (0.3, 2.0), (-2.0, 1.5), (150.0, 150.0), (1e6, 1e6), (1e10, 1e10))
+            H = Float64.(∇²f(a, x))
+            result = ForwardDiff.hessian!(DiffResults.HessianResult([a, x]), v -> SpecialFunctions.gamma_inc_dx(v[1], v[2]), [a, x])
+            # The gradient is computed with nested duals
+            @test DiffResults.gradient(result) ≈ setprecision(() -> Float64.(∇f(a, x)), BigFloat, 512) rtol = 8 * eps()
+            @test DiffResults.hessian(result) ≈ H rtol = 1e-14
+            @test @inferred(d2fdadx(a, x)) ≈ H[1, 2] rtol = 1e-14
+            @test d2fdxda(a, x) ≈ H[2, 1] rtol = 1e-14
+        end
+    else
+        # The x-derivative of `gamma_inc` falls back to the closed form
+        for (a, x) in ((2.5, 1.0), (30.0, 25.0))
+            @test ForwardDiff.derivative(x -> gamma_inc(a, x, 0)[1], x) === exp(-x) * x^(a - 1) / gamma(a)
+        end
+        @test ForwardDiff.derivative(x -> gamma_inc(150.0, x, 0)[1], 150.0) ≈ exp(149 * log(big(150)) - 150) / gamma(big(150)) rtol = 1e-12
+    end
+end
+
 @testset "Givens rotations: consistency with `LinearAlgebra.givensAlgorithm` for zero partials (no duals)" begin
     # Test different branches in `LinearAlgebra.givensAlgorithm`
     for f in [randexp(), -randexp()], g in [0.0, f / 2, 2f, -f / 2, -2f]