Add vectorize method for LKJCholesky

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.23.0"
+version = "0.23.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/DynamicPPL.jl

@@ -15,6 +15,8 @@ using Setfield: Setfield
 using ZygoteRules: ZygoteRules
 using LogDensityProblems: LogDensityProblems
 
+using LinearAlgebra: Cholesky
+
 using DocStringExtensions
 
 using Random: Random

src/abstract_varinfo.jl

@@ -570,7 +570,7 @@ end
 # NOTE: `reconstruct` is no-op if `val` is already of correct shape.
 """
     reconstruct_and_link(dist, val)
-    reconstruct_and_link(vi::AbstractVarInfo, vi::VarName, dist, val)
+    reconstruct_and_link(vi::AbstractVarInfo, vn::VarName, dist, val)
 
 Return linked `val` but reconstruct before linking, if necessary.
 

src/utils.jl

@@ -213,6 +213,7 @@
 vectorize(d::UnivariateDistribution, r::Real) = [r]
 vectorize(d::MultivariateDistribution, r::AbstractVector{<:Real}) = copy(r)
 vectorize(d::MatrixDistribution, r::AbstractMatrix{<:Real}) = copy(vec(r))
+vectorize(d::Distribution{CholeskyVariate}, r::Cholesky) = copy(vec(r.UL))
 
 # NOTE:
 # We cannot use reconstruct{T} because val is always Vector{Real} then T will be Real.
@@ -235,6 +236,13 @@
 reconstruct(::UnivariateDistribution, val::Real) = val
 reconstruct(::MultivariateDistribution, val::AbstractVector{<:Real}) = copy(val)
 reconstruct(::MatrixDistribution, val::AbstractMatrix{<:Real}) = copy(val)
+reconstruct(::Inverse{Bijectors.VecCorrBijector}, ::LKJ, val::AbstractVector) = copy(val)
+function reconstruct(
+    ::Inverse{Bijectors.VecCholeskyBijector}, ::LKJCholesky, val::AbstractVector
+)
+    return copy(val)
+end
+
 # TODO: Implement no-op `reconstruct` for general array variates.
 
 reconstruct(d::Distribution, val::AbstractVector) = reconstruct(size(d), val)
@@ -294,7 +302,12 @@
     sz = Bijectors.output_size(b, size(dist))
     return Bijectors.invlink(dist, randrealuni(rng, sz...))
 end
-
+function inittrans(rng, dist::Distribution{CholeskyVariate})
+    # Get the size of the unconstrained vector
+    b = link_transform(dist)
+    sz = Bijectors.output_size(b, size(dist))
+    return Bijectors.invlink(dist, randrealuni(rng, sz...))
+end
 ################################
 # Multi-sample initialisations #
 ################################

test/lkj.jl

@@ -0,0 +1,53 @@
+using Bijectors: pd_from_upper, pd_from_lower
+
+pd_from_triangular(X::AbstractMatrix, uplo::Char) = uplo == 'U' ? pd_from_upper(X) : pd_from_lower(X)
+
+@model lkj_prior_demo() = x ~ LKJ(2, 1)
+@model lkj_chol_prior_demo(uplo) = x ~ LKJCholesky(2, 1, uplo)
+
+# Same for both distributions
+target_mean = vec(Matrix{Float64}(I, 2, 2))
+
+_lkj_atol = 0.05
+
+@testset "Sample from x ~ LKJ(2, 1)" begin
+    model = lkj_prior_demo()
+    # `SampleFromPrior` will sample in constrained space.
+    @testset "SampleFromPrior" begin
+        samples = sample(model, SampleFromPrior(), 1_000)
+        @test mean(map(Base.Fix2(getindex, Colon()), samples)) ≈ target_mean atol =
+            _lkj_atol
+    end
+
+    # `SampleFromUniform` will sample in unconstrained space.
+    @testset "SampleFromUniform" begin
+        samples = sample(model, SampleFromUniform(), 1_000)
+        @test mean(map(Base.Fix2(getindex, Colon()), samples)) ≈ target_mean atol =
+            _lkj_atol
+    end
+end
+
+@testset "Sample from x ~ LKJCholesky(2, 1, $(uplo))" for uplo in ['U', 'L']
+    model = lkj_chol_prior_demo(uplo)
+    # `SampleFromPrior` will sample in unconstrained space.
+    @testset "SampleFromPrior" begin
+        samples = sample(model, SampleFromPrior(), 1_000)
+        # Build correlation matrix from factor
+        corr_matrices = map(samples) do s
+            M = reshape(s.metadata.vals, (2, 2))
+            pd_from_triangular(M, uplo)
+        end
+        @test vec(mean(corr_matrices)) ≈ target_mean atol = _lkj_atol
+    end
+
+    # `SampleFromUniform` will sample in unconstrained space.
+    @testset "SampleFromUniform" begin
+        samples = sample(model, SampleFromUniform(), 1_000)
+        # Build correlation matrix from factor
+        corr_matrices = map(samples) do s
+            M = reshape(s.metadata.vals, (2, 2))
+            pd_from_triangular(M, uplo)
+        end
+        @test vec(mean(corr_matrices)) ≈ target_mean atol = _lkj_atol
+    end
+end

test/runtests.jl

@@ -50,6 +50,8 @@ include("test_util.jl")
             include("serialization.jl")
 
             include("loglikelihoods.jl")
+
+            include("lkj.jl")
         end
 
         @testset "compat" begin

test/utils.jl

@@ -42,4 +42,10 @@
         @test getargs_tilde(:(@. x ~ Normal(μ, $(Expr(:$, :(sqrt(v))))))) === nothing
         @test getargs_tilde(:(@~ Normal.(μ, σ))) === nothing
     end
+
+    @testset "vectorize" begin
+        dist = LKJCholesky(2, 1)
+        x = rand(dist)
+        @test vectorize(dist, x) == vec(x.UL)
+    end
 end