diff --git a/CHANGELOG.md b/CHANGELOG.md index 049fefc..9aee259 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,51 @@ 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. +- `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 + 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` @@ -83,8 +128,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 2c8c4fc..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`, `Categorical`, `Bernoulli`, `Binomial`, `MvNormal`, -and `Multinomial`. +`Exponential`, `Uniform`, `Laplace`, `Cauchy`, `Gamma`, `InverseGamma`, +`Categorical`, `Bernoulli`, `Binomial`, `Poisson`, `MvNormal`, `Multinomial`, and +`Wishart`. ## Installation @@ -62,7 +63,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(α, θ)`, `InverseGamma(α, θ)`, +`Categorical(p)`, `Bernoulli(p)`, `Binomial(n, p)`, and `Poisson(λ)` each support: - `densityof` and `logdensityof` - `cdf`, `ccdf`, `logcdf`, and `logccdf` @@ -73,6 +75,21 @@ 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 `θ`. + +`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`: @@ -123,6 +140,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: @@ -204,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`, `Categorical`, `Bernoulli`, `Binomial`, `MvNormal`, and -`Multinomial`. 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/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/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 c6918cd..4b623db 100644 --- a/libs/ProbabilityMeasuresTest/src/implementations.jl +++ b/libs/ProbabilityMeasuresTest/src/implementations.jl @@ -54,6 +54,35 @@ 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) + +# 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]) ] @@ -238,3 +267,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 a96aec6..ce0044e 100644 --- a/src/ProbabilityMeasures.jl +++ b/src/ProbabilityMeasures.jl @@ -13,14 +13,16 @@ 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: 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/linalg.jl") include("core/support.jl") include("core/interface.jl") @@ -30,6 +32,8 @@ 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/continuous/inversegamma.jl") include("univariate/discrete/categorical.jl") include("univariate/discrete/bernoulli.jl") @@ -39,14 +43,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, @@ -57,11 +64,12 @@ export Support, RealInterval, IntegerRange, IntegerSimplex, - RealVectors + RealVectors, + PositiveDefiniteMatrices 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,11 +84,14 @@ export Exponential export Uniform export Laplace export Cauchy +export Gamma +export InverseGamma export Categorical export Bernoulli export Binomial export Poisson export MvNormal export Multinomial +export Wishart end 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/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/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 b4d02ad..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) @@ -66,3 +79,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/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/src/univariate/continuous/gamma.jl b/src/univariate/continuous/gamma.jl new file mode 100644 index 0000000..363fd3e --- /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() + +@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, 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). + +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(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, islower::Bool) where {T<:Number} + #= + 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 switch costs no precision. + =# + small = p <= one(T) / 2 + lower = islower == small + target = small ? logt(p) : log1p(-p) + + u = gammaquantile_start(a, target, 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 + +# `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 + +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, true) +end + +function Base.show(io::IO, d::Gamma) + return print(io, "Gamma(α=", d.α, ", θ=", d.θ, ")") +end 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/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..63c3671 --- /dev/null +++ b/test/test-gamma.jl @@ -0,0 +1,254 @@ +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 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") + 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 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 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