log-chol.R

# computes the log Cholesky decomposition
log_chol <- function (x) {
  L <- chol(x)
  diag(L) <- log(diag(L))
  L[upper.tri(L, TRUE)]
}

# computes the inverse of the log Cholesky decomposition
log_chol_inv <- function(x) {
  n <- round((sqrt(8 * length(x) + 1) - 1)/2)
  out <- matrix(0, n, n)
  out[upper.tri(out, TRUE)] <- x
  diag(out) <- exp(diag(out))
  crossprod(out)
}

# computes the Jacobian of z <- log_chol_inv(x); z[upper.tri(z), TRUE]
jac_log_chol_inv <- function(x){
  n <- round((sqrt(8 * length(x) + 1) - 1)/2)
  z <- chol(log_chol_inv(x))
  n <- NCOL(z)
  # TODO: matrixcalc::commutation.matrix is very slow. This can be done a lot
  #       smarter
  res <- kronecker(t(z), diag(n)) %*% matrixcalc::commutation.matrix(n) +
    kronecker(diag(n), t(z))

  mult <- rep(1., length(x))
  j <- 0L
  for(i in Reduce(`+`, 1:n, 0L, accumulate = TRUE)[-1]){
    j <- j + 1L
    mult[i] <- z[j, j]
  }

  res <- res[, upper.tri(z, TRUE)] * rep(mult, each = NROW(res))
  res[upper.tri(z, TRUE), ]
}

# test
set.seed(1234)
X <- drop(rWishart(1, 5, diag(5)))
stopifnot(all.equal(X, log_chol_inv(log_chol(X))))

num_jac <- numDeriv::jacobian(\(x) {
  res <- log_chol_inv(x)
  res[upper.tri(res, TRUE)]
}, log_chol(X))
stopifnot(all.equal(jac_log_chol_inv(log_chol(X)), num_jac))