---
title: "Maximum Diversity Clustering"

format:
  html:
    toc: true
    number-sections: true
    code-fold: false

execute:
  echo: true

vignette: >
  %\VignetteIndexEntry{Maximum Diversity Clustering}
  %\VignetteEngine{quarto::html}
  %\VignetteEncoding{UTF-8}
---

```{r setup}
#| echo: false
knitr::opts_chunk$set(
  fig.path = "figure/mdclust-"
)
```


# Introduction

The **mdclust** package provides an independent implementation of the *maximum diversity clustering* algorithm introduced by [Korsunsky *et al.* (2019)](https://www.nature.com/articles/s41592-019-0619-0) as the clustering component of the [Harmony](https://github.com/immunogenomics/harmony) framework for single-cell data integration. Harmony popularized maximum diversity clustering as an effective way of constructing clusters that are robust to batch effects. However, the clustering algorithm is tightly integrated into Harmony’s embedding correction workflow and is not readily reusable in other contexts. `mdclust` provides an independent implementation of the clustering algorithm, exposing all optimization parameters and intermediate quantities for standalone use and integration into custom workflows. `mdclust` supports user-defined centroid initialization, making it suitable for iterative optimization procedures requiring repeated diversity clustering.

This vignette introduces the basic workflow for using mdclust as both a standalone clustering method and as a reusable component in custom analysis pipelines.

# Background

The objective optimized by **mdclust** consists of three components:

- a spherical clustering objective that encourages observations to be close to their assigned centroids,
- an entropy term that allows observations to belong probabilistically to multiple clusters, and
- a diversity penalty that promotes balanced batch composition within each cluster.

Setting the diversity penalty (`theta = 0`) recovers ordinary soft spherical *k*-means, while increasing `theta` progressively favours clusters with improved batch mixing.

Throughout this vignette we use a small synthetic data set to illustrate the behaviour of the algorithm. The same workflow can be applied directly to embeddings obtained from common single-cell analysis pipelines, including principal component analysis, latent semantic indexing, variational autoencoders, or other dimensionality reduction methods.

# A complete example
We simulate two signal groups `A` and `B` separated but overlapping Gaussians in 2D, and randomly assign them to two batches `Batch 1` and `Batch 2`.
A strong batch effect is added by offsetting one of the batches along the y-axis.
```{r simulate-data}
# --- Simulate signal // groups A & B ---
n <- 600

signal <- rep(c("A", "B"), each = n / 2)
Z <- matrix(0, nrow = n, ncol = 2)
Z[signal == "A", ] <- cbind(rnorm(n / 2, -0.8, 0.8), rnorm(n / 2, 0.0, 0.5))
Z[signal == "B", ] <- cbind(rnorm(n / 2, 0.8, 0.8), rnorm(n / 2, 0.0, 0.5))

# --- Assign batches independently ---
batch <- sample(c("Batch 1", "Batch 2"), n, replace = TRUE)

# --- Add strong batch effect ---
shift <- c(0., 2.)
Z[batch == "Batch 2", ] <- sweep(Z[batch == "Batch 2", , drop = FALSE], 2, shift, "+")
```



We can plot the simulated data to confirm that, indeed, we have a strong batch separation of the biological signal. Ordinary clustering (e.g. k-means) tends to struggle with such data and overwhelmingly aligns with undesired batch variation.
```{r plot-data}
#| fig-width: 12
#| fig-height: 6
#| code-fold: true
library(ggplot2)
library(patchwork)

df <- data.frame(
  Dim1 = Z[, 1],
  Dim2 = Z[, 2],
  signal = signal,
  batch = batch
)

p_truth <- ggplot(df, aes(Dim1, Dim2)) +
  geom_point(
    aes(colour = signal),
    size = 2,
    alpha = 0.8
  ) +
  coord_equal() +
  labs(
    title = "Ground truth",
    colour = "Signal",
    x = "Dimension 1",
    y = "Dimension 2"
  ) +
  scale_colour_brewer(palette = "Dark2") +
  theme_classic(base_size = 12)

p_batch <- ggplot(df, aes(Dim1, Dim2)) +
  geom_point(
    aes(colour = batch),
    size = 2,
    alpha = 0.8
  ) +
  coord_equal() +
  labs(
    title = "Batch effect",
    colour = "Batch",
    x = "Dimension 1",
    y = "Dimension 2"
  ) +
  scale_colour_brewer(palette = "Set2") +
  theme_classic(base_size = 12)

p_truth + p_batch +
  plot_layout(widths = c(1, 1))


```


A minimal call to `mdclust` on this data looks as follows:

```{r minimal-mdclust}
library(mdclust)

set.seed(111)

mdc <- mdclust(t(Z), batch, K = 2)
mdc
```

Printing the fitted object confirms that the optimization converged. The resulting clustering is shown below.

```{r mdclust-result}
#| fig-width: 12
#| fig-height: 6
df$cluster <- factor(mdc$cluster)

p_mdclust <- ggplot(df, aes(Dim1, Dim2)) +
  geom_point(
    aes(colour = cluster),
    size = 2,
    alpha = 0.8
  ) +
  coord_equal() +
  labs(
    title = "Maximum diversity clustering",
    colour = "Cluster",
    x = "Dimension 1",
    y = "Dimension 2"
  ) +
  scale_colour_brewer(palette = "Set1") +
  theme_classic(base_size = 12)

p_truth + p_mdclust +
  plot_layout(widths = c(1, 1))
```

The recovered clusters closely follow the underlying biological groups despite the pronounced batch effect, illustrating the purpose of the diversity penalty.

Any run of `mdclust` requires at least these three arguments. Here, `t(Z)` is the low-dimensional embedding used for clustering.  Observations are in columns, so be mindful of orientation of your embedding. `t(Z)` should be of shape $d \times N$, and an error will be raised if there is a mismatch. Columns are internally L2-normalized per column for spherical k-means. `batch` encodes the label that is penalized during the optimization. It can be either factor, character or integer and only one label per observation is permitted. `K` is simply the number of clusters and completely analogous to ordinary k-means.

However, `mdclust` exposes most internals to the user and the algorithm is fully configurable. The following sections describe the available initialization methods and optimization parameters.

## Initialization

Maximum diversity clustering requires an initial matrix of centroid positions `Y_init`, and allows for three different initializations.
```r
mdclust(..., Y_init = "kmeans++")

mdclust(..., Y_init = "kmeans")

mdclust(..., Y_init = Y)
```
The default is `"kmeans++"`, which is an ordinary k-means optimization with an enhanced starting point selection. `"kmeans"` performs a _single_ k-means optimisation from a random starting point. In addition, `Y_init` can be assigned any valid centroid matrix `Y` of shape $d\times K$. This allows for custom initialisations, and especially warm starts or use of `mdclust` as a diversity clustering engine in iterative procedures.

## Clustering parameters

The behaviour of maximum diversity clustering is primarily controlled by four parameters. These are in principle analogous to their Harmony counterparts, hence more information on them is available in the Harmony documentation and the respective publication.


- `theta` controls the strength of the diversity penalty. Setting `theta = 0` disables the diversity objective entirely, reducing the algorithm to ordinary soft spherical *k*-means. Larger values increasingly favour clusters with balanced batch composition. The default (`theta = 2`) matches the Harmony implementation.

- `sigma` controls the softness of cluster assignments. Smaller values produce assignments closer to hard clustering, whereas larger values distribute observations across multiple clusters. The default is `sigma = 0.1`.

- `tau` protects very small batches from being over-penalized by reducing the effective diversity penalty for batches containing fewer than approximately `K × tau` observations. The default (`tau = 0`) disables this behaviour.

- `alpha` is a pseudocount used when computing observed-to-expected enrichment ratios for the diversity penalty. It improves numerical stability and rarely needs adjustment.

## Optimization parameters

The remaining parameters control the optimization procedure.

- `block_size` specifies the fraction of observations updated simultaneously. Smaller values provide more accurate online updates at the expense of runtime.
- `max_iter` sets the maximum number of optimization iterations.
- `epsilon` specifies the convergence threshold on the objective.

The k-means implementations have their own convergence controls. Convergence is declared once subsequent runs fall under the threshold `abs(new - old)/abs(old) < epsilon_kmeans`. If this condition is not met, k-means terminates after `max_iter_kmeans`. Both parameters are accessible with defaults `mdclust(..., epsilon_kmeans = 1e-3, max_iter_kmeans = 10L)`.

## Output 

An additional concise overview to printing of the finished clustering can be accessed via
```{r mdclust-summary}
summary(mdc)
```


The complete object contains almost all internal parameters of interest.
```{r mdclust-names}
str(mdc, max.level = 1)
```

Here

- `R`: Soft-assignment matrix of shape $K\times N$. Columns sum to one.
- `Y`: Final centroids.
- `cluster`: Hard-assignment via `max.col(t(mdc$R))`.
- `objective`: List containing the partial and full objective traces of the run.
- `O,E,logOE`: Expected and observed batch compositions. A useful diagnostic of the batch diversity objective.
- `converged`: Convergence flag.
- `batch_levels`: Batch levels used internally during optimization.



`mdclust` also supports rudimentary plots of the objective trace for quick visual diagnostics.
```{r mdclust-plot}
plot(mdc)
```


# Relationship to Harmony

This package is an independently developed, standalone implementation of Algorithm 2 from [Korsunsky *et al.* (2019)](https://www.nature.com/articles/s41592-019-0619-0). Compared with the current Harmony implementation (2.0.x), `mdclust` differs in a few key aspects.

Most notably, `mdclust` exposes all parameters of the underlying algorithm, making maximum diversity clustering available as a standalone method and facilitating its use in custom workflows.

The initialization procedure has also been extended. Whereas Harmony initializes diversity clustering using a k-means solution, `mdclust` additionally supports random initialization, k-means++, and user-supplied centroid matrices. Custom centroid initialization enables warm starts and simplifies the use of `mdclust` as a diversity clustering engine within iterative algorithms.

Finally, `mdclust` provides independent convergence criteria and optimization controls for both the initialization procedure and the diversity clustering optimization.


# Session information

```{r session-info}
sessionInfo()
```