From 96755292cbd0d96b8fcf7dbd994d23197b2f3199 Mon Sep 17 00:00:00 2001 From: simonsteiger Date: Wed, 2 Sep 2026 23:55:28 +0200 Subject: [PATCH 1/3] Add Gamma MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Gamma(α, θ)` takes a shape and a scale, so its mean is `α * θ`, and `Gamma(α)` sets the scale to one. The density is closed form and traces; the distribution functions are not. `loggammap` and `loggammaq` give the regularized incomplete gamma integrals in log space. Returning logarithms is what keeps `logcdf` and `logccdf` finite where the probabilities underflow, and costs nothing: each tail is already built from a logarithmic prefactor. They sum a series below `x = a + 1` and run a continued fraction above it, both in the type they are given, so `BigFloat` keeps its precision. `quantile` then inverts the tail with Newton's method on `log(x)`, which is what lets the deep lower tail return a subnormal rather than zero. All of them loop until their terms stop changing the result, so none can run in traced or device-side code. Relative accuracy sits at the rounding error of the argument type for shapes up to about 1000 and falls off roughly in proportion to the shape after that, since the prefactor's terms grow while their sum does not. Sampling uses Marsaglia and Tsang's rejection method, boosted by `Gamma(α, θ) = Gamma(α + 1, θ) · U^(1/α)` below a unit shape. The accept step reads `basevalue(α)`, a new interface function returning the plain floating-point value inside a wrapped number, with methods in the ForwardDiff and ReverseDiff extensions and an Enzyme inactivity rule. The loop therefore runs on plain numbers whatever type the parameters carry, and the accepted noise enters the draw through arithmetic on `α` and `θ`, which leaves the draw differentiable with respect to both. The README lists of implemented measures also gain `Poisson`, which they had been missing. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01HiwLngD494i4dSLJ2DGPBG --- CHANGELOG.md | 20 ++ README.md | 21 +- ext/ProbabilityMeasuresEnzymeCoreExt.jl | 3 + ext/ProbabilityMeasuresForwardDiffExt.jl | 7 +- ext/ProbabilityMeasuresReverseDiffExt.jl | 7 +- .../src/implementations.jl | 11 + src/ProbabilityMeasures.jl | 7 +- src/core/gammainc.jl | 99 +++++++ src/core/mathfuns.jl | 13 + src/univariate/continuous/gamma.jl | 208 +++++++++++++++ test/Project.toml | 1 + test/test-gamma.jl | 244 ++++++++++++++++++ 12 files changed, 631 insertions(+), 10 deletions(-) create mode 100644 src/core/gammainc.jl create mode 100644 src/univariate/continuous/gamma.jl create mode 100644 test/test-gamma.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 049fefc..bc3df98 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,26 @@ and this project adheres to [Semantic Versioning]. substitute fall back to a truncated sum that costs `O(λ)` and cannot run in traced or device-side code. Entropy always sums. Sampling uses CDF inversion to avoid underflow at large rates. These operations return `NaN` if the truncation bound exceeds `Int`. +- `Gamma(α, θ)`, taking a shape and a scale so that the mean is `α * θ`. Its density is + closed form, but its distribution functions are not: `cdf`, `ccdf`, `logcdf`, + `logccdf`, `quantile`, `median` and `entropy` sum a series, run a continued fraction, + or iterate Newton's method until the terms stop changing the result. They work in the + type they are given rather than dropping to `Float64`, so `BigFloat` keeps its + precision, and the price is that they cannot run in traced or device-side code. +- `loggammap(a, x)` and `loggammaq(a, x)`, the regularized incomplete gamma integrals in + log space. Returning logarithms is what keeps `Gamma`'s `logcdf` and `logccdf` finite + where the probabilities themselves underflow, and it costs nothing: each tail is + already computed from a logarithmic prefactor. Relative accuracy sits at the rounding + error of the argument type for shapes up to about `1000`, then falls off roughly in + proportion to the shape, since the prefactor's terms grow while their sum does not: + measured against `SpecialFunctions.gamma_inc` it is `5e-13` at shape `1000` and `2e-10` + at shape `10^5`. +- `basevalue(x)`, the plain floating-point value inside a wrapped number, with methods in + the ForwardDiff and ReverseDiff extensions and an Enzyme inactivity rule. `Gamma` is + the first measure whose sampler rejects, and its accept step reads `basevalue(α)`. The + loop therefore runs on plain numbers whatever type the parameters carry, and the + accepted noise enters the draw through arithmetic on `α` and `θ`, which is what leaves + the draw differentiable with respect to both. - `validateparams(d)`, which returns `d` or throws a `DomainError`, for the boundary where user-supplied parameters enter. It earns its place on `Categorical`, whose sum-to-one is the one invalid parameter a density cannot report: an unnormalized `p` diff --git a/README.md b/README.md index 2c8c4fc..5c321b0 100644 --- a/README.md +++ b/README.md @@ -11,8 +11,8 @@ density and sampling operations, and compatible with automatic differentiation, broadcasting on GPU arrays, and Reactant tracing. The package is experimental. At present it implements `Normal`, `LogNormal`, -`Exponential`, `Uniform`, `Laplace`, `Cauchy`, `Categorical`, `Bernoulli`, `Binomial`, `MvNormal`, -and `Multinomial`. +`Exponential`, `Uniform`, `Laplace`, `Cauchy`, `Gamma`, `Categorical`, `Bernoulli`, +`Binomial`, `Poisson`, `MvNormal`, and `Multinomial`. ## Installation @@ -62,7 +62,8 @@ than throwing. ## Available API `Normal(μ, σ)`, `LogNormal(μ, σ)`, `Exponential(θ)`, `Uniform(a, b)`, -`Laplace(μ, b)`, `Categorical(p)`, `Cauchy`, `Bernoulli(p)`, and `Binomial(n, p)` each support: +`Laplace(μ, b)`, `Cauchy(μ, σ)`, `Gamma(α, θ)`, `Categorical(p)`, `Bernoulli(p)`, +`Binomial(n, p)`, and `Poisson(λ)` each support: - `densityof` and `logdensityof` - `cdf`, `ccdf`, `logcdf`, and `logccdf` @@ -73,6 +74,14 @@ than throwing. `Cauchy(μ, σ)` has no finite mean or variance, so `mean`, `var` and `std` return `NaN`. `median` and `entropy` are exact. +`Gamma(α, θ)` takes a shape and a scale, so its mean is `α * θ`, and `Gamma(α)` sets the +scale to one. Its density is closed form, but its distribution functions are not: `cdf`, +`ccdf`, `logcdf`, `logccdf`, `quantile`, `median` and `entropy` sum a series or iterate +until the terms stop changing the result. They work in the type they are given, so +`BigFloat` keeps its precision, but they cannot run in traced or device-side code. +Sampling has no such limit: it uses rejection, and the accept step runs on plain +floating-point noise, which leaves the draw differentiable with respect to `α` and `θ`. + `Categorical(p)` assigns the probabilities in `p` to categories `1:length(p)`. Draws and quantiles use the promoted floating-point type of `p`: @@ -204,9 +213,9 @@ See the [contribution guide](docs/src/90-contributing.md) for contribution guide ## Current scope ProbabilityMeasures.jl currently contains `Normal`, `LogNormal`, `Exponential`, -`Uniform`, `Cauchy`, `Laplace`, `Categorical`, `Bernoulli`, `Binomial`, `MvNormal`, and -`Multinomial`. Transformed or composite measures and Distributions.jl interoperability -are not implemented yet. +`Uniform`, `Cauchy`, `Laplace`, `Gamma`, `Categorical`, `Bernoulli`, `Binomial`, +`Poisson`, `MvNormal`, and `Multinomial`. Transformed or composite measures and +Distributions.jl interoperability are not implemented yet. ## Citation diff --git a/ext/ProbabilityMeasuresEnzymeCoreExt.jl b/ext/ProbabilityMeasuresEnzymeCoreExt.jl index 799334e..3f1f30c 100644 --- a/ext/ProbabilityMeasuresEnzymeCoreExt.jl +++ b/ext/ProbabilityMeasuresEnzymeCoreExt.jl @@ -11,6 +11,9 @@ EnzymeRules.inactive(::typeof(ProbabilityMeasures.insupport), args...) = nothing EnzymeRules.inactive(::typeof(ProbabilityMeasures.noisetype), args...) = nothing EnzymeRules.inactive(::typeof(ProbabilityMeasures.basefloat), args...) = nothing +# A rejection sampler's accept step is a choice, not a value the draw depends on. +EnzymeRules.inactive(::typeof(ProbabilityMeasures.basevalue), args...) = nothing + # Supports contain no differentiable data. EnzymeRules.inactive_type(::Type{<:ProbabilityMeasures.Support}) = true diff --git a/ext/ProbabilityMeasuresForwardDiffExt.jl b/ext/ProbabilityMeasuresForwardDiffExt.jl index 408d89f..2ff367d 100644 --- a/ext/ProbabilityMeasuresForwardDiffExt.jl +++ b/ext/ProbabilityMeasuresForwardDiffExt.jl @@ -1,6 +1,6 @@ module ProbabilityMeasuresForwardDiffExt -using ForwardDiff: Dual +using ForwardDiff: Dual, value using ProbabilityMeasures: ProbabilityMeasures # Draw noise in the plain float type. Dual parameters still affect the returned sample. @@ -8,4 +8,9 @@ function ProbabilityMeasures.basefloat(::Type{<:Dual{T,V,N}}) where {T,V,N} return ProbabilityMeasures.basefloat(V) end +# Keep a rejection sampler's accept step off the derivative. +function ProbabilityMeasures.basevalue(x::Dual) + return ProbabilityMeasures.basevalue(value(x)) +end + end diff --git a/ext/ProbabilityMeasuresReverseDiffExt.jl b/ext/ProbabilityMeasuresReverseDiffExt.jl index 06ce7fe..ff3424c 100644 --- a/ext/ProbabilityMeasuresReverseDiffExt.jl +++ b/ext/ProbabilityMeasuresReverseDiffExt.jl @@ -1,11 +1,16 @@ module ProbabilityMeasuresReverseDiffExt using ProbabilityMeasures: ProbabilityMeasures -using ReverseDiff: TrackedReal +using ReverseDiff: TrackedReal, value # `float(TrackedReal)` is still wrapped, but random noise needs the plain type. function ProbabilityMeasures.basefloat(::Type{TrackedReal{V,D,O}}) where {V,D,O} return ProbabilityMeasures.basefloat(V) end +# Keep a rejection sampler's accept step off the tape. +function ProbabilityMeasures.basevalue(x::TrackedReal) + return ProbabilityMeasures.basevalue(value(x)) +end + end diff --git a/libs/ProbabilityMeasuresTest/src/implementations.jl b/libs/ProbabilityMeasuresTest/src/implementations.jl index c6918cd..06fdb70 100644 --- a/libs/ProbabilityMeasuresTest/src/implementations.jl +++ b/libs/ProbabilityMeasuresTest/src/implementations.jl @@ -54,6 +54,17 @@ function _invalids(::Cauchy) end _exactparams(::Cauchy) = Cauchy(0, 2) +@implements MeasureInterface{UNIVARIATE_OPTIONALS} Gamma [ + Gamma(2.0, 1.0), Gamma(0.5, 3.0), Gamma(4.5f0, 0.5f0) +] + +function _invalids(::Gamma) + return (Gamma(-1.0, 1.0), Gamma(0.0, 1.0), Gamma(1.0, -1.0), Gamma(Inf, 1.0)) +end + +# Use a shape and a scale that leave `loggamma(α)` and `log(θ)` non-zero. +_exactparams(::Gamma) = Gamma(3, 2) + @implements MeasureInterface{UNIVARIATE_OPTIONALS} Categorical [ Categorical([0.2, 0.3, 0.5]), Categorical([1.0]), Categorical(Float32[0.25, 0.75]) ] diff --git a/src/ProbabilityMeasures.jl b/src/ProbabilityMeasures.jl index a96aec6..794ceba 100644 --- a/src/ProbabilityMeasures.jl +++ b/src/ProbabilityMeasures.jl @@ -15,12 +15,13 @@ using DensityInterface: DensityInterface, densityof, logdensityof using IrrationalConstants: invsqrt2, log2π, logπ, logtwo, sqrt2 using LinearAlgebra: Diagonal, LowerTriangular, UniformScaling using Random: Random, AbstractRNG -using SpecialFunctions: erfc, erfcinv, gamma_inc, logerfc, loggamma +using SpecialFunctions: digamma, erfc, erfcinv, gamma_inc, logerfc, loggamma using Statistics: Statistics, cov, mean, median, quantile, std, var using StatsAPI: StatsAPI, params include("core/types.jl") include("core/mathfuns.jl") +include("core/gammainc.jl") include("core/support.jl") include("core/interface.jl") @@ -30,6 +31,7 @@ include("univariate/continuous/exponential.jl") include("univariate/continuous/uniform.jl") include("univariate/continuous/laplace.jl") include("univariate/continuous/cauchy.jl") +include("univariate/continuous/gamma.jl") include("univariate/discrete/categorical.jl") include("univariate/discrete/bernoulli.jl") @@ -61,7 +63,7 @@ export Support, export support, insupport # Interface -export checkparams, validateparams, noisetype, basefloat +export checkparams, validateparams, noisetype, basefloat, basevalue export cdf, ccdf, logcdf, logccdf, entropy # Re-export common operations from package dependencies. @@ -76,6 +78,7 @@ export Exponential export Uniform export Laplace export Cauchy +export Gamma export Categorical export Bernoulli export Binomial diff --git a/src/core/gammainc.jl b/src/core/gammainc.jl new file mode 100644 index 0000000..84166b7 --- /dev/null +++ b/src/core/gammainc.jl @@ -0,0 +1,99 @@ +#= + The regularized incomplete gamma integrals, in log space. Returning logarithms keeps + a tail readable where the probability itself underflows, and every step stays in the + type of the arguments, so `BigFloat` keeps its precision. +=# + +#= + Both expansions converge geometrically once `x` is on their own side of `a + 1`. At + the crossover they need about `8.5√a` terms in `Float64`, and more at higher + precision, so this bound covers `Float64` shapes up to roughly `10^6`. Past it the + remaining terms are dropped and the result loses accuracy. +=# +const GAMMAINC_MAXITER = 10_000 + +""" + loggammap(a, x) + +`log(P(a, x))`, where ``P`` is the lower incomplete gamma integral divided by +``\\Gamma(a)``: + +```math +P(a, x) = \\frac{1}{\\Gamma(a)} \\int_0^x t^{a-1} e^{-t} \\, \\mathrm{d}t. +``` + +`a` must be positive and `x` non-negative. Below `x = a + 1` this sums a series whose +terms are all positive; above it, it takes the complement of [`loggammaq`](@ref). Both +loop until the terms stop changing the result, which rules out tracing and device-side +evaluation. + +The prefactor is a difference of terms of order ``a \\log a``, so relative accuracy +holds to the rounding error of the argument type for shapes up to about `1000` and then +falls off roughly in proportion to the shape. + +See also [`loggammaq`](@ref). +""" +@inline function loggammap(a::T, x::T) where {T<:Number} + return x < a + one(T) ? gammap_series(a, x) : log1mexpt(gammaq_cf(a, x)) +end + +""" + loggammaq(a, x) + +`log(Q(a, x))`, where ``Q(a, x) = 1 - P(a, x)`` is the upper incomplete gamma integral +divided by ``\\Gamma(a)``. + +Above `x = a + 1` this evaluates a continued fraction; below it, it takes the +complement of [`loggammap`](@ref). +""" +@inline function loggammaq(a::T, x::T) where {T<:Number} + return x < a + one(T) ? log1mexpt(gammap_series(a, x)) : gammaq_cf(a, x) +end + +#= + `P(a, x) = x^a e^{-x} / Γ(a+1) · Σₙ xⁿ / ((a+1)⋯(a+n))`. Every term is positive, so + the sum loses nothing to cancellation. +=# +function gammap_series(a::T, x::T) where {T<:Number} + tol = eps(basefloat(T)) + term = one(T) + total = one(T) + n = 0 + while (n < GAMMAINC_MAXITER) & (abs(term) > tol * abs(total)) + n += 1 + term *= x / (a + n) + total += term + end + return muladd(a, logt(x), -x) - loggamma(a + one(T)) + logt(total) +end + +#= + `Q(a, x) = x^a e^{-x} / Γ(a) · CF`, with the continued fraction + + CF = 1/(b₀ - a₁/(b₁ - a₂/(b₂ - ⋯))), b₀ = x + 1 - a, bᵢ = b₀ + 2i, aᵢ = i(i - a) + + evaluated by Lentz's method, which builds the fraction from the front and so needs no + guess at where to truncate it. `tiny` replaces a denominator that rounds to zero, + which is what lets the recurrence step past a vanishing partial numerator. +=# +function gammaq_cf(a::T, x::T) where {T<:Number} + tol = eps(basefloat(T)) + tiny = convert(T, floatmin(basefloat(T))) / tol + b = x + one(T) - a + c = inv(tiny) + d = inv(b) + h = d + for i in 1:GAMMAINC_MAXITER + an = -i * (i - a) + b += 2 * one(T) + d = an * d + b + abs(d) < tiny && (d = tiny) + c = b + an / c + abs(c) < tiny && (c = tiny) + d = inv(d) + delta = d * c + h *= delta + abs(delta - one(T)) <= tol && break + end + return muladd(a, logt(x), -x) - loggamma(a) + logt(h) +end diff --git a/src/core/mathfuns.jl b/src/core/mathfuns.jl index b4d02ad..40f3cfa 100644 --- a/src/core/mathfuns.jl +++ b/src/core/mathfuns.jl @@ -66,3 +66,16 @@ basefloat(::Type{T}) where {T<:AbstractFloat} = T basefloat(::Type{T}) where {T<:Real} = float(T) basefloat(::Type{Bool}) = Float64 basefloat(::Type{<:Irrational}) = Float64 + +""" + basevalue(x) -> AbstractFloat + +The plain floating-point value inside `x`. + +A rejection sampler compares against it so that its accept step runs on plain numbers +and stays independent of the wrapped numeric types automatic differentiation and +tracing systems substitute. The accepted noise then enters the draw through arithmetic +on the parameters, which is what makes the draw differentiable. Package extensions add +methods for wrapped types. +""" +basevalue(x::Number) = float(x) diff --git a/src/univariate/continuous/gamma.jl b/src/univariate/continuous/gamma.jl new file mode 100644 index 0000000..4b7cd69 --- /dev/null +++ b/src/univariate/continuous/gamma.jl @@ -0,0 +1,208 @@ +""" + Gamma(α, θ) + Gamma(α) + +The gamma measure on ``(0, \\infty)`` with shape `α` and scale `θ`. Its density is + +```math +p(x) = \\frac{x^{\\alpha - 1} e^{-x/\\theta}}{\\Gamma(\\alpha)\\, \\theta^{\\alpha}} +``` + +The mean is ``\\alpha\\theta``, matching Distributions.jl. `Gamma(α)` sets the scale to +one. `Gamma(1, θ)` is `Exponential(θ)`. + +# Arguments + + - `α::Number`: the shape. + - `θ::Number`: the scale. + +The constructor does not check its arguments. Invalid parameters give a non-finite +density. Use [`checkparams`](@ref) to check them when needed. + +```julia +checkparams(Gamma(-1.0, 1.0)) # false +isnan(logdensityof(Gamma(-1.0, 1.0), 1.0)) # true +``` + +`logdensityof` is closed form, so it broadcasts on device arrays and traces. `cdf`, +`ccdf`, `logcdf`, `logccdf`, `quantile`, `median` and `entropy` have no closed form and +iterate until their terms stop changing the result, which rules out traced and +device-side evaluation; see [`loggammap`](@ref). + +Sampling uses Marsaglia and Tsang's rejection method. The accept step runs on plain +floating-point noise, so it costs the same whatever numeric type the parameters carry, +and the accepted noise enters the draw through arithmetic on `α` and `θ`, leaving the +draw differentiable with respect to both. +""" +struct Gamma{A<:Number,T<:Number} <: ContinuousUnivariateMeasure + α::A + θ::T +end + +Gamma(α::Number) = Gamma(α, one(α)) + +Base.eltype(::Type{Gamma{A,T}}) where {A,T} = float(promote_type(A, T)) + +function checkparams(d::Gamma) + return isfinite(d.α) & (d.α > zero(d.α)) & isfinite(d.θ) & (d.θ > zero(d.θ)) +end + +support(::Gamma) = PositiveReals() + +""" + valuetype(d::Gamma, x) + +The floating-point type of a density, tail probability, or quantile at `x`. + +It promotes the parameter types with the type of `x`, so exact parameters keep the +argument's precision. +""" +@inline function valuetype(d::Gamma, x::Number) + return float(promote_type(typeof(d.α), typeof(d.θ), typeof(x))) +end + +@inline function DensityInterface.logdensityof(d::Gamma, x::Number) + T = valuetype(d, x) + α, θ, y = convert(T, d.α), convert(T, d.θ), convert(T, x) + # `loggamma` throws for a non-positive argument, so an invalid shape takes `NaN`. + lg = select(α > zero(T), () -> loggamma(α), () -> convert(T, NaN)) + v = muladd(α - one(T), logt(y), -(y / θ)) - lg - α * logt(θ) + # Convert exact values to a float before returning `-Inf`. + return select(insupport(d, y), () -> v, () -> convert(T, -Inf)) +end + +""" + gammanoise(rng, a) + +The standard normal draw that Marsaglia and Tsang's method accepts for shape `a >= 1`. + +``(a - 1/3)(1 + z/\\sqrt{9a - 3})^3`` is then a draw from the unit-scale gamma measure. +""" +function gammanoise(rng::AbstractRNG, a::F) where {F<:AbstractFloat} + c = a - one(F) / 3 + w = inv(sqrt(9 * c)) + while true + z = randn(rng, F) + v = muladd(w, z, one(F)) + if v > zero(F) + u = rand(rng, F) + v3 = v^3 + # The squeeze accepts most draws without reaching the logarithms. + if u < muladd(-F(0.0331), z^4, one(F)) || + log(u) < z^2 / 2 + c * (one(F) - v3 + log(v3)) + return z + end + end + end +end + +#= + Marsaglia and Tsang's method needs a shape of at least one. Below that, the identity + `Gamma(α, θ) = Gamma(α + 1, θ) · U^(1/α)` with `U` uniform supplies the rest. +=# +@inline function Base.rand(rng::AbstractRNG, d::Gamma) + F = noisetype(d) + boost = basevalue(d.α) < one(F) + α = boost ? d.α + one(d.α) : d.α + z = gammanoise(rng, convert(F, basevalue(α))) + # Apply the parameters after drawing noise so automatic differentiation can follow + # them. This must repeat `gammanoise`'s arithmetic to land on the accepted draw. + c = α - one(α) / 3 + x = d.θ * c * muladd(inv(sqrt(9 * c)), z, one(c))^3 + return boost ? x * rand(rng, F)^inv(d.α) : x +end + +Statistics.mean(d::Gamma) = d.α * d.θ +Statistics.var(d::Gamma) = d.α * d.θ^2 + +function entropy(d::Gamma) + α = float(d.α) + # `loggamma` and `digamma` throw for a non-positive argument. + shape = select( + α > zero(α), () -> loggamma(α) + (one(α) - α) * digamma(α), () -> oftype(α, NaN) + ) + return α + logt(float(d.θ)) + shape +end + +#= + The four distribution functions differ only in the tail they take and in the two + values they hold outside `(0, ∞)`, where the answer is known without computing + anything. +=# +@inline function gammatail(tail, d::Gamma, x::Number, below, above) + T = valuetype(d, x) + checkparams(d) || return convert(T, NaN) + y = convert(T, x) / convert(T, d.θ) + isnan(y) && return convert(T, NaN) + y > zero(T) || return convert(T, below) + isfinite(y) || return convert(T, above) + return tail(convert(T, d.α), y) +end + +cdf(d::Gamma, x::Number) = gammatail((a, y) -> exp(loggammap(a, y)), d, x, 0, 1) +ccdf(d::Gamma, x::Number) = gammatail((a, y) -> exp(loggammaq(a, y)), d, x, 1, 0) +logcdf(d::Gamma, x::Number) = gammatail(loggammap, d, x, -Inf, 0) +logccdf(d::Gamma, x::Number) = gammatail(loggammaq, d, x, 0, -Inf) + +# Newton's method converges quadratically, so a starting point good to a few digits +# reaches `BigFloat` precision well inside this bound. +const GAMMAQUANTILE_MAXITER = 100 + +""" + gammaquantile(a, p) + +The `p`-quantile of the unit-scale gamma measure with shape `a`, for `0 < p < 1`. + +Newton's method on `log(x)` refines a closed-form starting point until the step stops +moving it. The logarithm is what keeps the deep lower tail, where the quantile itself +underflows, from collapsing onto zero at the first step. +""" +function gammaquantile(a::T, p::T) where {T<:Number} + #= + Solve on whichever tail holds the smaller probability. The other tail is near one, + where its logarithm is flat and Newton's method has almost no slope to descend. + `1 - p` is exact above one half, so the split costs no precision. + =# + lower = p <= one(T) / 2 + target = lower ? logt(p) : log1p(-p) + + u = gammaquantile_start(a, p, lower) + tol = eps(basefloat(T)) + for _ in 1:GAMMAQUANTILE_MAXITER + y = exp(u) + (isfinite(y) & (y > zero(T))) || break + # The log-density of the unit-scale measure, written in `u` so that it stays + # accurate where `y` underflows. + g = muladd(a - one(T), u, -y) - loggamma(a) + tail = lower ? loggammap(a, y) : loggammaq(a, y) + # The tails run in opposite directions, so their Newton steps do too. + step = (tail - target) * exp(tail - g - u) + u = lower ? u - step : u + step + abs(step) <= tol * max(abs(u), one(T)) && break + end + return exp(u) +end + +@inline function gammaquantile_start(a::T, p::T, lower::Bool) where {T<:Number} + # Small probabilities follow `P(a, y) ≈ y^a / Γ(a+1)`, which inverts directly. + logr = (logt(p) + loggamma(a + one(T))) / a + lower && logr < logt((one(T) + a) / 5) && return logr + # Elsewhere, Wilson and Hilferty's cube-root normal approximation. + z = -(sqrt2 * erfcinvt(2 * p)) + w = a * (one(T) - inv(9 * a) + z / (3 * sqrt(a)))^3 + return (isfinite(w) & (w > zero(T))) ? logt(w) : logr +end + +function Statistics.quantile(d::Gamma, p::Number) + T = valuetype(d, p) + checkparams(d) || return convert(T, NaN) + q = convert(T, p) + (isnan(q) | (q < zero(T)) | (q > one(T))) && return convert(T, NaN) + iszero(q) && return zero(T) + isone(q) && return convert(T, Inf) + return convert(T, d.θ) * gammaquantile(convert(T, d.α), q) +end + +function Base.show(io::IO, d::Gamma) + return print(io, "Gamma(α=", d.α, ", θ=", d.θ, ")") +end diff --git a/test/Project.toml b/test/Project.toml index e2629a5..6d25d1b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,6 +8,7 @@ ProbabilityMeasures = "d96443b0-2461-4412-bb6f-797ac5201ef9" ProbabilityMeasuresTest = "8c1f4a3e-27bd-4b0a-9f4d-6e2b1c5a7d90" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/test-gamma.jl b/test/test-gamma.jl new file mode 100644 index 0000000..4137b65 --- /dev/null +++ b/test/test-gamma.jl @@ -0,0 +1,244 @@ +using ProbabilityMeasures +using ProbabilityMeasuresTest: test_measure +using Distributions: Distributions +using ForwardDiff: ForwardDiff +using Random: Random, Xoshiro +using SpecialFunctions: digamma, gamma_inc +using Test + +@testset "conformance" begin + # Widen the parameters because Distributions.jl works at their own precision. + function reference_logpdf(m, x) + return Distributions.logpdf(Distributions.Gamma(Float64(m.α), Float64(m.θ)), x) + end + for d in (Gamma(2.0, 1.0), Gamma(0.5, 3.0), Gamma(4.5f0, 0.5f0), Gamma(3, 2)) + test_measure(d; name=string(d), reference_logpdf=reference_logpdf) + end +end + +@testset "traits" begin + d = Gamma(2.0, 3.0) + @test d isa AbstractProbabilityMeasure{Univariate,Continuous} + @test d isa ContinuousUnivariateMeasure + @test !(d isa DiscreteUnivariateMeasure) + @test string(d) == "Gamma(α=2.0, θ=3.0)" + @test params(d) === (α=2.0, θ=3.0) +end + +@testset "no promotion at construction" begin + dual = ForwardDiff.Dual(1.0, 1.0) + @test typeof(Gamma(dual, 1.0)) === Gamma{typeof(dual),Float64} + @test typeof(Gamma(2, 1.0)) === Gamma{Int,Float64} + @test typeof(Gamma(2.0f0, 1.0f0)) === Gamma{Float32,Float32} + + # The one-argument form sets the scale to one in the shape's own type. + @test Gamma(2) === Gamma(2, 1) + @test Gamma(2.0f0) === Gamma(2.0f0, 1.0f0) + + @test eltype(Gamma(2, 1)) === Float64 + @test eltype(Gamma(2.0f0, 1.0f0)) === Float32 + @test isbits(Gamma(2.0, 3.0)) +end + +@testset "precision follows the argument, not the parameters" begin + @test logdensityof(Gamma(3, 2), 1.0f0) isa Float32 + @test logdensityof(Gamma(3, 2), big"1.0") isa BigFloat + + # Integer parameters must not reduce `BigFloat` precision. + exact = logdensityof(Gamma(3, 2), big"1.0") + full = logdensityof(Gamma(big"3.0", big"2.0"), big"1.0") + @test abs(exact - full) < 1e-70 + + # Exact rational inputs must still return floating-point values. + @test logdensityof(Gamma(3, 2), 1//2) isa Float64 + @test logdensityof(Gamma(3//1, 2//1), 1//2) isa Float64 + @test logdensityof(Gamma(3, 2), -1//2) === -Inf +end + +@testset "construction never validates" begin + for d in (Gamma(-1.0, 1.0), Gamma(0.0, 1.0), Gamma(1.0, -1.0), Gamma(1.0, 0.0)) + @test !checkparams(d) + @test !isfinite(logdensityof(d, 1.0)) + end + @test !checkparams(Gamma(Inf, 1.0)) + @test !checkparams(Gamma(1.0, NaN)) + @test checkparams(Gamma(2.0, 3.0)) + + @test_throws DomainError validateparams(Gamma(-1.0, 1.0)) + @test validateparams(Gamma(2.0, 3.0)) === Gamma(2.0, 3.0) +end + +@testset "invalid parameters give NaN, not a partial answer" begin + for d in (Gamma(-1.0, 1.0), Gamma(2.0, -1.0), Gamma(NaN, 1.0)) + @test isnan(cdf(d, 1.0)) + @test isnan(ccdf(d, 1.0)) + @test isnan(logcdf(d, 1.0)) + @test isnan(logccdf(d, 1.0)) + @test isnan(quantile(d, 0.5)) + @test isnan(entropy(d)) + end +end + +@testset "support" begin + d = Gamma(2.0, 3.0) + @test support(d) === PositiveReals() + @test minimum(support(d)) === 0.0 + @test maximum(support(d)) === Inf + + @test insupport(d, 1e-300) + @test insupport(d, 1e300) + @test !insupport(d, 0.0) + @test !insupport(d, -1.0) + @test !insupport(d, Inf) + @test !insupport(d, NaN) +end + +@testset "density is total off the support" begin + # A shape below one has an infinite density at zero, one above it a zero density. + for d in (Gamma(0.5, 1.0), Gamma(1.0, 1.0), Gamma(2.0, 1.0)) + for x in (0.0, -1.0, -Inf, NaN, -floatmax(Float64)) + @test logdensityof(d, x) == -Inf + end + @test !isfinite(logdensityof(d, Inf)) + end +end + +@testset "a unit shape is the exponential measure" begin + for θ in (0.4, 1.0, 3.0), x in (0.2, 1.7, 8.0) + @test logdensityof(Gamma(1.0, θ), x) ≈ logdensityof(Exponential(θ), x) + @test cdf(Gamma(1.0, θ), x) ≈ cdf(Exponential(θ), x) + @test logccdf(Gamma(1.0, θ), x) ≈ logccdf(Exponential(θ), x) + end + @test entropy(Gamma(1.0, 2.5)) ≈ entropy(Exponential(2.5)) +end + +@testset "reference numerics against Distributions.jl" begin + ref(α, θ) = Distributions.Gamma(α, θ) + shapes = (0.1, 0.5, 1.0, 2.0, 7.5, 100.0) + for α in shapes, θ in (0.5, 1.0, 4.0) + d, r = Gamma(α, θ), ref(α, θ) + for p in (0.001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999) + x = Distributions.quantile(r, p) + @test logdensityof(d, x) ≈ Distributions.logpdf(r, x) + @test densityof(d, x) ≈ Distributions.pdf(r, x) + @test cdf(d, x) ≈ Distributions.cdf(r, x) + @test ccdf(d, x) ≈ Distributions.ccdf(r, x) + # Use an absolute tolerance near `log(1) == 0`. + @test logcdf(d, x) ≈ Distributions.logcdf(r, x) atol = 1e-12 + @test logccdf(d, x) ≈ Distributions.logccdf(r, x) atol = 1e-12 + @test quantile(d, p) ≈ x rtol = 1e-10 + end + @test mean(d) ≈ Distributions.mean(r) + @test var(d) ≈ Distributions.var(r) + @test std(d) ≈ Distributions.std(r) + @test median(d) ≈ Distributions.median(r) rtol = 1e-10 + @test entropy(d) ≈ Distributions.entropy(r) + end +end + +@testset "the regularized incomplete gamma matches SpecialFunctions" begin + for a in (0.1, 0.5, 1.0, 2.0, 7.5, 100.0, 1000.0), x in (0.01, 0.3, 1.0, 3.0, 20.0) + y = x * a + p, q = gamma_inc(a, y) + # Skip the tails that `gamma_inc` itself rounds to zero. + p > 0 && @test exp(ProbabilityMeasures.loggammap(a, y)) ≈ p rtol = 1e-11 + q > 0 && @test exp(ProbabilityMeasures.loggammaq(a, y)) ≈ q rtol = 1e-11 + end +end + +@testset "log tails stay finite where the probability underflows" begin + d = Gamma(2.0, 1.0) + # `Q(2, x) = e^{-x}(1 + x)` and `P(2, x) → x²/2` as `x → 0`. + @test ccdf(d, 1000.0) == 0.0 + @test logccdf(d, 1000.0) ≈ -1000 + log(1001) + @test cdf(d, 1e-200) == 0.0 + @test logcdf(d, 1e-200) ≈ 2 * log(1e-200) - log(2) + + # The quantile of a probability this small is representable, and inverts. + deep = quantile(d, 1e-300) + @test 0 < deep < 1e-100 + @test logcdf(d, deep) ≈ log(1e-300) +end + +@testset "quantile is total and inverts the CDF" begin + d = Gamma(2.5, 1.5) + for p in (-0.001, 1.001, -Inf, Inf, NaN) + @test isnan(quantile(d, p)) + end + @test quantile(d, 0.0) == 0.0 + @test quantile(d, 1.0) == Inf + + for α in (0.05, 0.5, 1.0, 3.0, 50.0), p in (1e-12, 1e-3, 0.1, 0.5, 0.9, 1 - 1e-9) + m = Gamma(α, 2.0) + @test cdf(m, quantile(m, p)) ≈ p rtol = 1e-10 + end +end + +@testset "the quantile keeps BigFloat precision" begin + setprecision(BigFloat, 256) do + d = Gamma(big"2.5", big"1.5") + for p in (big"1e-40", big"0.25", big"0.5", big"0.99") + x = quantile(d, p) + @test x isa BigFloat + @test abs(cdf(d, x) - p) < 1e-60 * p + end + end +end + +@testset "distribution functions keep their type" begin + for T in (Float32, Float64, BigFloat) + d = Gamma(T(2), T(3)) + @test cdf(d, T(1)) isa T + @test ccdf(d, T(1)) isa T + @test logcdf(d, T(1)) isa T + @test logccdf(d, T(1)) isa T + @test quantile(d, T(1) / 4) isa T + @test entropy(d) isa T + end +end + +@testset "log-density gradient with respect to the parameters" begin + for α in (0.5, 2.0, 7.5), θ in (0.5, 3.0), x in (0.3, 1.0, 9.0) + g = ForwardDiff.gradient([α, θ]) do p + logdensityof(Gamma(p[1], p[2]), x) + end + @test g[1] ≈ log(x) - digamma(α) - log(θ) + @test g[2] ≈ x / θ^2 - α / θ + end +end + +@testset "sampling" begin + d = Gamma(2.0, 1.5) + @test rand(Xoshiro(1), d) isa Float64 + @test rand(Xoshiro(1), Gamma(2.0f0, 1.5f0)) isa Float32 + @test size(rand(Xoshiro(1), d, 3, 4)) == (3, 4) + @test eltype(rand(Xoshiro(1), d, 5)) === Float64 + + v = zeros(4) + Random.rand!(Xoshiro(1), v, d) + @test all(x -> insupport(d, x), v) + + # Both the direct method and the boost below a unit shape. + for α in (0.2, 0.9, 1.0, 3.0, 40.0) + m = Gamma(α, 2.0) + draws = rand(Xoshiro(20250801), m, 200_000) + @test all(x -> insupport(m, x), draws) + @test mean(draws) ≈ mean(m) atol = 5 * std(m) / sqrt(200_000) + @test var(draws) ≈ var(m) rtol = 0.05 + end +end + +@testset "sample derivative follows the parameters" begin + # The accept step runs on plain noise, so the draw is a smooth function of both + # parameters at fixed noise. Compare with a central difference. + for α in (0.4, 2.0, 9.0), θ in (0.5, 2.0) + draw(p) = rand(Xoshiro(7), Gamma(p[1], p[2])) + g = ForwardDiff.gradient(draw, [α, θ]) + h = 1e-6 + dα = (draw([α + h, θ]) - draw([α - h, θ])) / 2h + dθ = (draw([α, θ + h]) - draw([α, θ - h])) / 2h + @test g[1] ≈ dα rtol = 1e-4 + @test g[2] ≈ dθ rtol = 1e-6 + @test !iszero(g[1]) + end +end From e3f74c32fa97286a07bd3b5460b3e83c27a366f7 Mon Sep 17 00:00:00 2001 From: simonsteiger Date: Thu, 3 Sep 2026 06:16:02 +0200 Subject: [PATCH 2/3] Add Wishart MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Wishart(ν, L)` is the first matrix-variate measure, so it brings the `Matrixvariate` variate form, the `ContinuousMatrixvariateMeasure` alias it dispatches on, and the `PositiveDefiniteMatrices` support. `L` is the lower-triangular factor of the scale matrix, the convention `MvNormal` already uses. It keeps `log|S|` a sum over a diagonal and the trace term a triangular solve rather than an inversion, and it makes Bartlett's decomposition, `L A A' L'`, plain arithmetic in the parameters. Sampling needs one chi-squared draw per dimension and so inherits `Gamma`'s sampler along with its derivative. Draws are symmetrized, which is what puts them exactly in the support rather than a rounding error away from it. `mean`, `var` and `std` take the shape of a draw. `cov` covers every pair of entries and so is indexed the way `vec` orders them; the conformance suite gains a `matrixsummaries` group for that shape. `src/core/linalg.jl` collects the triangular linear algebra the two factored measures share: `rowdot`, moved from `MvNormal`, alongside `rowsdot`, `forwardsolve`, `logdetdiag` and `cholfactor`. Each builds new arrays instead of writing into one, so reverse-mode backends, which reject array mutation, can follow them. `cholfactor` takes its pivots through the new `sqrtt`, so an indefinite argument yields a non-finite factor and the log-density reports `NaN` instead of throwing. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01HiwLngD494i4dSLJ2DGPBG --- CHANGELOG.md | 20 +- README.md | 25 +- .../src/ProbabilityMeasuresTest.jl | 2 +- .../src/implementations.jl | 48 ++++ libs/ProbabilityMeasuresTest/src/interface.jl | 12 + src/ProbabilityMeasures.jl | 12 +- src/core/linalg.jl | 98 ++++++++ src/core/mathfuns.jl | 13 ++ src/core/support.jl | 19 ++ src/core/types.jl | 8 +- src/matrixvariate/continuous/wishart.jl | 216 +++++++++++++++++ src/multivariate/continuous/mvnormal.jl | 23 +- test/test-wishart.jl | 220 ++++++++++++++++++ 13 files changed, 686 insertions(+), 30 deletions(-) create mode 100644 src/core/linalg.jl create mode 100644 src/matrixvariate/continuous/wishart.jl create mode 100644 test/test-wishart.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index bc3df98..00ca61d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -75,6 +75,21 @@ and this project adheres to [Semantic Versioning]. loop therefore runs on plain numbers whatever type the parameters carry, and the accepted noise enters the draw through arithmetic on `α` and `θ`, which is what leaves the draw differentiable with respect to both. +- `Wishart(ν, L)`, the first matrix-variate measure, with the `Matrixvariate` variate + form and the `ContinuousMatrixvariateMeasure` alias it dispatches on, and the + `PositiveDefiniteMatrices` support. `L` is the lower-triangular factor of the scale + matrix, the same convention `MvNormal` uses, which keeps ``\log|S|`` a sum over a + diagonal and the trace term a triangular solve rather than an inversion. `mean`, `var` + and `std` take the shape of a draw; `cov` covers every pair of entries and so is + indexed the way `vec` orders them. Sampling uses Bartlett's decomposition, one + chi-squared draw per dimension, and so inherits `Gamma`'s sampler and its derivative + with respect to the parameters. Draws are symmetrized, so they land exactly in the + support rather than a rounding error away from it. +- `cholfactor`, `forwardsolve`, `rowsdot` and `logdetdiag` in `src/core/linalg.jl`, with + `rowdot` moved there from `MvNormal`. Each builds new arrays instead of writing into + one, so reverse-mode backends, which reject array mutation, can follow them, and + `cholfactor` takes its pivots through `sqrtt`, so an indefinite argument gives a + non-finite factor rather than a `DomainError`. - `validateparams(d)`, which returns `d` or throws a `DomainError`, for the boundary where user-supplied parameters enter. It earns its place on `Categorical`, whose sum-to-one is the one invalid parameter a density cannot report: an unnormalized `p` @@ -103,8 +118,11 @@ and this project adheres to [Semantic Versioning]. - `MvNormal`'s `logdensityof` allocates, unlike the univariate measures'. Whitening needs a temporary, grown by `vcat` so that reverse-mode backends, which reject array mutation, can differentiate it. +- The conformance suite gained a `matrixsummaries` optional group for measures whose + draws are matrices. `mean` and `var` take the shape of a draw there, so the vector + form's check that `var` is the diagonal of `cov` needs a reshape. - The exported surface is intentionally minimal: every name is one a PPL is - expected to call. `mode`, `skewness`, `kurtosis`, `mgf`, `cf`, `Matrixvariate`, + expected to call. `mode`, `skewness`, `kurtosis`, `mgf`, `cf`, `variateform`/`valuesupport`, and the unused supports are omitted rather than shipped speculatively, since adding an export later is non-breaking and removing one is not. diff --git a/README.md b/README.md index 5c321b0..dc8f768 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ broadcasting on GPU arrays, and Reactant tracing. The package is experimental. At present it implements `Normal`, `LogNormal`, `Exponential`, `Uniform`, `Laplace`, `Cauchy`, `Gamma`, `Categorical`, `Bernoulli`, -`Binomial`, `Poisson`, `MvNormal`, and `Multinomial`. +`Binomial`, `Poisson`, `MvNormal`, `Multinomial`, and `Wishart`. ## Installation @@ -132,6 +132,27 @@ contains standard deviations, not variances. in `IntegerSimplex(n, length(p))`, and `var` and `std` return marginal values. As with `MvNormal`, multivariate `cdf`, `quantile`, and `median` are not provided. +`Wishart(ν, L)` draws symmetric positive-definite matrices. Like `MvNormal`, its second +argument is a lower-triangular factor, so the scale matrix is `L * L'` and the mean is +`ν * L * L'`. `ν` must exceed `size(L, 1) - 1`. + +```julia +using LinearAlgebra, ProbabilityMeasures + +S = [4.0 1.0; 1.0 2.5] +d = Wishart(5.0, Matrix(cholesky(S).L)) + +logdensityof(d, [6.0 1.0; 1.0 4.0]) +mean(d), var(d), std(d), entropy(d) +rand(d) +``` + +It supports `densityof`, `logdensityof`, `rand`, `mean`, `cov`, `var`, `std`, `entropy`, +`params`, `support`, `insupport`, and `checkparams`. `mean`, `var` and `std` take the +shape of a draw; `cov` covers every pair of entries and so is a `length(X)`-by-`length(X)` +matrix indexed the way `vec` orders them. Only the lower triangles of `L` and of the +argument are read. + The density result follows normal Julia promotion rules across the parameters and evaluation point: @@ -214,7 +235,7 @@ See the [contribution guide](docs/src/90-contributing.md) for contribution guide ProbabilityMeasures.jl currently contains `Normal`, `LogNormal`, `Exponential`, `Uniform`, `Cauchy`, `Laplace`, `Gamma`, `Categorical`, `Bernoulli`, `Binomial`, -`Poisson`, `MvNormal`, and `Multinomial`. Transformed or composite measures and +`Poisson`, `MvNormal`, `Multinomial`, and `Wishart`. Transformed or composite measures and Distributions.jl interoperability are not implemented yet. ## Citation diff --git a/libs/ProbabilityMeasuresTest/src/ProbabilityMeasuresTest.jl b/libs/ProbabilityMeasuresTest/src/ProbabilityMeasuresTest.jl index 9195f72..40226a0 100644 --- a/libs/ProbabilityMeasuresTest/src/ProbabilityMeasuresTest.jl +++ b/libs/ProbabilityMeasuresTest/src/ProbabilityMeasuresTest.jl @@ -15,7 +15,7 @@ using JLArrays: JLArray using Mooncake: Mooncake using ProbabilityMeasures using ProbabilityMeasures: ContinuousMeasure, DiscreteMeasure, UnivariateMeasure -using ProbabilityMeasures: DiagMvNormal, IsoMvNormal, unwhiten +using ProbabilityMeasures: DiagMvNormal, IsoMvNormal, scalematrix, unwhiten using QuadGK: quadgk using Random: Xoshiro using ReverseDiff: ReverseDiff diff --git a/libs/ProbabilityMeasuresTest/src/implementations.jl b/libs/ProbabilityMeasuresTest/src/implementations.jl index 06fdb70..c74835a 100644 --- a/libs/ProbabilityMeasuresTest/src/implementations.jl +++ b/libs/ProbabilityMeasuresTest/src/implementations.jl @@ -249,3 +249,51 @@ function _extremepoints(d::MvNormal) Float64[], ) end + +# Optional methods for matrix-variate measures. + +@implements MeasureInterface{(:matrixsummaries, :entropy)} Wishart [ + Wishart(5.0, [1.0 0.0; 0.0 1.0]), + Wishart(3.5, [2.0 0.0; 0.5 1.5]), + Wishart(4.0f0, Float32[1.0 0.0; -0.25 0.5]), +] + +# Keep invalid examples the same size as the measure under test. +function _invalids(d::Wishart) + p, T = size(d.L, 1), _elscalar(d) + singular, flipped = _identity(T, p), _identity(T, p) + singular[1, 1] = 0 + flipped[1, 1] = -1 + return ( + # A measure at `p - 1` degrees of freedom is singular and has no density. + Wishart(T(p) - one(T), _identity(T, p)), + Wishart(d.ν, singular), + Wishart(d.ν, flipped), + ) +end + +# Use a non-unit factor so the precision check includes a nonzero logarithm. +function _exactparams(d::Wishart) + p = size(d.L, 1) + return Wishart(2p, [i == j ? 2 : Int(i > j) for i in 1:p, j in 1:p]) +end + +# Multiples of the scale matrix are positive definite whatever the factor is. +function default_testpoints(d::Wishart) + S, T = scalematrix(d), _elscalar(d) + return [convert(T, c) .* S for c in (d.ν, 1, 2)] +end + +function _extremepoints(d::Wishart) + p = size(d.L, 1) + return ( + fill(Inf, p, p), + fill(-Inf, p, p), + fill(NaN, p, p), + fill(floatmax(Float64), p, p), + zeros(p, p), + -ones(p, p), + zeros(p + 1, p + 1), + zeros(0, 0), + ) +end diff --git a/libs/ProbabilityMeasuresTest/src/interface.jl b/libs/ProbabilityMeasuresTest/src/interface.jl index c0e6876..a5b69c2 100644 --- a/libs/ProbabilityMeasuresTest/src/interface.jl +++ b/libs/ProbabilityMeasuresTest/src/interface.jl @@ -64,6 +64,18 @@ testpoint(d, i::Int=1) = rand(Xoshiro(i), d) d -> var(d) ≈ [cov(d)[i, i] for i in axes(cov(d), 1)], "std is the elementwise square root of var" => d -> std(d) ≈ sqrt.(var(d)), ), + # Matrix-variate summaries. `mean` and `var` take the shape of a draw, while + # `cov` covers the pairs of its entries and so is indexed as `vec` orders them. + matrixsummaries=( + "mean is a matrix with the shape of a draw" => + d -> (mean(d) isa AbstractMatrix) && size(mean(d)) == size(testpoint(d)), + "cov is square, with the length of a draw" => + d -> size(cov(d)) == (length(testpoint(d)), length(testpoint(d))), + "cov is symmetric" => d -> cov(d) ≈ transpose(cov(d)), + "var is the diagonal of cov, in the shape of a draw" => + d -> vec(var(d)) ≈ [cov(d)[i, i] for i in axes(cov(d), 1)], + "std is the elementwise square root of var" => d -> std(d) ≈ sqrt.(var(d)), + ), ), ) """ Checks the required measure methods and any declared distribution functions or diff --git a/src/ProbabilityMeasures.jl b/src/ProbabilityMeasures.jl index 794ceba..c9c6d17 100644 --- a/src/ProbabilityMeasures.jl +++ b/src/ProbabilityMeasures.jl @@ -13,7 +13,7 @@ module ProbabilityMeasures using DensityInterface: DensityInterface, densityof, logdensityof using IrrationalConstants: invsqrt2, log2π, logπ, logtwo, sqrt2 -using LinearAlgebra: Diagonal, LowerTriangular, UniformScaling +using LinearAlgebra: Diagonal, LowerTriangular, UniformScaling, diag, issymmetric using Random: Random, AbstractRNG using SpecialFunctions: digamma, erfc, erfcinv, gamma_inc, logerfc, loggamma using Statistics: Statistics, cov, mean, median, quantile, std, var @@ -22,6 +22,7 @@ using StatsAPI: StatsAPI, params include("core/types.jl") include("core/mathfuns.jl") include("core/gammainc.jl") +include("core/linalg.jl") include("core/support.jl") include("core/interface.jl") @@ -41,14 +42,17 @@ include("univariate/discrete/poisson.jl") include("multivariate/continuous/mvnormal.jl") include("multivariate/discrete/multinomial.jl") +include("matrixvariate/continuous/wishart.jl") + # Export the operations commonly needed by probabilistic programs. # Core types export AbstractProbabilityMeasure -export VariateForm, Univariate, Multivariate +export VariateForm, Univariate, Multivariate, Matrixvariate export ValueSupport, Continuous, Discrete export ContinuousUnivariateMeasure, DiscreteUnivariateMeasure export ContinuousMultivariateMeasure, DiscreteMultivariateMeasure +export ContinuousMatrixvariateMeasure # Supports export Support, @@ -59,7 +63,8 @@ export Support, RealInterval, IntegerRange, IntegerSimplex, - RealVectors + RealVectors, + PositiveDefiniteMatrices export support, insupport # Interface @@ -85,5 +90,6 @@ export Binomial export Poisson export MvNormal export Multinomial +export Wishart end diff --git a/src/core/linalg.jl b/src/core/linalg.jl new file mode 100644 index 0000000..9043c79 --- /dev/null +++ b/src/core/linalg.jl @@ -0,0 +1,98 @@ +#= + Triangular linear algebra for the measures parameterized by a Cholesky factor. Each + routine builds new arrays instead of writing into one, so that reverse-mode + differentiation tools, which reject array mutation, can follow it. +=# + +""" + rowdot(L, v, i, k) + +The inner product of the first `k` entries of row `i` of `L` and `v`. +""" +@inline function rowdot(L::AbstractMatrix, v::AbstractVector, i::Integer, k::Integer) + # Zero times infinity is `NaN`, which must remain visible in the result. + acc = L[i, 1] * v[1] + for j in 2:k + acc = muladd(L[i, j], v[j], acc) + end + return acc +end + +""" + rowsdot(A, i, j, k) + +The inner product of the first `k` entries of rows `i` and `j` of `A`. +""" +@inline function rowsdot(A::AbstractMatrix, i::Integer, j::Integer, k::Integer) + acc = A[i, 1] * A[j, 1] + for m in 2:k + acc = muladd(A[i, m], A[j, m], acc) + end + return acc +end + +""" + forwardsolve(L, b) + +``L^{-1} b`` for lower-triangular `L`, by forward substitution. + +Only the lower triangle of `L` is read. +""" +function forwardsolve(L::AbstractMatrix{<:Number}, b::AbstractVector{<:Number}) + # The first entry sets the result type and keeps later inner products non-empty. + z = [b[1] / L[1, 1]] + for i in 2:length(b) + z = vcat(z, (b[i] - rowdot(L, z, i, i - 1)) / L[i, i]) + end + return z +end + +""" + logdetdiag(A, n, ::Type{T}) + +``\\sum_{i=1}^{n} \\log A_{ii}`` in `float(T)`. + +A non-positive diagonal entry gives a non-finite result instead of an error, which is +what makes the log-determinant of a triangular factor safe to evaluate on unvalidated +parameters. +""" +function logdetdiag(A, n::Integer, ::Type{T}) where {T} + R = float(T) + acc = zero(R) + for i in 1:n + acc += logt(convert(R, A[i, i])) + end + return acc +end + +""" + cholfactor(X) + +The lower-triangular Cholesky factor `C` of `X`, with `C * C' == X`. + +Only the lower triangle of `X` is read. Where `X` is not positive definite, the pivot +goes through [`sqrtt`](@ref) and the factor carries `NaN` from that column on, so a +caller gets a non-finite result rather than an error. +""" +function cholfactor(X::AbstractMatrix{<:Number}) + n = size(X, 1) + pivot = sqrtt(X[1, 1]) + C = reshape(map(i -> X[i, 1] / pivot, 1:n), n, 1) + for j in 2:n + C = hcat(C, cholcolumn(X, C, j, n)) + end + return C +end + +#= + Column `j` of the factor, given the `j - 1` columns before it: zero above the diagonal, + the pivot on it, and the rest of column `j` of `X` with those columns subtracted off. + + This is a function rather than a loop body because the pivot is captured by the closure + that fills the column. A captured variable that the enclosing scope reassigns is boxed, + and its contents are then opaque, which widens the whole factor to `Any`. +=# +function cholcolumn(X::AbstractMatrix, C::AbstractMatrix, j::Integer, n::Integer) + pivot = sqrtt(X[j, j] - rowsdot(C, j, j, j - 1)) + return map(i -> i < j ? zero(pivot) : (X[i, j] - rowsdot(C, i, j, j - 1)) / pivot, 1:n) +end diff --git a/src/core/mathfuns.jl b/src/core/mathfuns.jl index 40f3cfa..f70ffdc 100644 --- a/src/core/mathfuns.jl +++ b/src/core/mathfuns.jl @@ -20,6 +20,19 @@ Like `log`, but returns `NaN` instead of throwing for a negative input. return select(x < zero(x), () -> oftype(float(x), NaN), () -> log(x)) end +""" + sqrtt(x) + +Like `sqrt`, but returns `NaN` instead of throwing for a negative input. + +A Cholesky factorization takes the square root of a pivot, so this version is what lets +it report an indefinite matrix as a non-finite factor rather than an error. +""" +@inline function sqrtt(x::Number) + # Some tools evaluate both choices, so call `sqrt` only with a valid value. + return select(x < zero(x), () -> oftype(float(x), NaN), () -> sqrt(x)) +end + """ erfcinvt(y) diff --git a/src/core/support.jl b/src/core/support.jl index cbf9499..84a0ddf 100644 --- a/src/core/support.jl +++ b/src/core/support.jl @@ -121,3 +121,22 @@ end function insupport(s::RealVectors, x::AbstractVector{<:Number}) return (length(x) == s.n) & all(isfinite, x) end + +""" + PositiveDefiniteMatrices(n) + +The symmetric positive-definite `n`-by-`n` real matrices. + +Symmetry is exact rather than approximate, so a matrix assembled as a product has to be +symmetrized before it lands here. +""" +struct PositiveDefiniteMatrices <: Support + n::Int +end + +function insupport(s::PositiveDefiniteMatrices, x::AbstractMatrix{<:Number}) + (size(x) == (s.n, s.n)) && all(isfinite, x) && issymmetric(x) || return false + # Every pivot of the Cholesky factor is positive exactly when `x` is definite. + C = cholfactor(x) + return all(i -> isfinite(C[i, i]) & (C[i, i] > zero(C[i, i])), 1:(s.n)) +end diff --git a/src/core/types.jl b/src/core/types.jl index cf43604..4d34686 100644 --- a/src/core/types.jl +++ b/src/core/types.jl @@ -1,7 +1,8 @@ """ VariateForm -Describes the shape of one sample: [`Univariate`](@ref) or [`Multivariate`](@ref). +Describes the shape of one sample: [`Univariate`](@ref), [`Multivariate`](@ref), or +[`Matrixvariate`](@ref). """ abstract type VariateForm end @@ -11,6 +12,9 @@ struct Univariate <: VariateForm end "Draws are vectors." struct Multivariate <: VariateForm end +"Draws are matrices." +struct Matrixvariate <: VariateForm end + """ ValueSupport @@ -67,12 +71,14 @@ abstract type AbstractProbabilityMeasure{F<:VariateForm,S<:ValueSupport} end # Short names used by implementations and interface fallbacks. const UnivariateMeasure{S} = AbstractProbabilityMeasure{Univariate,S} const MultivariateMeasure{S} = AbstractProbabilityMeasure{Multivariate,S} +const MatrixvariateMeasure{S} = AbstractProbabilityMeasure{Matrixvariate,S} const ContinuousMeasure{F} = AbstractProbabilityMeasure{F,Continuous} const DiscreteMeasure{F} = AbstractProbabilityMeasure{F,Discrete} const ContinuousUnivariateMeasure = AbstractProbabilityMeasure{Univariate,Continuous} const DiscreteUnivariateMeasure = AbstractProbabilityMeasure{Univariate,Discrete} const ContinuousMultivariateMeasure = AbstractProbabilityMeasure{Multivariate,Continuous} const DiscreteMultivariateMeasure = AbstractProbabilityMeasure{Multivariate,Discrete} +const ContinuousMatrixvariateMeasure = AbstractProbabilityMeasure{Matrixvariate,Continuous} # Reuse the same measure for every value in a broadcast. Base.broadcastable(d::AbstractProbabilityMeasure) = Ref(d) diff --git a/src/matrixvariate/continuous/wishart.jl b/src/matrixvariate/continuous/wishart.jl new file mode 100644 index 0000000..382737c --- /dev/null +++ b/src/matrixvariate/continuous/wishart.jl @@ -0,0 +1,216 @@ +""" + Wishart(ν, L) + +The Wishart measure on the symmetric positive-definite ``p``-by-``p`` matrices, with `ν` +degrees of freedom and scale matrix ``S = L L'``. Its density is + +```math +p(X) = \\frac{|X|^{(\\nu - p - 1)/2} + \\exp\\left(-\\frac{1}{2}\\operatorname{tr}(S^{-1}X)\\right)} + {2^{\\nu p/2}\\, |S|^{\\nu/2}\\, \\Gamma_p(\\nu/2)} +``` + +where ``\\Gamma_p`` is the multivariate gamma function; see [`logmvgamma`](@ref). The +mean is ``\\nu S``. + +`L` is a lower-triangular scale factor, the same convention [`MvNormal`](@ref) uses for +its covariance. If you have a scale matrix `S`, factor it before construction: + +```julia +using LinearAlgebra +d = Wishart(5.0, Matrix(cholesky(S).L)) +``` + +`ν` must exceed ``p - 1``, below which the measure has no density. The constructor does +not check its arguments; use [`validateparams`](@ref) for user input. + +Only the lower triangles of `L` and of the argument are read, and `logdensityof` returns +a non-finite value for invalid parameters, for an argument of the wrong shape, and for +one that is not positive definite. Like [`MvNormal`](@ref), it allocates. + +Sampling uses Bartlett's decomposition, which needs one chi-squared draw per dimension +and so inherits [`Gamma`](@ref)'s sampler. The draw is symmetrized, so it lands exactly +in the support rather than a rounding error away from it. +""" +struct Wishart{N<:Number,M<:AbstractMatrix{<:Number}} <: ContinuousMatrixvariateMeasure + ν::N + L::M +end + +# Samples are matrices regardless of how the parameters are stored. +function Base.eltype(::Type{Wishart{N,M}}) where {N,M} + return Matrix{float(promote_type(N, eltype(M)))} +end + +""" + dimension(d::Wishart) + +The size of one side of a draw. +""" +@inline dimension(d::Wishart) = size(d.L, 1) + +""" + shapesmatch(d::Wishart, X) -> Bool + +Whether `X` and `L` are square matrices of the same size. +""" +@inline function shapesmatch(d::Wishart, X::AbstractMatrix) + p = dimension(d) + return (p >= 1) & (size(d.L, 2) == p) & (size(X, 1) == p) & (size(X, 2) == p) +end + +function checkparams(d::Wishart) + p = dimension(d) + (p >= 1) & (size(d.L, 2) == p) || return false + # Below `p - 1` degrees of freedom the measure is singular and has no density. + ok = isfinite(d.ν) & (d.ν > p - 1) + for i in 1:p + Lii = d.L[i, i] + ok &= isfinite(Lii) & (Lii > zero(Lii)) + for j in 1:(i - 1) + ok &= isfinite(d.L[i, j]) + end + end + return ok +end + +support(d::Wishart) = PositiveDefiniteMatrices(dimension(d)) + +""" + logmvgamma(p, a) + +`log(Γ_p(a))`, the logarithm of the multivariate gamma function + +```math +\\Gamma_p(a) = \\pi^{p(p-1)/4} \\prod_{j=1}^{p} + \\Gamma\\left(a + \\frac{1 - j}{2}\\right). +``` + +The product needs `2a > p - 1`, and returns `NaN` otherwise rather than throwing. +""" +function logmvgamma(p::Integer, a::T) where {T<:Number} + # `loggamma` throws for a non-positive argument, and the smallest one in the product + # is `a + (1 - p)/2`, so one test covers all of them. + valid = 2 * a > p - one(T) + safe = select(valid, () -> a, () -> convert(T, p)) + acc = convert(T, logπ) * ((p * (p - 1)) // 4) + for j in 1:p + acc += loggamma(safe + (one(T) - j) / 2) + end + # Copy the loop-assigned value to avoid boxing the closure capture. + total = acc + return select(valid, () -> total, () -> convert(T, NaN)) +end + +""" + mvdigamma(p, a) + +The derivative of [`logmvgamma`](@ref) with respect to `a`, ``\\sum_{j=1}^{p} +\\psi(a + (1 - j)/2)``. +""" +function mvdigamma(p::Integer, a::T) where {T<:Number} + valid = 2 * a > p - one(T) + safe = select(valid, () -> a, () -> convert(T, p)) + acc = zero(T) + for j in 1:p + acc += digamma(safe + (one(T) - j) / 2) + end + # Copy the loop-assigned value to avoid boxing the closure capture. + total = acc + return select(valid, () -> total, () -> convert(T, NaN)) +end + +# Result type of the log-density, promoting the parameters with the argument. +@inline function densitytype(d::Wishart, X::AbstractMatrix) + return float(promote_type(typeof(d.ν), eltype(d.L), eltype(X))) +end + +function DensityInterface.logdensityof(d::Wishart, X::AbstractMatrix{<:Number}) + R = densitytype(d, X) + shapesmatch(d, X) || return convert(R, NaN) + p = dimension(d) + ν = convert(R, d.ν) + C = cholfactor(X) + logdetX = 2 * logdetdiag(C, p, R) + # `tr(S⁻¹X)` is the squared Frobenius norm of `L⁻¹C`, taken one column at a time. + q = sum(j -> sum(abs2, forwardsolve(d.L, C[:, j])), 1:p) + return (ν - p - 1) / 2 * logdetX - q / 2 - ν * p / 2 * convert(R, logtwo) - + ν * logdetdiag(d.L, p, R) - logmvgamma(p, ν / 2) +end + +""" + bartlettfactor(rng, d::Wishart) + +A lower-triangular `A` such that ``L A A' L'`` is a draw from `d`. + +The strictly lower entries are standard normal and the diagonal holds the square roots +of chi-squared draws whose degrees of freedom count down from `ν`. +""" +function bartlettfactor(rng::AbstractRNG, d::Wishart) + p = dimension(d) + offdiagonal = randn(rng, noisetype(d), (p * (p - 1)) ÷ 2) + # `Gamma(k/2, 2)` is the chi-squared measure with `k` degrees of freedom. + diagonal = map(i -> sqrt(rand(rng, Gamma((d.ν - i + 1) / 2, 2))), 1:p) + T = eltype(diagonal) + # Column `j` holds its strictly lower entries at this offset. + offset(j) = (j - 1) * p - ((j - 1) * j) ÷ 2 - j + return map(CartesianIndices((p, p))) do I + i, j = Tuple(I) + i < j && return zero(T) + i == j && return diagonal[i] + return convert(T, offdiagonal[offset(j) + i]) + end +end + +function Base.rand(rng::AbstractRNG, d::Wishart) + M = LowerTriangular(d.L) * bartlettfactor(rng, d) + X = M * transpose(M) + # A product is only symmetric up to rounding, and the support is exact. + return (X + transpose(X)) / 2 +end + +""" + scalematrix(d::Wishart) + +The scale matrix ``S = L L'``. +""" +function scalematrix(d::Wishart) + F = LowerTriangular(d.L) + return F * F' +end + +Statistics.mean(d::Wishart) = d.ν * scalematrix(d) + +# `Var(X_ij) = ν (S_ij² + S_ii S_jj)`, in the shape of a draw. +function Statistics.var(d::Wishart) + S = scalematrix(d) + s = diag(S) + return d.ν .* (S .^ 2 .+ s * transpose(s)) +end + +Statistics.std(d::Wishart) = sqrt.(var(d)) + +# `Cov(X_ij, X_kl) = ν (S_ik S_jl + S_il S_jk)`, over the entries of a draw in the order +# `vec` takes them. +function Statistics.cov(d::Wishart) + S = scalematrix(d) + entries = vec(CartesianIndices((dimension(d), dimension(d)))) + return [ + d.ν * (S[a[1], b[1]] * S[a[2], b[2]] + S[a[1], b[2]] * S[a[2], b[1]]) for + a in entries, b in entries + ] +end + +function entropy(d::Wishart) + p = dimension(d) + R = float(promote_type(typeof(d.ν), eltype(d.L))) + ν = convert(R, d.ν) + logdetS = 2 * logdetdiag(d.L, p, R) + return (p + 1) / 2 * logdetS + + p * (p + 1) / 2 * convert(R, logtwo) + + logmvgamma(p, ν / 2) - (ν - p - 1) / 2 * mvdigamma(p, ν / 2) + ν * p / 2 +end + +function Base.show(io::IO, d::Wishart) + return print(io, "Wishart(ν=", d.ν, ", L=", d.L, ")") +end diff --git a/src/multivariate/continuous/mvnormal.jl b/src/multivariate/continuous/mvnormal.jl index 4de5780..79edf29 100644 --- a/src/multivariate/continuous/mvnormal.jl +++ b/src/multivariate/continuous/mvnormal.jl @@ -81,14 +81,7 @@ end Return ``\\log \\det L`` in `float(T)`. A non-positive diagonal entry gives a non-finite result instead of an error. """ -function logdetfactor(d::MvNormal, ::Type{T}) where {T} - R = float(T) - acc = zero(R) - for i in 1:length(d.μ) - acc += logt(convert(R, d.L[i, i])) - end - return acc -end +logdetfactor(d::MvNormal, ::Type{T}) where {T} = logdetdiag(d.L, length(d.μ), T) function logdetfactor(d::DiagMvNormal, ::Type{T}) where {T} R = float(T) @@ -135,20 +128,6 @@ end support(d::MvNormal) = RealVectors(length(d.μ)) -""" - rowdot(L, v, i, k) - -The inner product of the first `k` entries of row `i` of `L` and `v`. -""" -@inline function rowdot(L::AbstractMatrix, v::AbstractVector, i::Integer, k::Integer) - # Zero times infinity is `NaN`, which must remain visible in the result. - acc = L[i, 1] * v[1] - for j in 2:k - acc = muladd(L[i, j], v[j], acc) - end - return acc -end - """ whiten(d::MvNormal, x) diff --git a/test/test-wishart.jl b/test/test-wishart.jl new file mode 100644 index 0000000..13cfc7d --- /dev/null +++ b/test/test-wishart.jl @@ -0,0 +1,220 @@ +using ProbabilityMeasures +using ProbabilityMeasuresTest: test_measure +using Distributions: Distributions +using ForwardDiff: ForwardDiff +using LinearAlgebra: I, LowerTriangular +using Random: Random, Xoshiro +using SpecialFunctions: digamma, loggamma +using Test + +# A lower-triangular factor whose scale matrix has distinct, correlated entries. +function factor(p, T=Float64) + return T[i == j ? 1 + i//2 : (i > j ? (i - j)//4 : 0) for i in 1:p, j in 1:p] +end + +@testset "conformance" begin + for d in ( + Wishart(5.0, [1.0 0.0; 0.0 1.0]), + Wishart(3.5, [2.0 0.0; 0.5 1.5]), + Wishart(4.0f0, Float32[1.0 0.0; -0.25 0.5]), + Wishart(7.0, factor(3)), + ) + test_measure(d; name=string(d)) + end +end + +@testset "traits" begin + d = Wishart(5.0, [2.0 0.0; 0.5 1.5]) + @test d isa AbstractProbabilityMeasure{Matrixvariate,Continuous} + @test d isa ContinuousMatrixvariateMeasure + @test !(d isa ContinuousMultivariateMeasure) + @test string(d) == "Wishart(ν=5.0, L=[2.0 0.0; 0.5 1.5])" + @test keys(params(d)) === (:ν, :L) + @test params(d) == (ν=5.0, L=[2.0 0.0; 0.5 1.5]) +end + +@testset "no promotion at construction" begin + L = [2.0 0.0; 0.5 1.5] + @test typeof(Wishart(5, L)) === Wishart{Int,Matrix{Float64}} + @test typeof(Wishart(5.0f0, Float32.(L))) === Wishart{Float32,Matrix{Float32}} + + @test eltype(Wishart(5, [2 0; 1 2])) === Matrix{Float64} + @test eltype(Wishart(5.0f0, Float32.(L))) === Matrix{Float32} +end + +@testset "precision follows the argument, not the parameters" begin + exact = Wishart(5, [2 0; 1 2]) + X = [6.0 1.0; 1.0 4.0] + @test logdensityof(exact, Float32.(X)) isa Float32 + @test logdensityof(exact, big.(X)) isa BigFloat + + # Integer parameters must not reduce `BigFloat` precision. + widened = logdensityof(Wishart(big"5.0", big.([2 0; 1 2])), big.(X)) + @test abs(logdensityof(exact, big.(X)) - widened) < 1e-70 +end + +@testset "construction never validates" begin + p = 2 + identity2 = Matrix{Float64}(I, p, p) + # A measure at `p - 1` degrees of freedom is singular and has no density. + for d in ( + Wishart(1.0, identity2), + Wishart(0.5, identity2), + Wishart(Inf, identity2), + Wishart(5.0, [0.0 0.0; 0.5 1.5]), + Wishart(5.0, [-1.0 0.0; 0.5 1.5]), + Wishart(5.0, [NaN 0.0; 0.5 1.5]), + ) + @test !checkparams(d) + @test !isfinite(logdensityof(d, [6.0 1.0; 1.0 4.0])) + end + @test checkparams(Wishart(1.0001, identity2)) + @test !checkparams(Wishart(5.0, [1.0 0.0 0.0; 0.0 1.0 0.0])) + + @test_throws DomainError validateparams(Wishart(1.0, identity2)) + @test validateparams(Wishart(5.0, identity2)) isa Wishart +end + +@testset "support" begin + d = Wishart(5.0, [2.0 0.0; 0.5 1.5]) + @test support(d) === PositiveDefiniteMatrices(2) + + @test insupport(d, [1.0 0.0; 0.0 1.0]) + @test insupport(d, [4.0 1.0; 1.0 2.5]) + # Singular, indefinite, asymmetric, mis-shaped, and non-finite all fall outside. + @test !insupport(d, [1.0 1.0; 1.0 1.0]) + @test !insupport(d, [-1.0 0.0; 0.0 1.0]) + @test !insupport(d, [1.0 0.5; 0.4 1.0]) + @test !insupport(d, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]) + @test !insupport(d, [NaN 0.0; 0.0 1.0]) + + # Draws land exactly in the support, not a rounding error away from it. + for k in 1:20 + @test insupport(d, rand(Xoshiro(k), d)) + end +end + +@testset "density is total off the support" begin + d = Wishart(5.0, [2.0 0.0; 0.5 1.5]) + for X in ( + fill(Inf, 2, 2), + fill(-Inf, 2, 2), + fill(NaN, 2, 2), + fill(floatmax(Float64), 2, 2), + zeros(2, 2), + -ones(2, 2), + zeros(3, 3), + zeros(0, 0), + ) + @test !isfinite(logdensityof(d, X)) + end +end + +@testset "one dimension is a scaled chi-squared" begin + # `Wishart(ν, σ)` on 1-by-1 matrices is `Gamma(ν/2, 2σ²)`. + for ν in (1.5, 3.0, 8.0), σ in (0.5, 1.0, 2.0), x in (0.3, 1.0, 7.0) + d, g = Wishart(ν, fill(σ, 1, 1)), Gamma(ν / 2, 2σ^2) + @test logdensityof(d, fill(x, 1, 1)) ≈ logdensityof(g, x) + @test mean(d)[1] ≈ mean(g) + @test var(d)[1] ≈ var(g) + @test entropy(d) ≈ entropy(g) + end +end + +@testset "reference numerics against Distributions.jl" begin + for p in 1:4, ν in (p - 0.5, float(p), p + 0.5, p + 3.0, 20.0) + ν > p - 1 || continue + L = factor(p) + d, r = Wishart(ν, L), Distributions.Wishart(ν, Matrix(L * L')) + for k in 1:5 + X = rand(Xoshiro(k), d) + @test logdensityof(d, X) ≈ Distributions.logpdf(r, X) + @test densityof(d, X) ≈ Distributions.pdf(r, X) + end + @test mean(d) ≈ Distributions.mean(r) + @test var(d) ≈ Distributions.var(r) + @test cov(d) ≈ Distributions.cov(r) + @test std(d) ≈ sqrt.(Distributions.var(r)) + @test entropy(d) ≈ Distributions.entropy(r) + end +end + +@testset "the multivariate gamma function matches its product" begin + for p in 1:4, a in (p / 2 + 0.1, p / 2 + 1, 5.0, 12.5) + expected = (p * (p - 1) / 4) * log(π) + expected += sum(j -> loggamma(a + (1 - j) / 2), 1:p) + @test ProbabilityMeasures.logmvgamma(p, a) ≈ expected + @test ProbabilityMeasures.mvdigamma(p, a) ≈ sum(j -> digamma(a + (1 - j) / 2), 1:p) + end + # Below the domain of the product, both report `NaN` rather than throwing. + @test isnan(ProbabilityMeasures.logmvgamma(3, 1.0)) + @test isnan(ProbabilityMeasures.mvdigamma(3, 1.0)) +end + +@testset "Cholesky factorization is exact and total" begin + A = [4.0 1.0 0.5; 1.0 3.0 -0.25; 0.5 -0.25 2.0] + C = ProbabilityMeasures.cholfactor(A) + @test C ≈ LowerTriangular(C) + @test C * C' ≈ A + # Forward substitution inverts the factor. + b = [1.0, -2.0, 0.5] + @test C * ProbabilityMeasures.forwardsolve(C, b) ≈ b + + # An indefinite matrix gives a non-finite factor instead of an error. + @test any(isnan, ProbabilityMeasures.cholfactor([1.0 2.0; 2.0 1.0])) + @test any(!isfinite, ProbabilityMeasures.cholfactor(zeros(2, 2))) +end + +@testset "sampling" begin + d = Wishart(5.0, [2.0 0.0; 0.5 1.5]) + @test rand(Xoshiro(1), d) isa Matrix{Float64} + @test rand(Xoshiro(1), Wishart(5.0f0, Float32[2.0 0.0; 0.5 1.5])) isa Matrix{Float32} + @test size(rand(Xoshiro(1), d)) == (2, 2) + @test length(rand(Xoshiro(1), d, 3)) == 3 + @test eltype(rand(Xoshiro(1), d, 3)) === Matrix{Float64} + + # Degrees of freedom just above the singular bound take Gamma's boosted sampler. + for (ν, p) in ((1.2, 2), (5.0, 2), (3.1, 3), (12.0, 3)) + m = Wishart(ν, factor(p)) + draws = [rand(Xoshiro(k), m) for k in 1:100_000] + average = sum(draws) / length(draws) + spread = sum(X -> (X .- average) .^ 2, draws) / (length(draws) - 1) + @test all(X -> insupport(m, X), draws[1:100]) + @test average ≈ mean(m) rtol = 0.03 + @test spread ≈ var(m) rtol = 0.10 + end +end + +@testset "log-density gradient with respect to the parameters" begin + L = [2.0 0.0; 0.5 1.5] + X = [6.0 1.0; 1.0 4.0] + p0 = [5.0; vec(L)] + rebuild(v) = Wishart(v[1], reshape(v[2:end], 2, 2)) + g = ForwardDiff.gradient(v -> logdensityof(rebuild(v), X), p0) + + h = 1e-6 + for i in eachindex(p0) + step = [j == i ? h : 0.0 for j in eachindex(p0)] + expected = + (logdensityof(rebuild(p0 .+ step), X) - logdensityof(rebuild(p0 .- step), X)) / + 2h + @test g[i] ≈ expected rtol = 1e-5 atol = 1e-7 + end + # The upper triangle of the factor is never read. + @test iszero(g[4]) +end + +@testset "sample derivative follows the parameters" begin + L = [2.0 0.0; 0.5 1.5] + p0 = [5.0; vec(L)] + rebuild(v) = Wishart(v[1], reshape(v[2:end], 2, 2)) + draw(v) = sum(rand(Xoshiro(7), rebuild(v))) + g = ForwardDiff.gradient(draw, p0) + + h = 1e-6 + for i in eachindex(p0) + step = [j == i ? h : 0.0 for j in eachindex(p0)] + @test g[i] ≈ (draw(p0 .+ step) - draw(p0 .- step)) / 2h rtol = 1e-4 atol = 1e-6 + end + @test !iszero(g[1]) +end From 8313982d29d32abc4592d987cbcc33e6dcf01539 Mon Sep 17 00:00:00 2001 From: simonsteiger Date: Thu, 3 Sep 2026 06:58:46 +0200 Subject: [PATCH 3/3] Add InverseGamma MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `InverseGamma(α, θ)` is the measure of `1/X` for `X` distributed as `Gamma(α, 1/θ)`. Its mean is infinite at or below a unit shape and its variance at or below a shape of two, both of which it reports rather than returning a negative number. Its distribution functions read the upper incomplete gamma integral at `θ/x`, which puts its lower tail on the integral's upper one. Inverting that through the lower tail would mean forming `1 - p`, and a probability of `1e-300` has no complement in floating point. So `gammaquantile` now takes which tail its probability measures, and `quantile` inverts the upper one directly, keeping small probabilities down to the point where the quantile itself underflows. `valuetype(d, x)` moves to `core/interface.jl`, where both gamma measures reach it, and `masstype` becomes its discrete case rather than a second copy of the same promotion. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01HiwLngD494i4dSLJ2DGPBG --- CHANGELOG.md | 10 + README.md | 22 +- .../src/implementations.jl | 18 ++ src/ProbabilityMeasures.jl | 2 + src/core/interface.jl | 21 +- src/univariate/continuous/gamma.jl | 58 ++-- src/univariate/continuous/inversegamma.jl | 124 +++++++++ test/test-gamma.jl | 10 + test/test-inversegamma.jl | 259 ++++++++++++++++++ 9 files changed, 482 insertions(+), 42 deletions(-) create mode 100644 src/univariate/continuous/inversegamma.jl create mode 100644 test/test-inversegamma.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 00ca61d..9aee259 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -75,6 +75,16 @@ and this project adheres to [Semantic Versioning]. loop therefore runs on plain numbers whatever type the parameters carry, and the accepted noise enters the draw through arithmetic on `α` and `θ`, which is what leaves the draw differentiable with respect to both. +- `InverseGamma(α, θ)`, the measure of `1/X` for `X` distributed as `Gamma(α, 1/θ)`. Its + mean is infinite at or below a unit shape and its variance at or below a shape of two, + both of which it reports rather than returning a negative number. Its distribution + functions read the upper incomplete gamma integral at `θ/x`, and `quantile` inverts + that same tail: `gammaquantile` now takes which tail its probability measures, so a + probability as small as `1e-300` gives a positive quantile instead of the zero a + detour through `1 - p` would give. +- `valuetype(d, x)`, the promoted floating-point type of a density, tail probability, or + quantile, in `core/interface.jl`. `masstype` becomes its discrete case rather than a + second copy of the same promotion. - `Wishart(ν, L)`, the first matrix-variate measure, with the `Matrixvariate` variate form and the `ContinuousMatrixvariateMeasure` alias it dispatches on, and the `PositiveDefiniteMatrices` support. `L` is the lower-triangular factor of the scale diff --git a/README.md b/README.md index dc8f768..9420d01 100644 --- a/README.md +++ b/README.md @@ -11,8 +11,9 @@ density and sampling operations, and compatible with automatic differentiation, broadcasting on GPU arrays, and Reactant tracing. The package is experimental. At present it implements `Normal`, `LogNormal`, -`Exponential`, `Uniform`, `Laplace`, `Cauchy`, `Gamma`, `Categorical`, `Bernoulli`, -`Binomial`, `Poisson`, `MvNormal`, `Multinomial`, and `Wishart`. +`Exponential`, `Uniform`, `Laplace`, `Cauchy`, `Gamma`, `InverseGamma`, +`Categorical`, `Bernoulli`, `Binomial`, `Poisson`, `MvNormal`, `Multinomial`, and +`Wishart`. ## Installation @@ -62,8 +63,8 @@ than throwing. ## Available API `Normal(μ, σ)`, `LogNormal(μ, σ)`, `Exponential(θ)`, `Uniform(a, b)`, -`Laplace(μ, b)`, `Cauchy(μ, σ)`, `Gamma(α, θ)`, `Categorical(p)`, `Bernoulli(p)`, -`Binomial(n, p)`, and `Poisson(λ)` each support: +`Laplace(μ, b)`, `Cauchy(μ, σ)`, `Gamma(α, θ)`, `InverseGamma(α, θ)`, +`Categorical(p)`, `Bernoulli(p)`, `Binomial(n, p)`, and `Poisson(λ)` each support: - `densityof` and `logdensityof` - `cdf`, `ccdf`, `logcdf`, and `logccdf` @@ -82,6 +83,13 @@ until the terms stop changing the result. They work in the type they are given, Sampling has no such limit: it uses rejection, and the accept step runs on plain floating-point noise, which leaves the draw differentiable with respect to `α` and `θ`. +`InverseGamma(α, θ)` is the measure of `1/X` for `X` distributed as `Gamma(α, 1/θ)`. Its +mean is `θ/(α - 1)` and is infinite at or below a unit shape; the variance needs a shape +above two. Its distribution functions read the *upper* incomplete gamma integral, and +`quantile` inverts that same tail, so `quantile(d, 1e-300)` is a positive number rather +than the zero a detour through `1 - p` would give. It shares `Gamma`'s limits otherwise: +a closed-form density, and iterative distribution functions. + `Categorical(p)` assigns the probabilities in `p` to categories `1:length(p)`. Draws and quantiles use the promoted floating-point type of `p`: @@ -234,9 +242,9 @@ See the [contribution guide](docs/src/90-contributing.md) for contribution guide ## Current scope ProbabilityMeasures.jl currently contains `Normal`, `LogNormal`, `Exponential`, -`Uniform`, `Cauchy`, `Laplace`, `Gamma`, `Categorical`, `Bernoulli`, `Binomial`, -`Poisson`, `MvNormal`, `Multinomial`, and `Wishart`. Transformed or composite measures and -Distributions.jl interoperability are not implemented yet. +`Uniform`, `Cauchy`, `Laplace`, `Gamma`, `InverseGamma`, `Categorical`, `Bernoulli`, +`Binomial`, `Poisson`, `MvNormal`, `Multinomial`, and `Wishart`. Transformed or composite +measures and Distributions.jl interoperability are not implemented yet. ## Citation diff --git a/libs/ProbabilityMeasuresTest/src/implementations.jl b/libs/ProbabilityMeasuresTest/src/implementations.jl index c74835a..4b623db 100644 --- a/libs/ProbabilityMeasuresTest/src/implementations.jl +++ b/libs/ProbabilityMeasuresTest/src/implementations.jl @@ -65,6 +65,24 @@ end # Use a shape and a scale that leave `loggamma(α)` and `log(θ)` non-zero. _exactparams(::Gamma) = Gamma(3, 2) +# Keep the shape above four so the sampled variance in the moment check has a finite +# variance of its own. +@implements MeasureInterface{UNIVARIATE_OPTIONALS} InverseGamma [ + InverseGamma(6.0, 1.0), InverseGamma(8.0, 2.5), InverseGamma(7.0f0, 0.5f0) +] + +function _invalids(::InverseGamma) + return ( + InverseGamma(-1.0, 1.0), + InverseGamma(0.0, 1.0), + InverseGamma(1.0, -1.0), + InverseGamma(Inf, 1.0), + ) +end + +# Use a shape and a scale that leave `loggamma(α)` and `log(θ)` non-zero. +_exactparams(::InverseGamma) = InverseGamma(3, 2) + @implements MeasureInterface{UNIVARIATE_OPTIONALS} Categorical [ Categorical([0.2, 0.3, 0.5]), Categorical([1.0]), Categorical(Float32[0.25, 0.75]) ] diff --git a/src/ProbabilityMeasures.jl b/src/ProbabilityMeasures.jl index c9c6d17..ce0044e 100644 --- a/src/ProbabilityMeasures.jl +++ b/src/ProbabilityMeasures.jl @@ -33,6 +33,7 @@ include("univariate/continuous/uniform.jl") include("univariate/continuous/laplace.jl") include("univariate/continuous/cauchy.jl") include("univariate/continuous/gamma.jl") +include("univariate/continuous/inversegamma.jl") include("univariate/discrete/categorical.jl") include("univariate/discrete/bernoulli.jl") @@ -84,6 +85,7 @@ export Uniform export Laplace export Cauchy export Gamma +export InverseGamma export Categorical export Bernoulli export Binomial diff --git a/src/core/interface.jl b/src/core/interface.jl index 3ce1ab8..18ee60f 100644 --- a/src/core/interface.jl +++ b/src/core/interface.jl @@ -46,7 +46,7 @@ validateparams(Categorical([2.0, 2.0])) # DomainError: does not sum to one """ function validateparams(d::AbstractProbabilityMeasure) checkparams(d) && return d - throw(DomainError(d, "invalid parameters; see `checkparams`")) + return throw(DomainError(d, "invalid parameters; see `checkparams`")) end """ @@ -63,17 +63,26 @@ end end """ - masstype(d, x) + valuetype(d, x) -The floating-point type for the probability of `x` under a discrete `d`. +The floating-point type of a density, tail probability, or quantile of `d` at `x`. -It promotes the parameter types with the type of `x`, so a `BigFloat` argument keeps -its precision. +It promotes the parameter types with the type of `x`, so a `BigFloat` argument keeps its +precision and exact parameters do not reduce it. """ -@inline function masstype(::D, x::Number) where {D<:DiscreteMeasure} +@inline function valuetype(::D, x::Number) where {D<:AbstractProbabilityMeasure} return float(promote_type(_promoted_paramtype(D), typeof(x))) end +""" + masstype(d, x) + +The floating-point type for the probability of `x` under a discrete `d`. + +See [`valuetype`](@ref), of which this is the discrete case. +""" +@inline masstype(d::DiscreteMeasure, x::Number) = valuetype(d, x) + function Random.rand( rng::AbstractRNG, sp::Random.SamplerTrivial{<:AbstractProbabilityMeasure} ) diff --git a/src/univariate/continuous/gamma.jl b/src/univariate/continuous/gamma.jl index 4b7cd69..363fd3e 100644 --- a/src/univariate/continuous/gamma.jl +++ b/src/univariate/continuous/gamma.jl @@ -49,18 +49,6 @@ end support(::Gamma) = PositiveReals() -""" - valuetype(d::Gamma, x) - -The floating-point type of a density, tail probability, or quantile at `x`. - -It promotes the parameter types with the type of `x`, so exact parameters keep the -argument's precision. -""" -@inline function valuetype(d::Gamma, x::Number) - return float(promote_type(typeof(d.α), typeof(d.θ), typeof(x))) -end - @inline function DensityInterface.logdensityof(d::Gamma, x::Number) T = valuetype(d, x) α, θ, y = convert(T, d.α), convert(T, d.θ), convert(T, x) @@ -149,24 +137,31 @@ logccdf(d::Gamma, x::Number) = gammatail(loggammaq, d, x, 0, -Inf) const GAMMAQUANTILE_MAXITER = 100 """ - gammaquantile(a, p) + gammaquantile(a, p, islower) + +The point `y` where the unit-scale gamma measure with shape `a` puts probability `p` on +one tail, for `0 < p < 1`. `islower` selects the tail: `true` solves +[`loggammap`](@ref)`(a, y) == log(p)` and `false` solves [`loggammaq`](@ref). -The `p`-quantile of the unit-scale gamma measure with shape `a`, for `0 < p < 1`. +Taking the tail rather than a lower-tail probability is what lets a measure built on the +upper tail, such as [`InverseGamma`](@ref), invert it without first forming `1 - p` and +losing the small tail to rounding. -Newton's method on `log(x)` refines a closed-form starting point until the step stops +Newton's method on `log(y)` refines a closed-form starting point until the step stops moving it. The logarithm is what keeps the deep lower tail, where the quantile itself underflows, from collapsing onto zero at the first step. """ -function gammaquantile(a::T, p::T) where {T<:Number} +function gammaquantile(a::T, p::T, islower::Bool) where {T<:Number} #= - Solve on whichever tail holds the smaller probability. The other tail is near one, + Solve on whichever tail holds the smaller probability. The other one is near one, where its logarithm is flat and Newton's method has almost no slope to descend. - `1 - p` is exact above one half, so the split costs no precision. + `1 - p` is exact above one half, so the switch costs no precision. =# - lower = p <= one(T) / 2 - target = lower ? logt(p) : log1p(-p) + small = p <= one(T) / 2 + lower = islower == small + target = small ? logt(p) : log1p(-p) - u = gammaquantile_start(a, p, lower) + u = gammaquantile_start(a, target, lower) tol = eps(basefloat(T)) for _ in 1:GAMMAQUANTILE_MAXITER y = exp(u) @@ -183,13 +178,18 @@ function gammaquantile(a::T, p::T) where {T<:Number} return exp(u) end -@inline function gammaquantile_start(a::T, p::T, lower::Bool) where {T<:Number} - # Small probabilities follow `P(a, y) ≈ y^a / Γ(a+1)`, which inverts directly. - logr = (logt(p) + loggamma(a + one(T))) / a - lower && logr < logt((one(T) + a) / 5) && return logr - # Elsewhere, Wilson and Hilferty's cube-root normal approximation. - z = -(sqrt2 * erfcinvt(2 * p)) - w = a * (one(T) - inv(9 * a) + z / (3 * sqrt(a)))^3 +# `target` is the logarithm of the tail being solved, which always holds at most half +# the mass, so the lower-tail probability is `target` itself or its complement. +@inline function gammaquantile_start(a::T, target::T, lower::Bool) where {T<:Number} + tail = exp(target) + loglower = lower ? target : log1p(-tail) + # A small lower tail follows `P(a, y) ≈ y^a / Γ(a+1)`, which inverts directly. + logr = (loglower + loggamma(a + one(T))) / a + logr < logt((one(T) + a) / 5) && return logr + # Elsewhere, Wilson and Hilferty's cube-root normal approximation. The normal + # quantile of the tail being solved changes sign with the tail. + z = sqrt2 * erfcinvt(2 * tail) + w = a * (one(T) - inv(9 * a) + (lower ? -z : z) / (3 * sqrt(a)))^3 return (isfinite(w) & (w > zero(T))) ? logt(w) : logr end @@ -200,7 +200,7 @@ function Statistics.quantile(d::Gamma, p::Number) (isnan(q) | (q < zero(T)) | (q > one(T))) && return convert(T, NaN) iszero(q) && return zero(T) isone(q) && return convert(T, Inf) - return convert(T, d.θ) * gammaquantile(convert(T, d.α), q) + return convert(T, d.θ) * gammaquantile(convert(T, d.α), q, true) end function Base.show(io::IO, d::Gamma) diff --git a/src/univariate/continuous/inversegamma.jl b/src/univariate/continuous/inversegamma.jl new file mode 100644 index 0000000..6a42561 --- /dev/null +++ b/src/univariate/continuous/inversegamma.jl @@ -0,0 +1,124 @@ +""" + InverseGamma(α, θ) + InverseGamma(α) + +The inverse-gamma measure on ``(0, \\infty)`` with shape `α` and scale `θ`. If ``X`` +follows `InverseGamma(α, θ)`, then ``1/X`` follows `Gamma(α, 1/θ)`. Its density is + +```math +p(x) = \\frac{\\theta^{\\alpha}}{\\Gamma(\\alpha)}\\, x^{-\\alpha - 1} e^{-\\theta/x} +``` + +The mean is ``\\theta/(\\alpha - 1)`` for ``\\alpha > 1`` and infinite otherwise; the +variance needs ``\\alpha > 2``. `InverseGamma(α)` sets the scale to one. + +# Arguments + + - `α::Number`: the shape. + - `θ::Number`: the scale. + +The constructor does not check its arguments. Invalid parameters give a non-finite +density. Use [`checkparams`](@ref) to check them when needed. + +`logdensityof` is closed form, so it broadcasts on device arrays and traces. `cdf`, +`ccdf`, `logcdf`, `logccdf`, `quantile`, `median` and `entropy` come from the incomplete +gamma integrals and iterate until their terms stop changing the result, which rules out +traced and device-side evaluation; see [`loggammap`](@ref). + +The distribution functions read the *upper* incomplete gamma integral at ``\\theta/x``, +and `quantile` inverts that same tail, so a small probability never has to be recovered +from `1 - p`. Sampling inverts a [`Gamma`](@ref) draw and inherits its derivative with +respect to the parameters. +""" +struct InverseGamma{A<:Number,T<:Number} <: ContinuousUnivariateMeasure + α::A + θ::T +end + +InverseGamma(α::Number) = InverseGamma(α, one(α)) + +Base.eltype(::Type{InverseGamma{A,T}}) where {A,T} = float(promote_type(A, T)) + +function checkparams(d::InverseGamma) + return isfinite(d.α) & (d.α > zero(d.α)) & isfinite(d.θ) & (d.θ > zero(d.θ)) +end + +support(::InverseGamma) = PositiveReals() + +@inline function DensityInterface.logdensityof(d::InverseGamma, x::Number) + T = valuetype(d, x) + α, θ, y = convert(T, d.α), convert(T, d.θ), convert(T, x) + # `loggamma` throws for a non-positive argument, so an invalid shape takes `NaN`. + lg = select(α > zero(T), () -> loggamma(α), () -> convert(T, NaN)) + v = muladd(α, logt(θ), -lg) - (α + one(T)) * logt(y) - θ / y + # Convert exact values to a float before returning `-Inf`. + return select(insupport(d, y), () -> v, () -> convert(T, -Inf)) +end + +# The scale carries the numeric type of the unit-scale gamma draw, so the quotient lands +# in `eltype(d)` whatever types the two parameters were given separately. +@inline function Base.rand(rng::AbstractRNG, d::InverseGamma) + return d.θ / rand(rng, Gamma(d.α, one(d.θ))) +end + +function Statistics.mean(d::InverseGamma) + E = eltype(d) + α, θ = convert(E, d.α), convert(E, d.θ) + return select(α > one(E), () -> θ / (α - one(E)), () -> convert(E, Inf)) +end + +function Statistics.var(d::InverseGamma) + E = eltype(d) + α, θ = convert(E, d.α), convert(E, d.θ) + below = (α - one(E))^2 * (α - 2 * one(E)) + return select(α > 2 * one(E), () -> θ^2 / below, () -> convert(E, Inf)) +end + +function entropy(d::InverseGamma) + α = float(d.α) + # `loggamma` and `digamma` throw for a non-positive argument. + shape = select( + α > zero(α), () -> loggamma(α) - (one(α) + α) * digamma(α), () -> oftype(α, NaN) + ) + return α + logt(float(d.θ)) + shape +end + +#= + `P(X ≤ x) = Q(α, θ/x)`, so the lower tail of this measure is the upper tail of the + incomplete gamma integral, and the two swap places throughout. As with `Gamma`, the + four functions differ only in the tail they take and in the two values they hold + outside `(0, ∞)`. +=# +@inline function invgammatail(tail, d::InverseGamma, x::Number, below, above) + T = valuetype(d, x) + checkparams(d) || return convert(T, NaN) + y = convert(T, x) + isnan(y) && return convert(T, NaN) + y > zero(T) || return convert(T, below) + isfinite(y) || return convert(T, above) + s = convert(T, d.θ) / y + # An argument small enough to overflow the ratio is below everything the tail can + # resolve, which is the same answer as an argument at zero. + isfinite(s) || return convert(T, below) + return tail(convert(T, d.α), s) +end + +cdf(d::InverseGamma, x::Number) = invgammatail((a, s) -> exp(loggammaq(a, s)), d, x, 0, 1) +ccdf(d::InverseGamma, x::Number) = invgammatail((a, s) -> exp(loggammap(a, s)), d, x, 1, 0) +logcdf(d::InverseGamma, x::Number) = invgammatail(loggammaq, d, x, -Inf, 0) +logccdf(d::InverseGamma, x::Number) = invgammatail(loggammap, d, x, 0, -Inf) + +function Statistics.quantile(d::InverseGamma, p::Number) + T = valuetype(d, p) + checkparams(d) || return convert(T, NaN) + q = convert(T, p) + (isnan(q) | (q < zero(T)) | (q > one(T))) && return convert(T, NaN) + iszero(q) && return zero(T) + isone(q) && return convert(T, Inf) + # `p` measures the upper tail of the incomplete gamma integral, so invert that one. + return convert(T, d.θ) / gammaquantile(convert(T, d.α), q, false) +end + +function Base.show(io::IO, d::InverseGamma) + return print(io, "InverseGamma(α=", d.α, ", θ=", d.θ, ")") +end diff --git a/test/test-gamma.jl b/test/test-gamma.jl index 4137b65..63c3671 100644 --- a/test/test-gamma.jl +++ b/test/test-gamma.jl @@ -174,6 +174,16 @@ end end end +@testset "the quantile solver inverts either tail" begin + # `InverseGamma` reads the upper tail, so the solver takes which tail `p` measures. + for a in (0.5, 2.0, 9.0), p in (1e-8, 0.1, 0.5, 0.9, 1 - 1e-9) + lower = ProbabilityMeasures.gammaquantile(a, p, true) + upper = ProbabilityMeasures.gammaquantile(a, p, false) + @test exp(ProbabilityMeasures.loggammap(a, lower)) ≈ p rtol = 1e-10 + @test exp(ProbabilityMeasures.loggammaq(a, upper)) ≈ p rtol = 1e-10 + end +end + @testset "the quantile keeps BigFloat precision" begin setprecision(BigFloat, 256) do d = Gamma(big"2.5", big"1.5") diff --git a/test/test-inversegamma.jl b/test/test-inversegamma.jl new file mode 100644 index 0000000..4906e1a --- /dev/null +++ b/test/test-inversegamma.jl @@ -0,0 +1,259 @@ +using ProbabilityMeasures +using ProbabilityMeasuresTest: test_measure +using Distributions: Distributions +using ForwardDiff: ForwardDiff +using Random: Random, Xoshiro +using SpecialFunctions: digamma +using Test + +@testset "conformance" begin + # Widen the parameters because Distributions.jl works at their own precision. + function reference_logpdf(m, x) + r = Distributions.InverseGamma(Float64(m.α), Float64(m.θ)) + return Distributions.logpdf(r, x) + end + for d in ( + InverseGamma(6.0, 1.0), + InverseGamma(8.0, 2.5), + InverseGamma(7.0f0, 0.5f0), + InverseGamma(5, 2), + ) + test_measure(d; name=string(d), reference_logpdf=reference_logpdf) + end + # A shape at or below two leaves the variance infinite, so the sampled moments + # have nothing to converge to. + for d in (InverseGamma(0.5, 1.0), InverseGamma(1.5, 2.0)) + test_measure(d; name=string(d), check_moments=false) + end +end + +@testset "traits" begin + d = InverseGamma(3.0, 2.0) + @test d isa AbstractProbabilityMeasure{Univariate,Continuous} + @test d isa ContinuousUnivariateMeasure + @test string(d) == "InverseGamma(α=3.0, θ=2.0)" + @test params(d) === (α=3.0, θ=2.0) +end + +@testset "no promotion at construction" begin + dual = ForwardDiff.Dual(1.0, 1.0) + @test typeof(InverseGamma(dual, 1.0)) === InverseGamma{typeof(dual),Float64} + @test typeof(InverseGamma(3, 1.0)) === InverseGamma{Int,Float64} + @test typeof(InverseGamma(3.0f0, 1.0f0)) === InverseGamma{Float32,Float32} + + # The one-argument form sets the scale to one in the shape's own type. + @test InverseGamma(3) === InverseGamma(3, 1) + + @test eltype(InverseGamma(3, 1)) === Float64 + @test eltype(InverseGamma(3.0f0, 1.0f0)) === Float32 + @test isbits(InverseGamma(3.0, 2.0)) +end + +@testset "precision follows the argument, not the parameters" begin + @test logdensityof(InverseGamma(3, 2), 1.0f0) isa Float32 + @test logdensityof(InverseGamma(3, 2), big"1.0") isa BigFloat + + exact = logdensityof(InverseGamma(3, 2), big"1.0") + full = logdensityof(InverseGamma(big"3.0", big"2.0"), big"1.0") + @test abs(exact - full) < 1e-70 + + @test logdensityof(InverseGamma(3, 2), 1//2) isa Float64 + @test logdensityof(InverseGamma(3, 2), -1//2) === -Inf +end + +@testset "construction never validates" begin + for d in ( + InverseGamma(-1.0, 1.0), + InverseGamma(0.0, 1.0), + InverseGamma(1.0, -1.0), + InverseGamma(1.0, 0.0), + ) + @test !checkparams(d) + @test !isfinite(logdensityof(d, 1.0)) + end + @test !checkparams(InverseGamma(Inf, 1.0)) + @test checkparams(InverseGamma(3.0, 2.0)) + + @test_throws DomainError validateparams(InverseGamma(-1.0, 1.0)) + @test validateparams(InverseGamma(3.0, 2.0)) === InverseGamma(3.0, 2.0) +end + +@testset "invalid parameters give NaN, not a partial answer" begin + for d in (InverseGamma(-1.0, 1.0), InverseGamma(3.0, -1.0), InverseGamma(NaN, 1.0)) + @test isnan(cdf(d, 1.0)) + @test isnan(ccdf(d, 1.0)) + @test isnan(logcdf(d, 1.0)) + @test isnan(logccdf(d, 1.0)) + @test isnan(quantile(d, 0.5)) + @test isnan(entropy(d)) + end +end + +@testset "support" begin + d = InverseGamma(3.0, 2.0) + @test support(d) === PositiveReals() + @test insupport(d, 1e-300) + @test insupport(d, 1e300) + @test !insupport(d, 0.0) + @test !insupport(d, -1.0) + @test !insupport(d, Inf) + @test !insupport(d, NaN) +end + +@testset "density is total off the support" begin + for d in (InverseGamma(0.5, 1.0), InverseGamma(3.0, 2.0)) + for x in (0.0, -1.0, -Inf, NaN, -floatmax(Float64)) + @test logdensityof(d, x) == -Inf + end + @test !isfinite(logdensityof(d, Inf)) + end +end + +@testset "the reciprocal of a gamma draw" begin + # If `X` follows `InverseGamma(α, θ)` then `1/X` follows `Gamma(α, 1/θ)`, so the + # densities differ by the Jacobian `1/x²`. + for α in (0.5, 1.0, 3.0, 9.0), θ in (0.5, 2.0), x in (0.1, 0.7, 3.0, 20.0) + g = Gamma(α, 1 / θ) + @test logdensityof(InverseGamma(α, θ), x) ≈ logdensityof(g, 1 / x) - 2 * log(x) + @test cdf(InverseGamma(α, θ), x) ≈ ccdf(g, 1 / x) + @test logccdf(InverseGamma(α, θ), x) ≈ logcdf(g, 1 / x) + end +end + +@testset "reference numerics against Distributions.jl" begin + ref(α, θ) = Distributions.InverseGamma(α, θ) + for α in (0.1, 0.5, 1.0, 2.0, 5.0, 20.0, 200.0), θ in (0.5, 1.0, 4.0) + d, r = InverseGamma(α, θ), ref(α, θ) + for p in (0.001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999) + x = Distributions.quantile(r, p) + @test logdensityof(d, x) ≈ Distributions.logpdf(r, x) + @test densityof(d, x) ≈ Distributions.pdf(r, x) + @test cdf(d, x) ≈ Distributions.cdf(r, x) + @test ccdf(d, x) ≈ Distributions.ccdf(r, x) + # Use an absolute tolerance near `log(1) == 0`. + @test logcdf(d, x) ≈ Distributions.logcdf(r, x) atol = 1e-12 + @test logccdf(d, x) ≈ Distributions.logccdf(r, x) atol = 1e-12 + @test quantile(d, p) ≈ x rtol = 1e-10 + end + @test mean(d) ≈ Distributions.mean(r) + @test var(d) ≈ Distributions.var(r) + @test median(d) ≈ Distributions.median(r) rtol = 1e-10 + @test entropy(d) ≈ Distributions.entropy(r) + end +end + +@testset "undefined moments are infinite" begin + @test mean(InverseGamma(0.5, 2.0)) == Inf + @test mean(InverseGamma(1.0, 2.0)) == Inf + @test mean(InverseGamma(2.0, 2.0)) == 2.0 + @test var(InverseGamma(1.5, 2.0)) == Inf + @test var(InverseGamma(2.0, 2.0)) == Inf + @test var(InverseGamma(3.0, 2.0)) == 1.0 + @test std(InverseGamma(1.5, 2.0)) == Inf +end + +@testset "log tails stay finite where the probability underflows" begin + d = InverseGamma(3.0, 2.0) + # `P(X > x) = P(3, 2/x)`, which behaves like `(2/x)³/6` as `x` grows. + @test ccdf(d, 1e120) == 0.0 + @test logccdf(d, 1e120) ≈ 3 * log(2e-120) - log(6.0) rtol = 1e-6 + # The lower tail decays like `e^{-2/x}`, which underflows well before its logarithm. + @test cdf(d, 1e-3) == 0.0 + @test isfinite(logcdf(d, 1e-3)) +end + +@testset "the lower tail inverts without forming 1 - p" begin + #= + A probability this small has no complement in floating point: `1 - p` rounds to + one. Inverting the upper incomplete gamma integral directly keeps it. + =# + for α in (0.5, 3.0, 9.0), θ in (1.0, 2.5) + d = InverseGamma(α, θ) + for p in (1e-300, 1e-100, 1e-20, 1e-8) + x = quantile(d, p) + @test 0 < x < Inf + @test cdf(d, x) ≈ p rtol = 1e-10 + end + end +end + +@testset "quantile is total and inverts the CDF" begin + d = InverseGamma(3.0, 2.0) + for p in (-0.001, 1.001, -Inf, Inf, NaN) + @test isnan(quantile(d, p)) + end + @test quantile(d, 0.0) == 0.0 + @test quantile(d, 1.0) == Inf + + for α in (0.05, 0.5, 1.0, 3.0, 50.0), p in (1e-12, 1e-3, 0.1, 0.5, 0.9, 1 - 1e-9) + m = InverseGamma(α, 2.0) + @test cdf(m, quantile(m, p)) ≈ p rtol = 1e-10 + end +end + +@testset "the quantile keeps BigFloat precision" begin + setprecision(BigFloat, 256) do + d = InverseGamma(big"2.5", big"1.5") + for p in (big"1e-40", big"0.25", big"0.5", big"0.99") + x = quantile(d, p) + @test x isa BigFloat + @test abs(cdf(d, x) - p) < 1e-60 * p + end + end +end + +@testset "distribution functions keep their type" begin + for T in (Float32, Float64, BigFloat) + d = InverseGamma(T(3), T(2)) + @test cdf(d, T(1)) isa T + @test ccdf(d, T(1)) isa T + @test logcdf(d, T(1)) isa T + @test logccdf(d, T(1)) isa T + @test quantile(d, T(1) / 4) isa T + @test entropy(d) isa T + end +end + +@testset "log-density gradient with respect to the parameters" begin + for α in (0.5, 3.0, 9.0), θ in (0.5, 2.0), x in (0.3, 1.0, 9.0) + g = ForwardDiff.gradient([α, θ]) do p + logdensityof(InverseGamma(p[1], p[2]), x) + end + @test g[1] ≈ log(θ) - digamma(α) - log(x) + @test g[2] ≈ α / θ - 1 / x + end +end + +@testset "sampling" begin + d = InverseGamma(3.0, 2.0) + @test rand(Xoshiro(1), d) isa Float64 + @test rand(Xoshiro(1), InverseGamma(3.0f0, 2.0f0)) isa Float32 + @test size(rand(Xoshiro(1), d, 3, 4)) == (3, 4) + @test eltype(rand(Xoshiro(1), d, 5)) === Float64 + + v = zeros(4) + Random.rand!(Xoshiro(1), v, d) + @test all(x -> insupport(d, x), v) + + # Both the direct sampler and Gamma's boost below a unit shape. + for α in (0.2, 0.9, 6.0, 40.0) + m = InverseGamma(α, 2.0) + draws = rand(Xoshiro(20250801), m, 200_000) + @test all(x -> insupport(m, x), draws) + # Compare medians: the mean is undefined for the smaller shapes. + @test median(draws) ≈ median(m) rtol = 0.02 + end +end + +@testset "sample derivative follows the parameters" begin + for α in (0.4, 3.0, 9.0), θ in (0.5, 2.0) + draw(p) = rand(Xoshiro(7), InverseGamma(p[1], p[2])) + g = ForwardDiff.gradient(draw, [α, θ]) + h = 1e-6 + dα = (draw([α + h, θ]) - draw([α - h, θ])) / 2h + dθ = (draw([α, θ + h]) - draw([α, θ - h])) / 2h + @test g[1] ≈ dα rtol = 1e-4 + @test g[2] ≈ dθ rtol = 1e-6 + @test !iszero(g[1]) + end +end