--- title: "An introduction to mbkmeans" author: "Yuwei Ni, Davide Risso, Stephanie Hicks, and Elizabeth Purdom" date: "Last modified: November 7, 2020; Compiled: `r format(Sys.time(), '%B %d, %Y')`" bibliography: biblio.bib output: BiocStyle::html_document: toc: true vignette: > %\VignetteEncoding{UTF-8} --- # Installation To install the package, please use the following. ```{r, eval=FALSE} if (!requireNamespace("BiocManager")) install.packages("BiocManager") BiocManager::install("mbkmeans") ``` # Introduction This vignette provides an introductory example on how to work with the `mbkmeans` package, which contains an implementation of the mini-batch k-means algorithm proposed in [@sculley2010web] for large single-cell sequencing data. This algorithm runs the _k_means iterations on small, random subsamples of data (``mini batches'') rather than the entire dataset. This both speeds up computation and reduces the amount of data that must be contained in memory. The main function to be used by the users is `mbkmeans`. This is implemented as an S4 generic and methods are implemented for `matrix`, `Matrix`, `HDF5Matrix`, `DelayedMatrix`, `SummarizedExperiment`, and `SingleCellExperiment`. Most of this work was inspired by the `MiniBatchKmeans()` function implemented in the `ClusterR` R package and we re-use many of the C++ functions implemented there. Our main contribution here is to provide an interface to the `DelayedArray` and `HDF5Array` packages so that `mbkmeans` algorithm only brings into memory at any time the "mini-batches" needed for the current iteration, and never requires loading the entire dataset into memory. This allows the user to run the mini-batch _k_-means algorithm on very large datasets that do not fit entirely in memory. `mbkmeans` also runs seamlessly on in-memory matrices provided by the user, and doing so significantly increases the speed of the algorithm compared to the standard _k_-means algorithm ([@Hicks2020]). A complete comparison of `mbkmeans` with other implementations can be found in [@Hicks2020]. The motivation for this work is the clustering of large single-cell RNA-sequencing (scRNA-seq) datasets, and hence the main focus is on Bioconductor's `SingleCellExperiment`and `SummarizedExperiment` data container. For this reason, `mbkmeans` assumes a data representation typical of genomic data, in which genes (variables) are in the rows and cells (observations) are in the column. This is contrary to most other statistical applications, and notably to the `stats::kmeans()` and `ClusterR::MiniBatchKmeans()` functions that assume observations in rows. We provide a lower-level `mini_batch()` function that expects observations in rows and is expected to be a direct replacement of `ClusterR::MiniBatchKmeans()` for on-disk data representations such as `HDF5` files. The rest of this document shows the typical use case through the `mbkmeans()` interface; users interested in the `mini_batch()` function should refer to its manual page. ## Example dataset ```{r options, include=FALSE, message=FALSE, warning=FALSE} knitr::opts_chunk$set(cache=FALSE, error=FALSE, message=FALSE, warning=FALSE) ``` To illustrate a typical use case, we use the `pbmc4k` dataset of the [`TENxPBMCData` package](https://bioconductor.org/packages/release/data/experiment/html/TENxPBMCData.html). This dataset contains a set of about 4,000 cells from peripheral blood from a healthy donor and is expected to contain many types or clusters of cell. We would note that `mbkmeans` is designed for very large datasets, in particular for datasets large enough that either the full data matrix can't be held in memory, or the computations necessary cannot be computed. This vignette works with a small dataset to allow the user to quickly follow along and understand the commands. [@Hicks2020] shows the example of using `mbkmeans` with a dataset with 1.3 million cells, which is a more appropriate size for observing improved memory usage and improved speed. First, we load the needed packages. ```{r packages, message=FALSE, warning=FALSE} library(TENxPBMCData) library(scater) library(SingleCellExperiment) library(mbkmeans) library(DelayedMatrixStats) ``` Note that in this vignette, we do not aim at identifying biologically meaningful clusters (that would entail a more sophisticated normalization of data and dimensionality reduction), but instead we only aim to show how to run mini-batch _k_-means on a HDF5-backed matrix. We normalize the data simply by scaling for the total number of counts using `scater` and select the 1,000 most variable genes and a random set of 100 cells to speed-up computations. ```{r} tenx_pbmc4k <- TENxPBMCData(dataset = "pbmc4k") set.seed(1034) idx <- sample(seq_len(NCOL(tenx_pbmc4k)), 100) sce <- tenx_pbmc4k[, idx] #normalization sce <- logNormCounts(sce) vars <- rowVars(logcounts(sce)) names(vars) <- rownames(sce) vars <- sort(vars, decreasing = TRUE) sce1000 <- sce[names(vars)[1:1000],] sce1000 ``` # `mbkmeans` The main function, `mbkmeans()`, returns a list object including `centroids`, `WCSS_per_cluster` (where WCSS stands for within-cluster-sum-of-squares), `best_initialization`, `iters_per_initiazation` and `Clusters`. It takes any matrix-like object as input, such as `SummarizedExperiment`, `SingleCellExperiment`, `matrix`, `DelayedMatrix` and `HDF5Matrix`. In this example, the input is a `SingleCellExperiment` object. ```{r} res <- mbkmeans(sce1000, clusters = 5, reduceMethod = NA, whichAssay = "logcounts") ``` The number of clusters (such as _k_ in the _k_-means algorithm) is set through the `clusters` argument. In this case, we set `clusters = 5` for no particular reason. For `SingleCellExperiment` objects, the function provides the `reduceMethod` and `whichAssay` arguments. The `reduceMethod` argument should specify the dimensionality reduction slot to use for the clustering, and the default is "PCA". Note that this *does not perform* PCA but only looks at a slot called "PCA" already stored in the object. Alternatively, one can specify `whichAssay` as the assay to use as input to mini-batch _k_-means. This is used only when `reduceMethod` option is `NA`. See `?mbkmeans` for more details. ## Choice of arguments There are additional arguements in `mbkmeans()` that make the function more flexible and suitable for more situations. ### Batch size The size of the mini batches is set through the `batch_size` argument. This argument gives the size of the samples of data brought into memory at any one time, and thus constrains the memory usage of the alogrithm. The `blocksize()` function can be used to set the batch size to the maximum allowed by the available memory. It considers both the number of columns in the dataset and the amount of RAM on the current matchine to calculate as large of a batch size as reasonable for the RAM available to the session. The calculation uses `get_ram()` function in `benchmarkme` package. See the `benchmarkme` vignette for more details. ```{r} batchsize <- blocksize(sce1000) batchsize ``` In this simple example, because the whole data fits in memory, the default batch size would be a single batch of size `r ncol(sce1000)`. [@Hicks2020] provides a comparison of the effect of batch size for large datasets (i.e. where the full data cannot be stored in memory). Their results show the algorithm to be robust to a large range of batch sizes (1,000-10,000 cells): accuracy of the results were equivalent if the batch size was above 500-1000 cells, with no noticable advantage in memory usage or speed as long as the batches were smaller than roughly 10,000 cells (corresponding to a subsample of the full data matrix whose memory foot print is not noticeable on most modern computers). For this reason, we would err on the larger range of batch size suggested in [@Hicks2020] -- 10,000 cells. Furthermore, since the initialization uses by default the same number of cells as batch size, this gives a robust sample for determining the initial start point. We would note that it is better to make choices on batch size based on choosing the *absolute* number of cells in a batch rather than choosing batch size based on a percentage of the overall sample size. If batch size is chosen based on the percentage of cells, very large datasets will use very large batches of cells, negating the memory and speed gains. Moreover such large batches are not needed for good performance of the algorithm. The continual resampling of the batches works like stochastic gradient descent and allows for (local) minimization of the objective function [@sculley2010web] with relatively small batch sizes, as demonstrated on single-cell data in the work of [@Hicks2020]. If anything, the size of the batch should be governed more by the complexity of the problem (e.g. the number of underlying true clusters or subtypes), rather than the number of overall cells in the dataset. ### Initialization The performance of mini-batch _k_-means greatly depends on the process of initialization. We implemented two different initialization methods: 1. Random initialization, as in regular _k_-means; 2. `kmeans++`, as proposed in [@arthur2007k]. The default is "kmeans++". The percentage of data to use for the initialization centroids is set through the `init_fraction` argument, which should be larger than 0 and less than 1. The default value is given so that the proportion matches the value of `batch_size`, converted to a proportion of the data. ```{r} res_random <- mbkmeans(sce1000, clusters = 5, reduceMethod = NA, whichAssay = "logcounts", initializer = "random") table(res$Clusters, res_random$Clusters) ``` Note that large values of `init_fraction` will result in that proportion of the data being held in memory and used in the initialization of the `kmeans++` algorithm. Therefore, large values of `init_fraction` will require a great deal of memory -- potentially more than that used by the actual clustering algorithm. Indeed keeping the fraction used constant for larger and larger datasets will result in commiserate numbers of cells being used in the initialization and increasing the memory usage (similar to the issues discussed above for batch size). Therefore, we recommend that users keep the default and make changes to the `batch_size` parameter as desired, ensuring consistent memory usage across both the initialization stage and the actual clustering stage of the algorithm. As mentioned in the discussion above on `batch_size`, for this reason we would recommend `batch_size` on the larger range suggested as reasonable in [@Hicks2020] so as to improve the initialization. ## Running `mbkmeans` with multiple values of $k$ The main parameter to set in $k$-means and its variants is the number of clusters $k$. `mbkmeans` is quick enough to rerun the clustering algorithm with different number of clusters even in very large datasets. Here, we apply `mbkmeans` with $k$ from 5 to 15 and select the number of clusters using the elbow method (i.e., the value that corresponds to the point of inflection on the curve). We note that this is just a rule of thumb for selecting the number of clusters, and many different methods exist to decide the appropriate value of $k$. To speed up the computations, we will cluster on the top 20 PCs. ```{r} sce1000 <- runPCA(sce1000, ncomponents=20) ks <- seq(5, 15) res <- lapply(ks, function(k) { mbkmeans(sce1000, clusters = k, reduceMethod = "PCA", calc_wcss = TRUE, num_init=10) }) wcss <- sapply(res, function(x) sum(x$WCSS_per_cluster)) plot(ks, wcss, type = "b") ``` From the plot, it seems that `k = 12` could be a reasonable value. # Comparison with _k_-means Note that if we set `init_fraction = 1`, `initializer = "random"`, and `batch_size = ncol(x)`, we recover the classic _k_-means algorithm. ```{r} res_full <- mbkmeans(sce1000, clusters = 5, reduceMethod = NA, whichAssay = "logcounts", initializer = "random", batch_size = ncol(sce1000)) res_classic <- kmeans(t(logcounts(sce1000)), centers = 5) table(res_full$Clusters, res_classic$cluster) ``` Note however, that since the two algorithms start from different random initializations, they will not always converge to the same solution. Furthremore, if the algorithm uses batches smaller than the full dataset (`batch_size