Merge pull request #755 from njtierney/add-warmup-print-i652

Add warmup print i652

Author: njtierney

NEWS.md

@@ -3,6 +3,7 @@
 ## Changes
 
 - `log.greta_array()` function warns if user uses the `base` arg, as it was unused, (#597).
+- Add warmup information to MCMC print method (#652, resolved by #755).
 
 # greta 0.5.0
 

R/greta_mcmc_list.R

@@ -62,13 +62,15 @@ window.greta_mcmc_list <- function(x, start, end, thin, ...) {
 #' @export
 print.greta_mcmc_list <- function(x, ..., n = 5){
 
+  n_warmup <- n_warmup(x)
   n_chain <- coda::nchain(x)
   n_iter <- coda::niter(x)
   n_thin <- coda::thin(x)
   cli::cli_h1("MCMC draws from {.pkg greta}")
   cli::cli_bullets(
     c(
       "*" = "Iterations = {n_iter}",
+      "*" = "Warmup = {n_warmup}",
       "*" = "Chains  = {n_chain}",
       "*" = "Thinning  = {n_thin}"
     )

R/inference.R

@@ -394,7 +394,10 @@ run_samplers <- function(samplers,
     greta_stash$trace_log_files <- trace_log_files
     greta_stash$percentage_log_files <- percentage_log_files
     greta_stash$progress_bar_log_files <- progress_bar_log_files
-    greta_stash$mcmc_info <- list(n_samples = n_samples)
+    greta_stash$mcmc_info <- list(
+      n_samples = n_samples,
+      warmup = warmup
+      )
   }
 
   if (plan_is$parallel) {
@@ -514,7 +517,8 @@ stashed_samples <- function() {
       model_info <- list(
         raw_draws = free_state_draws,
         samplers = samplers,
-        model = samplers[[1]]$model
+        model = samplers[[1]]$model,
+        warmup = samplers[[1]]$warmup
       )
 
       values_draws <- as_greta_mcmc_list(values_draws, model_info)

R/inference_class.R

@@ -244,6 +244,7 @@ sampler <- R6Class(
     n_chains = 1,
     numerical_rejections = 0,
     thin = 1,
+    warmup = 1,
 
     # tuning information
     mean_accept_stat = 0.5,
@@ -336,6 +337,7 @@ sampler <- R6Class(
                          one_by_one, plan_is, n_cores, float_type,
                          trace_batch_size,
                          from_scratch = TRUE) {
+      self$warmup <- warmup
       self$thin <- thin
       dag <- self$model$dag
 

R/utils.R

@@ -1174,3 +1174,8 @@ are_initials <- function(x){
     FUN.VALUE = logical(1)
     )
 }
+
+n_warmup <- function(x){
+  x_info <- attr(x, "model_info")
+  x_info$warmup
+}

tests/testthat/_snaps/greta_mcmc_list_class.md

@@ -6,6 +6,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 10
+      * Warmup = 10
       * Chains = 4
       * Thinning = 1
       
@@ -36,6 +37,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 20
+      * Warmup = 20
       * Chains = 4
       * Thinning = 1
       
@@ -66,6 +68,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 20
+      * Warmup = 20
       * Chains = 4
       * Thinning = 1
       
@@ -109,6 +112,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 20
+      * Warmup = 20
       * Chains = 4
       * Thinning = 1
       
@@ -153,6 +157,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 20
+      * Warmup = 20
       * Chains = 4
       * Thinning = 1
       
@@ -196,6 +201,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 2
+      * Warmup = 2
       * Chains = 4
       * Thinning = 1
       
@@ -221,6 +227,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 2
+      * Warmup = 2
       * Chains = 4
       * Thinning = 1
       
@@ -247,6 +254,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 2
+      * Warmup = 2
       * Chains = 4
       * Thinning = 1
       

tests/testthat/_snaps/print_calculate.md

@@ -28,6 +28,7 @@
       
       -- MCMC draws from greta -------------------------------------------------------
       * Iterations = 10
+      * Warmup = 10
       * Chains = 4
       * Thinning = 1