Implement parallel processing in Horizontal.Integration.Assemble and multimodal integration

R/DAVINCHI.R

@@ -14,7 +14,10 @@ pkgs <- c("Rcpp",
 "cluster",
 "mclust",
 "dendextend",
-"BiocNeighbors")
+"BiocNeighbors",
+"future",
+"future.apply",
+"progressr",)
 
 
 .onAttach <- function(libname, pkgname){

R/import.R

@@ -16,5 +16,8 @@
 #' @importFrom stats wilcox.test
 #' @importFrom utils head
 #' @importFrom utils packageVersion
+#' @importFrom future plan multicore
+#' @importFrom future.apply future_lapply
+#' @importFrom progressr handlers handler_txtprogressbar with_progress
 ## usethis namespace: end
 NULL

R/new.R

@@ -1,3 +1,7 @@
+library(future)
+library(future.apply) 
+library(progressr)
+
 #' Adaptive Leiden/Louvain clustering
 #'
 #' When the number of clusters is set, run louvain/leiden clustering in an adaptive manner.
@@ -585,9 +589,6 @@ Proximity.Dependency.scatter <- function(gene.name,
 
 
 
-
-
-
 #Other parameters
 #save.label <- "rna" #rna, peak as input        #preprocess
 #normaliza.version <- "gene" #should be used when processing the data     #preprocess
@@ -873,6 +874,322 @@ Horizontal.Integration.Assemble <- function(
 
 
 
+#Other parameters
+#save.label <- "rna" #rna, peak as input        #preprocess
+#normaliza.version <- "gene" #should be used when processing the data     #preprocess
+# - rna, adt -> gene
+# - atac -> gene, LSI_1, LSI_2, LSI_3
+
+#clustering input parameters
+# - L2.opt <- "Yes" #Yes, No - whether L2norm before finding the nearest neighbors or mclust
+# - k.opt <- 40 #neighboring paramters for louvain
+
+#' Cross-sample (Horizontal) integration wrapper
+#'
+#' Cross-sample (Horizontal) integration
+#' 
+#' @export
+Horizontal.Integration.Assemble.Parallel <- function(
+    dataset.opts = NULL,
+    Y.list,
+    L.list, 
+    coor.list,
+    k.arg.list = c(4, 6, 8, 10, 11, 12, 13, 15, 16, 18, 20, 22, 24), #the parameters for self-contrastive learning step
+    L4.arg = 50, #parameters
+    input.opt = "FineTune", #OnlyDeco, FineTune
+    h.opt = "first",   #default, first
+    mod.opt = "all",   #all, common
+    L2.in = "default", #default, L2norm, L2norm.joint
+    random.seed = 123,
+    smooth = F,
+    workers = 2, # number of cores/workers for parallelization
+    ram = 15, # in GB max RAM usage 
+    sequential = F,
+    multi = "multicore" # multicore (mac and linux), multisession (for windows)
+    ){ 
+
+
+  #sanity check
+  #########################################################
+  if (is.null(dataset.opts)){
+    stop("Please include a vector of character names corresponding to the slices.")
+  }#if
+
+  if (sequential){
+    #sequential execution
+    future::plan(future::sequential)
+  }else{
+    #parallel execution
+    if (multi == "multisession"){
+      future::plan(future::multisession, workers = workers)
+    }else{
+      future::plan(future::multicore, workers = workers)
+    }
+    pkgs_to_load <- if (multi == "multisession") c("DaVinci", "progressr") else NULL
+  }
+
+  options(future.globals.maxSize = ram * 1024^3)
+
+  # progress bar settings 
+  progressr::handlers("progress")
+
+  options(
+    progressr.enable = TRUE,
+    progressr.enable_after = 0,      
+    progressr.delay_stdout = FALSE   
+  )
+
+  progressr::handlers(progressr::handler_txtprogressbar(
+    file = stdout(),
+    enable_after = 0
+  ))
+
+  #self-contrastive learning step
+  #########################################################
+
+  exhaustive.list <- progressr::with_progress({
+
+    p <- progressr::progressor(steps = length(Y.list) * length(k.arg.list))
+
+    future.apply::future_lapply(seq_along(Y.list), function(ii) {
+
+      message(paste0("Working on Sample ", ii, " / ", length(Y.list)))
+      proj <- list()
+      s0 <- NULL  
+
+
+      count <- 0
+      for (k.arg in k.arg.list){
+        message("Working on k=", k.arg)
+
+        if (is.null(s0)){
+
+          ICAp.res <- manifoldDecomp_adaptive(Y.list[[ii]], 
+                                              L.list[[ii]], 
+                                              k = k.arg, 
+                                              L4 = L4.arg, 
+                                              L4_adaptive = 2, 
+                                              to_drop = T, 
+                                              save.complete = T,
+                                              verbose = F, 
+                                              random.seed = random.seed)  
+        }else{
+
+          ICAp.res <- manifoldDecomp_adaptive(Y.list[[ii]], 
+                                              L.list[[ii]], 
+                                              k = k.arg, 
+                                              L4 = L4.arg, 
+                                              L4_adaptive = 2, 
+                                              to_drop = T, 
+                                              save.complete = T, 
+                                              shur0 = s0, 
+                                              verbose = F, 
+                                              random.seed = random.seed)  
+        }#else
+
+        count <- count+1
+        proj[[count]] <- ICAp.res
+
+        p(sprintf("Sample %d: k=%d", ii, k.arg))
+
+        if (is.null(s0)){
+          s0 <- ICAp.res$shur0
+        }#if
+
+      }#for k.arg
+
+      return(proj)
+    }, future.seed = TRUE, future.packages = pkgs_to_load)
+
+  })
+
+
+    #self contrastive learning step
+    #######################################################################################
+  if (length(k.arg.list) == 1){
+
+      #create the object needed for the downstream analysis
+      embed.list <- list()
+      embed.list.finetune <- list()
+      dav.res.list <- list()
+
+      for (ii in 1:length(Y.list)){
+          embed.list[[ii]] <- exhaustive.list[[ii]][[1]]  
+          embed.list.finetune[[ii]] <- exhaustive.list[[ii]][[1]]  
+          dav.res.list[[ii]] <- exhaustive.list[[ii]][[1]]
+      }#for ii
+
+  }else{
+
+    #self contrastive learning
+    #############################
+
+    embed.list <- future.apply::future_lapply(seq_along(exhaustive.list), function(ii) {
+      self_deco(exhaustive.list[[ii]], 
+                LVs.filter.thr = 0.9, 
+                freq = 1, 
+                opt = "B")
+    }, future.seed = TRUE, future.packages = pkgs_to_load)
+
+
+    #run again to finetune
+    #############################
+
+    embed.list.finetune <- future.apply::future_lapply(seq_along(embed.list), function(ii) {
+      manifoldDecomp_adaptive(Y.list[[ii]], 
+                              L.list[[ii]], 
+                              k = nrow(embed.list[[ii]]$LVs), 
+                              B = embed.list[[ii]]$LVs, 
+                              L4 = L4.arg, 
+                              L4_adaptive = 2, 
+                              to_drop = T, 
+                              save.complete = T, 
+                              shur0 = exhaustive.list[[ii]][[1]]$shur0, 
+                              verbose = F, 
+                              random.seed = random.seed)
+    }, future.seed = TRUE, future.packages = pkgs_to_load)
+
+
+    #Finetune after self-contrastive learning
+    if (input.opt == "OnlyDeco"){
+
+    dav.res.list <- future.apply::future_lapply(seq_along(embed.list.finetune), function(ii) {
+      ICAp.res <- embed.list.finetune[[ii]]
+      #construct the pseudo ICAp.res objects
+      list(Z = t(embed$LVs.pair), 
+           B = embed$LVs, 
+           L4 = ICAp.res$L4, 
+           shur0 = ICAp.res$shur0, 
+           L1 = ICAp.res$L1, 
+           L2 = ICAp.res$L2)
+    }, future.seed = TRUE, future.packages = pkgs_to_load)
+
+    }else if (input.opt == "FineTune"){
+
+      dav.res.list <- embed.list.finetune
+
+    }#else if
+  }#else
+
+
+
+  if (length(dataset.opts)==1){
+
+      #no need to integrate here, directly return here
+      mat <- t(dav.res.list[[1]]$B)
+      mat <- as.matrix(mat)
+      mat.slice.id <- rep(dataset.opts[1], nrow(mat))
+      #no integration.res
+      return(list(Y.list = Y.list, 
+                  coor.list = coor.list, 
+                  L.list = L.list, 
+                  embed.list = embed.list, 
+                  embed.list.finetune = embed.list.finetune, 
+                  dataset.opts = dataset.opts, 
+                  mat.slice.id = mat.slice.id, 
+                  mat = mat,
+                  exhaustive.list = exhaustive.list))
+
+  }else{
+    #Horizontal integration
+    ########################################################
+
+    if (h.opt == "default"){
+
+      integration.res <- Horizontal.Integration(Y.list, 
+                                                L.list, 
+                                                coor.list, 
+                                                dav.res.list = dav.res.list, 
+                                                LVs.filter.thr = 0.8, 
+                                                mod = mod.opt, 
+                                                remove.LV1 = F, 
+                                                L2.option = L2.in)
+
+    }else if (h.opt == "first"){
+
+      integration.res <- Horizontal.Integration.first(Y.list, 
+                                                      L.list, 
+                                                      coor.list, 
+                                                      dav.res.list = dav.res.list, 
+                                                      LVs.filter.thr = 0.8, 
+                                                      mod = mod.opt, 
+                                                      remove.LV1 = F, 
+                                                      L2.option = L2.in)
+
+    }#else
+
+
+    #cleanup the output
+    #########################################################
+    embed <- integration.res$LVs_embeddings
+    #make sure the sample names can align
+    sids <- unlist(lapply(strsplit( rownames(embed), "_"), function(x){paste0(x[-1], collapse = "_")}))
+    mat.slice.id <- unlist(lapply(strsplit(rownames(embed),"_"), function(x){x[1]}))
+    mat.slice.id <- dataset.opts[as.numeric(mat.slice.id)]
+
+    rownames(embed) <- sids             
+    mat <- as.matrix(embed)
+
+
+
+    #smoothing first no matter used or not
+    #########################################################
+    if (!smooth){
+        return(list(Y.list = Y.list, 
+                    coor.list = coor.list, 
+                    L.list = L.list, 
+                    embed.list = embed.list, 
+                    embed.list.finetune = embed.list.finetune, 
+                    dataset.opts = dataset.opts, 
+                    mat.slice.id = mat.slice.id, 
+                    integration.res = integration.res, 
+                    mat = mat,
+                    exhaustive.list = exhaustive.list))
+    }else{
+      mat.smooth <- mat
+
+    smooth_updates <- future.apply::future_lapply(seq_along(dataset.opts), function(ii) {
+      print(ii)
+
+      subpart.index <- which(mat.slice.id==dataset.opts[ii])
+      subpart <- mat[subpart.index,]
+      subpart.smooth <- refinement.batch(subpart, 
+                                         as.matrix(coor.list[[ii]])[rownames(subpart),], 
+                                         neighbor.option = "KNN", 
+                                         neighbor.arg = 8, 
+                                         tasks = "continuous")
+
+      return(list(idx = subpart.index, 
+                  val = (subpart+t(subpart.smooth))/2))
+    }, future.seed = TRUE, future.packages = pkgs_to_load)
+
+    for (update in smooth_updates) {
+      mat.smooth[update$idx, ] <- update$val
+    }
+
+      return(list(Y.list = Y.list, 
+                  coor.list = coor.list, 
+                  L.list = L.list, 
+                  embed.list = embed.list, 
+                  embed.list.finetune = embed.list.finetune, 
+                  dataset.opts = dataset.opts, 
+                  mat.slice.id = mat.slice.id, 
+                  integration.res = integration.res, 
+                  mat = mat, 
+                  mat.smooth = mat.smooth,
+                  exhaustive.list = exhaustive.list))
+
+    }#else
+
+  }#else
+
+  future::plan(sequential) # Explicitly close multisession workers by switching plan
+
+}#Horizontal.Integration.Assemble
+
+
+
+