switch K to Cholesky LL^T to simplify multiple steps

Author: MikeDMorgan

NAMESPACE

@@ -97,6 +97,7 @@ importFrom(BiocParallel,bpstopOnError)
 importFrom(BiocParallel,bptry)
 importFrom(BiocSingular,bsparam)
 importFrom(Matrix,Matrix)
+importFrom(Matrix,chol)
 importFrom(Matrix,colSums)
 importFrom(Matrix,crossprod)
 importFrom(Matrix,diag)

R/glmm.R

@@ -76,7 +76,7 @@
 #' @name fitGLMM
 #'
 #' @importMethodsFrom Matrix %*%
-#' @importFrom Matrix Matrix solve crossprod kronecker
+#' @importFrom Matrix Matrix solve crossprod kronecker chol
 #' @importFrom stats runif var
 #' @importFrom BiocParallel bpstopOnError
 #' @export
@@ -176,9 +176,13 @@ fitGLMM <- function(X, Z, y, offsets, init.theta=NULL, Kin=NULL,
 
         ## add the genetic components
         ## augment Z with I
-        geno.I <- diag(nrow(full.Z))
+        ## Use K=LD^1/2(LD^1/2)^T instead of ZZ^T for GRM
+        ## chol return L^T
+        geno.I <- t(chol(Kin))
+        # geno.I <- diag(nrow(full.Z))
         colnames(geno.I) <- paste0("CovarMat", seq_len(ncol(geno.I)))
         full.Z <- do.call(cbind, list(full.Z, geno.I))
+
         # add a genetic variance component
         sigma_g <- Matrix(runif(1, 0, 1), ncol=1, nrow=1, sparse=TRUE)
         rownames(sigma_g) <- "CovarMat"
@@ -201,7 +205,10 @@ fitGLMM <- function(X, Z, y, offsets, init.theta=NULL, Kin=NULL,
 
         # if we only have a GRM then Z _is_ full.Z?
         # full.Z <- initializeFullZ(Z, cluster_levels=random.levels)
-        full.Z <- Z
+        ## Use K=LD^1/2(LD^1/2)^T instead of ZZ^T for GRM
+        ## chol return L^T
+        full.Z <- t(chol(Kin))
+        # full.Z <- Z
         # should this be the matrix R?
         colnames(full.Z) <- paste0(names(random.levels), seq_len(ncol(full.Z)))
 
@@ -456,7 +463,7 @@ initialiseG <- function(cluster_levels, sigmas, Kin=NULL){
             diag(G[c(i:(i+x.q-1)), c(i:(i+x.q-1)), drop=FALSE]) <- sigmas[x, ] # is this sufficient to transform the sigma to the model scale?
         } else{
             if(rownames(sigmas[x, , drop=FALSE]) %in% c("Genetic")){
-                diag(G[c(i:(i+x.q-1)), c(i:(i+x.q-1)), drop=FALSE]) <- sigmas[x, ] * Kin
+                diag(G[c(i:(i+x.q-1)), c(i:(i+x.q-1)), drop=FALSE]) <- sigmas[x, ] #* Kin
             }else{
                 diag(G[c(i:(i+x.q-1)), c(i:(i+x.q-1)), drop=FALSE]) <- sigmas[x, ] # is this sufficient to transform the sigma to the model scale?
             }

src/paramEst.cpp

@@ -700,7 +700,12 @@ arma::mat vectoriseZGenetic(const arma::mat& Z, const Rcpp::List& u_indices,
 
         // always set the last component to the genetic variance if there is a kinship matrix
         if(i == c-1){
-            arma::vec _vecZ = Kin(lower_indices);
+            // arma::vec _vecZ = Kin(lower_indices);
+            // Cholesky LL^T of kinship == ZZ^T
+            arma::mat _ZZT = PZ.cols(u_idx - 1) * Z.cols(u_idx - 1).t() * P.t(); // REML projection
+
+            // vectorise
+            arma::vec _vecZ = _ZZT(lower_indices);
             vecMat.col(i+1) = _vecZ;
             // vecMat.col(i) = _vecZ;
         } else{
@@ -741,7 +746,12 @@ arma::mat vectoriseZGeneticML(const arma::mat& Z, const Rcpp::List& u_indices,
 
         // always set the last component to the genetic variance if there is a kinship matrix
         if(i == c-1){
-            arma::vec _vecZ = Kin(lower_indices);
+            // arma::vec _vecZ = Kin(lower_indices);
+            // Cholesky LL^T of kinship == ZZ^T
+            arma::mat _ZZT = Z.cols(u_idx - 1) * Z.cols(u_idx - 1).t();
+
+            // vectorise
+            arma::vec _vecZ = _ZZT(lower_indices);
             vecMat.col(i+1) = _vecZ;
             // vecMat.col(i) = _vecZ;
         } else{

src/pseudovarPartial.cpp

@@ -75,7 +75,8 @@ List computePZList_G(const List& u_indices, const arma::mat& PZ, const arma::mat
         arma::mat _pzz(n, n);
 
         if(i == c - 1){
-            _pzz = PZ.cols(u_idx-1) * K * Z.cols(u_idx-1).t(); // convert 1-based to 0-based
+            // _pzz = PZ.cols(u_idx-1) * K * Z.cols(u_idx-1).t(); // convert 1-based to 0-based
+            _pzz = PZ.cols(u_idx-1) * Z.cols(u_idx-1).t(); // convert 1-based to 0-based
         } else{
             _pzz = PZ.cols(u_idx-1) * Z.cols(u_idx-1).t(); // convert 1-based to 0-based
         }
@@ -160,7 +161,8 @@ List pseudovarPartial_VG(const List& u_indices, const arma::mat& Z, const arma::
         arma::mat Zcols = Z.cols(u_idx-1).t();
 
         if(i == c - 1){
-            omat = VstarZ.cols(u_idx-1) * K * Zcols;
+            //omat = VstarZ.cols(u_idx-1) * K * Zcols;
+            omat = VstarZ.cols(u_idx-1) * Zcols;
         } else{
             omat = VstarZ.cols(u_idx-1) * Zcols;
         }
@@ -181,7 +183,8 @@ List pseudovarPartial_G(arma::mat Z, const arma::mat& K, List u_indices){
     for(unsigned int i = 0; i < items; i++){
         if(i == items - 1){
             arma::uvec icols = u_indices[i];
-            arma::mat _omat(Z.cols(icols - 1) * K * Z.cols(icols - 1).t());
+            // arma::mat _omat(Z.cols(icols - 1) * K * Z.cols(icols - 1).t());
+            arma::mat _omat(K);
             outlist[i] = _omat; // K is equivalent to ZZ^T
         } else{
             arma::uvec icols = u_indices[i];