Accelerate ExpandNeighbours() with (R)cpp code

Author: fspecque

NAMESPACE

@@ -111,6 +111,7 @@ importFrom(Matrix,t)
 importFrom(Rcpp,sourceCpp)
 importFrom(RhpcBLASctl,blas_get_num_procs)
 importFrom(RhpcBLASctl,blas_set_num_threads)
+importFrom(RhpcBLASctl,omp_get_max_threads)
 importFrom(RhpcBLASctl,omp_get_num_procs)
 importFrom(RhpcBLASctl,omp_set_num_threads)
 importFrom(Seurat,"DefaultAssay<-")

NEWS.md

@@ -1,5 +1,8 @@
 # SeuratIntegrate (development version)
 
+* Speed up Dijkstra's algorithm-like used in `ExpandNeighbours` with a new c++
+implementation
+
 * The most suited matrix format is automatically chosen for corrected counts
 output by integration methods (should be dense matrix most of the time)
 

R/RcppExports.R

@@ -1,6 +1,10 @@
 # Generated by using Rcpp::compileAttributes() -> do not edit by hand
 # Generator token: 10BE3573-1514-4C36-9D1C-5A225CD40393
 
+dijkstra_cpp <- function(m, k) {
+    .Call(`_SeuratIntegrate_dijkstra_cpp`, m, k)
+}
+
 n_zeros_dense_mat <- function(mat) {
     .Call(`_SeuratIntegrate_n_zeros_dense_mat`, mat)
 }

R/expand_knn_graph.R

@@ -20,13 +20,13 @@ NULL
 #' See \strong{Details} section for further explanations
 #' @param algo One of "dijkstra" or "diffusion". "diffusion" is suited for
 #' connectivity matrices only
-#' @param which.dijkstra one of "igraph", "fast" or "slow". "auto" (default)
-#' chooses for you. See \strong{Details} section
+#' @param which.dijkstra one of "cpp", "igraph", "fast" or "slow". "auto"
+#' (default) chooses for you. See \strong{Details} section
 #' @param dijkstra.ncores number of cores to use for Dijkstra's algorithm.
 #' Ignored when \code{which.dijkstra = "igraph"}
 #' @param dijkstra.tol number of sequential iterations with identical best
 #' neighbours found to consider that Dijkstra's algorithm should be stopped.
-#' Ignored when \code{which.dijkstra = "igraph"}
+#' Ignored when \code{which.dijkstra = "igraph" | "cpp"}
 #' @param diffusion.iter maximum number of iterations to reach \code{k.target}
 #' @param verbose whether to print progress messages
 #'
@@ -43,10 +43,29 @@ NULL
 #'
 #' One can choose to keep the graph as it is and consider it as a directed graph
 #' (\code{do.symmetrize = FALSE}).
+#'
 #' The alternative solution is to use all computed distances to extend the knn
 #' graph by making the matrix symmetric. Note that connectivity graphs are
 #' already symmetric, so the argument value should have no effect on the result.
 #'
+#' The choice of Dijkstra's algorithm implementation (\code{which.dijkstra}) is
+#' handled automatically by default. The igraph implementation is only preferred
+#' when the number of cells is low (1,000 at most). However, for a large
+#' \code{k.target} value (from approximately 1,000), igraph is faster on a single
+#' thread and thus can be imposed by setting \code{which.dijkstra = 'igraph'}.
+#'
+#' "fast" and "slow" are pure R implementations. "fast" is faster than "slow"
+#' but is not accurate on symmetric graph because it is suited for knn graphs
+#' with a constant k value. "slow" is slower but more accurate. No matter,
+#' \strong{"slow" and "fast" are both much slower than "cpp" and are deprecated}.
+#' They remain available in the even that some users have trouble running the
+#' "cpp" implementation.
+#'
+#' @note
+#' igraph implementation of Dijkstra's algorithm is single threaded only.
+#'
+#' "cpp" implementation is parallelized with OpenMP
+#'
 #' @importFrom SeuratObject DefaultAssay Cells as.Graph
 #'
 #' @export
@@ -56,7 +75,7 @@ setGeneric("ExpandNeighbours",
                     graph.type = c("distances", "connectivities"),
                     k.target = 90L, do.symmetrize = FALSE,
                     algo = c("dijkstra", "diffusion"),
-                    which.dijkstra = c("auto", "igraph", "fast", "slow"),
+                    which.dijkstra = c("auto", "cpp", "igraph", "fast", "slow"),
                     dijkstra.ncores = 1L, dijkstra.tol = 1L, diffusion.iter = 26L,
                     assay = NULL, verbose = TRUE)
              standardGeneric("ExpandNeighbours"))
@@ -68,7 +87,7 @@ setMethod("ExpandNeighbours", "Seurat",
                    graph.type = c("distances", "connectivities"),
                    k.target = 90L, do.symmetrize = FALSE,
                    algo = c("dijkstra", "diffusion"),
-                   which.dijkstra = c("auto", "igraph", "fast", "slow"),
+                   which.dijkstra = c("auto", "cpp", "igraph", "fast", "slow"),
                    dijkstra.ncores = 1L, dijkstra.tol = 1L, diffusion.iter = 26L,
                    assay = NULL, verbose = TRUE) {
             assay <- assay %||% DefaultAssay(object)
@@ -118,18 +137,19 @@ setMethod("ExpandNeighbours", "Seurat",
 setGeneric("expand_neighbours_dijkstra",
            function(object, graph.type = c("distances", "connectivities"),
                     k.target = 90L, do.symmetrize = FALSE,
-                    which.dijkstra = c("auto", "igraph", "fast", "slow"),
+                    which.dijkstra = c("auto", "cpp", "igraph", "fast", "slow"),
                     ncores = 1L, tol = 1L, verbose = TRUE)
              standardGeneric("expand_neighbours_dijkstra"))
 
 #' @importFrom SeuratObject as.Neighbor
-#' @importFrom Matrix sparseMatrix drop0
+#' @importFrom Matrix sparseMatrix t drop0
+#' @importFrom RhpcBLASctl omp_get_max_threads omp_set_num_threads
 #' @keywords internal
 #' @noRd
 setMethod("expand_neighbours_dijkstra", "Matrix",
           function(object, graph.type = c("distances", "connectivities"),
                    k.target = 90L, do.symmetrize = FALSE,
-                   which.dijkstra = c("auto", "igraph", "fast", "slow"),
+                   which.dijkstra = c("auto", "cpp", "igraph", "fast", "slow"),
                    ncores = 1L, tol = 1L, verbose = TRUE) {
             n <- ncol(object)
             const.k <- is.kconstant(object)
@@ -142,14 +162,17 @@ setMethod("expand_neighbours_dijkstra", "Matrix",
             }
 
             if (which.dijkstra == "auto") {
-              if (n <= 1e4) {
-                which.dijkstra <- "igraph"
-              } else if (const.k) {
-                which.dijkstra <- "fast"
+              which.dijkstra <- if (n <= 1e4) {
+                "igraph"
               } else {
-                which.dijkstra <- "slow"
+                "cpp"
               }
             }
+            if (which.dijkstra == "cpp") {
+              oomp <- omp_get_max_threads()
+              omp_set_num_threads(ncores)
+              on.exit(omp_set_num_threads(oomp))
+            }
             object.symmetry <- const.k.symmetry <- NULL
             igraph.mode <- "directed"
             if (do.symmetrize && !isSymmetric(object)) {
@@ -189,7 +212,7 @@ setMethod("expand_neighbours_dijkstra", "Matrix",
                 warning("When all cells have the same number k of nearest ",
                         "neighbors, ", sQuote("fast"), " implementation of ",
                         "Dijkstra's algorithm is recommended")
-              } else {
+              } else if (which.dijkstra == "fast"){
                 object <- as.Neighbor(x = object)
               }
             } else if (which.dijkstra == "fast") {
@@ -204,6 +227,7 @@ setMethod("expand_neighbours_dijkstra", "Matrix",
             message(msg[verbose], appendLF = F)
             beginning <- Sys.time()
             res <- switch (which.dijkstra,
+              cpp = dijkstra_cpp(m = t(drop0(object)), k = k.target),
               igraph = dijkstra.igraph(knnmat = object, k.target = k.target,
                                        mode = igraph.mode, weighted = T,
                                        diag = F),
@@ -225,17 +249,31 @@ setMethod("expand_neighbours_dijkstra", "Matrix",
             i <- rep(1:nrow(res$knn.idx), k.target)
             j <- as.vector(res$knn.idx)
             x <- as.vector(res$knn.dist)
-            # correct igraph output when not enough neighbours
-            infs <- which(is.infinite(x))
-            if (length(infs) > 0) {
-              x[is.infinite(x)] <- 0
-              warning('Dijkstra (igraph) : could not find enough neighbours',
-                      ' for ', length(infs), ' cell(s) ',
-                      paste0(infs, collapse = ', '),
-                      call. = F, immediate. = F)
+            if (which.dijkstra == "igraph") {
+              # correct igraph output when not enough neighbours
+              infs <- which(is.infinite(x))
+              if (length(infs) > 0) {
+                less_neighbours <- sort(unique(i[infs]))
+                i <- i[-infs]
+                j <- j[-infs]
+                x <- x[-infs]
+                warning('Dijkstra (igraph): could not find enough neighbours',
+                        ' for ', length(less_neighbours), ' cell(s) ',
+                        paste0(less_neighbours, collapse = ', '),
+                        call. = F, immediate. = F)
+              }
             }
+
             expanded.mat <- sparseMatrix(i = i, j = j, x = x, dims = rep(n, 2))
-            expanded.mat <- drop0(expanded.mat)
+            if (which.dijkstra == "cpp") {
+              less_neighbours <- which(get.k(expanded.mat, 'all') < k.target)
+              if (length(less_neighbours) > 0) {
+                warning('Dijkstra (cpp): could not find enough neighbours',
+                        ' for ', length(less_neighbours), ' cell(s) ',
+                        paste0(less_neighbours, collapse = ', '),
+                        call. = F, immediate. = F)
+              }
+            }
             if(do.symmetrize && which.dijkstra == "fast") {
               expanded.mat <- SymmetrizeKnn(expanded.mat, use.max = FALSE)
             }
@@ -249,7 +287,7 @@ setMethod("expand_neighbours_dijkstra", "Matrix",
 setMethod("expand_neighbours_dijkstra", "Neighbor",
           function(object, graph.type = c("distances", "connectivities"),
                    k.target = 90L, do.symmetrize = FALSE,
-                   which.dijkstra = c("auto", "igraph", "fast", "slow"),
+                   which.dijkstra = c("auto", "cpp", "igraph", "fast", "slow"),
                    ncores = 1L, tol = 1L, verbose = TRUE) {
 
             return(expand_neighbours_dijkstra(as.Graph(object),

man/ExpandNeighbours.Rd

@@ -11,6 +11,18 @@ Rcpp::Rostream<true>&  Rcpp::Rcout = Rcpp::Rcpp_cout_get();
 Rcpp::Rostream<false>& Rcpp::Rcerr = Rcpp::Rcpp_cerr_get();
 #endif
 
+// dijkstra_cpp
+List dijkstra_cpp(const arma::sp_mat& m, const int k);
+RcppExport SEXP _SeuratIntegrate_dijkstra_cpp(SEXP mSEXP, SEXP kSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const arma::sp_mat& >::type m(mSEXP);
+    Rcpp::traits::input_parameter< const int >::type k(kSEXP);
+    rcpp_result_gen = Rcpp::wrap(dijkstra_cpp(m, k));
+    return rcpp_result_gen;
+END_RCPP
+}
 // n_zeros_dense_mat
 int64_t n_zeros_dense_mat(const arma::mat& mat);
 RcppExport SEXP _SeuratIntegrate_n_zeros_dense_mat(SEXP matSEXP) {
@@ -35,6 +47,7 @@ END_RCPP
 }
 
 static const R_CallMethodDef CallEntries[] = {
+    {"_SeuratIntegrate_dijkstra_cpp", (DL_FUNC) &_SeuratIntegrate_dijkstra_cpp, 2},
     {"_SeuratIntegrate_n_zeros_dense_mat", (DL_FUNC) &_SeuratIntegrate_n_zeros_dense_mat, 1},
     {"_SeuratIntegrate_n_zeros_sparse_mat", (DL_FUNC) &_SeuratIntegrate_n_zeros_sparse_mat, 1},
     {NULL, NULL, 0}

src/RcppExports.cpp

@@ -0,0 +1,57 @@
+#include <RcppArmadillo.h>
+#include <queue>
+#include <unordered_set>
+
+using namespace Rcpp ;
+
+typedef std::pair<double, int> e;
+
+// [[Rcpp::plugins(openmp)]]
+
+// [[Rcpp::export]]
+List dijkstra_cpp(const arma::sp_mat& m, const int k) {
+
+  const int n_rows = m.n_rows;
+  arma::uvec seq_idx = arma::regspace<arma::uvec>(0, 1, n_rows - 1);
+  arma::umat knn_idx = arma::repmat(seq_idx, 1, k);
+  arma::dmat knn_dist(n_rows, k, arma::fill::zeros);
+
+  #ifdef _OPENMP
+  #pragma omp parallel num_threads(omp_get_max_threads())
+  #pragma omp for
+  #endif
+  for (int root = 0; root < n_rows; root++) {
+    arma::sp_mat::const_iterator it;
+    std::priority_queue<e, std::vector<e>, std::greater<e> > pq;
+    std::unordered_set<int> reached;
+    reached.insert(root);
+    arma::sp_mat col(m.col(root));
+    for (it = col.begin(); it != col.end(); ++it){
+      pq.push(std::make_pair((*it), it.row()));
+    }
+
+    int i = 1;
+    while(i < k) {
+      e current = pq.top();
+      int source = current.second;
+      pq.pop();
+      auto found = reached.find(source);
+      if (found == reached.end()) {
+        knn_idx(root,  i) = source;
+        knn_dist(root, i) = current.first;
+        i++;
+        reached.insert(source);
+
+        col = m.col(source);
+        for (it = col.begin(); it != col.end(); ++it){
+          pq.push(std::make_pair((*it) + current.first, it.row()));
+        }
+      } else if (pq.size() == 0) {
+        break;
+      }
+    }
+  }
+
+  return(List::create(Named("knn.idx") = knn_idx + 1,
+                      Named("knn.dist") = knn_dist));
+}