Inverse gamma

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.

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
 

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
 

ext/ProbabilityMeasuresForwardDiffExt.jl

@@ -1,11 +1,16 @@
 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.
 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

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

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

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

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

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