Add LKJ

Project.toml

@@ -31,7 +31,9 @@ FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Calculus", "Distributed", "FiniteDifferences", "ForwardDiff", "JSON", "StaticArrays", "Test"]
+test = ["Calculus", "Distributed", "FiniteDifferences", "ForwardDiff", "JSON",
+        "StaticArrays", "HypothesisTests", "Test"]

docs/src/matrix.md

@@ -38,6 +38,7 @@ MatrixNormal
 MatrixTDist
 MatrixBeta
 MatrixFDist
+LKJ
 ```
 
 ## Internal Methods (for creating your own matrix-variate distributions)

src/Distributions.jl

@@ -106,6 +106,7 @@ export
     KSOneSided,
     Laplace,
     Levy,
+    LKJ,
     LocationScale,
     Logistic,
     LogNormal,

src/matrix/lkj.jl

@@ -0,0 +1,219 @@
+"""
+    LKJ(d, η)
+```julia
+d::Int   dimension
+η::Real  positive scale
+```
+The [LKJ](https://doi.org/10.1016/j.jmva.2009.04.008) distribution is a distribution over
+``d\\times d`` real correlation matrices (positive-definite matrices with ones on the diagonal).
+If ``\\mathbf{R}\\sim \\textrm{LKJ}_{d}(\\eta)``, then its probability density function is
+
+```math
+f(\\mathbf{R};\\eta) = c_0 |\\mathbf{R}|^{\\eta-1},
+```
+
+where (among many equivalent expressions)
+
+```math
+c_0 = \\prod_{k=1}^{d-1}\\pi^{\\frac{k}{2}}
+\\frac{\\Gamma\\left(\\eta+\\frac{d-1-k}{2}\\right)}{\\Gamma\\left(\\eta+\\frac{d-1}{2}\\right)}.
+```
+
+If ``\\eta = 1``, then the LKJ distribution is uniform over the space of correlation matrices.
+"""
+struct LKJ{T <: Real, D <: Integer} <: ContinuousMatrixDistribution
+    d::D
+    η::T
+    logc0::T
+end
+
+#  -----------------------------------------------------------------------------
+#  Constructors
+#  -----------------------------------------------------------------------------
+
+function LKJ(d::Integer, η::Real)
+    d > 0 || throw(ArgumentError("Matrix dimension must be positive."))
+    η > 0 || throw(ArgumentError("Scale parameter must be positive."))
+    logc0 = lkj_logc0(d, η)
+    T = Base.promote_eltype(η, logc0)
+    LKJ{T, typeof(d)}(d, T(η), T(logc0))
+end
+
+#  -----------------------------------------------------------------------------
+#  REPL display
+#  -----------------------------------------------------------------------------
+
+show(io::IO, d::LKJ) = show_multline(io, d, [(:d, d.d), (:η, d.η)])
+
+#  -----------------------------------------------------------------------------
+#  Conversion
+#  -----------------------------------------------------------------------------
+
+function convert(::Type{LKJ{T}}, d::LKJ) where T <: Real
+    LKJ{T, typeof(d.d)}(d.d, T(d.η), T(d.logc0))
+end
+
+function convert(::Type{LKJ{T}}, d::Integer, η, logc0) where T <: Real
+    LKJ{T, typeof(d)}(d, T(η), T(logc0))
+end
+
+#  -----------------------------------------------------------------------------
+#  Properties
+#  -----------------------------------------------------------------------------
+
+dim(d::LKJ) = d.d
+
+size(d::LKJ) = (dim(d), dim(d))
+
+rank(d::LKJ) = dim(d)
+
+insupport(d::LKJ, R::AbstractMatrix) = isreal(R) && size(R) == size(d) && isone(Diagonal(R)) && isposdef(R)
+
+mean(d::LKJ) = Matrix{partype(d)}(I, dim(d), dim(d))
+
+mode(d::LKJ) = mean(d)
+
+var(d::LKJ) = (σ² = var(_marginal(d)); σ² * (ones(size(d)) - I))
+
+params(d::LKJ) = d.η
+
+@inline partype(d::LKJ{T}) where {T <: Real} = T
+
+#  -----------------------------------------------------------------------------
+#  Evaluation
+#  -----------------------------------------------------------------------------
+
+function lkj_logc0(d::Integer, η::Real)
+    if isone(η)
+        if iseven(d)
+            logc0 = lkj_onion_logc0_uniform_even(d)
+        else
+            logc0 = lkj_onion_logc0_uniform_odd(d)
+        end
+    else
+        logc0 = lkj_onion_logc0(d, η)
+    end
+    return logc0
+end
+
+logkernel(d::LKJ, R::AbstractMatrix) = (d.η - 1) * logdet(R)
+
+_logpdf(d::LKJ, R::AbstractMatrix) = logkernel(d, R) + d.logc0
+
+#  -----------------------------------------------------------------------------
+#  Sampling
+#  -----------------------------------------------------------------------------
+
+function _rand!(rng::AbstractRNG, d::LKJ, R::AbstractMatrix)
+    R .= _lkj_onion_sampler(d.d, d.η, rng)
+end
+
+function _lkj_onion_sampler(d::Integer, η::Real, rng::AbstractRNG = Random.GLOBAL_RNG)
+    #  Section 3.2 in LKJ (2009 JMA)
+    #  1. Initialization
+    R = ones(typeof(η), d, d)
+    β = η + 0.5d - 1
+    u = rand(rng, Beta(β, β))
+    R[1, 2] = 2u - 1
+    R[2, 1] = R[1, 2]
+    #  2.
+    for k in 2:d - 1
+        #  (a)
+        β -= 0.5
+        #  (b)
+        y = rand(rng, Beta(k / 2, β))
+        #  (c)
+        u = randn(rng, k)
+        u = u / norm(u)
+        #  (d)
+        w = sqrt(y) * u
+        A = cholesky(R[1:k, 1:k]).L
+        z = A * w
+        #  (e)
+        R[1:k, k + 1] = z
+        R[k + 1, 1:k] = z'
+    end
+    #  3.
+    return R
+end
+
+#  -----------------------------------------------------------------------------
+#  The free elements of an LKJ matrix each have the same marginal distribution
+#  -----------------------------------------------------------------------------
+
+function _marginal(lkj::LKJ)
+    d = lkj.d
+    η = lkj.η
+    α = η + 0.5d - 1
+    LocationScale(-1, 2, Beta(α, α))
+end
+
+#  -----------------------------------------------------------------------------
+#  Several redundant implementations of the integrating constant
+#  used for unit testing
+#  -----------------------------------------------------------------------------
+
+function lkj_onion_logc0(d::Integer, η::Real)
+    #  Equation (17) in LKJ (2009 JMA)
+    sumlogs = zero(η)
+    for k in 2:d - 1
+        sumlogs += 0.5k*logπ + loggamma(η + 0.5(d - 1 - k))
+    end
+    α = η + 0.5d - 1
+    logc0 = (2η + d - 3)*logtwo + logbeta(α, α) + sumlogs - (d - 2) * loggamma(η + 0.5(d - 1))
+    return logc0
+end
+
+function lkj_onion_logc0_uniform_odd(d::Integer)
+    #  Theorem 5 in LKJ (2009 JMA)
+    sumlogs = 0.0
+    for k in 1:div(d - 1, 2)
+        sumlogs += loggamma(2k)
+    end
+    logc0 = 0.25(d^2 - 1)*logπ + sumlogs - 0.25(d - 1)^2*logtwo - (d - 1)*loggamma(0.5(d + 1))
+    return logc0
+end
+
+function lkj_onion_logc0_uniform_even(d::Integer)
+    #  Theorem 5 in LKJ (2009 JMA)
+    sumlogs = 0.0
+    for k in 1:div(d - 2, 2)
+        sumlogs += loggamma(2k)
+    end
+    logc0 = 0.25d*(d - 2)*logπ + 0.25(3d^2 - 4d)*logtwo + d*loggamma(0.5d) + sumlogs - (d - 1)*loggamma(d)
+end
+
+function lkj_vine_logc0(d::Integer, η::Real)
+    #  Equation (16) in LKJ (2009 JMA)
+    expsum = zero(η)
+    betasum = zero(η)
+    for k in 1:d - 1
+        α = η + 0.5(d - k - 1)
+        expsum += (2η - 2 + d - k) * (d - k)
+        betasum += (d - k) * logbeta(α, α)
+    end
+    logc0 = expsum * logtwo + betasum
+    return logc0
+end
+
+function lkj_vine_logc0_uniform(d::Integer)
+    #  Equation after (16) in LKJ (2009 JMA)
+    expsum = 0.0
+    betasum = 0.0
+    for k in 1:d - 1
+        α = (k + 1) / 2
+        expsum += k ^ 2
+        betasum += k * logbeta(α, α)
+    end
+    logc0 = expsum * logtwo + betasum
+    return logc0
+end
+
+function lkj_logc0_alt(d::Integer, η::Real)
+    #  Third line in first proof of Section 3.3 in LKJ (2009 JMA)
+    logc0 = zero(η)
+    for k in 1:d - 1
+        logc0 += 0.5k*logπ + loggamma(η + 0.5(d - 1 - k)) - loggamma(η + 0.5(d - 1))
+    end
+    return logc0
+end

src/matrixvariates.jl

@@ -213,6 +213,7 @@ _logpdf(d::MatrixDistribution, x::AbstractArray)
 ##### Specific distributions #####
 
 for fname in ["wishart.jl", "inversewishart.jl", "matrixnormal.jl",
-              "matrixtdist.jl", "matrixbeta.jl", "matrixfdist.jl"]
+              "matrixtdist.jl", "matrixbeta.jl", "matrixfdist.jl",
+              "lkj.jl"]
     include(joinpath("matrix", fname))
 end

test/lkj.jl

@@ -0,0 +1,123 @@
+using Distributions, Random
+using Test, LinearAlgebra, PDMats, Statistics, HypothesisTests
+
+
+d = 4
+η = abs(3randn())
+η₀ = 1
+
+G = LKJ(d, η)
+F = LKJ(d, η₀)
+
+@testset "LKJ construction errors" begin
+    @test_throws ArgumentError LKJ(-1, η)
+    @test_throws ArgumentError LKJ(d, -1)
+end
+
+@testset "LKJ params" begin
+    η̃ = params(G)
+    η̃₀ = params(F)
+    @test η̃ == η
+    @test η̃₀ == η₀
+end
+
+@testset "LKJ dim" begin
+    @test dim(G) == d
+    @test dim(F) == d
+end
+
+@testset "LKJ size" begin
+    @test size(G) == (d, d)
+    @test size(F) == (d, d)
+end
+
+@testset "LKJ rank" begin
+    @test rank(G) == d
+    @test rank(F) == d
+    @test rank(G) == rank(rand(G))
+    @test rank(F) == rank(rand(F))
+end
+
+@testset "LKJ insupport" begin
+    @test insupport(G, rand(G))
+    @test insupport(F, rand(F))
+
+    @test !insupport(G, rand(G) + rand(G) * im)
+    @test !insupport(G, randn(d, d + 1))
+    @test !insupport(G, randn(d, d))
+end
+
+@testset "LKJ sample moments" begin
+    @test isapprox(mean(rand(G, 100000)), mean(G) , atol = 0.1)
+    @test isapprox(var(rand(G, 100000)), var(G) , atol = 0.1)
+end
+
+@testset "LKJ marginals" begin
+    M = 10000
+    α = 0.05
+    L = sum(1:(d - 1))
+    ρ = Distributions._marginal(G)
+    mymats = zeros(d, d, M)
+    for m in 1:M
+        mymats[:, :, m] = rand(G)
+    end
+    for i in 1:d
+        for j in 1:i-1
+            kstest = ExactOneSampleKSTest(mymats[i, j, :], ρ)
+            @test pvalue(kstest) >= α / L  #  multiple comparisons
+        end
+    end
+end
+
+@testset "LKJ conversion" for elty in (Float32, Float64, BigFloat)
+    Gel1 = convert(LKJ{elty}, G)
+    Gel2 = convert(LKJ{elty}, G.d, G.η, G.logc0)
+
+    @test Gel1 isa LKJ{elty, typeof(d)}
+    @test Gel2 isa LKJ{elty, typeof(d)}
+    @test partype(Gel1) == elty
+    @test partype(Gel2) == elty
+end
+
+@testset "LKJ integrating constant" begin
+    #  =============
+    #  odd non-uniform
+    #  =============
+    d = 5
+    η = 2.3
+    lkj = LKJ(d, η)
+    @test Distributions.lkj_vine_logc0(d, η) ≈ Distributions.lkj_onion_logc0(d, η)
+    @test Distributions.lkj_onion_logc0(d, η) ≈ Distributions.lkj_logc0_alt(d, η)
+    @test lkj.logc0 == Distributions.lkj_onion_logc0(d, η)
+    #  =============
+    #  odd uniform
+    #  =============
+    d = 5
+    η = 1.0
+    lkj = LKJ(d, η)
+    @test Distributions.lkj_vine_logc0(d, η) ≈ Distributions.lkj_onion_logc0(d, η)
+    @test Distributions.lkj_onion_logc0(d, η) ≈ Distributions.lkj_onion_logc0_uniform_odd(d)
+    @test Distributions.lkj_vine_logc0(d, η) ≈ Distributions.lkj_vine_logc0_uniform(d)
+    @test Distributions.lkj_onion_logc0(d, η) ≈ Distributions.lkj_logc0_alt(d, η)
+    @test lkj.logc0 == Distributions.lkj_onion_logc0_uniform_odd(d)
+    #  =============
+    #  even non-uniform
+    #  =============
+    d = 6
+    η = 2.3
+    lkj = LKJ(d, η)
+    @test Distributions.lkj_vine_logc0(d, η) ≈ Distributions.lkj_onion_logc0(d, η)
+    @test Distributions.lkj_onion_logc0(d, η) ≈ Distributions.lkj_logc0_alt(d, η)
+    @test lkj.logc0 == Distributions.lkj_onion_logc0(d, η)
+    #  =============
+    #  even uniform
+    #  =============
+    d = 6
+    η = 1.0
+    lkj = LKJ(d, η)
+    @test Distributions.lkj_vine_logc0(d, η) ≈ Distributions.lkj_onion_logc0(d, η)
+    @test Distributions.lkj_onion_logc0(d, η) ≈ Distributions.lkj_onion_logc0_uniform_even(d)
+    @test Distributions.lkj_vine_logc0(d, η) ≈ Distributions.lkj_vine_logc0_uniform(d)
+    @test Distributions.lkj_onion_logc0(d, η) ≈ Distributions.lkj_logc0_alt(d, η)
+    @test lkj.logc0 == Distributions.lkj_onion_logc0_uniform_even(d)
+end

test/runtests.jl

@@ -5,6 +5,7 @@ using Distributed
 using Random
 using StatsBase
 using LinearAlgebra
+using HypothesisTests
 
 import JSON
 import ForwardDiff
@@ -39,6 +40,7 @@ const tests = [
     "matrixfdist",
     "matrixnormal",
     "matrixtdist",
+    "lkj",
     "vonmisesfisher",
     "conversion",
     "convolution",