Optimising rgeneric using caching

haziqj · May 28, 2024 · 56d4e7d · 56d4e7d
1 parent b6a1763
commit 56d4e7d
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 54 deletions.
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: INLAvaan
 Type: Package
 Title: Bayesian Structural Equation Modelling with INLA
-Version: 0.1.0.9007
+Version: 0.1.0.9008
 Authors@R: 
     c(person("Haziq", "Jamil", 
              email = "[email protected]",

diff --git a/R/20-rgeneric.R b/R/20-rgeneric.R
@@ -9,14 +9,66 @@ inla_sem <- function(
   # - init (initial values)
   # - partable
 
-  interpret.theta <- function() {
+  # CACHE
+  envir <- parent.env(environment())
+  if (!exists("INLAvaan_SEM_cache", envir = envir)) {
+
+    # FIXME: These are indices for free parameters theta. But sometimes, with
+    # fixed parameter values we still need the indices...
 
+    # Indices
     idx_lam   <- partable$free[partable$mat == "lambda" & partable$free > 0]
     idx_beta  <- partable$free[partable$mat == "beta" & partable$free > 0]
     idx_theta <- partable$free[partable$mat == "theta" & partable$free > 0]
     idx_rho   <- partable$free[partable$mat == "rho" & partable$free > 0]
     idx_psi   <- partable$free[partable$mat == "psi" & partable$free > 0]
     idx_lvrho <- partable$free[partable$mat == "lvrho" & partable$free > 0]
+    assign("idx_lam", idx_lam, envir = envir)
+    assign("idx_beta", idx_beta, envir = envir)
+    assign("idx_theta", idx_theta, envir = envir)
+    assign("idx_rho", idx_rho, envir = envir)
+    assign("idx_psi", idx_psi, envir = envir)
+    assign("idx_lvrho", idx_lvrho, envir = envir)
+
+    # Lambda matrix
+    Lam_df <- partable[partable$mat == "lambda", ]
+    Lambda <- matrix(0, nrow = max(Lam_df$row), ncol = max(Lam_df$col))
+    Lambda[cbind(Lam_df$row, Lam_df$col)] <- Lam_df$est
+    LAM_IDX <- as.matrix(Lam_df[Lam_df$free > 0, c("row", "col")])
+    assign("Lambda", Lambda, envir = envir)
+    assign("LAM_IDX", LAM_IDX, envir = envir)
+
+    # (I-B) matrix
+    B_df <- partable[partable$mat == "beta", ]
+    B_IDX <- as.matrix(B_df[B_df$free > 0, c("row", "col")])
+    IminB <- 1
+    if (length(idx_beta) > 0) {
+      IminB <- diag(q)
+      IminB[B_IDX] <- -B_df$est
+    }
+    assign("IminB", IminB, envir = envir)
+    assign("B_IDX", B_IDX, envir = envir)
+
+    # Rho and Theta matrix
+    Rho_df <- partable[partable$mat == "rho", ]
+    Rho <- Theta <- diag(p)
+    RHO_IDX <- cbind(Rho_df$row, Rho_df$col)
+    assign("Rho", Rho, envir = envir)
+    assign("RHO_IDX", RHO_IDX, envir = envir)
+    assign("Theta", Theta, envir = envir)
+
+    # LVRho and Psi matrix
+    LVRho_df <- partable[partable$mat == "lvrho", ]
+    LVRho <- Psi <- diag(q)
+    LVRHO_IDX <- cbind(LVRho_df$row, LVRho_df$col)
+    assign("LVRho", LVRho, envir = envir)
+    assign("LVRHO_IDX", LVRHO_IDX, envir = envir)
+    assign("Psi", Psi, envir = envir)
+
+    assign("INLAvaan_SEM_cache", TRUE, envir = envir)
+  }
+
+  interpret.theta <- function() {
 
     lambda <- theta[idx_lam]
     beta <- theta[idx_beta]
@@ -38,52 +90,30 @@ inla_sem <- function(
   Q <- function() {
     params <- interpret.theta()
 
-    # Lambda matrix
-    Lam_df <- partable[partable$mat == "lambda", ]
-    Lam_df$est[Lam_df$free > 0] <- params$lambda
-    Lambda <- matrix(0, nrow = max(Lam_df$row), ncol = max(Lam_df$col))
-    Lambda[cbind(Lam_df$row, Lam_df$col)] <- Lam_df$est
-
-    # B matrix
-    B_df <- partable[partable$mat == "beta", ]
-    B_df$est[B_df$free > 0] <- params$beta
-    if (length(params$beta) > 0) {
-      sizeB <- max(B_df$row, B_df$col)
-      B <- matrix(0, nrow = sizeB, ncol = sizeB)
-      B[cbind(B_df$row, B_df$col)] <- B_df$est
-    } else {
-      B <- 0
-    }
+    # The matrices are cached and updated here
+    Lambda[LAM_IDX] <- params$lambda
+    if (length(idx_beta) > 0) IminB[B_IDX] <- -params$beta
 
     # Theta matrix
-    if (length(params$rho) > 0) {
-      SD <- diag(x = params$sd_e)
-      Rho_df <- partable[partable$mat == "rho", ]
-      Rho_df$est[Rho_df$free > 0] <- params$rho
-      sizeRho <- length(params$sd_e)
-      Rho <- matrix(0, nrow = sizeRho, ncol = sizeRho)
-      Rho[cbind(Rho_df$row, Rho_df$col)] <- Rho_df$est
-      Rho <- Rho + t(Rho)
-      diag(Rho) <- 1
-      Theta <- SD %*% Rho %*% SD
-    } else {
-      Theta <- diag(params$sd_e ^ 2)
+    diag(Theta) <- params$sd_e ^ 2
+    if (length(idx_rho) > 0) {
+      for (k in seq_along(nrow(RHO_IDX))) {
+        i <- RHO_IDX[k, 1]
+        j <- RHO_IDX[k, 2]
+        Theta[i, j] <- Theta[j, i] <-
+          params$rho[k] * params$sd_e[i] * params$sd_e[j]
+      }
     }
 
-
     # Psi matrix
-    if (length(params$lvrho) > 0) {
-      SD <- diag(x = params$sd_z)
-      LVRho_df <- partable[partable$mat == "lvrho", ]
-      LVRho_df$est[LVRho_df$free > 0] <- params$lvrho
-      sizeLVRho <- length(params$sd_z)
-      LVRho <- matrix(0, nrow = sizeLVRho, ncol = sizeLVRho)
-      LVRho[cbind(LVRho_df$row, LVRho_df$col)] <- LVRho_df$est
-      LVRho <- LVRho + t(LVRho)
-      diag(LVRho) <- 1
-      Psi <- SD %*% LVRho %*% SD
-    } else {
-      Psi <- diag(params$sd_z ^ 2)
+    diag(Psi) <- params$sd_z ^ 2
+    if (length(idx_psi) > 0) {
+      for (k in seq_along(nrow(LVRHO_IDX))) {
+        i <- LVRHO_IDX[k, 1]
+        j <- LVRHO_IDX[k, 2]
+        Psi[i, j] <- Psi[j, i] <-
+          params$lvrho[k] * params$sd_z[i] * params$sd_z[j]
+      }
     }
 
     # Using Woodbury Matrix identity
@@ -96,26 +126,25 @@ inla_sem <- function(
     # Sigma_inv <- A_inv - A_inv %*% U %*% solve(int_mat) %*% t(U) %*% A_inv
 
     # Regular way
-    front <- Lambda %*% solve(diag(1, nrow = q) - B)
-    Sigma <- front %*% Psi %*% t(front) + Theta
+    if (IminB == 1) {
+      front <- Lambda
+    } else {
+      front <- Lambda %*% solve(IminB)
+    }
+    Sigma <- front %*% tcrossprod(Psi, front) + Theta
     solve(Sigma)
   }
 
-  mu <- function() {
-    return(numeric(0))
-  }
+  mu <- function() { numeric(0) }
 
-  log.norm.const = function() { numeric(0) }
+  log.norm.const <- function() { numeric(0) }
 
   log.prior = function() {
     params <- interpret.theta()
 
     # If sigma ~ gamma(shape,rate) then this is the logpdf transform of theta = log(sigma^2)
     log_pdf_scale <- function(x, shape = 1, rate = 0.5) {
-      sigma <- sqrt(exp(x))
-      gamma_density <- dgamma(sigma, shape, rate, log = TRUE)
-      log_jacobian <- x
-      gamma_density + log_jacobian
+      dgamma(sqrt(exp(x)), shape, rate, log = TRUE) + x
     }
 
     # If rho ~ beta(a,b) then this is the logpdf of theta = log(rho / (1 - rho))

diff --git a/R/40-lav_export_INLA.R b/R/40-lav_export_INLA.R
@@ -1302,6 +1302,7 @@ lav2inla <- function(
     q = pta$nfac[[1]],
     init = inlastart,
     partable = partable
+    # optimize = TRUE
   )
 
   control_fixed <- INLA::control.fixed()

diff --git a/inst/demo.R b/inst/demo.R
@@ -17,11 +17,13 @@ true_model <- "
   y5 ~~ 0.1*y5
   y6 ~~ 0.1*y6
 "
-dat <- lavaan::simulateData(true_model, sample.nobs = 1000)
+dat <- lavaan::simulateData(true_model, sample.nobs = 500)
 
 mod <- "
   eta1 =~ y1 + y2 + y3
   eta2 =~ y4 + y5 + y6
+  # eta2 ~ eta1
+  y1 ~~ y4
 "
 fit <- isem(model = mod, data = dat, meanstructure = FALSE, verbose = TRUE)