Pearson residuals function in R, from Lause, Berens, and Kobak

pearson_residuals.R

pearson_residuals <- function(counts, theta = 100) {
  counts_sum0 <- as(matrix(Matrix::colSums(counts),nrow=1),"dgCMatrix")
  counts_sum1 <- as(matrix(Matrix::rowSums(counts),ncol=1),"dgCMatrix")
  counts_sum <- sum(counts@x)
  
  #get residuals
  mu <- (counts_sum1 %*% counts_sum0) / counts_sum
  z <- (counts - mu) / sqrt(mu + mu^2/theta)
  
  #clip to sqrt(n)
  n <- ncol(counts)
  z[z > sqrt(n)] <- sqrt(n)
  z[z < -sqrt(n)] <- -sqrt(n)
  
  z
}

sparse_pearson_residuals <- function(counts, theta = 100) {
  counts_sum0 <- Matrix::colSums(counts)
  counts_sum1 <- Matrix::rowSums(counts)
  counts_sum <- sum(counts@x)
  
  #get residuals sparsely
  mu <- Matrix::sparseMatrix(x = counts_sum1[counts@i + 1] * rep(counts_sum0, diff(counts@p)) / counts_sum,
                             i = counts@i,
                             p = counts@p,
                             dims = dim(counts),
                             index1 = FALSE)
  
  z <- counts - mu
  mu2 <- sqrt(mu + mu^2/theta)
  z@x <- z@x / mu2@x
  
  # clip to sqrt(n)
  n <- ncol(counts)
  z@x[z@x > sqrt(n)] <- sqrt(n)
  z@x[z@x < -sqrt(n)] <- -sqrt(n)
  
  z
}

Based on Lause, Berens, and Kobak on bioRxiv:
https://www.biorxiv.org/content/10.1101/2020.12.01.405886v1.full

Original Python implementation:
https://github.com/berenslab/umi-normalization/blob/main/tools.py