diff --git a/docs/src/index.md b/docs/src/index.md index 69f76d8..61c3d2c 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -98,6 +98,7 @@ as𝕀 ```@docs UnitVector CorrCholeskyFactor +PosDefCholeskyFactor ``` # Defining custom transformations diff --git a/src/special_arrays.jl b/src/special_arrays.jl index 653d2a0..8bf9b72 100644 --- a/src/special_arrays.jl +++ b/src/special_arrays.jl @@ -1,4 +1,4 @@ -export UnitVector, CorrCholeskyFactor +export UnitVector, CorrCholeskyFactor, PosDefCholeskyFactor """ (y, r, ℓ) = $SIGNATURES @@ -129,3 +129,63 @@ function inverse!(x::RealVector, t::CorrCholeskyFactor, U::UpperTriangular) end x end + + + +""" + PosDefCholeskyFactor(n) + +Cholesky factor of a symmetric positive-definite matrix of size `n`. + +Transforms ``n×(n+1)/2`` real numbers to an ``n×n`` upper-triangular matrix `Ω`, such that +`Ω'*Ω` is a positive definite matrix. +""" +@calltrans struct PosDefCholeskyFactor <: VectorTransform + n::Int + function PosDefCholeskyFactor(n) + @argcheck n ≥ 1 "Dimension should be positive." + new(n) + end +end + +dimension(t::PosDefCholeskyFactor) = (t.n*(t.n+1)) ÷ 2 + +function transform_with(flag::LogJacFlag, t::PosDefCholeskyFactor, x::RealVector) + @unpack n = t + T = extended_eltype(x) + ℓ = logjac_zero(flag, T) + U = Matrix{T}(undef, n, n) + index = firstindex(x) + @inbounds for col in 1:n + for row in 1:(col-1) + U[row, col] = x[index] + index += 1 + end + if flag isa NoLogJac + U[col, col] = transform(asℝ₊, x[index]) + else + U[col, col], ℓi = transform_and_logjac(asℝ₊, x[index]) + ℓ += ℓi + end + index += 1 + end + UpperTriangular(U), ℓ +end + +inverse_eltype(t::PosDefCholeskyFactor, U::UpperTriangular) = extended_eltype(U) + +function inverse!(x::RealVector, t::PosDefCholeskyFactor, U::UpperTriangular) + @unpack n = t + @argcheck size(U, 1) == n + @argcheck length(x) == dimension(t) + index = firstindex(x) + @inbounds for col in 1:n + for row in 1:(col-1) + x[index] = U[row,col] + index += 1 + end + x[index] = inverse(asℝ₊, U[col, col]) + index += 1 + end + x +end diff --git a/test/runtests.jl b/test/runtests.jl index ddac45b..fa0c6ba 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -339,3 +339,14 @@ end @test_nowarn @inferred transform(t, x) @test_nowarn @inferred transform_and_logjac(t, x) end + +@testset "positive definite cholesky factor" begin + t = PosDefCholeskyFactor(4) + d = dimension(t) + + v = randn(d) + U = transform(t, v) + @test U isa UpperTriangular + @test size(U) = (4,4) + @test inverse(t,U) ≈ v +end