Better handling of output printing

andrjohns · Dec 25, 2023 · f31abc3 · f31abc3
1 parent 5d8ed46
commit f31abc3
Show file tree

Hide file tree

Showing 15 changed files with 77 additions and 11 deletions.
diff --git a/NAMESPACE b/NAMESPACE
@@ -9,6 +9,7 @@ export(stan_variational)
 exportMethods(summary)
 importFrom(Rcpp,RcppLdFlags)
 importFrom(RcppParallel,RcppParallelLibs)
+importFrom(callr,r_bg)
 importFrom(methods,new)
 importFrom(posterior,as_draws_df)
 importFrom(stats,setNames)

diff --git a/R/cpp_exports.R b/R/cpp_exports.R
@@ -1,4 +1,4 @@
-call_stan <- function(options_vector, ll_fun, grad_fun) {
+call_stan_impl <- function(options_vector, ll_fun, grad_fun) {
   sinkfile <- tempfile()
   sink(file = file(sinkfile, open = "wt"), type = "message")
   status <- .Call(`call_stan_`, options_vector, ll_fun, grad_fun)

diff --git a/R/diagnose.R b/R/diagnose.R
@@ -10,6 +10,7 @@
 #' @param upper Upper bound constraint(s) for parameters
 #' @param seed Random seed
 #' @param refresh Number of iterations for printing
+#' @param quiet (logical) Whether to suppress Stan's output
 #' @param output_dir Directory to store outputs
 #' @param output_basename Basename to use for output files
 #' @param sig_figs Number of significant digits to use for printing
@@ -19,6 +20,7 @@ stan_diagnose <- function(fn, par_inits, additional_args = list(),
                              grad_fun = NULL, lower = -Inf, upper = Inf,
                               seed = NULL,
                               refresh = NULL,
+                              quiet = FALSE,
                               output_dir = NULL,
                               output_basename = NULL,
                               sig_figs = NULL) {
@@ -37,5 +39,5 @@ stan_diagnose <- function(fn, par_inits, additional_args = list(),
                           init = inputs$init_filepath,
                           seed = seed,
                           output_args = output)
-  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function)
+  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function, quiet)
 }
diff --git a/R/laplace.R b/R/laplace.R
@@ -35,6 +35,7 @@ setMethod("summary", "StanLaplace", function(object, ...) {
 #' @param upper Upper bound constraint(s) for parameters
 #' @param seed Random seed
 #' @param refresh Number of iterations for printing
+#' @param quiet (logical) Whether to suppress Stan's output
 #' @param output_dir Directory to store outputs
 #' @param output_basename Basename to use for output files
 #' @param sig_figs Number of significant digits to use for printing
@@ -52,6 +53,7 @@ stan_laplace <- function(fn, par_inits, additional_args = list(),
                              grad_fun = NULL, lower = -Inf, upper = Inf,
                               seed = NULL,
                               refresh = NULL,
+                              quiet = FALSE,
                               output_dir = NULL,
                               output_basename = NULL,
                               sig_figs = NULL,
@@ -112,7 +114,7 @@ stan_laplace <- function(fn, par_inits, additional_args = list(),
                           seed = seed,
                           output_args = output)
 
-  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function)
+  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function, quiet)
 
   parsed <- parse_csv(inputs$output_filepath)
   estimates <- setNames(data.frame(parsed$samples), parsed$header)

diff --git a/R/optimize.R b/R/optimize.R
@@ -36,6 +36,7 @@ setMethod("summary", "StanOptimize", function(object, ...) {
 #' @param upper Upper bound constraint(s) for parameters
 #' @param seed Random seed
 #' @param refresh Number of iterations for printing
+#' @param quiet (logical) Whether to suppress Stan's output
 #' @param output_dir Directory to store outputs
 #' @param output_basename Basename to use for output files
 #' @param sig_figs Number of significant digits to use for printing
@@ -59,6 +60,7 @@ stan_optimize <- function(fn, par_inits, additional_args = list(), algorithm = "
                           grad_fun = NULL, lower = -Inf, upper = Inf,
                           seed = NULL,
                           refresh = NULL,
+                          quiet = FALSE,
                           output_dir = NULL,
                           output_basename = NULL,
                           sig_figs = NULL,
@@ -104,7 +106,7 @@ stan_optimize <- function(fn, par_inits, additional_args = list(), algorithm = "
                           init = inputs$init_filepath,
                           seed = seed,
                           output_args = output)
-  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function)
+  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function, quiet)
 
   parsed <- parse_csv(inputs$output_filepath)
 

diff --git a/R/pathfinder.R b/R/pathfinder.R
@@ -34,6 +34,7 @@ setMethod("summary", "StanPathfinder", function(object, ...) {
 #' @param upper Upper bound constraint(s) for parameters
 #' @param seed Random seed
 #' @param refresh Number of iterations for printing
+#' @param quiet (logical) Whether to suppress Stan's output
 #' @param output_dir Directory to store outputs
 #' @param output_basename Basename to use for output files
 #' @param sig_figs Number of significant digits to use for printing
@@ -61,6 +62,7 @@ stan_pathfinder <- function(fn, par_inits, additional_args = list(), grad_fun =
                           lower = -Inf, upper = Inf,
                           seed = NULL,
                           refresh = NULL,
+                          quiet = FALSE,
                           output_dir = NULL,
                           output_basename = NULL,
                           sig_figs = NULL,
@@ -104,7 +106,7 @@ stan_pathfinder <- function(fn, par_inits, additional_args = list(), grad_fun =
                           seed = seed,
                           output_args = output)
 
-  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function)
+  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function, quiet)
 
   parsed <- parse_csv(inputs$output_filepath)
 

diff --git a/R/sample.R b/R/sample.R
@@ -42,6 +42,7 @@ setMethod("summary", "StanMCMC", function(object, ...) {
 #' @param upper Upper bound constraint(s) for parameters
 #' @param seed Random seed
 #' @param refresh Number of iterations for printing
+#' @param quiet (logical) Whether to suppress Stan's output
 #' @param output_dir Directory to store outputs
 #' @param output_basename Basename to use for output files
 #' @param sig_figs Number of significant digits to use for printing
@@ -89,6 +90,7 @@ stan_sample <- function(fn, par_inits, additional_args = list(),
                           grad_fun = NULL, lower = -Inf, upper = Inf,
                           seed = NULL,
                           refresh = NULL,
+                          quiet = FALSE,
                           output_dir = NULL,
                           output_basename = NULL,
                           sig_figs = NULL,
@@ -163,26 +165,37 @@ stan_sample <- function(fn, par_inits, additional_args = list(),
   r_bg_procs <- lapply(seq_len(parallel_procs), function(chain) {
     list(
       chain_id = chain,
-      proc = callr::r_bg(call_stan, args = chain_calls[[chain]], package = "StanEstimators")
+      proc = callr::r_bg(call_stan_impl, args = chain_calls[[chain]], package = "StanEstimators", supervise = TRUE)
     )
   })
 
   chains_alive <- parallel_procs
   chains_to_run <- num_chains - parallel_procs
+  finished_metadata <- rep(FALSE, parallel_procs)
   while(chains_alive > 0) {
     for (chain in seq_len(parallel_procs)) {
       if (r_bg_procs[[chain]]$proc$is_alive()) {
         r_bg_procs[[chain]]$proc$wait(0.1)
         r_bg_procs[[chain]]$proc$poll_io(0)
-        lines <- r_bg_procs[[chain]]$proc$read_output_lines()
-        if (length(lines) > 0) {
-          cat(paste0("Chain ", r_bg_procs[[chain]]$chain_id, ": ", lines), sep = "\n")
+        if (!quiet) {
+          lines <- r_bg_procs[[chain]]$proc$read_output_lines()
+          if (length(lines) > 0) {
+            for (line in lines) {
+              if (finished_metadata[chain] && line != "") {
+                cat(paste0("Chain ", r_bg_procs[[chain]]$chain_id, ": ", line, "\n"))
+              }
+              if (grepl("num_threads", line)) {
+                finished_metadata[chain] <- TRUE
+              }
+            }
+          }
         }
       } else if (chains_to_run > 0) {
         r_bg_procs[[chain]] <- list(
           chain_id = num_chains - chains_to_run + 1,
-          proc = callr::r_bg(call_stan, args = chain_calls[[num_chains - chains_to_run + 1]], package = "StanEstimators")
+          proc = callr::r_bg(call_stan_impl, args = chain_calls[[num_chains - chains_to_run + 1]], package = "StanEstimators", supervise = TRUE)
         )
+        finished_metadata[chain] <- FALSE
         chains_to_run <- chains_to_run - 1
       }
     }

diff --git a/R/util.R b/R/util.R
@@ -297,3 +297,27 @@ build_stan_call <- function(method, method_args, data_file, init, seed, output_a
   args <- unlist(c(method, method_string, data_string, init_string, random_string, output_string))
   args[args != ""]
 }
+
+call_stan <- function(args_list, ll_fun, grad_fun, quiet) {
+  finished_metadata <- FALSE
+  proc <- callr::r_bg(call_stan_impl, args = list(args_list, ll_fun, grad_fun),
+                      supervise = TRUE,
+                      package = "StanEstimators")
+  while (proc$is_alive()) {
+    proc$wait(0.1)
+    proc$poll_io(0)
+    if (!quiet) {
+      lines <- proc$read_output_lines()
+      if (length(lines) > 0) {
+        for (line in lines) {
+          if (finished_metadata && line != "") {
+            cat(line, "\n")
+          }
+          if (grepl("num_threads", line)) {
+            finished_metadata <- TRUE
+          }
+        }
+      }
+    }
+  }
+}
diff --git a/R/variational.R b/R/variational.R
@@ -36,6 +36,7 @@ setMethod("summary", "StanVariational", function(object, ...) {
 #' @param upper Upper bound constraint(s) for parameters
 #' @param seed Random seed
 #' @param refresh Number of iterations for printing
+#' @param quiet (logical) Whether to suppress Stan's output
 #' @param output_dir Directory to store outputs
 #' @param output_basename Basename to use for output files
 #' @param sig_figs Number of significant digits to use for printing
@@ -60,6 +61,7 @@ stan_variational <- function(fn, par_inits, additional_args = list(), algorithm
                              grad_fun = NULL, lower = -Inf, upper = Inf,
                               seed = NULL,
                               refresh = NULL,
+                              quiet = FALSE,
                               output_dir = NULL,
                               output_basename = NULL,
                               sig_figs = NULL,
@@ -100,7 +102,7 @@ stan_variational <- function(fn, par_inits, additional_args = list(), algorithm
                           seed = seed,
                           output_args = output)
 
-  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function)
+  call_stan(args, ll_fun = inputs$ll_function, grad_fun = inputs$grad_function, quiet)
 
   parsed <- parse_csv(inputs$output_filepath)
   estimates <- setNames(data.frame(parsed$samples), parsed$header)

diff --git a/man/stan_diagnose.Rd b/man/stan_diagnose.Rd
diff --git a/man/stan_laplace.Rd b/man/stan_laplace.Rd
diff --git a/man/stan_optimize.Rd b/man/stan_optimize.Rd
diff --git a/man/stan_pathfinder.Rd b/man/stan_pathfinder.Rd
diff --git a/man/stan_sample.Rd b/man/stan_sample.Rd
diff --git a/man/stan_variational.Rd b/man/stan_variational.Rd