From c8d21e6560673aaddcf465ddf8fd9c190710e987 Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Sun, 1 Feb 2026 13:29:41 -0500 Subject: [PATCH 01/13] [feat] Initial implementation of stochastic GTM for atlas-level training. --- NAMESPACE | 1 + R/GTM.R | 54 +++++++++++- R/RcppExports.R | 4 + man/sGTM.Rd | 57 ++++++++++++ src/CellMap.cpp | 34 +++++++ src/CellMap.h | 11 +++ src/GTM.cpp | 4 +- src/GTM.h | 2 +- src/RcppExports.cpp | 29 ++++++ src/sGTM.cpp | 176 +++++++++++++++++++++++++++++++++++++ src/sGTM.h | 33 +++++++ tests/testthat/test-GTM.R | 2 + tests/testthat/test-sGTM.R | 45 ++++++++++ 13 files changed, 447 insertions(+), 5 deletions(-) create mode 100644 man/sGTM.Rd create mode 100644 src/sGTM.cpp create mode 100644 src/sGTM.h create mode 100644 tests/testthat/test-sGTM.R diff --git a/NAMESPACE b/NAMESPACE index 1e1c65b..6c49782 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -3,6 +3,7 @@ export(BuildDiseaseAlpha) export(GTM) export(RTM) +export(sGTM) export(STM) export(STM_Torch) export(STM_Torch_burnin) diff --git a/R/GTM.R b/R/GTM.R index 642ea11..1a4925d 100644 --- a/R/GTM.R +++ b/R/GTM.R @@ -42,9 +42,57 @@ 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, + 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, + shuffle) + return(scte) +} #### Prediction #### #' Topic Inference using Pretrained Model diff --git a/R/RcppExports.R b/R/RcppExports.R index c20d8e9..e7ac955 100644 --- a/R/RcppExports.R +++ b/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, 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, shuffle)) +} + diff --git a/man/sGTM.Rd b/man/sGTM.Rd new file mode 100644 index 0000000..938ac31 --- /dev/null +++ b/man/sGTM.Rd @@ -0,0 +1,57 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/GTM.R +\name{sGTM} +\alias{sGTM} +\title{Stochastic Guided Topic Model (sGTM)} +\usage{ +sGTM( + scte, + K, + D, + batch_size = 1024, + num_threads = 1, + maxiter = 100, + verbal = TRUE, + zero_gamma = FALSE, + rand_gamma = TRUE, + thresh = 1e-08, + burnin = 1, + lr = 0.1, + shuffle = TRUE +) +} +\arguments{ +\item{scte}{A `SingleCellTopicExperiment` or `SpatialTopicExperiment` object containing count data and prior matrices.} + +\item{K}{Integer, the number of topics to infer.} + +\item{D}{Integer, the number of cells (documents) in the dataset.} + +\item{batch_size}{Integer, the number of cells per mini-batch. Default is 1024.} + +\item{num_threads}{Integer, the number of threads to use for parallel computation. Default is 1.} + +\item{maxiter}{Integer, the maximum number of epochs for training. Default is 100.} + +\item{verbal}{Logical, whether to print progress messages. Default is TRUE.} + +\item{zero_gamma}{Logical, whether to initialize gamma values to zero. Default is FALSE.} + +\item{rand_gamma}{Logical, whether to initialize gamma values randomly. Default is TRUE.} + +\item{thresh}{Convergence threshold based on the change in `nwk`. Default is 1e-8.} + +\item{burnin}{Maximum number of iterations to run during the GTM E-step for each sample (default is 1).} + +\item{lr}{Numeric, the SVI learning rate for global updates. Default is 0.1.} + +\item{shuffle}{Logical, whether to shuffle cells each epoch. Default is TRUE.} +} +\value{ +A `SingleCellTopicExperiment` or `SpatialTopicExperiment` object with updated topic distributions. +} +\description{ +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. +} diff --git a/src/CellMap.cpp b/src/CellMap.cpp index 48bcda1..0c0fdd7 100644 --- a/src/CellMap.cpp +++ b/src/CellMap.cpp @@ -106,6 +106,40 @@ std::unordered_map build_Cell_Map(arma::sp_mat& counts, return CellMap; } +std::unordered_map 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 CellMap; + for (arma::uword idx = 0; idx < batch_ids.n_elem; idx++){ + 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); + + CellMap[cell_id] = Cell(cell_id,mtx,alpha,beta,K, + zero_gamma, + rand_gamma); + } + + return CellMap; +} + std::unordered_map build_Predict_Cell_Map(arma::sp_mat& counts, arma::vec& celltypes, arma::vec& genes, diff --git a/src/CellMap.h b/src/CellMap.h index 8c2df90..5aa75ac 100644 --- a/src/CellMap.h +++ b/src/CellMap.h @@ -44,6 +44,17 @@ std::unordered_map build_Cell_Map(arma::sp_mat& counts, bool verbal, int num_threads); +std::unordered_map 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 build_Predict_Cell_Map(arma::sp_mat& counts, arma::vec& celltypes, arma::vec& genes, diff --git a/src/GTM.cpp b/src/GTM.cpp index 954d009..31be039 100644 --- a/src/GTM.cpp +++ b/src/GTM.cpp @@ -3,7 +3,6 @@ #include #include -#include // [[Rcpp::depends(RcppArmadillo)]] using namespace Rcpp; using namespace std; @@ -48,6 +47,7 @@ void build_ndk(arma::mat& n_dk, } + // [[Rcpp::export]] arma::mat get_theta(const arma::mat& n_dk, const arma::mat& alpha){ @@ -195,6 +195,7 @@ void run_epoch(std::unordered_map& CellMap,const arma::mat& alpha, } + /* Full model */ @@ -304,6 +305,7 @@ void train_gtm(arma::sp_mat& counts, } + /// PREDICTION void predict_epoch(std::unordered_map& CellMap,const double& alpha, int K,int D,int M,arma::mat& n_dk, diff --git a/src/GTM.h b/src/GTM.h index 208e650..8efaba9 100644 --- a/src/GTM.h +++ b/src/GTM.h @@ -6,7 +6,6 @@ #include #include #include -#include // [[Rcpp::depends(RcppArmadillo)]] using namespace Rcpp; using namespace std; @@ -73,6 +72,7 @@ void train_gtm(arma::sp_mat& counts, + // Prediction void predict_epoch(std::unordered_map& CellMap,const double& alpha, int K,int D,int M,arma::mat& n_dk, diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp index f1ffac5..889c93c 100644 --- a/src/RcppExports.cpp +++ b/src/RcppExports.cpp @@ -336,6 +336,34 @@ BEGIN_RCPP return rcpp_result_gen; END_RCPP } +// train_sgtm +void train_sgtm(arma::sp_mat& counts, arma::vec& celltypes, arma::vec& genes, const arma::mat& alpha, const arma::mat& beta, int K, int D, arma::mat& n_dk, arma::mat& n_wk, int batch_size, int num_threads, int maxiter, bool verbal, bool zero_gamma, bool rand_gamma, double thresh, int burnin, double lr, bool shuffle); +RcppExport SEXP _SpaTM_train_sgtm(SEXP countsSEXP, SEXP celltypesSEXP, SEXP genesSEXP, SEXP alphaSEXP, SEXP betaSEXP, SEXP KSEXP, SEXP DSEXP, SEXP n_dkSEXP, SEXP n_wkSEXP, SEXP batch_sizeSEXP, SEXP num_threadsSEXP, SEXP maxiterSEXP, SEXP verbalSEXP, SEXP zero_gammaSEXP, SEXP rand_gammaSEXP, SEXP threshSEXP, SEXP burninSEXP, SEXP lrSEXP, SEXP shuffleSEXP) { +BEGIN_RCPP + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< arma::sp_mat& >::type counts(countsSEXP); + Rcpp::traits::input_parameter< arma::vec& >::type celltypes(celltypesSEXP); + Rcpp::traits::input_parameter< arma::vec& >::type genes(genesSEXP); + Rcpp::traits::input_parameter< const arma::mat& >::type alpha(alphaSEXP); + Rcpp::traits::input_parameter< const arma::mat& >::type beta(betaSEXP); + Rcpp::traits::input_parameter< int >::type K(KSEXP); + Rcpp::traits::input_parameter< int >::type D(DSEXP); + Rcpp::traits::input_parameter< arma::mat& >::type n_dk(n_dkSEXP); + Rcpp::traits::input_parameter< arma::mat& >::type n_wk(n_wkSEXP); + Rcpp::traits::input_parameter< int >::type batch_size(batch_sizeSEXP); + Rcpp::traits::input_parameter< int >::type num_threads(num_threadsSEXP); + Rcpp::traits::input_parameter< int >::type maxiter(maxiterSEXP); + Rcpp::traits::input_parameter< bool >::type verbal(verbalSEXP); + Rcpp::traits::input_parameter< bool >::type zero_gamma(zero_gammaSEXP); + Rcpp::traits::input_parameter< bool >::type rand_gamma(rand_gammaSEXP); + Rcpp::traits::input_parameter< double >::type thresh(threshSEXP); + Rcpp::traits::input_parameter< int >::type burnin(burninSEXP); + Rcpp::traits::input_parameter< double >::type lr(lrSEXP); + Rcpp::traits::input_parameter< bool >::type shuffle(shuffleSEXP); + 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, shuffle); + return R_NilValue; +END_RCPP +} static const R_CallMethodDef CallEntries[] = { {"_SpaTM_test_cell_map", (DL_FUNC) &_SpaTM_test_cell_map, 10}, @@ -356,6 +384,7 @@ static const R_CallMethodDef CallEntries[] = { {"_SpaTM_progress_bar", (DL_FUNC) &_SpaTM_progress_bar, 5}, {"_SpaTM_stm_torch_estep", (DL_FUNC) &_SpaTM_stm_torch_estep, 20}, {"_SpaTM_test_mlp_forward", (DL_FUNC) &_SpaTM_test_mlp_forward, 4}, + {"_SpaTM_train_sgtm", (DL_FUNC) &_SpaTM_train_sgtm, 19}, {NULL, NULL, 0} }; diff --git a/src/sGTM.cpp b/src/sGTM.cpp new file mode 100644 index 0000000..638b0a4 --- /dev/null +++ b/src/sGTM.cpp @@ -0,0 +1,176 @@ +#include "CellMap.h" +#include "sGTM.h" +#include +#include + +#include +// [[Rcpp::depends(RcppArmadillo)]] +using namespace Rcpp; +using namespace std; + +static void build_nwk_batch(arma::mat& n_wk_batch, + std::unordered_map& CellMap, + const arma::uvec& batch_ids){ + int tokens; + n_wk_batch.zeros(); + for (arma::uword b = 0; b < batch_ids.n_elem; b++){ + int cell_id = batch_ids(b); + tokens = CellMap[cell_id].cell_mtx.n_rows; + for (int i = 0;i < tokens; i++){ + int gene = CellMap[cell_id].cell_mtx(i,2)-1; + n_wk_batch.row(gene) += CellMap[cell_id].cell_gamma.row(i) * CellMap[cell_id].cell_mtx(i,0); + } + } +} + +static void run_epoch_batch(std::unordered_map& CellMap,const arma::mat& alpha, + const arma::mat& beta,int K,arma::mat& n_dk,arma::mat& n_wk, + const arma::uvec& batch_ids, + int num_threads = 1, + int burnin = 1){ + omp_set_num_threads(num_threads); + + + arma::rowvec beta_sum = arma::sum(beta,0); + arma::rowvec nwk_sum = arma::sum(n_wk,0); + #pragma omp parallel for + for (arma::uword b = 0; b < batch_ids.n_elem; b++){ + int cell_id = batch_ids(b); + arma::mat cur_gamma = CellMap[cell_id].cell_gamma; + arma::mat m = CellMap[cell_id].cell_mtx; + int ndk_id = cell_id - 1; + int tokens = m.n_rows; + double diff = 1; + int max_update = burnin; + int u = 0; + arma::rowvec cur_ndk = n_dk.row(ndk_id); + while (diff > 0.01 && u < max_update){ + u++; + for (int token = 0; token < tokens; token++){ + int counts = m(token,0); + int gene = m(token,2)-1; + arma::rowvec cur_counts = cur_gamma.row(token)*counts; + arma::rowvec gamma_k = (alpha.row(ndk_id) + n_dk.row(ndk_id) - cur_counts) % + (beta.row(gene)+n_wk.row(gene)- cur_counts) / + (beta_sum+nwk_sum- cur_counts); + gamma_k = gamma_k/sum(gamma_k); + if (arma::any(gamma_k < 0)){ + gamma_k = (alpha.row(ndk_id) + n_dk.row(ndk_id)) % (beta.row(gene)+ + n_wk.row(gene)) / + (beta_sum+nwk_sum); + } + cur_gamma.row(token) = gamma_k; + } + + n_dk.row(ndk_id).zeros(); + for (int token = 0; token < tokens; token++){ + n_dk.row(ndk_id) += cur_gamma.row(token)*m(token,0); + } + + diff = abs(arma::accu(cur_ndk - n_dk.row(ndk_id))); + cur_ndk = n_dk.row(ndk_id); + } + CellMap[cell_id].cell_gamma = cur_gamma; + } +} + +// [[Rcpp::export]] +void train_sgtm(arma::sp_mat& counts, + arma::vec& celltypes, + arma::vec& genes, + const arma::mat& alpha, const arma::mat& beta, + int K,int D,arma::mat& n_dk,arma::mat& n_wk, + int batch_size = 1024, + int num_threads = 1,int maxiter = 100, + bool verbal = true, + bool zero_gamma = false, + bool rand_gamma = true, + double thresh = 0.00001, + int burnin = 1, + double lr = 0.1, + bool shuffle = true){ + if (batch_size <= 0){ + batch_size = D; + } + if (batch_size > D){ + batch_size = D; + } + + int M = genes.n_elem; + int num_batches = (D + batch_size - 1) / batch_size; + double prog = 0; + int prog_width = 50; + int pos = 0; + arma::mat old_nwk = n_wk; + + for (int iter = 0; iter < maxiter; iter++){ + arma::uvec order = shuffle ? (arma::randperm(D) + 1) : arma::regspace(1,D); + + for (int b = 0; b < num_batches; b++){ + arma::uword start = static_cast(b * batch_size); + arma::uword end = static_cast(std::min(D, static_cast(start) + batch_size)); + arma::uvec batch_ids = order.subvec(start, end - 1); + + std::unordered_map CellMap = build_Cell_Map_batch(counts, + celltypes, + genes, + alpha, + beta, + batch_ids, + K, + zero_gamma, + rand_gamma, + num_threads); + + run_epoch_batch(CellMap,alpha,beta,K,n_dk,n_wk, + batch_ids,num_threads,burnin); + + arma::mat batch_nwk = arma::zeros(M,K); + build_nwk_batch(batch_nwk,CellMap,batch_ids); + + double scale = static_cast(D) / static_cast(batch_ids.n_elem); + n_wk = (1.0 - lr) * n_wk + lr * (scale * batch_nwk); + CellMap.clear(); + } + + double diff = arma::max(arma::max(arma::abs(n_wk - old_nwk))); + if (verbal){ + Rcout.flush(); + prog += 1.0/maxiter; + Rcout << "["; + pos = int(prog_width * prog); + for(int p = 0; p < prog_width; p++){ + if (p < pos){ Rcout << "=" ;} + else { Rcout << " " ;} + } + Rcout << "] " << int(prog * 100.0) << "% || Iter: " << iter << " || Delta: " << diff <<" \r"; + } + + if (diff < thresh){ + if (verbal){ + Rcout.flush(); + Rcout << "["; + for(int p = 0; p < prog_width; p++){ + Rcout << "=" ; + } + Rcout << "] " << "100% || Iter: " << iter << " || Delta: " << diff << std::endl; + Rcout << "Model Converged at iteration : " << iter << std::endl; + } + return; + } + + old_nwk = n_wk; + } + + if (verbal){ + Rcout.flush(); + Rcout << "["; + for(int p = 0; p < prog_width; p++){ + Rcout << "=" ; + } + Rcout << "] " << "100% || Iter: " << maxiter << " || Delta: " << arma::max(arma::max(arma::abs(n_wk - old_nwk))) << std::endl; + Rcout << "Max Iteration Reached" << std::endl; + } + return; + +} diff --git a/src/sGTM.h b/src/sGTM.h new file mode 100644 index 0000000..967c396 --- /dev/null +++ b/src/sGTM.h @@ -0,0 +1,33 @@ +#ifndef SGTM_H_ +#define SGTM_H_ + +#include "CellMap.h" + +#include +#include +#include +// [[Rcpp::depends(RcppArmadillo)]] +using namespace Rcpp; +using namespace std; +using namespace arma; + +void train_sgtm(arma::sp_mat& counts, + arma::vec& celltypes, + arma::vec& genes, + const arma::mat& alpha, + const arma::mat& beta, + int K,int D, + arma::mat& n_dk, + arma::mat& n_wk, + int batch_size, + int num_threads, + int maxiter, + bool verbal, + bool zero_gamma, + bool rand_gamma, + double thresh, + int burnin, + double lr, + bool shuffle); + +#endif diff --git a/tests/testthat/test-GTM.R b/tests/testthat/test-GTM.R index 259abd5..7ea0bad 100644 --- a/tests/testthat/test-GTM.R +++ b/tests/testthat/test-GTM.R @@ -66,3 +66,5 @@ test_that("inferTopics should only update ndk and theta matrices with correct su }) + + diff --git a/tests/testthat/test-sGTM.R b/tests/testthat/test-sGTM.R new file mode 100644 index 0000000..4e1f02a --- /dev/null +++ b/tests/testthat/test-sGTM.R @@ -0,0 +1,45 @@ +library(testthat) +library(SingleCellExperiment) +library(Matrix) +library(SpaTM) + +## Logical helper to check approximate equality +near_match <- function(target, test, tol = 1e-8) { + all(abs(target - test) < tol) +} + +set.seed(42) + +test_sce_sgtm <- SingleCellExperiment( + assays = list(counts = matrix(rpois(80 * 40, lambda = 60), nrow = 80, ncol = 40)) +) +counts(test_sce_sgtm) <- as(counts(test_sce_sgtm),'dgCMatrix') + +test_that("sGTM updates ndk and produces valid nwk", { + scte <- SingleCellTopicExperiment(test_sce_sgtm, K = 4) + + scte_trained <- sGTM(scte, K = 4, D = ncol(scte), batch_size = 8, + num_threads = 1, maxiter = 3, verbal = FALSE, + lr = 0.5, shuffle = FALSE) + + expect_false(all(ndk(scte_trained) == 0), info = "ndk should not be a zero matrix") + + row_sums_ndk <- rowSums(ndk(scte_trained)) + colsums_counts <- Matrix::colSums(counts(scte_trained)) + expect_true(near_match(row_sums_ndk,colsums_counts), info = "Row sums of ndk should match col sums of counts matrix") + + expect_false(all(nwk(scte_trained) == 0), info = "nwk should not be a zero matrix") + expect_true(all(is.finite(nwk(scte_trained))), info = "nwk should contain finite values") +}) + +test_that("sGTM handles full-batch setting", { + scte <- SingleCellTopicExperiment(test_sce_sgtm, K = 3) + + scte_trained <- sGTM(scte, K = 3, D = ncol(scte), batch_size = ncol(scte), + num_threads = 1, maxiter = 2, verbal = FALSE, + lr = 0.3, shuffle = FALSE) + + row_sums_ndk <- rowSums(ndk(scte_trained)) + colsums_counts <- Matrix::colSums(counts(scte_trained)) + expect_true(near_match(row_sums_ndk,colsums_counts), info = "Full batch ndk sums should match counts") +}) From 065c77b299c785e8cbf3f83f9bb2049d1fcc3eec Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Mon, 2 Feb 2026 08:53:49 -0500 Subject: [PATCH 02/13] [fix] Suppressed startup messages when loading packages in unit tests. --- tests/testthat.R | 6 ++++-- tests/testthat/test-GTM.R | 10 ++++++---- tests/testthat/test-Get_Neighbors.R | 8 +++++--- tests/testthat/test-RTM.R | 10 ++++++---- tests/testthat/test-STM.R | 14 ++++++++------ tests/testthat/test-SingleCellTopicExperiment.R | 12 +++++++----- tests/testthat/test-SpatialTopicExperiment.R | 12 +++++++----- tests/testthat/test-TopicExperiment.utils.R | 8 +++++--- tests/testthat/test-sGTM.R | 10 ++++++---- 9 files changed, 54 insertions(+), 36 deletions(-) diff --git a/tests/testthat.R b/tests/testthat.R index f6584ab..1db0d42 100644 --- a/tests/testthat.R +++ b/tests/testthat.R @@ -6,7 +6,9 @@ # * https://r-pkgs.org/testing-design.html#sec-tests-files-overview # * https://testthat.r-lib.org/articles/special-files.html -library(testthat) -library(SpaTM) +suppressPackageStartupMessages({ + library(testthat) + library(SpaTM) +}) test_check("SpaTM") diff --git a/tests/testthat/test-GTM.R b/tests/testthat/test-GTM.R index 7ea0bad..5ebbd65 100644 --- a/tests/testthat/test-GTM.R +++ b/tests/testthat/test-GTM.R @@ -1,7 +1,9 @@ -library(testthat) -library(SingleCellExperiment) -library(Matrix) -library(SpaTM) +suppressPackageStartupMessages({ + library(testthat) + library(SingleCellExperiment) + library(Matrix) + library(SpaTM) +}) ##Logical Function to check that outputs are 'nearly' identical diff --git a/tests/testthat/test-Get_Neighbors.R b/tests/testthat/test-Get_Neighbors.R index 617c4c9..70e9bbd 100644 --- a/tests/testthat/test-Get_Neighbors.R +++ b/tests/testthat/test-Get_Neighbors.R @@ -1,6 +1,8 @@ -library(testthat) -library(SpatialExperiment) -library(SpaTM) +suppressPackageStartupMessages({ + library(testthat) + library(SpatialExperiment) + library(SpaTM) +}) # Helper function to create a fixed mock SpatialExperiment object mock_spe <- function() { diff --git a/tests/testthat/test-RTM.R b/tests/testthat/test-RTM.R index 0db11f4..d98f4ee 100644 --- a/tests/testthat/test-RTM.R +++ b/tests/testthat/test-RTM.R @@ -1,7 +1,9 @@ -library(testthat) -library(SpatialExperiment) -library(SpaTM) -library(bluster) +suppressPackageStartupMessages({ + library(testthat) + library(SpatialExperiment) + library(SpaTM) + library(bluster) +}) #TODO need tests for the helper functions # Helper function to create a mock SpatialTopicExperiment object diff --git a/tests/testthat/test-STM.R b/tests/testthat/test-STM.R index e1c0f28..1b9ad73 100644 --- a/tests/testthat/test-STM.R +++ b/tests/testthat/test-STM.R @@ -1,9 +1,11 @@ -library(testthat) -library(SpaTM) -library(SpatialExperiment) -library(Matrix) -library(torch) -library(dplyr) +suppressPackageStartupMessages({ + library(testthat) + library(SpaTM) + library(SpatialExperiment) + library(Matrix) + library(torch) + library(dplyr) +}) #TODO need to make sure pred_mlp works on GLM and NN #TODO need to work on how the MLP is save if desired diff --git a/tests/testthat/test-SingleCellTopicExperiment.R b/tests/testthat/test-SingleCellTopicExperiment.R index b376996..bad6c42 100644 --- a/tests/testthat/test-SingleCellTopicExperiment.R +++ b/tests/testthat/test-SingleCellTopicExperiment.R @@ -1,13 +1,15 @@ -library(testthat) -library(SingleCellExperiment) -library(Matrix) -library(SpaTM) # Ensure package is loaded +suppressPackageStartupMessages({ + library(testthat) + library(SingleCellExperiment) + library(Matrix) + library(SpaTM) # Ensure package is loaded +}) # TODO proofread tests to ensure correctness test_that("SingleCellTopicExperiment class is correctly defined", { - library(SpaTM) # Ensure package is loaded + suppressPackageStartupMessages(library(SpaTM)) # Ensure package is loaded expect_true(!is.null(getClassDef("SingleCellTopicExperiment"))) diff --git a/tests/testthat/test-SpatialTopicExperiment.R b/tests/testthat/test-SpatialTopicExperiment.R index ec93df4..044de14 100644 --- a/tests/testthat/test-SpatialTopicExperiment.R +++ b/tests/testthat/test-SpatialTopicExperiment.R @@ -1,10 +1,12 @@ -library(testthat) -library(SpatialExperiment) -library(Matrix) -library(SpaTM) # Ensure package is loaded +suppressPackageStartupMessages({ + library(testthat) + library(SpatialExperiment) + library(Matrix) + library(SpaTM) # Ensure package is loaded +}) test_that("SpatialTopicExperiment class is correctly defined", { - library(SpaTM) # Ensure package is loaded + suppressPackageStartupMessages(library(SpaTM)) # Ensure package is loaded expect_true(!is.null(getClassDef("SpatialTopicExperiment"))) diff --git a/tests/testthat/test-TopicExperiment.utils.R b/tests/testthat/test-TopicExperiment.utils.R index 9d77fd5..b4baa5f 100644 --- a/tests/testthat/test-TopicExperiment.utils.R +++ b/tests/testthat/test-TopicExperiment.utils.R @@ -1,6 +1,8 @@ -library(testthat) -library(SpatialExperiment) -library(SpaTM) +suppressPackageStartupMessages({ + library(testthat) + library(SpatialExperiment) + library(SpaTM) +}) # Mock TopicExperiment objects test_spe <- SpatialExperiment( assays = list(counts = matrix(rpois(100 * 50, lambda = 10), nrow = 100, ncol = 50)) diff --git a/tests/testthat/test-sGTM.R b/tests/testthat/test-sGTM.R index 4e1f02a..69b1c6e 100644 --- a/tests/testthat/test-sGTM.R +++ b/tests/testthat/test-sGTM.R @@ -1,7 +1,9 @@ -library(testthat) -library(SingleCellExperiment) -library(Matrix) -library(SpaTM) +suppressPackageStartupMessages({ + library(testthat) + library(SingleCellExperiment) + library(Matrix) + library(SpaTM) +}) ## Logical helper to check approximate equality near_match <- function(target, test, tol = 1e-8) { From e17c2452ca8e37161613e7611de472e3a6d4c289 Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Mon, 2 Feb 2026 08:54:23 -0500 Subject: [PATCH 03/13] [feat] Added batch-based progress bar to sGTM --- src/sGTM.cpp | 36 +++++++++++++++++++++++++----------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/src/sGTM.cpp b/src/sGTM.cpp index 638b0a4..ff5f996 100644 --- a/src/sGTM.cpp +++ b/src/sGTM.cpp @@ -131,19 +131,29 @@ void train_sgtm(arma::sp_mat& counts, double scale = static_cast(D) / static_cast(batch_ids.n_elem); n_wk = (1.0 - lr) * n_wk + lr * (scale * batch_nwk); CellMap.clear(); + + if (verbal){ + double batch_prog = (static_cast(iter) + + (static_cast(b + 1) / static_cast(num_batches))) / + static_cast(maxiter); + Rcout.flush(); + prog = batch_prog; + Rcout << "["; + pos = int(prog_width * prog); + for(int p = 0; p < prog_width; p++){ + if (p < pos){ Rcout << "=" ;} + else { Rcout << " " ;} + } + double diff = arma::max(arma::max(arma::abs(n_wk - old_nwk))); + Rcout << "] " << int(prog * 100.0) << "% || Iter: " << iter + << " || Batches: " << (b + 1) << "/" << num_batches + << " || Delta: " << diff <<" \r"; + } } double diff = arma::max(arma::max(arma::abs(n_wk - old_nwk))); if (verbal){ - Rcout.flush(); - prog += 1.0/maxiter; - Rcout << "["; - pos = int(prog_width * prog); - for(int p = 0; p < prog_width; p++){ - if (p < pos){ Rcout << "=" ;} - else { Rcout << " " ;} - } - Rcout << "] " << int(prog * 100.0) << "% || Iter: " << iter << " || Delta: " << diff <<" \r"; + prog = (static_cast(iter + 1)) / static_cast(maxiter); } if (diff < thresh){ @@ -153,7 +163,9 @@ void train_sgtm(arma::sp_mat& counts, for(int p = 0; p < prog_width; p++){ Rcout << "=" ; } - Rcout << "] " << "100% || Iter: " << iter << " || Delta: " << diff << std::endl; + Rcout << "] " << "100% || Iter: " << iter + << " || Batches: " << num_batches << "/" << num_batches + << " || Delta: " << diff << std::endl; Rcout << "Model Converged at iteration : " << iter << std::endl; } return; @@ -168,7 +180,9 @@ void train_sgtm(arma::sp_mat& counts, for(int p = 0; p < prog_width; p++){ Rcout << "=" ; } - Rcout << "] " << "100% || Iter: " << maxiter << " || Delta: " << arma::max(arma::max(arma::abs(n_wk - old_nwk))) << std::endl; + Rcout << "] " << "100% || Iter: " << maxiter + << " || Batches: " << num_batches << "/" << num_batches + << " || Delta: " << arma::max(arma::max(arma::abs(n_wk - old_nwk))) << std::endl; Rcout << "Max Iteration Reached" << std::endl; } return; From c1a7ed426756a46c3985b2b17c2c45aa6faad9bd Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Mon, 2 Feb 2026 08:59:51 -0500 Subject: [PATCH 04/13] [fix] Made changes to the inferTopics algorithm to be less memory hungry and address openMP-based flaws. --- src/GTM.cpp | 151 ++++++++++++++++++++++++++------------------ src/GTM.h | 12 ++-- src/RcppExports.cpp | 16 ++--- 3 files changed, 102 insertions(+), 77 deletions(-) diff --git a/src/GTM.cpp b/src/GTM.cpp index 31be039..8d15910 100644 --- a/src/GTM.cpp +++ b/src/GTM.cpp @@ -1,6 +1,7 @@ #include "CellMap.h" #include "GTM.h" #include +#include #include // [[Rcpp::depends(RcppArmadillo)]] @@ -307,86 +308,110 @@ void train_gtm(arma::sp_mat& counts, /// PREDICTION -void predict_epoch(std::unordered_map& CellMap,const double& alpha, - int K,int D,int M,arma::mat& n_dk, +struct PredictCell { + arma::vec counts; + arma::uvec genes; + arma::mat gamma; + + PredictCell() {} + PredictCell(const arma::vec& counts_in, const arma::uvec& genes_in, int K) + : counts(counts_in), genes(genes_in), gamma(counts_in.n_elem, K) { + if (counts_in.n_elem > 0){ + gamma.fill(1.0 / K); + } + } +}; + +static void init_predict_cells(const arma::sp_mat& counts, + const arma::vec& genes, + int K, int D, + std::vector& cells, + arma::mat& n_dk){ + cells.clear(); + cells.reserve(D); + n_dk.set_size(D, K); + n_dk.zeros(); + for (int i = 0; i < D; i++){ + arma::uvec geneidx = arma::find(counts.col(i)); + arma::vec token_counts = arma::nonzeros(counts.col(i)); + arma::uvec gene_ids = arma::conv_to::from(genes(geneidx)); + if (gene_ids.n_elem > 0){ + gene_ids -= 1; + } + cells.emplace_back(token_counts, gene_ids, K); + if (token_counts.n_elem > 0){ + double total = arma::accu(token_counts); + n_dk.row(i).fill(total / K); + } + } +} + +static void predict_epoch(std::vector& cells,const double& alpha, + int K,int D,arma::mat& n_dk, const arma::mat& phi, int num_threads = 1, int burnin = 1){ omp_set_num_threads(num_threads); - int gene,counts,ndk_id; - arma::rowvec gamma_k = arma::zeros(K); - arma::rowvec cur_counts = arma::zeros(K); #pragma omp parallel for - for (int i = 1; i <= D; i++){ + for (int i = 0; i < D; i++){ + PredictCell& cell = cells[i]; + int tokens = cell.counts.n_elem; + if (tokens == 0){ + n_dk.row(i).zeros(); + continue; + } - arma::mat cur_gamma = CellMap[i].cell_gamma; //TODO - arma::mat m = CellMap[i].cell_mtx; //TODO - ndk_id = i-1; //To accomodate row indices for matrix - int tokens = m.n_rows; + arma::rowvec gamma_k(K, arma::fill::zeros); + arma::rowvec cur_counts(K, arma::fill::zeros); double diff = 1; int max_update = burnin; int u = 0; - arma::rowvec cur_ndk = n_dk.row(ndk_id); - //Gamma burn-in + arma::rowvec cur_ndk = n_dk.row(i); while (diff > 0.01 && u < max_update){ u++; for (int token = 0; token < tokens; token++){ - //Rcout << "1" << std::endl; - counts = m(token,0); - gene = m(token,2)-1; - cur_counts = cur_gamma.row(token)*counts; //TODO - // Rcout << "2" << std::endl; - gamma_k = (alpha + n_dk.row(ndk_id) - cur_counts) % phi.row(gene); + int gene = cell.genes(token); + double counts = cell.counts(token); + cur_counts = cell.gamma.row(token)*counts; + gamma_k = (alpha + n_dk.row(i) - cur_counts) % phi.row(gene); - //sequential par if (arma::any(gamma_k < 0)){ - gamma_k = (alpha + n_dk.row(ndk_id)) % phi.row(gene); + gamma_k = (alpha + n_dk.row(i)) % phi.row(gene); } gamma_k = gamma_k/sum(gamma_k); - //Rcout << "3" << std::endl; - cur_gamma.row(token) = gamma_k; //TODO - - //update params - //Rcout << "4" << std::endl; + cell.gamma.row(token) = gamma_k; } - n_dk.row(ndk_id).zeros(); - for (int token = 0; token < tokens; token++){ //TODO vectorise - n_dk.row(ndk_id) += cur_gamma.row(token)*m(token,0); + n_dk.row(i).zeros(); + for (int token = 0; token < tokens; token++){ + n_dk.row(i) += cell.gamma.row(token)*cell.counts(token); } - diff = abs(arma::accu(cur_ndk - n_dk.row(ndk_id))); - cur_ndk = n_dk.row(ndk_id); + diff = abs(arma::accu(cur_ndk - n_dk.row(i))); + cur_ndk = n_dk.row(i); } - CellMap[i].cell_gamma = cur_gamma; //TODO - - } } // [[Rcpp::export]] -void infer_topics_cpp(arma::sp_mat& counts, - arma::vec& celltypes, - arma::vec& genes, +void infer_topics_cpp(const arma::sp_mat& counts, + const arma::vec& celltypes, + const arma::vec& genes, double& alpha, int K,int D,int M,arma::mat& n_dk, const arma::mat& phi, int num_threads = 1,int maxiter = 100, bool verbal = true,int burnin = 1){ - std::unordered_map CellMap = build_Predict_Cell_Map(counts, - celltypes, - genes, - D,K); - - build_ndk(n_dk,CellMap,D,K); + static_cast(celltypes); + static_cast(M); + std::vector cells; + init_predict_cells(counts, genes, K, D, cells, n_dk); double prog = 0; int prog_width = 50; int pos = 0; for (int i = 0; i < maxiter; i++){ - predict_epoch(CellMap,alpha,K,D,M,n_dk,phi,num_threads,burnin); - build_ndk(n_dk,CellMap,D,K); + predict_epoch(cells,alpha,K,D,n_dk,phi,num_threads,burnin); if (verbal){ - Rcout.flush(); prog += 1.0/maxiter; @@ -408,31 +433,28 @@ void infer_topics_cpp(arma::sp_mat& counts, Rcout << "] " << "100% || Iter: " << maxiter << std::endl; Rcout << "Max Iteration Reached" << std::endl; } - CellMap.clear(); return; } // [[Rcpp::export]] -List infer_gex_cpp(arma::sp_mat& counts, - arma::vec& celltypes, - arma::vec& genes, +List infer_gex_cpp(const arma::sp_mat& counts, + const arma::vec& celltypes, + const arma::vec& genes, double& alpha, int K, int D, int M, arma::mat& n_dk, const arma::mat& phi, int num_threads = 1, int maxiter = 100, bool verbal = true, int burnin = 1 ) { -std::unordered_map CellMap = build_Predict_Cell_Map(counts, - celltypes, - genes, - D, K); -build_ndk(n_dk, CellMap, D, K); +static_cast(M); +std::vector cells; +init_predict_cells(counts, genes, K, D, cells, n_dk); double prog = 0; int prog_width = 50; int pos = 0; for (int i = 0; i < maxiter; i++) { -predict_epoch(CellMap, alpha, K, D, M, n_dk, phi, num_threads,burnin); +predict_epoch(cells, alpha, K, D, n_dk, phi, num_threads,burnin); if (verbal) { Rcout.flush(); prog += 1.0 / maxiter; @@ -458,15 +480,22 @@ Rcout << "] " << "100% || Iter: " << maxiter << std::endl; Rcout << "Max Iteration Reached" << std::endl; } -// Collect gamma and mtx matrices from CellMap +// Collect gamma and mtx matrices from predicted cells List result(D); -for (int i = 1; i <= D; i++) { +for (int i = 0; i < D; i++) { +PredictCell& cell = cells[i]; +arma::mat mtx(cell.counts.n_elem,3); +mtx.col(0) = cell.counts; +mtx.col(1).fill(celltypes(i)); +if (cell.genes.n_elem > 0){ + arma::vec gene_vals = arma::conv_to::from(cell.genes) + 1; + mtx.col(2) = gene_vals; +} List cell_data = List::create( -Named("cell_gamma") = CellMap[i].cell_gamma, -Named("cell_mtx") = CellMap[i].cell_mtx +Named("cell_gamma") = cell.gamma, +Named("cell_mtx") = mtx ); -result[i - 1] = cell_data; +result[i] = cell_data; } -CellMap.clear(); return result; } diff --git a/src/GTM.h b/src/GTM.h index 8efaba9..ec25c57 100644 --- a/src/GTM.h +++ b/src/GTM.h @@ -74,16 +74,12 @@ void train_gtm(arma::sp_mat& counts, // Prediction -void predict_epoch(std::unordered_map& CellMap,const double& alpha, - int K,int D,int M,arma::mat& n_dk, - const arma::mat& phi, int num_threads); - -void infer_topics_cpp(arma::sp_mat& counts, - arma::vec& celltypes, - arma::vec& genes, +void infer_topics_cpp(const arma::sp_mat& counts, + const arma::vec& celltypes, + const arma::vec& genes, double& alpha, int K,int D,int M,arma::mat& n_dk, const arma::mat& phi, int num_threads,int maxiter, - bool verbal); + bool verbal,int burnin); #endif diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp index 889c93c..d2b8704 100644 --- a/src/RcppExports.cpp +++ b/src/RcppExports.cpp @@ -92,13 +92,13 @@ BEGIN_RCPP END_RCPP } // infer_topics_cpp -void infer_topics_cpp(arma::sp_mat& counts, arma::vec& celltypes, arma::vec& genes, double& alpha, int K, int D, int M, arma::mat& n_dk, const arma::mat& phi, int num_threads, int maxiter, bool verbal, int burnin); +void infer_topics_cpp(const arma::sp_mat& counts, const arma::vec& celltypes, const arma::vec& genes, double& alpha, int K, int D, int M, arma::mat& n_dk, const arma::mat& phi, int num_threads, int maxiter, bool verbal, int burnin); RcppExport SEXP _SpaTM_infer_topics_cpp(SEXP countsSEXP, SEXP celltypesSEXP, SEXP genesSEXP, SEXP alphaSEXP, SEXP KSEXP, SEXP DSEXP, SEXP MSEXP, SEXP n_dkSEXP, SEXP phiSEXP, SEXP num_threadsSEXP, SEXP maxiterSEXP, SEXP verbalSEXP, SEXP burninSEXP) { BEGIN_RCPP Rcpp::RNGScope rcpp_rngScope_gen; - Rcpp::traits::input_parameter< arma::sp_mat& >::type counts(countsSEXP); - Rcpp::traits::input_parameter< arma::vec& >::type celltypes(celltypesSEXP); - Rcpp::traits::input_parameter< arma::vec& >::type genes(genesSEXP); + Rcpp::traits::input_parameter< const arma::sp_mat& >::type counts(countsSEXP); + Rcpp::traits::input_parameter< const arma::vec& >::type celltypes(celltypesSEXP); + Rcpp::traits::input_parameter< const arma::vec& >::type genes(genesSEXP); Rcpp::traits::input_parameter< double& >::type alpha(alphaSEXP); Rcpp::traits::input_parameter< int >::type K(KSEXP); Rcpp::traits::input_parameter< int >::type D(DSEXP); @@ -114,14 +114,14 @@ BEGIN_RCPP END_RCPP } // infer_gex_cpp -List infer_gex_cpp(arma::sp_mat& counts, arma::vec& celltypes, arma::vec& genes, double& alpha, int K, int D, int M, arma::mat& n_dk, const arma::mat& phi, int num_threads, int maxiter, bool verbal, int burnin); +List infer_gex_cpp(const arma::sp_mat& counts, const arma::vec& celltypes, const arma::vec& genes, double& alpha, int K, int D, int M, arma::mat& n_dk, const arma::mat& phi, int num_threads, int maxiter, bool verbal, int burnin); RcppExport SEXP _SpaTM_infer_gex_cpp(SEXP countsSEXP, SEXP celltypesSEXP, SEXP genesSEXP, SEXP alphaSEXP, SEXP KSEXP, SEXP DSEXP, SEXP MSEXP, SEXP n_dkSEXP, SEXP phiSEXP, SEXP num_threadsSEXP, SEXP maxiterSEXP, SEXP verbalSEXP, SEXP burninSEXP) { BEGIN_RCPP Rcpp::RObject rcpp_result_gen; Rcpp::RNGScope rcpp_rngScope_gen; - Rcpp::traits::input_parameter< arma::sp_mat& >::type counts(countsSEXP); - Rcpp::traits::input_parameter< arma::vec& >::type celltypes(celltypesSEXP); - Rcpp::traits::input_parameter< arma::vec& >::type genes(genesSEXP); + Rcpp::traits::input_parameter< const arma::sp_mat& >::type counts(countsSEXP); + Rcpp::traits::input_parameter< const arma::vec& >::type celltypes(celltypesSEXP); + Rcpp::traits::input_parameter< const arma::vec& >::type genes(genesSEXP); Rcpp::traits::input_parameter< double& >::type alpha(alphaSEXP); Rcpp::traits::input_parameter< int >::type K(KSEXP); Rcpp::traits::input_parameter< int >::type D(DSEXP); From 52687d0bbc5517d9b0167db5602962402398cc06 Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Tue, 3 Feb 2026 14:42:19 -0500 Subject: [PATCH 05/13] [fix] Edited code to parallelize CellMap construction during sGTM inference. --- src/CellMap.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/CellMap.cpp b/src/CellMap.cpp index 0c0fdd7..7d0d065 100644 --- a/src/CellMap.cpp +++ b/src/CellMap.cpp @@ -119,6 +119,8 @@ std::unordered_map build_Cell_Map_batch(arma::sp_mat& counts, omp_set_num_threads(num_threads); std::unordered_map CellMap; + + #pragma omp parallel for for (arma::uword idx = 0; idx < batch_ids.n_elem; idx++){ int cell_id = batch_ids(idx); int col_id = cell_id - 1; From cda883039cd1c44e72d29ac93f635152bf641f91 Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Tue, 3 Feb 2026 15:13:08 -0500 Subject: [PATCH 06/13] [fix] Fixed possible data racing issues with parallelized build_CellMap function. Included simple tests to ensure multi-threading works and a warning for systems with access to a single thread. --- src/CellMap.cpp | 34 ++++++++++++++++++++++++++-------- tests/testthat/test-GTM.R | 19 +++++++++++++++++++ tests/testthat/test-RTM.R | 21 +++++++++++++++++++++ tests/testthat/test-sGTM.R | 22 ++++++++++++++++++++++ 4 files changed, 88 insertions(+), 8 deletions(-) diff --git a/src/CellMap.cpp b/src/CellMap.cpp index 7d0d065..fdf7a4c 100644 --- a/src/CellMap.cpp +++ b/src/CellMap.cpp @@ -1,6 +1,7 @@ #include "CellMap.h" #include +#include // [[Rcpp::depends(RcppArmadillo)]] @@ -88,9 +89,10 @@ std::unordered_map build_Cell_Map(arma::sp_mat& counts, omp_set_num_threads(num_threads); std::unordered_map CellMap; - //int cell_start,cell_end; + std::vector> 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); @@ -98,9 +100,17 @@ std::unordered_map build_Cell_Map(arma::sp_mat& 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; @@ -119,9 +129,10 @@ std::unordered_map build_Cell_Map_batch(arma::sp_mat& counts, omp_set_num_threads(num_threads); std::unordered_map CellMap; - + std::vector> 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){ @@ -134,9 +145,16 @@ std::unordered_map build_Cell_Map_batch(arma::sp_mat& counts, mtx.col(1).fill(celltypes(col_id)); mtx.col(2) = genes(geneidx); - CellMap[cell_id] = Cell(cell_id,mtx,alpha,beta,K, - zero_gamma, - rand_gamma); + 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; diff --git a/tests/testthat/test-GTM.R b/tests/testthat/test-GTM.R index 5ebbd65..66489d7 100644 --- a/tests/testthat/test-GTM.R +++ b/tests/testthat/test-GTM.R @@ -12,6 +12,17 @@ near_match <- function(target, test, tol = 1e-8) { all(abs(target - test) < tol) } +check_multithread_available <- function() { + if (!requireNamespace("parallel", quietly = TRUE)) { + install.packages("parallel") + } + cores <- parallel::detectCores(logical = FALSE) + if (is.na(cores) || cores < 2) { + warning("Only 1 CPU core detected. Multi-threaded tests need at least 2 cores; consider running on a multi-core machine.") + } + invisible(cores) +} + test_sce <- SingleCellExperiment( assays = list(counts = matrix(rpois(100 * 50, lambda = 80), nrow = 100, ncol = 50)) ) @@ -44,6 +55,14 @@ test_that("ndk and nwk are updated correctly after training", { }) +test_that("GTM runs with num_threads = 2", { + check_multithread_available() + scte <- SingleCellTopicExperiment(test_sce, K = 4) + expect_silent( + GTM(scte, K = 4, D = ncol(scte), num_threads = 2, maxiter = 2, verbal = FALSE) + ) +}) + # Test 2: Check that phi and theta have colSums and rowSums equal to 1 test_that("inferTopics should only update ndk and theta matrices with correct sums", { scte <- SingleCellTopicExperiment(test_sce, K = 5) diff --git a/tests/testthat/test-RTM.R b/tests/testthat/test-RTM.R index d98f4ee..1adb40c 100644 --- a/tests/testthat/test-RTM.R +++ b/tests/testthat/test-RTM.R @@ -33,6 +33,17 @@ create_test_spatial_topic_experiment <- function() { SpatialTopicExperiment(spe,K = 10) } +check_multithread_available <- function() { + if (!requireNamespace("parallel", quietly = TRUE)) { + install.packages("parallel") + } + cores <- parallel::detectCores(logical = FALSE) + if (is.na(cores) || cores < 2) { + warning("Only 1 CPU core detected. Multi-threaded tests need at least 2 cores; consider running on a multi-core machine.") + } + invisible(cores) +} + test_that("RTM function returns expected results for valid input", { ste <- create_test_spatial_topic_experiment() nbr_list <- get_nbrs(ste, samples = "sample_id", cell_ids = "int_cell", dist = 5) @@ -86,6 +97,16 @@ test_that("RTM throws error for invalid input", { expect_error(RTM(invalid_input, K = 5, nbr_list = nbr_list, num_threads = 1, maxiter = 5, verbal = FALSE)) }) +test_that("RTM runs with num_threads = 2", { + check_multithread_available() + ste <- create_test_spatial_topic_experiment() + nbr_list <- get_nbrs(ste, samples = "sample_id", cell_ids = "int_cell", dist = 5) + K <- ncol(alphaPrior(ste)) + expect_silent( + RTM(ste, K = K, nbr_list = nbr_list, num_threads = 2, maxiter = 2, verbal = FALSE) + ) +}) + #### new tests diff --git a/tests/testthat/test-sGTM.R b/tests/testthat/test-sGTM.R index 69b1c6e..b857b26 100644 --- a/tests/testthat/test-sGTM.R +++ b/tests/testthat/test-sGTM.R @@ -10,6 +10,17 @@ near_match <- function(target, test, tol = 1e-8) { all(abs(target - test) < tol) } +check_multithread_available <- function() { + if (!requireNamespace("parallel", quietly = TRUE)) { + install.packages("parallel") + } + cores <- parallel::detectCores(logical = FALSE) + if (is.na(cores) || cores < 2) { + warning("Only 1 CPU core detected. Multi-threaded tests need at least 2 cores; consider running on a multi-core machine.") + } + invisible(cores) +} + set.seed(42) test_sce_sgtm <- SingleCellExperiment( @@ -45,3 +56,14 @@ test_that("sGTM handles full-batch setting", { colsums_counts <- Matrix::colSums(counts(scte_trained)) expect_true(near_match(row_sums_ndk,colsums_counts), info = "Full batch ndk sums should match counts") }) + +test_that("sGTM runs with num_threads = 2", { + check_multithread_available() + scte <- SingleCellTopicExperiment(test_sce_sgtm, K = 3) + + expect_silent( + sGTM(scte, K = 3, D = ncol(scte), batch_size = 10, + num_threads = 2, maxiter = 2, verbal = FALSE, + lr = 0.5, shuffle = FALSE) + ) +}) From bd0465b9915ef505a0f6455d4b51e4cfba18e250 Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Tue, 3 Feb 2026 16:01:15 -0500 Subject: [PATCH 07/13] [fix] Added learning rate scheduler and catches for invalid calculations. --- src/sGTM.cpp | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/src/sGTM.cpp b/src/sGTM.cpp index ff5f996..35dca5c 100644 --- a/src/sGTM.cpp +++ b/src/sGTM.cpp @@ -30,7 +30,7 @@ static void run_epoch_batch(std::unordered_map& CellMap,const arma::ma int burnin = 1){ omp_set_num_threads(num_threads); - + const double eps = 1e-12; arma::rowvec beta_sum = arma::sum(beta,0); arma::rowvec nwk_sum = arma::sum(n_wk,0); #pragma omp parallel for @@ -50,15 +50,22 @@ static void run_epoch_batch(std::unordered_map& CellMap,const arma::ma int counts = m(token,0); int gene = m(token,2)-1; arma::rowvec cur_counts = cur_gamma.row(token)*counts; + arma::rowvec denom = beta_sum + nwk_sum - cur_counts; + denom = arma::clamp(denom, eps, arma::datum::inf); arma::rowvec gamma_k = (alpha.row(ndk_id) + n_dk.row(ndk_id) - cur_counts) % (beta.row(gene)+n_wk.row(gene)- cur_counts) / - (beta_sum+nwk_sum- cur_counts); - gamma_k = gamma_k/sum(gamma_k); - if (arma::any(gamma_k < 0)){ + denom; + double gamma_sum = arma::accu(gamma_k); + if (!gamma_k.is_finite() || gamma_sum <= eps || arma::any(gamma_k < 0)){ gamma_k = (alpha.row(ndk_id) + n_dk.row(ndk_id)) % (beta.row(gene)+ n_wk.row(gene)) / (beta_sum+nwk_sum); + gamma_sum = arma::accu(gamma_k); + if (!gamma_k.is_finite() || gamma_sum <= eps){ + gamma_k.fill(1.0 / static_cast(K)); + } } + gamma_k = gamma_k/arma::accu(gamma_k); cur_gamma.row(token) = gamma_k; } @@ -102,6 +109,9 @@ void train_sgtm(arma::sp_mat& counts, int prog_width = 50; int pos = 0; arma::mat old_nwk = n_wk; + const double tau0 = 1.0; + const double kappa = 0.7; + double step = 0.0; for (int iter = 0; iter < maxiter; iter++){ arma::uvec order = shuffle ? (arma::randperm(D) + 1) : arma::regspace(1,D); @@ -129,7 +139,10 @@ void train_sgtm(arma::sp_mat& counts, build_nwk_batch(batch_nwk,CellMap,batch_ids); double scale = static_cast(D) / static_cast(batch_ids.n_elem); - n_wk = (1.0 - lr) * n_wk + lr * (scale * batch_nwk); + double lr_t = lr * std::pow(tau0 + step, -kappa); + lr_t = std::min(1.0, std::max(lr_t, 1e-6)); + n_wk = (1.0 - lr_t) * n_wk + lr_t * (scale * batch_nwk); + step += 1.0; CellMap.clear(); if (verbal){ From 82ad4eba9c18f7613eef244e3727b9e9b6c7eef8 Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Tue, 3 Feb 2026 18:21:06 -0500 Subject: [PATCH 08/13] [fix] Switched n_wk diff to be against scaled batch n_wk instead of old. Updated DESCRIPTION to use biocViews for dependencies. --- DESCRIPTION | 1 + src/sGTM.cpp | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index 4154c37..ef4e6d2 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -7,6 +7,7 @@ Author: Adrien Osakwe Maintainer: Adrien Osakwe Description: Topic Modelling framework for Spatial, Bulk, and Single cell analyses. License: MIT + file LICENSE +biocViews: Imports: Rcpp (>= 1.0.10), SingleCellExperiment, S4Vectors, diff --git a/src/sGTM.cpp b/src/sGTM.cpp index 35dca5c..f128f37 100644 --- a/src/sGTM.cpp +++ b/src/sGTM.cpp @@ -157,7 +157,10 @@ void train_sgtm(arma::sp_mat& counts, if (p < pos){ Rcout << "=" ;} else { Rcout << " " ;} } - double diff = arma::max(arma::max(arma::abs(n_wk - old_nwk))); + //double diff = arma::max(arma::max(arma::abs(n_wk - old_nwk))); + //Switched to scaled current batch to better assess convergence + double diff = arma::max(arma::max(arma::abs(n_wk - (scale * batch_nwk)))); + Rcout << "] " << int(prog * 100.0) << "% || Iter: " << iter << " || Batches: " << (b + 1) << "/" << num_batches << " || Delta: " << diff <<" \r"; From 7d7a3f75ea5ecbc67094f20881b6096ab2245329 Mon Sep 17 00:00:00 2001 From: Adrien Osakwe Date: Thu, 5 Feb 2026 14:58:21 -0500 Subject: [PATCH 09/13] [feat] Added scheduler parameters for sGTM to the wrapper function arguments. --- R/GTM.R | 4 ++++ R/RcppExports.R | 4 ++-- src/RcppExports.cpp | 10 ++++++---- src/sGTM.cpp | 5 +++-- 4 files changed, 15 insertions(+), 8 deletions(-) diff --git a/R/GTM.R b/R/GTM.R index 1a4925d..35494af 100644 --- a/R/GTM.R +++ b/R/GTM.R @@ -71,6 +71,8 @@ sGTM <- function(scte,K,D,batch_size = 1024,num_threads = 1,maxiter = 100,verbal 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, @@ -90,6 +92,8 @@ sGTM <- function(scte,K,D,batch_size = 1024,num_threads = 1,maxiter = 100,verbal thresh, burnin, lr, + tau0, + kappa, shuffle) return(scte) } diff --git a/R/RcppExports.R b/R/RcppExports.R index e7ac955..59ed9aa 100644 --- a/R/RcppExports.R +++ b/R/RcppExports.R @@ -87,7 +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, 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, shuffle)) +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)) } diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp index d2b8704..aee8599 100644 --- a/src/RcppExports.cpp +++ b/src/RcppExports.cpp @@ -337,8 +337,8 @@ BEGIN_RCPP END_RCPP } // train_sgtm -void train_sgtm(arma::sp_mat& counts, arma::vec& celltypes, arma::vec& genes, const arma::mat& alpha, const arma::mat& beta, int K, int D, arma::mat& n_dk, arma::mat& n_wk, int batch_size, int num_threads, int maxiter, bool verbal, bool zero_gamma, bool rand_gamma, double thresh, int burnin, double lr, bool shuffle); -RcppExport SEXP _SpaTM_train_sgtm(SEXP countsSEXP, SEXP celltypesSEXP, SEXP genesSEXP, SEXP alphaSEXP, SEXP betaSEXP, SEXP KSEXP, SEXP DSEXP, SEXP n_dkSEXP, SEXP n_wkSEXP, SEXP batch_sizeSEXP, SEXP num_threadsSEXP, SEXP maxiterSEXP, SEXP verbalSEXP, SEXP zero_gammaSEXP, SEXP rand_gammaSEXP, SEXP threshSEXP, SEXP burninSEXP, SEXP lrSEXP, SEXP shuffleSEXP) { +void train_sgtm(arma::sp_mat& counts, arma::vec& celltypes, arma::vec& genes, const arma::mat& alpha, const arma::mat& beta, int K, int D, arma::mat& n_dk, arma::mat& n_wk, int batch_size, int num_threads, int maxiter, bool verbal, bool zero_gamma, bool rand_gamma, double thresh, int burnin, double lr, const double tau0, const double kappa, bool shuffle); +RcppExport SEXP _SpaTM_train_sgtm(SEXP countsSEXP, SEXP celltypesSEXP, SEXP genesSEXP, SEXP alphaSEXP, SEXP betaSEXP, SEXP KSEXP, SEXP DSEXP, SEXP n_dkSEXP, SEXP n_wkSEXP, SEXP batch_sizeSEXP, SEXP num_threadsSEXP, SEXP maxiterSEXP, SEXP verbalSEXP, SEXP zero_gammaSEXP, SEXP rand_gammaSEXP, SEXP threshSEXP, SEXP burninSEXP, SEXP lrSEXP, SEXP tau0SEXP, SEXP kappaSEXP, SEXP shuffleSEXP) { BEGIN_RCPP Rcpp::RNGScope rcpp_rngScope_gen; Rcpp::traits::input_parameter< arma::sp_mat& >::type counts(countsSEXP); @@ -359,8 +359,10 @@ BEGIN_RCPP Rcpp::traits::input_parameter< double >::type thresh(threshSEXP); Rcpp::traits::input_parameter< int >::type burnin(burninSEXP); Rcpp::traits::input_parameter< double >::type lr(lrSEXP); + Rcpp::traits::input_parameter< const double >::type tau0(tau0SEXP); + Rcpp::traits::input_parameter< const double >::type kappa(kappaSEXP); Rcpp::traits::input_parameter< bool >::type shuffle(shuffleSEXP); - 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, shuffle); + 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); return R_NilValue; END_RCPP } @@ -384,7 +386,7 @@ static const R_CallMethodDef CallEntries[] = { {"_SpaTM_progress_bar", (DL_FUNC) &_SpaTM_progress_bar, 5}, {"_SpaTM_stm_torch_estep", (DL_FUNC) &_SpaTM_stm_torch_estep, 20}, {"_SpaTM_test_mlp_forward", (DL_FUNC) &_SpaTM_test_mlp_forward, 4}, - {"_SpaTM_train_sgtm", (DL_FUNC) &_SpaTM_train_sgtm, 19}, + {"_SpaTM_train_sgtm", (DL_FUNC) &_SpaTM_train_sgtm, 21}, {NULL, NULL, 0} }; diff --git a/src/sGTM.cpp b/src/sGTM.cpp index f128f37..cd3bff1 100644 --- a/src/sGTM.cpp +++ b/src/sGTM.cpp @@ -95,6 +95,8 @@ void train_sgtm(arma::sp_mat& counts, double thresh = 0.00001, int burnin = 1, double lr = 0.1, + const double tau0 = 1.0, + const double kappa = 0.7, bool shuffle = true){ if (batch_size <= 0){ batch_size = D; @@ -109,8 +111,7 @@ void train_sgtm(arma::sp_mat& counts, int prog_width = 50; int pos = 0; arma::mat old_nwk = n_wk; - const double tau0 = 1.0; - const double kappa = 0.7; + double step = 0.0; for (int iter = 0; iter < maxiter; iter++){ From 79514e4e44fec326284af65bbdc0e0993ec4ea24 Mon Sep 17 00:00:00 2001 From: aosakwe Date: Tue, 2 Jun 2026 16:57:45 -0400 Subject: [PATCH 10/13] [fix] removed hard-coded single-thread run for infertopics. --- R/GTM.R | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/R/GTM.R b/R/GTM.R index 35494af..dc6eb1d 100644 --- a/R/GTM.R +++ b/R/GTM.R @@ -108,9 +108,10 @@ sGTM <- function(scte,K,D,batch_size = 1024,num_threads = 1,maxiter = 100,verbal #' @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, @@ -119,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) } @@ -146,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) } From 9ec6dad6714fcc766a7854a79fb2337989a29e96 Mon Sep 17 00:00:00 2001 From: aosakwe Date: Tue, 2 Jun 2026 16:58:56 -0400 Subject: [PATCH 11/13] added .gitignore to avoid commiting gcc outputs --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..aba3460 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +*.o +*.so +data/ From f5e1c22d17fa9cb1671ffe40781178092a12ad7f Mon Sep 17 00:00:00 2001 From: aosakwe Date: Tue, 2 Jun 2026 17:30:52 -0400 Subject: [PATCH 12/13] Tweaked DESCRIPTION to automatically install bioconductor dependencies --- DESCRIPTION | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/DESCRIPTION b/DESCRIPTION index ef4e6d2..2de8363 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -7,20 +7,21 @@ Author: Adrien Osakwe Maintainer: Adrien Osakwe Description: Topic Modelling framework for Spatial, Bulk, and Single cell analyses. License: MIT + file LICENSE -biocViews: -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, @@ -31,3 +32,4 @@ Encoding: UTF-8 Config/testthat/edition: 3 VignetteBuilder: knitr URL: https://li-lab-mcgill.github.io/SpaTM/ +biocViews: Software, Transcriptomics From fbb689f552d31b9a4dc4aa4a127959d2a75cfc6f Mon Sep 17 00:00:00 2001 From: aosakwe Date: Tue, 2 Jun 2026 17:31:27 -0400 Subject: [PATCH 13/13] added comments to GTM rcpp code --- src/GTM.cpp | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/GTM.cpp b/src/GTM.cpp index 8d15910..db5cd17 100644 --- a/src/GTM.cpp +++ b/src/GTM.cpp @@ -308,10 +308,15 @@ void train_gtm(arma::sp_mat& counts, /// PREDICTION +/// +/// n_dk inference: trained gene-topic matrix phi is held constant; only +/// per-cell sufficient statistics n_dk (R slot ndk) are updated +/// After infer_topics_cpp, compute cell-level theta on the R side (buildTheta). +// Per-cell token store for prediction (0-based cell/gene indices; training struct PredictCell { - arma::vec counts; - arma::uvec genes; - arma::mat gamma; + arma::vec counts; // nonzero counts per token for this cell + arma::uvec genes; // 0-based row indices into phi + arma::mat gamma; // tokens x K topic responsibilities per token PredictCell() {} PredictCell(const arma::vec& counts_in, const arma::uvec& genes_in, int K) @@ -322,6 +327,7 @@ struct PredictCell { } }; +// One PredictCell per counts column; initial n_dk row is total_counts / K. static void init_predict_cells(const arma::sp_mat& counts, const arma::vec& genes, int K, int D, @@ -346,6 +352,9 @@ static void init_predict_cells(const arma::sp_mat& counts, } } +// Parallel E-step/M-step per cell: gamma_k propto (alpha + n_dk - cur_counts) % phi_g, +// with leave-one-out cur_counts; then n_dk.row(i) = sum_t gamma_t * count_t. +// Inner loop runs up to burnin times until n_dk change < 0.01. static void predict_epoch(std::vector& cells,const double& alpha, int K,int D,arma::mat& n_dk, const arma::mat& phi, int num_threads = 1, @@ -392,6 +401,10 @@ static void predict_epoch(std::vector& cells,const double& alpha, } +// Infer cell-topic counts n_dk in place (ndk on the R object). Repeats +// predict_epoch maxiter times; no ELBO early stop. celltypes and M are unused +// (API symmetry with training). Per-token gamma stays in C++ only; use +// infer_gex_cpp to export gamma and cell_mtx. // [[Rcpp::export]] void infer_topics_cpp(const arma::sp_mat& counts, const arma::vec& celltypes,