Speed up match_quantiles function for ComBat_seq functon

[lrf2019]

match_quantiles <- function(counts_sub, old_mu, old_phi, new_mu, new_phi){
new_counts_sub <- matrix(NA, nrow=nrow(counts_sub), ncol=ncol(counts_sub))
for(a in 1:nrow(counts_sub)){
for(b in 1:ncol(counts_sub)){
if(counts_sub[a, b] <= 1){
new_counts_sub[a,b] <- counts_sub[a, b]
}else{
tmp_p <- pnbinom(counts_sub[a, b]-1, mu=old_mu[a, b], size=1/old_phi[a])
if(abs(tmp_p-1)<1e-4){
new_counts_sub[a,b] <- counts_sub[a, b]
# for outlier count, if p==1, will return Inf values -> use original count instead
}else{
new_counts_sub[a,b] <- 1+qnbinom(tmp_p, mu=new_mu[a, b], size=1/new_phi[a])
}
}
}
}
return(new_counts_sub)
}

new code for speed up this function

1. use which to transfer matrix boolean calculation to array boolean calculation to speed up

2. use matrix calculation to replace two for cycles to speed up the new matrix generation

match_quantiles_new <- function(counts_sub, old_mu, old_phi, new_mu, new_phi){
new_counts_sub <- counts_sub
condition1 <- which( counts_sub > 1) # array boolean calculation
colid <- floor(condition1/nrow(counts_sub)) + 1 # calculate the colid for rowid calculation
rowid <- condition1 - (colid - 1)*nrow(counts_sub) # calculate rowid for old_phi and new_phi value selection
tmp_p <- pnbinom(counts_sub[condition1] - 1, mu = old_mu[condition1], size = 1/old_phi[rowid])
tmp_q <- 1 + qnbinom(tmp_p[condition1], mu = new_mu[condition1], size = 1/new_phi[rowid])
condition2 <- which((abs(tmp_p-1) >= 1e-4))

new_counts_sub[ condition1[condition2] ] = tmp_q[condition2]
return(new_counts_sub)
}