Most scalable pseudotime ordering algorithm

Question

What algorithms for linear Pseudotime trajectory construction (diffusion-based) are the most scalable to large datasets?

I'm currently using Slingshot based on the recommendation in this manuscript: https://www.biorxiv.org/content/biorxiv/early/2018/03/05/276907.full.pdf

Slingshot is nice and easy to use, but it turns out it is terrible at scaling to big datasets.

For 2000 highest variable genes and a random sampling of a variable number of cells in my ~80% sparse expression matrix I get the following computation times in R (96 GB RAM):

#cells / Time (min)
100 / 0.7
1000 / 4.3
2000 / 12.3
5000 / 71.1

At this rate it would take nearly 1.5 days to process 200,000 cells and 4.5 months to process 2 million cells based on 2,000 HVG.

I feel like Slingshot is limited by a fundamentally unnecessary comparison of diffusion for all cells against all cells, rather than preclustering and then finer resolution of pseudotime across cluster edges.

What algorithms are better, but still implement a similar realization of pseudotime ordering? I just need a single linear trajectory.

Answer 1

As it turns out, pseudotime algorithms are exceedingly difficult to parallelize. Pseudotime does not scale linearly to a subset of data compared to the full dataset, and as such a course-grained approach to parallelizing pseudotime calculations on random bins of the data, followed by bin merging, does not work. These values will never converge on the actual full-dataset solution.

I put in a pretty good effort trying to parallelize slingshot, but with no success. I realize my code may be of limited utility without more explanation, but I'm not putting more time into this. FWIW:

This code:

1. Finds computationally optimal bin size for slingshot based on a number of genes and features
2. Randomly splits the SingleCellExperiment object into that number of bins and runs slingshot on each bin
3. Orients each bin so that a cell with a pseudotime of 1 in the first bin is similar to a pseudotime of 1 in every bin. Pseudotime can get flipped depending on the sampled subset.
4. Merges all bins into a single matrix so each cell has a assigned pseudotime based on randomly sampled bins.
5. Repeats steps #2 to #5 with a different random seed for binning every time.
6. Measures convergence towards actual solution at three pseudotime values, where convergence is defined as correlation of cell indices along the trajectory.

I abandoned the project after observing that more iterations to not change the convergence towards the actual solution. While pseudotime values close to 0 recapitulated the actual solution quite well, pseudotime values towards 100 failed to correlate at all with the actual solution. There was no convergence with more iterations.

Maybe this will be helpful to somebody, which is why I'm posting.

# Parallel Iterative Sampling with Slingshot
# pissshot is a course-grained estimator of the actual slingshot solution, and iteratively converges towards the actual solution
# pissshot randomly partitions cells in a large database into bins, bin size is determined by finding the computationally optimal number of cells
# pissshot is run independently on each bin, pseudotime vectors are aligned across all bins, and then all bins are merged
# pissshot repeats this process and the average of pseudotime assignments for each cell is taken at the end of each iteration. 
# The convergence of the blingshot model at the end of each iteration is calculated by running slingshot on a small sample of adjacent cells at a random pseudotime interval and measuring correlation with the predicted model
# When satisfactory convergence is reached, blingshot returns a singlecellobject, just as slingshot would

pissshot <- function(sce, num_iterations = 20){
    cat("Running blingshot on", dim(sce)[2],"cells and",dim(sce)[1],"features\n")
    cat("  Step 1/5: Finding computationally optimal increment sizes\n   ")
#    increment_size <- FindOptimalIncrementSize(sce)
    increment_size <- 688
    num_increments <- ceiling(dim(sce)[2]/increment_size)-1
    cat("\n  ...Increment size of",increment_size,"cells is computationally optimal\n")
    cat("  Step 2/5: Course-grained pseudotime assignment across",num_increments,"increments")
    ps <- list()
    convergenceArr <- c()
    for(iteration in seq(from=1,to=num_iterations,by=1)){
        cat("\n   ... iteration",iteration,"...\n")
        cat("         randomly assigning cells to",num_increments,"bins\n         0%.")
        bins <- PartitionBins(sce, increment_size = increment_size, seed = 7*iteration)
        cat("100%\n         running slingshot on each bin\n         0%")
        cell_pseudotimes <- c()
        for(b in seq(from=1,to=length(bins),by=1)){
            bin_with_pseudotime <- suppressMessages(slingshot(bins[[b]]))
            # Find right cell and left cell in pseudotime array (cell with highest/lowest pseudotime values)
            pcells <- as.data.frame(colData(bin_with_pseudotime)$slingPseudotime_1)
            rownames(pcells) <- colnames(bin_with_pseudotime)
            pcells <- data.frame(lapply(pcells, function(x) as.numeric(as.character(x))), check.names = F, row.names = rownames(pcells))
            pcells.sorted <- pcells[order(-pcells[,1]), , drop = FALSE]
            right_cell_ID <- rownames(pcells.sorted)[dim(pcells.sorted)[1]]
            left_cell_ID <- rownames(pcells.sorted)[1]
            right_cell_pos <- match(right_cell_ID,colnames(bins[[b]]))
            left_cell_pos <- match(left_cell_ID,colnames(bins[[b]]))
            right_cell_logcounts <- assays(bins[[b]])$logcounts[,right_cell_pos]
            left_cell_logcounts <- assays(bins[[b]])$logcounts[,left_cell_pos]
            if(b==1&&iteration==1){
                ref_right_cell_logcounts <- right_cell_logcounts
                ref_left_cell_logcounts <- left_cell_logcounts
            }
            # Figure out whether this pseudotime in this bin has flipped relative to reference (first bin), and if so, flip pseudotime values
            cor_opposite_ends <- mean(cor(right_cell_logcounts,ref_left_cell_logcounts),cor(left_cell_logcounts,ref_right_cell_logcounts))
            cor_same_ends <- mean(cor(right_cell_logcounts,ref_right_cell_logcounts),cor(left_cell_logcounts,ref_left_cell_logcounts))
            flipped <- FALSE
            if(cor_opposite_ends > cor_same_ends){
                flipped <- TRUE
            }
            if(flipped==TRUE){
                # find max pseudotime in colData(bin_with_pseudotime)$slingPseudotime
                # Recalculate pseudotime in bin_with_pseudotime as max-value
                pseudotimes <- colData(bin_with_pseudotime)$slingPseudotime_1
                colData(bin_with_pseudotime)$slingPseudotime_1 <- max(pseudotimes)-pseudotimes 
            }
            cells_with_pseudotime <- as.data.frame(colData(bin_with_pseudotime)$slingPseudotime_1,colnames(bin_with_pseudotime))
            cell_pseudotimes <- rbind(cell_pseudotimes, cells_with_pseudotime)
            cat(".")
        }
        cat("100%\n")

        ps[[iteration]] <- cell_pseudotimes[order(row.names(cell_pseudotimes)), , drop = FALSE]

        pss <- data.frame(ps)
        if(iteration>1){
            convergence10 <- measureConvergence(pss, sce = sce, ptime = 10, increment_size = increment_size)
            convergence50 <- measureConvergence(pss, sce = sce, ptime = 50, increment_size = increment_size)
            convergence99 <- measureConvergence(pss, sce = sce, ptime = 99, increment_size = increment_size)
            cat("   convergence at 10:",convergence10,"... at 50:",convergence50,"...at 100:",convergence99,"\n")
            convergenceArr <- rbind(convergenceArr,c(iteration,convergence10,convergence50,convergence99))
        }
        # Iterate until the model converges towards the actual slingshot solution, within an indicated fraction (i.e. convergence = 0.02)
        # The convergence measure compares how similar the model is to the actual slingshot solution, i.e. a convergence of 0.02 means that cell pseudotime values are within 2% of the actual solution
        # Assess this by selecting a pseudotime at random, analyzing the following increment_size cells, and comparing accuracy of pseudotime measurement across that trajectory with the composite of the course-grained analysis

        # need to study how convergence changes over pseudotime
    }

    # remove cells with really high standard deviations

    return(pss)
#    return(convergenceArr)
}


measureConvergence <- function(mat, sce = sce, ptime = 45, increment_size = 500){
    # this line is throwing an error
    mat[,"avg"] <- apply(mat[,1:dim(mat)[2]],1,mean)
    mat <- mat[order(mat[,"avg"]), , drop = FALSE]
    # select a random pseudotime increment
    mat <- mat[mat[, "avg"] > ptime,]
    mat <- mat[1:increment_size,]
    cell_IDs <- rownames(mat)
    # pull out this list of cell_IDs from the sce object and run slingshot
    sce2 <- suppressMessages(slingshot(sce[,cell_IDs]))
    cpcells <- as.data.frame(colData(sce2)$slingPseudotime_1)
    rownames(cpcells) <- colnames(sce2)
    cpcells <- data.frame(lapply(cpcells, function(x) as.numeric(as.character(x))), check.names = F, row.names = rownames(cpcells))
    # sort cpcells by increasing pseudotime
    cpcells <- cpcells[order(cpcells[,1]), , drop = FALSE]
    ordered_cells <- cbind(rownames(mat),rownames(cpcells))
    # get the index of each cell in the dataframe, run a correlation on how well the indices line up
    cell_indices <- c()
    for(cell in cell_IDs){
        new_indices <- c(which(ordered_cells[,1] == cell), which(ordered_cells[,2] == cell))
        cell_indices <- rbind(cell_indices,new_indices)
    }
    convergence_measure <- cor(cell_indices[,1],cell_indices[,2])
    plot(cell_indices[,1],cell_indices[,2])
    return(convergence_measure)
}

PartitionBins <- function(sce, increment_size = 500, seed = 123){
    bins <- c()
    while(dim(sce)[2]>=increment_size){
        if(dim(sce)[2]<increment_size*2){
            bins <- c(bins, sce)
            sce <- sce[,1]
        }
        else{
            set.seed(seed)
            train_ind <- sample(seq_len(ncol(sce)), size = increment_size)
            bins <- c(bins, sce[,train_ind])
            sce <- sce[,-train_ind]
            cat(".")
        }
    }
    return(bins)
}

FindOptimalIncrementSize <- function(sce){
    # Find the computationally optimal increment_size for grainy alignments using 5 increment points (100, 500, 1000, 2500)
    res <- c(100,250,500,750,1000,1500,2500,5000)
    runtimes <- c()
    for(i in res){
        runtimes <- c(runtimes, suppressMessages(system.time(slingshot(sce[,0:i])))["elapsed"])
    }
    mat <- as.data.frame(t(rbind(res,runtimes)))
    c <- summary(lm(mat$runtimes ~ poly(mat$res, 2, raw=TRUE)))$coefficients[,"Estimate"]
    predicted_runtimes <- c()
    for(i in seq(from=100, to=5000, by=1)){
        predicted_runtime <- (c[1] + c[2]*i + c[3]*i^2)*(dim(sce)[2]/i)
        predicted_runtimes <- rbind(predicted_runtimes,c(i,predicted_runtime))
        cat(".")
    }
    increment_size <- predicted_runtimes[which.min(predicted_runtimes[,2]),1]
    return(increment_size)
}

Answer 2

There is a recent publication that compares methods for trajectory inference. In Figure 2 in terms of scalability PAGA, PAGA Tree and Component 1 get the highest rating. The authors offer a shiny app to help select the best method for your dataset.