refactor(imputation): Update AFT imputation method to linear time complexity

DESCRIPTION

@@ -40,7 +40,8 @@ Imports:
     parallel,
     rlang,
     matter,
-    BiocParallel
+    BiocParallel,
+    RhpcBLASctl
 Suggests: 
     BiocStyle,
     knitr,
@@ -62,4 +63,4 @@ Packaged: 2017-10-20 02:13:12 UTC; meenachoi
 LinkingTo: 
     Rcpp,
     RcppArmadillo
-Config/roxygen2/version: 8.0.0
+Config/roxygen2/version: 8.1.0

NAMESPACE

@@ -80,6 +80,7 @@ importFrom(MSstatsConvert,ProgenesistoMSstatsFormat)
 importFrom(MSstatsConvert,SkylinetoMSstatsFormat)
 importFrom(MSstatsConvert,SpectronauttoMSstatsFormat)
 importFrom(Rcpp,sourceCpp)
+importFrom(RhpcBLASctl,blas_set_num_threads)
 importFrom(data.table,as.data.table)
 importFrom(data.table,data.table)
 importFrom(data.table,fifelse)
@@ -122,14 +123,20 @@ importFrom(plotly,subplot)
 importFrom(preprocessCore,normalize.quantiles)
 importFrom(rlang,.data)
 importFrom(stats,dist)
+importFrom(stats,dnorm)
 importFrom(stats,fitted)
 importFrom(stats,formula)
 importFrom(stats,hclust)
 importFrom(stats,lm)
+importFrom(stats,lm.fit)
 importFrom(stats,loess)
 importFrom(stats,median)
+importFrom(stats,model.frame)
+importFrom(stats,model.matrix)
+importFrom(stats,model.response)
 importFrom(stats,na.omit)
 importFrom(stats,p.adjust)
+importFrom(stats,pnorm)
 importFrom(stats,predict)
 importFrom(stats,qbinom)
 importFrom(stats,qnorm)

R/MSstatsSummarizeWithMultipleCores.R

@@ -278,7 +278,7 @@
 #' @noRd
 .build_summarize_worker <- function(
         use_TMP, impute, censored_symbol, remove50missing,
-        aft_iterations, equal_variance
+        aft_iterations, equal_variance, aft_solver, aft_verbose
 ) {
     unpack_fn        <- .unpack_protein_slot
     use_TMP_         <- use_TMP
@@ -287,19 +287,22 @@
     remove50missing_ <- remove50missing
     aft_iterations_  <- aft_iterations
     equal_variance_  <- equal_variance
+    aft_solver_ <- aft_solver
+    aft_verbose_ <- aft_verbose
 
     function(record) {
         meta   <- record$meta
         protein_dt <- unpack_fn(record$packed, meta)
         result <- if (use_TMP_) {
             MSstatsSummarizeSingleTMP(
                 protein_dt, impute_, censored_symbol_,
-                remove50missing_, aft_iterations_)
+                remove50missing_, aft_iterations_, aft_solver_, aft_verbose_)
         } else {
             MSstatsSummarizeSingleLinear(
                 protein_dt, impute_, censored_symbol_,
                 remove50missing_, aft_iterations_,
-                equal_variances = equal_variance_)
+                equal_variances = equal_variance_, 
+                aft_solver = aft_solver_, aft_verbose = aft_verbose_)
         }
         result
     }
@@ -313,9 +316,11 @@
     list(worker = i, pid = Sys.getpid(), max_rss_mb = .max_rss_mb())
 }
 
+#' @importFrom RhpcBLASctl blas_set_num_threads
 .warmup_worker <- function(i) {
     library(MSstats, quietly = TRUE, warn.conflicts = FALSE)
     data.table::setDTthreads(1)
+    RhpcBLASctl::blas_set_num_threads(1)
     NULL
 }
 
@@ -330,6 +335,14 @@
 #' @param equal_variance only for method = "linear"; assume equal variance among feature intensities
 #' @param numberOfCores number of cores for parallel processing (Linux/Mac only)
 #' @param aft_iterations number of AFT model iterations
+#' @param aft_solver only used when impute = TRUE; linear solve
+#'   used in the AFT imputation model's Newton-Raphson step: "cholesky"
+#'   (default, via \code{survival::survreg}), "cg" (conjugate gradient), or
+#'   "pcg" (conjugate gradient with a Jacobi/inverse-diagonal preconditioner).
+#'   "cg"/"pcg" are experimental.
+#' @param aft_verbose if \code{TRUE}, \code{message()} AFT fitting diagnostics
+#'   (problem size, elapsed time, and for "cg"/"pcg" per-Newton-iteration
+#'   conjugate-gradient counts) for every protein fit. Default \code{FALSE}.
 #' @param verbose whether to print verbose output
 #' @param BPPARAM optional \code{BiocParallelParam} instance
 #' @param track_memory whether to report per-worker maximum RSS memory usage.
@@ -361,15 +374,18 @@ MSstatsSummarizeWithMultipleCores <- function(
         equal_variance,
         numberOfCores  = 1L,
         aft_iterations = 90L,
+        aft_solver = "cholesky", 
+        aft_verbose = FALSE,
         verbose        = FALSE,
         BPPARAM        = NULL,
         track_memory   = FALSE,
         max_proteins_per_worker = 50L
 ) {
+    .checkAFTSolver(aft_solver)
     if (numberOfCores <= 1L && is.null(BPPARAM)) {
         return(MSstatsSummarizeWithSingleCore(
             input, method, impute, censored_symbol,
-            remove50missing, equal_variance, aft_iterations))
+            remove50missing, equal_variance, aft_iterations, aft_solver, aft_verbose))
     }
 
     start_time <- proc.time()[["elapsed"]]
@@ -419,7 +435,7 @@ MSstatsSummarizeWithMultipleCores <- function(
 
     worker_fn <- .build_summarize_worker(
         use_TMP, impute, censored_symbol, remove50missing,
-        aft_iterations, equal_variance)
+        aft_iterations, equal_variance, aft_solver, aft_verbose)
 
     if (is.null(BPPARAM)) {
         tasks <- if (max_proteins_per_worker > 0L) {

R/dataProcess.R

@@ -62,6 +62,20 @@
 #' a logfile named `MSstats_dataProcess_log_progress.log` is created to 
 #' track progress. Only works for Linux & Mac OS. Default is 1.
 #' @param aft_iterations Number of iterations for AFT model fitting. Default is 90.
+#' @param aft_solver Which linear solve to use for the AFT imputation
+#' model's Newton-Raphson step: "cholesky" (default) delegates to
+#' \code{survival::survreg}, which solves it via Cholesky factorization.
+#' "cg" solves the same Newton step with a vendored conjugate-gradient
+#' routine instead; "pcg" is the same conjugate-gradient routine with a
+#' Jacobi (inverse-diagonal) preconditioner, which can reduce the number
+#' of conjugate-gradient iterations needed. "cg"/"pcg" are experimental
+#' alternatives, currently opt-in only.
+#' @param aft_verbose If \code{TRUE}, \code{message()} diagnostics for
+#' every protein fit: problem size and elapsed fitting time for all
+#' solvers, plus per-Newton-iteration conjugate-gradient iteration counts
+#' and timing when \code{aft_solver} is "cg" or "pcg" - useful for
+#' evaluating solver time complexity, but produces one block of output
+#' per protein, so leave at the default \code{FALSE} for routine runs.
 #' @inheritParams .documentFunction
 #' 
 #' @importFrom utils sessionInfo
@@ -130,7 +144,8 @@ dataProcess = function(
     equalFeatureVar = TRUE, censoredInt = "NA", MBimpute = TRUE, 
     remove50missing = FALSE, fix_missing = NULL, maxQuantileforCensored = 0.999, 
     use_log_file = TRUE, append = FALSE, verbose = TRUE, log_file_path = NULL,
-    numberOfCores = 1, aft_iterations=90
+    numberOfCores = 1, aft_iterations=90, aft_solver = "cholesky",
+    aft_verbose = FALSE
 ) {
     MSstatsConvert::MSstatsLogsSettings(use_log_file, append, verbose, 
                                         log_file_path,
@@ -143,6 +158,7 @@ dataProcess = function(
         list(method = summaryMethod, equal_var = equalFeatureVar),
         list(symbol = censoredInt, MB = MBimpute),
         colnames(raw))
+    .checkAFTSolver(aft_solver)
 
     peptides_dict = makePeptidesDictionary(as.data.table(unclass(raw)), normalization)
     input = MSstatsPrepareForDataProcess(raw, logTrans, fix_missing)
@@ -164,9 +180,10 @@ dataProcess = function(
     input = MSstatsPrepareForSummarization(input, summaryMethod, MBimpute, censoredInt,
                                            remove_uninformative_feature_outlier)
     summarized = tryCatch(MSstatsSummarizeWithMultipleCores(input, summaryMethod,
-                                           MBimpute, censoredInt, 
-                                           remove50missing, equalFeatureVar, 
-                                           numberOfCores, aft_iterations),
+                                           MBimpute, censoredInt,
+                                           remove50missing, equalFeatureVar,
+                                           numberOfCores, aft_iterations,
+                                           aft_solver, aft_verbose),
                           error = function(e) {
                               print(e)
                               NULL
@@ -211,8 +228,9 @@ dataProcess = function(
 #' head(summarized[[1]][[1]]) # run-level summary
 #' 
 MSstatsSummarizeWithSingleCore = function(input, method, impute, censored_symbol,
-                            remove50missing, equal_variance, aft_iterations = 90) {
-
+                            remove50missing, equal_variance, aft_iterations = 90,
+                            aft_solver = "cholesky", aft_verbose = FALSE) {
+    .checkAFTSolver(aft_solver)
 
     is_labeled_reference = "is_labeled_ref" %in% colnames(input) && any(input$is_labeled_ref, na.rm = TRUE)
     if (is_labeled_reference) {
@@ -227,8 +245,9 @@ MSstatsSummarizeWithSingleCore = function(input, method, impute, censored_symbol
         for (protein_id in seq_len(num_proteins)) {
             single_protein = input[protein_indices[[protein_id]],]
             summarized_results[[protein_id]] = MSstatsSummarizeSingleTMP(
-                single_protein, impute, censored_symbol, remove50missing, 
-                aft_iterations)
+                single_protein, impute, censored_symbol, remove50missing,
+                aft_iterations, aft_solver = aft_solver,
+                aft_verbose = aft_verbose)
             setTxtProgressBar(pb, protein_id)
         }
         close(pb)
@@ -237,8 +256,9 @@ MSstatsSummarizeWithSingleCore = function(input, method, impute, censored_symbol
         for (protein_id in seq_len(num_proteins)) {
             single_protein = input[protein_indices[[protein_id]],]
             summarized_result = MSstatsSummarizeSingleLinear(
-                single_protein, impute, censored_symbol, 
-              remove50missing, aft_iterations)
+                single_protein, impute, censored_symbol,
+              remove50missing, aft_iterations, aft_solver = aft_solver,
+              aft_verbose = aft_verbose)
 
             summarized_results[[protein_id]] = summarized_result
             setTxtProgressBar(pb, protein_id)
@@ -256,9 +276,16 @@ MSstatsSummarizeWithSingleCore = function(input, method, impute, censored_symbol
 #' @param remove50missing if TRUE, proteins with more than 50\% missing values in each run are removed
 #' @param aft_iterations number of iterations for AFT model fitting
 #' @param equal_variances if TRUE, observation are assumed to be homoskedastic
-#' 
+#' @param aft_solver Which linear solve to use for the AFT imputation
+#' model's Newton-Raphson step: "cholesky" (default, via
+#' \code{survival::survreg}), "cg" (conjugate gradient), or "pcg"
+#' (conjugate gradient with a Jacobi/inverse-diagonal preconditioner).
+#' @param aft_verbose If \code{TRUE}, log AFT fitting diagnostics for
+#' every protein fit. See \code{.fitSurvival}'s and
+#' \code{.fitSurvivalCG}'s \code{verbose}.
+#'
 #' @return list with protein-level data
-#' 
+#'
 #' @importFrom stats xtabs
 #' 
 #' @export
@@ -286,8 +313,11 @@ MSstatsSummarizeSingleLinear = function(single_protein,
                                         censored_symbol,
                                         remove50missing,
                                         aft_iterations = 90,
-                                        equal_variances = TRUE) {
+                                        equal_variances = TRUE,
+                                        aft_solver = "cholesky",
+                                        aft_verbose = FALSE) {
     ABUNDANCE = RUN = FEATURE = PROTEIN = LogIntensities = NULL
+    .checkAFTSolver(aft_solver)
 
     cols = intersect(
       colnames(single_protein),
@@ -315,7 +345,8 @@ MSstatsSummarizeSingleLinear = function(single_protein,
         } else {
             single_protein[, cols, with = FALSE]
         }
-        survival_fit = .fitSurvival(fit_data, aft_iterations)
+        survival_fit = .fitAFTModel(fit_data, aft_iterations, aft_solver,
+                                    aft_verbose)
         sigma2 = survival_fit$scale^2
 
         single_protein[, c("predicted", "imputation_var") := {
@@ -437,9 +468,12 @@ MSstatsSummarizeSingleLinear = function(single_protein,
 #' head(single_protein_summary[[1]])
 #' 
 MSstatsSummarizeSingleTMP = function(single_protein, impute, censored_symbol,
-                                     remove50missing, aft_iterations = 90) {
+                                     remove50missing, aft_iterations = 90,
+                                     aft_solver = "cholesky",
+                                     aft_verbose = FALSE) {
     newABUNDANCE = n_obs = n_obs_run = RUN = FEATURE = LABEL = NULL
     predicted = censored = NULL
+    .checkAFTSolver(aft_solver)
     cols = intersect(colnames(single_protein), c("newABUNDANCE", "cen", "RUN",
                                                  "FEATURE", "ref_covariate"))
     is_labeled_reference = "is_labeled_ref" %in% colnames(single_protein) &&
@@ -455,6 +489,7 @@ MSstatsSummarizeSingleTMP = function(single_protein, impute, censored_symbol,
 
         # Flag to track convergence warning
         converged = TRUE
+        convergence_messages = character(0)
 
         fit_data = if (is_labeled_reference) {
             single_protein[(!is_labeled_ref), cols, with = FALSE]
@@ -464,14 +499,30 @@ MSstatsSummarizeSingleTMP = function(single_protein, impute, censored_symbol,
 
         # Try to fit survival model and catch convergence warnings
         survival_fit = withCallingHandlers({
-            .fitSurvival(fit_data, aft_iterations)
+            .fitAFTModel(fit_data, aft_iterations, aft_solver, aft_verbose)
         }, warning = function(w) {
-            if (grepl("converge", conditionMessage(w), ignore.case = TRUE)) {
-                message("Convergence warning caught: ", conditionMessage(w))
+            warning_message = conditionMessage(w)
+            if (grepl("converge", warning_message, ignore.case = TRUE)) {
+                convergence_messages <<- c(convergence_messages,
+                                           warning_message)
                 converged <<- FALSE
+                invokeRestart("muffleWarning")
             }
         })
 
+        protein_name = as.character(unique(single_protein$PROTEIN))[1]
+        log_fun = getOption("MSstatsLog")
+        if (!converged) {
+            msg = paste0("CONVERGENCE WARNING for protein: ", protein_name,
+                         " (", length(convergence_messages),
+                         " warning(s)) - ",
+                         paste(unique(convergence_messages), collapse = " | "))
+            message(msg)
+            if (is.function(log_fun)) {
+                log_fun("INFO", msg)
+            }
+        }
+
         if (converged) {
             single_protein[, predicted := predict(survival_fit, newdata = .SD)]
         } else {