diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..aba3460 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +*.o +*.so +data/ diff --git a/DESCRIPTION b/DESCRIPTION index 4154c37..2de8363 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -7,19 +7,21 @@ Author: Adrien Osakwe Maintainer: Adrien Osakwe 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 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..dc6eb1d 100644 --- a/R/GTM.R +++ b/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) } diff --git a/R/RcppExports.R b/R/RcppExports.R index c20d8e9..59ed9aa 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, 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/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..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,61 @@ 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; +} + +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; + 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){ + 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; 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..db5cd17 100644 --- a/src/GTM.cpp +++ b/src/GTM.cpp @@ -1,9 +1,9 @@ #include "CellMap.h" #include "GTM.h" #include +#include #include -#include // [[Rcpp::depends(RcppArmadillo)]] using namespace Rcpp; using namespace std; @@ -48,6 +48,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 +196,7 @@ void run_epoch(std::unordered_map& CellMap,const arma::mat& alpha, } + /* Full model */ @@ -304,87 +306,125 @@ 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, +/// +/// 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; // 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) + : counts(counts_in), genes(genes_in), gamma(counts_in.n_elem, K) { + if (counts_in.n_elem > 0){ + gamma.fill(1.0 / K); + } + } +}; + +// 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, + 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); + } + } +} + +// 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, 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 - - } } +// 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(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; @@ -406,31 +446,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; @@ -456,15 +493,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 208e650..ec25c57 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,17 +72,14 @@ 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, +// Prediction +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 f1ffac5..aee8599 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); @@ -336,6 +336,36 @@ 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, 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); + 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< 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, tau0, kappa, shuffle); + return R_NilValue; +END_RCPP +} static const R_CallMethodDef CallEntries[] = { {"_SpaTM_test_cell_map", (DL_FUNC) &_SpaTM_test_cell_map, 10}, @@ -356,6 +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, 21}, {NULL, NULL, 0} }; diff --git a/src/sGTM.cpp b/src/sGTM.cpp new file mode 100644 index 0000000..cd3bff1 --- /dev/null +++ b/src/sGTM.cpp @@ -0,0 +1,207 @@ +#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); + + 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 + 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 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) / + 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; + } + + 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, + const double tau0 = 1.0, + const double kappa = 0.7, + 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; + + double step = 0.0; + + 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); + 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){ + 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))); + //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"; + } + } + + double diff = arma::max(arma::max(arma::abs(n_wk - old_nwk))); + if (verbal){ + prog = (static_cast(iter + 1)) / static_cast(maxiter); + } + + if (diff < thresh){ + if (verbal){ + Rcout.flush(); + Rcout << "["; + for(int p = 0; p < prog_width; p++){ + Rcout << "=" ; + } + Rcout << "] " << "100% || Iter: " << iter + << " || Batches: " << num_batches << "/" << num_batches + << " || 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 + << " || 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; + +} 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.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 259abd5..66489d7 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 @@ -10,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)) ) @@ -42,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) @@ -66,3 +87,5 @@ test_that("inferTopics should only update ndk and theta matrices with correct su }) + + 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..1adb40c 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 @@ -31,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) @@ -84,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-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 new file mode 100644 index 0000000..b857b26 --- /dev/null +++ b/tests/testthat/test-sGTM.R @@ -0,0 +1,69 @@ +suppressPackageStartupMessages({ + 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) +} + +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( + 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") +}) + +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) + ) +})