From 322e01034d4ff785ca89645876786b380a62f6ab Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Sun, 1 Dec 2024 19:23:46 +0000
Subject: [PATCH] Loop in reverse direction to improve performance

---
 src/bijectors/corr.jl | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/bijectors/corr.jl b/src/bijectors/corr.jl
index ffb8e3e3..845408da 100644
--- a/src/bijectors/corr.jl
+++ b/src/bijectors/corr.jl
@@ -293,20 +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
-        for i in 1:(j - 1)
-            remainder_norm = norm(W[i:end, j])
-            z = W[i, j] / remainder_norm
+        remainder_sq = W[j, j]^2
+        for i in (j - 1):-1:1
+            remainder_sq += W[i, j]^2
+            z = W[i, j] / sqrt(remainder_sq)
             y[i, j] = atanh(z)
         end
         for i in j:K
             y[i, j] = 0
         end
     end
-
     return y
 end
 
@@ -316,15 +316,17 @@ 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
-        for i in 2:(j - 1)
-            remainder_norm = norm(W[i:end, j])
-            z = W[i, j] / remainder_norm
+        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
+            remainder_sq += W[i, j]^2
+            z = W[i, j] / sqrt(remainder_sq)
             y[idx] = atanh(z)
-            idx += 1
+            starting_idx += 1
         end
     end