R/run_nmf.R

In zdebruine/scNMF: Fast Non-negative matrix factorization toolkit for single cell data

#' Run NMF on cluster centers calculated from a Seurat object ident class, and project cells onto NMF coordinates
#'
#' @param object Seurat object
#' @param k Decomposition rank, integer
#' @param rel.tol Stop criterion, defined as the relative tolerance between two successive iterations: |e2-e1|/avg(e1,e2). Default 1e-3.
#' @param verbose 0 = no tracking, 1 = progress bar relative to max.iter, 2 = iteration info. Default 2.
#' @param n.threads Number of threads/CPUs to use. Default to 0 (all cores).
#' @param ident Ident to average across (default is the active Idents()). Specify any valid discrete \code{object\@meta.data$ slot}.
#' @param assay Assay to average across (default is the active assay). Specify any valid assay name.
#' @param slot Slot to average, default is counts. Specify any valid counts slot within the assay.
#' @param project.block.size Number of cells to project in a single block during the projection phase (default 1000, should not need adjustment).
#' @param return Seurat object with a dimensional reduction NMF slot
run.nmf <- function(object, k = NULL, ident = "active.ident", assay = "default", slot = "counts", verbose = 1, rel.tol = 1e-4, n.threads = 0, project.block.size) {

  A <- average.expression(object, ident = ident, assay = assay, slot = slot, verbose = verbose)
  if (assay == "default") assay <- DefaultAssay(object)

  if (verbose > 0) message("\n\nRunning NMF on cluster centers:")
  nmf.mod <- nnmf(A, rel.tol = rel.tol, n.threads = n.threads, verbose = verbose, k = k)

  # now project cells 1000 at a time
  data <- GetAssayData(object, assay = assay)
  message("\n\nProjecting cells onto NMF factor coordinates")
  H <- nnls.project(data, W.model = nmf.mod$W, n.threads = n.threads, block.size = project.block.size)

  # create a new dimensional reduction object
  nmf.dr <- CreateDimReducObject(embeddings = t(H), loadings = nmf.mod$W, key = "NMF_", assay = assay)

  # add the dimensional reduction object to Seurat
  object[["nmf"]] <- nmf.dr
  return(object)
}