Adding SVI functionality for atlas-level training

.gitignore

@@ -0,0 +1,3 @@
+*.o
+*.so
+data/

DESCRIPTION

@@ -7,19 +7,21 @@ Author: Adrien Osakwe
 Maintainer: Adrien Osakwe <adrien.osakwe@mail.mcgill.ca>
 Description: Topic Modelling framework for Spatial, Bulk, and Single cell analyses.
 License: MIT + file LICENSE
-Imports: Rcpp (>= 1.0.10),
-         SingleCellExperiment,
-         S4Vectors,
-         SpatialExperiment,
-         scran,
-         dplyr,
-         Matrix,
-         RcppArmadillo,
-         bluster,
-         methods,
-         scuttle,
-         torch,
-         coro
+SystemRequirements: gsfonts (or fonts-urw-base35), libglpk-dev
+Imports: 
+    Rcpp (>= 1.0.10),
+    dplyr,
+    Matrix,
+    RcppArmadillo,
+    methods,
+    torch,
+    coro,
+    SingleCellExperiment,
+    SpatialExperiment,
+    S4Vectors,
+    scran,
+    bluster,
+    scuttle
 Suggests: 
     knitr,
     rmarkdown,
@@ -30,3 +32,4 @@ Encoding: UTF-8
 Config/testthat/edition: 3
 VignetteBuilder: knitr
 URL: https://li-lab-mcgill.github.io/SpaTM/
+biocViews: Software, Transcriptomics

NAMESPACE

@@ -3,6 +3,7 @@
 export(BuildDiseaseAlpha)
 export(GTM)
 export(RTM)
+export(sGTM)
 export(STM)
 export(STM_Torch)
 export(STM_Torch_burnin)

R/GTM.R

@@ -42,9 +42,61 @@ GTM <- function(scte,K,D,num_threads = 1,maxiter = 100,verbal = TRUE,
   return(scte)
 }
 
-# full-batch
-
-# mini-batch
+#' Stochastic Guided Topic Model (sGTM)
+#'
+#' This function trains a stochastic variational inference (SVI) version of the
+#' Guided Topic Model (GTM) using mini-batches of cells. It builds the `CellMap`
+#' object one batch at a time to reduce memory usage during inference.
+#'
+#' @param scte A `SingleCellTopicExperiment` or `SpatialTopicExperiment` object containing count data and prior matrices.
+#' @param K Integer, the number of topics to infer.
+#' @param D Integer, the number of cells (documents) in the dataset.
+#' @param batch_size Integer, the number of cells per mini-batch. Default is 1024.
+#' @param num_threads Integer, the number of threads to use for parallel computation. Default is 1.
+#' @param maxiter Integer, the maximum number of epochs for training. Default is 100.
+#' @param verbal Logical, whether to print progress messages. Default is TRUE.
+#' @param zero_gamma Logical, whether to initialize gamma values to zero. Default is FALSE.
+#' @param rand_gamma Logical, whether to initialize gamma values randomly. Default is TRUE.
+#' @param thresh Convergence threshold based on the change in `nwk`. Default is 1e-8.
+#' @param burnin Maximum number of iterations to run during the GTM E-step for each sample (default is 1).
+#' @param lr Numeric, the SVI learning rate for global updates. Default is 0.1.
+#' @param shuffle Logical, whether to shuffle cells each epoch. Default is TRUE.
+#'
+#' @return A `SingleCellTopicExperiment` or `SpatialTopicExperiment` object with updated topic distributions.
+#'
+#' @import SingleCellExperiment
+#' @export
+sGTM <- function(scte,K,D,batch_size = 1024,num_threads = 1,maxiter = 100,verbal = TRUE,
+                  zero_gamma = FALSE,
+                  rand_gamma = TRUE,
+                  thresh = 1e-8,burnin = 1,
+                  lr = 0.1,
+                  tau0 = 1.0,
+                  kappa = 0.7,
+                  shuffle = TRUE){
+  train_sgtm(counts(scte),
+              celltypes = scte$int_cell,
+              genes = rowData(scte)$gene_ints,
+              alpha = alphaPrior(scte),
+              beta = betaPrior(scte),
+              K,
+              D,
+              ndk(scte),
+              nwk(scte),
+              batch_size,
+              num_threads,
+              maxiter,
+              verbal,
+              zero_gamma,
+              rand_gamma,
+              thresh,
+              burnin,
+              lr,
+              tau0,
+              kappa,
+              shuffle)
+  return(scte)
+}
 
 #### Prediction ####
 #' Topic Inference using Pretrained Model
@@ -56,9 +108,10 @@ GTM <- function(scte,K,D,num_threads = 1,maxiter = 100,verbal = TRUE,
 #' @param maxiter Integer, the maximum number of iterations for inference. Default is 50.
 #' @param verbal Logical, whether to print progress messages. Default is TRUE.
 #' @param phi A matrix representing the topic-gene distribution from a trained model.
-#' @param burnin Maximum number of iterations to run during the STM LDA step for each sample (default is 0)
+#' @param burnin Maximum number of iterations to run during the LDA E-step for each sample (default is 1).
 #'
-#' @return A `SingleCellTopicExperiment` or `SpatialTopicExperiment` object with updated inferred topic distributions.
+#' @return A `SingleCellTopicExperiment` or `SpatialTopicExperiment` object with an updated
+#'   `ndk` matrix. Call `buildTheta()` to refresh cell-level topic proportions `theta`.
 #'
 #' @export
 inferTopics <- function(scte,num_threads = 1,maxiter = 50,verbal = TRUE,
@@ -67,7 +120,8 @@ inferTopics <- function(scte,num_threads = 1,maxiter = 50,verbal = TRUE,
   infer_topics_cpp(counts(scte),
                celltypes = scte$int_cell,
                genes = rowData(scte)$gene_ints,
-               alpha = 1,ncol(phi),ncol(scte),nrow(scte),ndk(scte),phi,1,maxiter,verbal,burnin)
+               alpha = 1,ncol(phi),ncol(scte),nrow(scte),ndk(scte),phi,
+               num_threads,maxiter,verbal,burnin)
 
   return(scte)
 }
@@ -94,7 +148,8 @@ inferGEx <- function(scte,num_threads = 1,maxiter = 50,verbal = TRUE,
   gex_res <- infer_gex_cpp(counts(scte),
                    celltypes = scte$int_cell,
                    genes = rowData(scte)$gene_ints,
-                   alpha = 1,ncol(phi),ncol(scte),nrow(scte),ndk(scte),phi,1,maxiter,verbal)
+                   alpha = 1,ncol(phi),ncol(scte),nrow(scte),ndk(scte),phi,
+                   num_threads,maxiter,verbal)
 
   return(gex_res)
 }

R/RcppExports.R

@@ -87,3 +87,7 @@ test_mlp_forward <- function(X, layers, weights, dummy_topic = FALSE) {
     .Call(`_SpaTM_test_mlp_forward`, X, layers, weights, dummy_topic)
 }
 
+train_sgtm <- function(counts, celltypes, genes, alpha, beta, K, D, n_dk, n_wk, batch_size = 1024L, num_threads = 1L, maxiter = 100L, verbal = TRUE, zero_gamma = FALSE, rand_gamma = TRUE, thresh = 0.00001, burnin = 1L, lr = 0.1, tau0 = 1.0, kappa = 0.7, shuffle = TRUE) {
+    invisible(.Call(`_SpaTM_train_sgtm`, counts, celltypes, genes, alpha, beta, K, D, n_dk, n_wk, batch_size, num_threads, maxiter, verbal, zero_gamma, rand_gamma, thresh, burnin, lr, tau0, kappa, shuffle))
+}
+

src/CellMap.cpp

@@ -1,6 +1,7 @@
 #include "CellMap.h"
 
 #include <RcppArmadillo.h>
+#include <vector>
 
 
 // [[Rcpp::depends(RcppArmadillo)]]
@@ -88,19 +89,72 @@ std::unordered_map<int,Cell> build_Cell_Map(arma::sp_mat& counts,
   omp_set_num_threads(num_threads);
 
   std::unordered_map<int,Cell> CellMap;
-  //int cell_start,cell_end;
+  std::vector<std::unordered_map<int,Cell>> thread_maps(num_threads);
   #pragma omp parallel for
   for (int i = 0; i < D; i++){
+    int tid = omp_get_thread_num();
     arma::uvec geneidx = arma::find(counts.col(i));
     arma::mat mtx(geneidx.n_elem,3);
 
     mtx.col(0) = arma::nonzeros(counts.col(i));//counts(nonzero_counts);
     mtx.col(1).fill(celltypes(i));
     mtx.col(2) = genes(geneidx);
 
-    CellMap[i+1] = Cell(i+1,mtx,alpha,beta,K,
-                        zero_gamma,
-                        rand_gamma);
+    int cell_id = i + 1;
+    thread_maps[tid].emplace(cell_id, Cell(cell_id,mtx,alpha,beta,K,
+                                           zero_gamma,
+                                           rand_gamma));
+  }
+
+  CellMap.reserve(D);
+  for (int t = 0; t < num_threads; t++){
+    for (auto& kv : thread_maps[t]){
+      CellMap.emplace(kv.first, std::move(kv.second));
+    }
+  }
+
+  return CellMap;
+}
+
+std::unordered_map<int,Cell> build_Cell_Map_batch(arma::sp_mat& counts,
+                                            arma::vec& celltypes,
+                                            arma::vec& genes,
+                                            const arma::mat& alpha,
+                                            const arma::mat& beta,
+                                            const arma::uvec& batch_ids,
+                                            int K,
+                                            bool zero_gamma,
+                                            bool rand_gamma = true,
+                                            int num_threads = 1){
+  omp_set_num_threads(num_threads);
+
+  std::unordered_map<int,Cell> CellMap;
+  std::vector<std::unordered_map<int,Cell>> thread_maps(num_threads);
+  #pragma omp parallel for
+  for (arma::uword idx = 0; idx < batch_ids.n_elem; idx++){
+    int tid = omp_get_thread_num();
+    int cell_id = batch_ids(idx);
+    int col_id = cell_id - 1;
+    if (col_id < 0 || col_id >= counts.n_cols){
+      continue;
+    }
+    arma::uvec geneidx = arma::find(counts.col(col_id));
+    arma::mat mtx(geneidx.n_elem,3);
+
+    mtx.col(0) = arma::nonzeros(counts.col(col_id));
+    mtx.col(1).fill(celltypes(col_id));
+    mtx.col(2) = genes(geneidx);
+
+    thread_maps[tid].emplace(cell_id, Cell(cell_id,mtx,alpha,beta,K,
+                                           zero_gamma,
+                                           rand_gamma));
+  }
+
+  CellMap.reserve(batch_ids.n_elem);
+  for (int t = 0; t < num_threads; t++){
+    for (auto& kv : thread_maps[t]){
+      CellMap.emplace(kv.first, std::move(kv.second));
+    }
   }
 
   return CellMap;

src/CellMap.h

@@ -44,6 +44,17 @@ std::unordered_map<int,Cell> build_Cell_Map(arma::sp_mat& counts,
                                             bool verbal,
                                             int num_threads);
 
+std::unordered_map<int,Cell> build_Cell_Map_batch(arma::sp_mat& counts,
+                                            arma::vec& celltypes,
+                                            arma::vec& genes,
+                                            const arma::mat& alpha,
+                                            const arma::mat& beta,
+                                            const arma::uvec& batch_ids,
+                                            int K,
+                                            bool zero_gamma,
+                                            bool rand_gamma,
+                                            int num_threads);
+
 std::unordered_map<int,Cell> build_Predict_Cell_Map(arma::sp_mat& counts,
                                                     arma::vec& celltypes,
                                                     arma::vec& genes,