Apply varimax rotation to PCA dimensionality reduction in a Seurat object to aid interpretability.

RunVarimax.R

## Typical usage
#
# library(Seurat)
# obj <- CreateSeuratObject(mat)
# obj <- NormalizeData(obj)
# obj <- FindVariableFeatures(obj)
# obj <- ScaleData(obj)
# obj <- RunPCA(obj)
#
# source("RunVarimax.R")
# obj <- RunVarimax(obj, dims=1:10)
#
# print(obj$vm)
#
# library(langevitour)
# langevitour(obj$vm@cell.embdeddings)
#


RunVarimax  <- function(obj, dims, reduction="pca") {
    red <- obj@reductions[[reduction]]
    loadings <- red@feature.loadings[,dims,drop=FALSE]
    scores <- red@cell.embeddings[,dims,drop=FALSE]

    # Find varimax rotation
    rotation <- varimax(loadings, normalize=FALSE)$rotmat

    # Optional step to make it a little nicer: 
    # Flip components so loadings have positive skew
    flips <- ifelse(colSums((loadings %*% rotation) ** 3) < 0, -1, 1)
    rotation <- sweep(rotation, 2, flips, '*')

    # Optional step to make it a little nicer: 
    # Order by mean absolute scores
    reorder <- (scores %*% rotation) |> abs() |> colMeans() |> order() |> rev()
    rotation <- rotation[, reorder, drop=FALSE]
    
    colnames(rotation) <- paste0("VM_",seq_len(ncol(rotation)))
    
    # Apply rotations to get new loadings and scores
    obj$vm <- CreateDimReducObject(
        embeddings=scores %*% rotation, 
        loadings=loadings %*% rotation,
        assay=red@assay.used, 
        key="VM_")
    
    obj
}