Cholesky numerical stability: Forward transform

src/bijectors/corr.jl

@@ -293,21 +293,20 @@ which is the above implementation.
 function _link_chol_lkj(W::AbstractMatrix)
     K = LinearAlgebra.checksquare(W)
 
-    y = similar(W) # z is also UpperTriangular. 
+    y = similar(W) # W is upper triangular.
     # Some zero filling can be avoided. Though diagnoal is still needed to be filled with zero.
 
     @inbounds for j in 1:K
-        remainder_sq = one(eltype(W))
-        for i in 1:(j - 1)
+        remainder_sq = W[j, j]^2
+        for i in (j - 1):-1:1
             z = W[i, j] / sqrt(remainder_sq)
-            y[i, j] = atanh(z)
-            remainder_sq -= W[i, j]^2
+            y[i, j] = asinh(z)
+            remainder_sq += W[i, j]^2
         end
         for i in j:K
             y[i, j] = 0
         end
     end
-
     return y
 end
 
@@ -317,17 +316,18 @@ function _link_chol_lkj_from_upper(W::AbstractMatrix)
 
     y = similar(W, N)
 
-    idx = 1
+    starting_idx = 1
     @inbounds for j in 2:K
-        y[idx] = atanh(W[1, j])
-        idx += 1
-        remainder_sq = 1 - W[1, j]^2
-        for i in 2:(j - 1)
+        y[starting_idx] = atanh(W[1, j])
+        starting_idx += 1
+        remainder_sq = W[j, j]^2
+        for i in (j - 1):-1:2
+            idx = starting_idx + i - 2
             z = W[i, j] / sqrt(remainder_sq)
-            y[idx] = atanh(z)
-            remainder_sq -= W[i, j]^2
-            idx += 1
+            y[idx] = asinh(z)
+            remainder_sq += W[i, j]^2
         end
+        starting_idx += length((j - 1):-1:2)
     end
 
     return y

test/transform.jl

@@ -237,18 +237,46 @@ end
 end
 
 @testset "LKJCholesky" begin
+    # Convert Cholesky factor to its free parameters, i.e. its off-diagonal elements
+    function chol_3by3_to_free_params(x::Cholesky)
+        if x.uplo == :U
+            return [x.U[1, 2], x.U[1, 3], x.U[2, 3]]
+        else
+            return [x.L[2, 1], x.L[3, 1], x.L[3, 2]]
+        end
+        # TODO: Generalise to arbitrary dimension using this code:
+        # inds = [
+        #     LinearIndices(size(x))[I] for I in CartesianIndices(size(x)) if
+        #     (uplo === :L && I[2] < I[1]) || (uplo === :U && I[2] > I[1])
+        # ]
+    end
+
+    # Reconstruct Cholesky factor from its free parameters
+    # Note that x[i, i] is always positive so we don't need to worry about the sign
+    function free_params_to_chol_3by3(free_params::AbstractVector, uplo::Symbol)
+        x = UpperTriangular(zeros(eltype(free_params), 3, 3))
+        x[1, 1] = 1
+        x[1, 2] = free_params[1]
+        x[1, 3] = free_params[2]
+        x[2, 2] = sqrt(1 - free_params[1]^2)
+        x[2, 3] = free_params[3]
+        x[3, 3] = sqrt(1 - free_params[2]^2 - free_params[3]^2)
+        if uplo == :U
+            return Cholesky(x)
+        else
+            return Cholesky(transpose(x))
+        end
+    end
+
     @testset "uplo: $uplo" for uplo in [:L, :U]
         dist = LKJCholesky(3, 1, uplo)
         single_sample_tests(dist)
 
         x = rand(dist)
-
-        inds = [
-            LinearIndices(size(x))[I] for I in CartesianIndices(size(x)) if
-            (uplo === :L && I[2] < I[1]) || (uplo === :U && I[2] > I[1])
-        ]
-        J = ForwardDiff.jacobian(z -> link(dist, Cholesky(z, x.uplo, x.info)), x.UL)
-        J = J[:, inds]
+        # Here, we need to pass ForwardDiff only the free parameters of the
+        # Cholesky factor so that we get a square Jacobian matrix
+        free_params = chol_3by3_to_free_params(x)
+        J = ForwardDiff.jacobian(z -> link(dist, free_params_to_chol_3by3(z, uplo)), free_params)
         logpdf_turing = logpdf_with_trans(dist, x, true)
         @test logpdf(dist, x) - _logabsdet(J) ≈ logpdf_turing
     end