3.1 L_2 scaling to induce cosine distance

A key preprocessing step of Harmony clustering is L2 normalization. As shown in Haghverdi et al 2018, scaling each cell to have L2 norm equal to 1 induces a special property: Euclidean distance of the scaled cells is equivalent to cosine distance of the unscaled cells. Cosine distance is a considerably more robust measure of cell-to-cell similarity (CITE Martin and Vlad). Moreover, it has been used in clustering analysis of high dimensional text datasets (CITE NLP spherical kmeans).

\(L_2\) Normalization of cell \(i\):

TL;DR Harmony clustering uses cosine distance. By normalizing each cell to have unit length, we can directly visualize the cosine distances between cells (right). These relationships are not obvious in Euclidean space (left).

In the Harmony object, we now have 2 copies of the cell embeddings. The first, \(Z_{orig}\) is the original PCA matrix (PCs by cells). The second, \(Z_{cos}\) is the new \(L_2\) scaled matrix. Since this scaling projects cells into a unit hypersphere, cells appear pushed away from the origin (0,0).

do_scatter(t(harmonyObj$Z_orig), meta_data, 'dataset', no_guides = TRUE, do_labels = TRUE) + 
    labs(title = 'Z_orig', subtitle = 'Euclidean distance', x = 'PC1', y = 'PC2') +
do_scatter(t(harmonyObj$Z_cos), meta_data, 'dataset', no_guides = TRUE, do_labels = TRUE) + 
    labs(title = 'Z_cos', subtitle = 'Induced Cosine distance', x = 'PC1', y = 'PC2') +
NULL

In the \(Z_{cos}\) scatterplot (right), cells that are nearby have a high cosine similarity. Although it is not obvious in this example, cells closeby in Euclidean space do not always have a high cosine similarity!

Above, we only visualize the first two PCs. In this simple example with cell lines, this is sufficient to visualize most of the variation. Note, however, that all clustering and correction in Harmony uses all the PCs. For completeness, we can visualize the quantiles of PCA embeddings for all 20 PCs, colored by original dataset.

harmonyObj$Z_cos %>% t %>% data.frame() %>% 
    cbind(meta_data) %>% 
    tidyr::gather(key, val, X1:X20) %>% 
    ggplot(aes(reorder(gsub('X', 'PC', key), as.integer(gsub('X', '', key))), val)) + 
        geom_boxplot(aes(color = dataset)) + 
        scale_color_manual(values = colors_use) + 
        labs(x = 'PC number', y = 'PC embedding value', title = 'Z_cos (unit scaled PCA embeddings) for all 20 PCs') + 
        theme_tufte(base_size = 14) + geom_rangeframe() + 
        theme(axis.text.x = element_text(angle = 45, hjust = 1)) + 
        NULL

3.2 Initial clustering

Initializing the Harmony object also triggered initialization of all the clustering data structures. Harmony currently uses regular kmeans, with 10 random restarts, to find initial locations for the cluster centroids. Let’s visualize these centroids directly! We can do this by accessing the Y matrix in the Harmony object. This is a matrix with \(d=20\) rows and \(K=5\) columns, so each column represents one 20-dimensional centroid.

Remember that we set the number of clusters to 5 above, so there are now 5 clusters below.

cluster_centroids <- harmonyObj$Y

do_scatter(t(harmonyObj$Z_cos), meta_data, 'dataset', no_guides = FALSE, do_labels = FALSE) + 
    labs(title = 'Initial kmeans cluster centroids', subtitle = '', x = 'PC1', y = 'PC2') +
    geom_point(
        data = data.frame(t(cluster_centroids)), 
        color = 'black', fill = 'black', alpha = .8,
        shape = 21, size = 6
    ) +
NULL

Based on these cluster centroids, we also assigned probabilistic cluster memberships to each cell. In the algorithm, this is done using the formula below.

Above, \(R_{ki}\) is a value from \(0\) to \(1\) and denotes the probability that cell \(i\) is assigned to cluster \(k\). Accordingly, the squared distance \(||Z_i - Y_k||^2_2\) is the distance between cell \(i\) and the centroid of cluster \(k\). Because we’re using cosine distance (i.e. cells and centroids have unit length), we can simplify the distance computation:

Finally, the \(\propto\) symbol means that we will normalize R to form a proper probability distribution for each cell:

Let’s take a look at these initial cluster assignments. We can find these assignments in the \(K\) row by \(N\) column matrix \(R\).

cluster_assignment_matrix <- harmonyObj$R

The plots below color each cell by cluster membership, from 0 (grey) to 1 (blue). For clarity, each column is a different dataset. Each row is one of the 5 clusters.

t(harmonyObj$Z_cos) %>% data.frame() %>%
    cbind(meta_data) %>% 
    tibble::rowid_to_column('id') %>% 
    dplyr::inner_join(
        cluster_assignment_matrix %>% t() %>% data.table() %>% 
            tibble::rowid_to_column('id') %>%
            tidyr::gather(cluster, r, -id) %>% 
            dplyr::mutate(cluster = gsub('V', 'Cluster ', cluster)), 
        by = 'id'
    ) %>% 
    dplyr::sample_frac(1L) %>% 
    ggplot(aes(X1, X2, color = r)) + 
        geom_point(shape = '.') + 
        theme_tufte(base_size = 12) + theme(panel.background = element_rect()) + 
        facet_grid(cluster ~ dataset) + 
        scale_color_gradient(low = 'lightgrey', breaks = seq(0, 1, .1)) + 
        labs(x = 'Scaled PC1', y = 'Scaled PC2', title = 'Initial probabilistic cluster assignments')

3.3 Evaluating initial cluster diversity

A key part of clustering in Harmony is diversity. We can evaluate the initial diversity of clustering by aggregating the number of cells from each batch assigned to each cluster. For this, we need two data structures:

1. \(\phi\) (B rows, N columns): the one-hot encoded design matrix.

2. \(R\) (K rows, N columns): the cluster assignment matrix.

The cross product \(R\phi^T\) gives us a matrix of the number of cells from batch b (columns) that are in cluster k (rows). Note that since cluster assignment is probabilistic, the observed counts don’t have to be integer valued. For simplicity, we round the values to their closest integers.

observed_counts <- with(harmonyObj, R %*% t(Phi))
round(observed_counts)

##      [,1] [,2] [,3]
## [1,]  900    0 1574
## [2,]   14 1645    0
## [3,]  679    0 1285
## [4,] 1730   23    0
## [5,]   41 1587    0

In fact, this information is already stored in the Harmony model object! The observed cluster by batch counts are stored in the \(O\) matrix. The expected counts are in the \(E\) matrix. We can check that the observed counts matrix has exactly the same values we computed above.

## observed counts
round(harmonyObj$O)

##      [,1] [,2] [,3]
## [1,]  900    0 1574
## [2,]   14 1645    0
## [3,]  679    0 1285
## [4,] 1730   23    0
## [5,]   41 1587    0

## observed counts
round(harmonyObj$E)

##      [,1] [,2] [,3]
## [1,]  878  850  746
## [2,]  589  570  501
## [3,]  697  674  592
## [4,]  622  602  529
## [5,]  578  559  491

It looks like clusters 2, 4, and 5 are not very diverse, with most cells coming from a single dataset. However, clusters 1 and 3 look pretty well mixed already! Cluster 1 has 900 cells from batch \(b=1\) (half dataset) and 1574 cells from batch \(b=3\) (t293 dataset). As we move into the maximum diversity clustering, we should see the clusters getting more and more mixed!

In this benchmark, we also have some ground truth cell types. In the same way that we evaluated the cluster diversity, we can evaluate the cluster accuracy. Since we didn’t tell Harmony what the ground truth cell types are, we need to first construct a cell-type design matrix (shown below). We want these columns to be as mutually exclusive as possible. It looks like the initial clustering is fairly accurate. The only mistakes are the \(n=2\) jurkat cells clustered with the \(n=2472\) 293t cells in cluster \(k=1\) and \(n=12\) jurkat cells clustered with \(n=1952\) t293 cells in cluster \(k=3\).

phi_celltype <- harmony:::onehot(meta_data$cell_type) %>% t
observed_cell_counts <- harmonyObj$R %*% t(phi_celltype)
round(observed_cell_counts)

##      jurkat t293
## [1,]      2 2472
## [2,]   1659    0
## [3,]     12 1952
## [4,]   1753    0
## [5,]   1628    0

4.1 Diverse cluster assignment

Now let’s re-assign cells to cluster centroids. We did this above, when we assigned cells during the Harmony initialization step. The difference is that we want to assign cells to clusters that are both nearby and will increase diversity.

In the algorithm, this assignment is defined by

Let’s see what this looks like in code. Then we’ll break down the formula to see what it does.

with(harmonyObj, {
    distance_matrix <- 2 * (1 - t(Y) %*% Z_cos)
    distance_score <- exp(-distance_matrix / as.numeric(sigma))
    diversity_score <- sweep(E / O, 2, theta, '/') %*% Phi

    ## new assignments are based on distance and diversity
    R_new <- distance_score * diversity_score
    
    ## normalize R so each cell sums to 1
    R_new <- prop.table(R_new, 2)
    
})

So how does the formula we used above help to create more diverse cluster assignment?

The diversity penalty is encoded in the new term: \((\frac{E}{O})^\theta \phi\). This has some familiar data structures: \(O\) for observed counts, \(E\) for expected counts, and \(\phi\) for the design matrix. \(\theta\) is a new term. \(\theta\) decides how much weight to give diversity versus distance to cluster centroids.

With \(\theta=0\), there is no penalty and each cluster gets a score of 1.

## with theta = 0
with(harmonyObj, {
    (E / O) ^ 0
})

##      [,1] [,2] [,3]
## [1,]    1    1    1
## [2,]    1    1    1
## [3,]    1    1    1
## [4,]    1    1    1
## [5,]    1    1    1

As we increase \(\theta\), let’s see what happens (shown below). Recall that in cluster \(k=1\), batches 1 and 3 were well represented. Below, note that in that cluster (\(k=1\)), the penalties for batches 1 and 3 are relatively low (0.98 and 0.47). On the other hand, batch 2 gets a penalty score of 30914. This means that cells from batches 1 and 3 will be encouraged to move into cluster \(k=1\). On the other hand, cluster \(k=2\) is replete with batch 2. The penalty for batch 2 in cluster \(k=2\) is relatively low, and noticeably smaller than the penalty score for batch 2 in cluster \(k=1\). Thus, cells from batch 2 will be discouraged from moving into cluster \(k=1\), as this cluster has a higher penalty score for cells from batch 2 compared to other clusters (such as \(k=1\)).

## with theta = 1
with(harmonyObj, {
    round((E / O) ^ 1, 2)
})

##      [,1]     [,2]     [,3]
## [1,] 0.98 30641.11     0.47
## [2,] 7.03     0.36 65405.27
## [3,] 1.02  2342.28     0.46
## [4,] 0.39     4.26  2319.43
## [5,] 5.13     0.37 42089.15

We should always be wary of setting \(\theta\) too high, since the diversity scores can go to \(\infty\). Below, we set \(\theta\) to 1 million. We do not recommend setting \(\theta\) to 1 million!

## as theta approach infinity
with(harmonyObj, {
    round((E / O) ^ 1e6, 2)
})

##      [,1] [,2] [,3]
## [1,]    0  Inf    0
## [2,]  Inf    0  Inf
## [3,]  Inf  Inf    0
## [4,]    0  Inf  Inf
## [5,]  Inf    0  Inf

Finally, it is important to note that we cannot re-assign cells independently as we did above. Why not? As soon as we re-assign one cell, the diversity counts in the \(O\) and \(E\) matrices change. Thus, the assignment formula for all other cells is different! For this reason, we need to assign one cell at a time and update the \(O\) and \(E\) as we go. In practice, we can update some chunk of cells (e.g. 5%), update the \(O\) and \(E\) matrices, and update another chunk of cells.

4.2 Cluster centroid estimation

In the previous step, we re-assigned cells to maximize diversity within the clusters. With this new assignment, we need to update the cluster centroids. In this step, we’ll use the cell positions \(Z_{cos}\) and the cluster assignments (\(R\)) to re-position cluster centroids to be close to their assigned cells.

Y_unscaled <- with(harmonyObj, Z_cos %*% t(R))

We then scale Y to make each centroid unit length:

Y_new <- cosine_normalize(Y_unscaled, 2)

5.1 Mixture of Experts model

The theory behind this algorithm is based on the Mixture of Experts model. This is a natural extension of linear modeling, in which each cluster is deemed an expert and is assigned its own linear model.

We model each PC coordinate with a combination of linear factors.

In the model above, each cluster gets 4 \(\beta\) terms: \(\beta_{0,k}\) is the intercept term. This term is independent of which dataset a cell comes from. Therefore, it represents the contribution of cell type or cell state to the PC score. The other three \(\beta\) terms are accompanied by an indicator variable. This means that a cell from dataset half will have \(\mathbb{1}_{(dataset = half)}\) equal to 1 and the rest 0.

Following this cell from dataset half half, we can write rewrite the MoE equation above as

5.2 Estimate MoE model parameters

We estimate the matrix of linear regression terms using the formula described in the manuscript:

The matrix above contains linear regression terms for the the intercept \(W_k[0] = \beta_{0,k}\) and the batch terms:

\(W_k[1] = \beta_{1, k} \mbox{ (for dataset half)}\)

\(W_k[2] = \beta_{2, k} \mbox{ (for dataset jurkat)}\)

\(W_k[3] = \beta_{3, k} \mbox{ (for dataset 293T)}\)

W <- list()
with(harmonyObj, {
    for (k in 1:K) {
        W[[k]] <<- solve(Phi_moe %*% diag(R[k, ]) %*% t(Phi_moe) + lambda) %*% (Phi_moe %*% diag(R[k, ])) %*% t(Z_orig)
    }    
})

Let’s take a look at how these regression terms relate to the data. Recall that the mixture of experts model is trying to estimate the contribution of intercept and batch to cell’s positions in space. So first we’ll take a look at the positions of each batch and each cluster in the original PCA embeddings. The color below represents soft cluster membership learned using the maximum diversity clustering above.

cluster_assignment_matrix <- harmonyObj$R

t(harmonyObj$Z_orig) %>% data.frame() %>%
    cbind(meta_data) %>% 
    tibble::rowid_to_column('id') %>% 
    dplyr::inner_join(
        cluster_assignment_matrix %>% t() %>% data.table() %>% 
            tibble::rowid_to_column('id') %>%
            tidyr::gather(cluster, r, -id) %>% 
            dplyr::mutate(cluster = gsub('V', 'Cluster ', cluster)), 
        by = 'id'
    ) %>% 
    dplyr::sample_frac(1L) %>% 
    ggplot(aes(X1, X2, color = r)) + 
        geom_point(shape = '.') + 
        theme_tufte(base_size = 14) + theme(panel.background = element_rect()) + 
        facet_grid(cluster ~ dataset) + 
        scale_color_gradient(low = 'lightgrey', breaks = seq(0, 1, .2)) + 
        labs(x = 'PC1', y = 'PC2', title = 'Cluster assigned in original PCA space (Z_orig)')

Now let’s draw the \(\beta\) terms into this space. For each cluster, we expect the sum of the intercept plus the batch terms to land squarely in the center of each batch:cluster. The arrows below represent the intercept (in black) and batch (colored) offsets.

plt_list <- lapply(1:harmonyObj$K, function(k) {
    plt_df <- W[[k]] %>% data.frame() %>% 
        dplyr::select(X1, X2)

    plt_df <- plt_df %>% 
        cbind(
            data.frame(t(matrix(unlist(c(c(0, 0), rep(plt_df[1, ], 3))), nrow = 2))) %>% 
                dplyr::rename(x0 = X1, y0 = X2) 
        ) %>%
        cbind(type = c('intercept', unique(meta_data$dataset))) 

    plt <- plt_df %>% 
        ggplot() + 
            geom_point(aes(X1, X2),
                data = t(harmonyObj$Z_orig) %>% data.frame(),
                shape = '.', 
                color = 'lightgrey'
            ) + 
            geom_segment(
                aes(x = x0, y = y0, xend = X1 + x0, yend = X2 + y0, color = type)
                # ,
                # arrow = arrow(length = unit(0.3, "cm"), type = 'closed')
            ) + 
            scale_color_manual(values = c('intercept' = 'black', colors_use)) + 
            theme_tufte(base_size = 14) + theme(panel.background = element_rect()) + 
            labs(x = 'PC 1', y = 'PC 2', title = sprintf('Cluster %d', k)) + 
            NULL     
    plt <- plt + guides(color = guide_legend(override.aes = list(stroke = 1, alpha = 1, shape = 16, size = 4))) 
    
    # if (k == harmonyObj$K) {
    # } else {
    #     plt <- plt + guides(color = FALSE)
    # }
    plt
})

Reduce(`+`, plt_list) + 
  patchwork::plot_annotation(title = 'Mixture of experts beta terms before correction (Z_orig)') + 
  plot_layout(ncol = 1)

After correction, we remove the batch specific terms (colored arrows above). We can see the result in the corrected linear space (\(Z_{corr}\)). Notice that now, the cells are centered around the tips of the black arrows, which represent the intercept term. This is because we’ve removed the effect of the batch terms (colored arrows).

plt_list <- lapply(1:harmonyObj$K, function(k) {
    plt_df <- W[[k]] %>% data.frame() %>% 
        dplyr::select(X1, X2)

    plt_df <- plt_df %>% 
        cbind(
            data.frame(t(matrix(unlist(c(c(0, 0), rep(plt_df[1, ], 3))), nrow = 2))) %>% 
                dplyr::rename(x0 = X1, y0 = X2) 
        ) %>%
        cbind(type = c('intercept', unique(meta_data$dataset))) 

    plt <- plt_df %>% 
        ggplot() + 
            geom_point(aes(X1, X2),
                data = t(harmonyObj$Z_corr) %>% data.frame(),
                shape = '.', 
                color = 'lightgrey'
            ) + 
            geom_segment(
                aes(x = x0, y = y0, xend = X1 + x0, yend = X2 + y0, color = type, size = 2)
                # ,
                # arrow = arrow(length = unit(0.1, "npc"), type = 'closed')
            ) + 
            scale_color_manual(values = c('intercept' = 'black', colors_use)) + 
            theme_tufte(base_size = 14) + theme(panel.background = element_rect()) + 
            labs(x = 'PC 1', y = 'PC 2', title = sprintf('Cluster %d', k)) + 
            NULL     
    
    plt <- plt + guides(color = guide_legend(override.aes = list(stroke = 1, alpha = 1, shape = 16, size = 4)))
    # if (k == harmonyObj$K) {
    #     plt <- plt + guides(color = guide_legend(override.aes = list(stroke = 1, alpha = 1, shape = 16, size = 4))) 
    # } else {
    #     plt <- plt + guides(color = FALSE)
    # }
    plt
})

Reduce(`+`, plt_list) + 
  patchwork::plot_annotation(title = 'Mixture of experts beta terms after correction (Z_corr)') + 
  plot_layout(ncol = 1)

5.3 Cell specific corrections

How does one cell get its correction factor?

Recall from above that each cell \(i\) is now modeled with intercept and batch specific terms:

Z_i <- harmonyObj$Z_orig[, 5]
Z_i_pred <- Reduce(`+`, lapply(1:harmonyObj$K, function(k) {
    W[[k]] * harmonyObj$Phi_moe[, 5] * harmonyObj$R[k, 5]
})) %>% colSums

The plot below shows the observed and predicted values of all 20 PCs for cell 5.

data.table(obs = Z_i, pred = Z_i_pred) %>% 
    tibble::rowid_to_column('PC') %>% 
    ggplot(aes(obs, pred)) + 
        geom_point(shape = 21) + 
        geom_label_repel(aes(label = PC)) + 
        geom_abline(slope = 1, intercept = 0) + 
        theme_tufte() + geom_rangeframe() + 
        labs(x = 'Observed PC score', 'Predicted PC score', title = 'Observed and predicted values of PC scores\nfor cell 5') + 
        NULL

Now that we’ve modeled all these contributions to PCs, we can remove the batch specific terms from cell \(i\) to get its corrected position (\(\hat{Z}_{i}\)) in \(Z_{corr}\):

delta <- Reduce(`+`, lapply(1:harmonyObj$K, function(k) {
    W[[k]][2:4, ] * harmonyObj$Phi[, 5] * harmonyObj$R[k, 5]
})) %>% colSums

Z_corrected <- harmonyObj$Z_orig[, 5] - delta

Let’s see where this one cell moves in the original embeddings. Cell 5 is highlighted in red. It’s individual correction factor is shown with the red arrow.

harmonyObj$Z_orig %>% t %>% data.frame() %>% 
    ggplot(aes(X1, X2)) + 
        geom_point(shape = '.') + 
        geom_point(
            data = data.frame(t(harmonyObj$Z_orig[, 5, drop = FALSE])), 
            color = 'red'
        ) + 
        geom_segment(
            data = data.table(x0 = harmonyObj$Z_orig[1, 5], 
                              y0 = harmonyObj$Z_orig[2, 5], 
                              x1 = Z_corrected[1],
                              y1 = Z_corrected[2]), 
            aes(x = x0, y = y0, xend = x1, yend = y1), 
            color = 'red', 
            arrow = arrow(length = unit(0.05, "npc"), type = 'closed')            
        ) + 
        theme_tufte(base_size = 14) + geom_rangeframe() + 
        labs(x = 'PC1', y = 'PC2', title = 'Correction of cell #5')