Fast implementation of incomplete extended gamma function, QuadGK?

mleprovost
Thu, 28 Mar 2024 16:35:17 +0000
Numerics

Hello,

I am interested in developing a fast routine to compute the following extended incomplete gamma function for various pairs (x, q).

\gamma(\alpha, \kappa, p, q)(x) = \int_{0}^x t^{\alpha-1} \exp( -\frac{t^\kappa}{p} - \frac{q}{t}) \mathrm{d}t with \alpha \in \mathbb{R}, p>0, \kappa \in \mathbb{R}_{\neq 0}, q \geq 0

I would also be interested in computing its inverse \gamma^{-1}(\alpha, \kappa, p, q)(z) for z \geq 0.

I attach my current implementation based on QuadGK. I would appreciate any feedback to accelerate computations or reduce allocations.

Can it be beneficial to define the function \gamma as the solution of the ODE \frac{\mathrm{d} \gamma(\alpha, \kappa, p, q)(x) }{\mathrm{d} x} = x^{\alpha-1} \exp( -\frac{x^\kappa}{p} - \frac{q}{x}) , then pass this ODE to DifferentialEquations.jl.
Then we can compute its inverse by using a root-finding algorithm combined with the interpolation for free tools from DifferentialEquations

unnorm_pdf!(out, α::Float64, κ::Float64, p::Float64, q::Float64, t) = out .= t^(α-1)*exp(-t^κ/p - q/t)

# Routine to compute the un-normalized probability density function
function unnorm_pdf(α::Float64, κ::Float64, p::Float64, q::Float64, t)
    out = zeros(eltype(t),1)
    unnorm_pdf!(out, α, κ, p, q, t)
    return out
end

# Routine to compute the normalizing constant of the probability density function
function Z!(out, α::Float64, κ::Float64, p::Float64, q::Float64)
    quadgk!((b,a)->unnorm_pdf!(b,α, κ, p, q, a),out, 0.0, Inf)
    nothing
end

# Routine to compute the normalizing constant of the probability density function
function Z(α::Float64, κ::Float64, p::Float64, q::Float64)
    out = zeros(eltype(t),1)
    Z!(out, α, κ, p, q)
    return out
end

struct ExtendedGamma
    α::Float64
    κ::Float64
    p::Float64
    q::MVector{1, Float64}
    # Normalizing constant
    Z::MVector{1, Float64}
    cache::MVector{1, Float64}
end

function ExtendedGamma(α::Float64, κ::Float64, p::Float64, q::Float64)
    cache = MVector{1, Float64}(0.0)
    # Compute normalizing constant
    Z!(cache,α, κ, p, q)
    ZΓ = copy(cache)
    return ExtendedGamma(α, κ, p, MVector{1, Float64}(q), ZΓ, cache)
end

unnorm_pdf!(out, Γ::ExtendedGamma, t::Real) = unnorm_pdf!(out, Γ.α, Γ.κ, Γ.p, Γ.q[1], t)

function unnorm_pdf(Γ::ExtendedGamma, t::Real)
    unnorm_pdf!(Γ.cache, Γ, t)
    return Γ.cache
end

function pdf!(out, Γ::ExtendedGamma, t::Real)
    unnorm_pdf!(out, Γ, t)
    out ./= Γ.Z[1]
    nothing
end

function pdf(Γ::ExtendedGamma, t::Real)
    pdf!(Γ.cache, Γ, t)
    return Γ.cache
end

Z!(out, Γ::ExtendedGamma) = Z!(out, Γ.α, Γ.κ, Γ.p, Γ.q[1])

function Z(Γ::ExtendedGamma)
    Z!(Γ.cache, Γ)
    return Γ.cache
end

update_Z!(Γ::ExtendedGamma) = Z!(Γ.Z, Γ)

function cdf!(out, Γ::ExtendedGamma, x)
    quadgk!((b,a)-> unnorm_pdf!(b, Γ, a),out, 0.0, x)[1]
    out ./= Γ.Z[1]
    nothing
end

function cdf(Γ::ExtendedGamma, x)
    cdf!(Γ.cache, Γ, x)
    return Γ.cache
end

# Complementary of the CDF to 1
function ccdf!(out, Γ::ExtendedGamma, x)
    quadgk!((b,a)-> unnorm_pdf!(b, Γ, a),out, x, Inf)[1]
    out ./= Γ.Z[1]
    nothing
end

# Complementary of the CDF to 1
function ccdf(Γ::ExtendedGamma, x)
    ccdf!(Γ.cache, Γ, x)
    return Γ.cache
end

Example

using LinearAlgebra
using QuadGK
using StaticArrays
using BenchmarkTools

α = 2.0
κ = 2.0
p = 1.0
q = 0.8
Γstruct = ExtendedGamma(α, κ, p, q)

# Time evaluation of the incomplete expanded Gamma function
@btime begin
    cdf!(Γstruct.cache, Γstruct, 1.0)
end

Output

 2.185 μs (18 allocations: 672 bytes)
1-element MVector{1, Float64} with indices SOneTo(1):
 0.40523030232449925

5 posts - 3 participants