Keep mcmc()'s one_by_one advice when a numerical error comes before the first draw (#577)

NEWS.md

@@ -65,8 +65,10 @@
 * `%*%` now multiplies a data matrix used directly in the product (an R matrix, or a greta array from `as_data()` or `as_data_mutable()`) once, instead of once per chain, so models such as `X %*% beta` run faster in `mcmc()` with several chains (#858).
 * `[<-` on a greta array fills a matrix or array replacement column by column, as R does (#844).
 * `mcmc()` and `extra_samples()` keep `n_samples %/% thin` draws, regardless of values of `pb_update` and `one_by_one`. A burst shorter than `thin` used to error, as in a final partial burst, `one_by_one = TRUE` with `thin` above 1, or `extra_samples()` with `pb_update` at or below `thin`. They run all `n_samples` iterations even when `thin` does not divide it, `one_by_one = TRUE` rejects only the proposal that caused a numerical error, and a `thin` larger than `n_samples` is an informative error (#241, #318, #567, #609).
+* `mcmc()` gives its error suggesting `one_by_one = TRUE` when a numerical error stops sampling before its first draw, which with `verbose = FALSE` is any numerical error, rather than "subscript out of bounds" (#577).
 * `mcmc()`'s progress bar reaches its total when `pb_update` does not divide the number of iterations, and fills in step with the iterations run; a bar could read 1010/1010 while still empty (#318).
 * `opt()`'s `convergence` comes from the optimiser: `bfgs()` and `nelder_mead()` report it, and a Keras optimiser such as `adam()` has converged when its last step changed the objective by no more than `tolerance`, rather than being guessed from whether it stopped before `max_iterations - 1` (#569).
+* `stashed_samples()` returns no draws when `mcmc()` with more than one chain stopped before its first draw, rather than erroring with "length of 'dimnames' [2] not equal to array extent" (#577).
 
 ## Deprecated and defunct
 

R/inference_class.R

@@ -43,7 +43,7 @@ inference <- R6Class(
       self$set_initial_values(initial_values)
       # maybe inefficient-  potentially speed up?
 
-      self$n_traced <- length(model$dag$trace_values(self$free_state))
+      self$n_traced <- ncol(model$dag$trace_values(self$free_state))
       self$seed <- seed
     },
 

R/sampler_class.R

@@ -516,13 +516,21 @@ sampler <- R6Class(
           "This will slow down the sampler slightly.",
           "The error encountered can be recovered and viewed with:",
           "{.code greta_notes_tf_num_error()}"
-        )
+        ),
+        # the methods this is raised from mean nothing to the user
+        call = NULL
       )
     },
 
     # convert traced free state to the traced values, accounting for
     # chain dimension
     trace_values = function(trace_batch_size) {
+      # dag$trace_values() fails on a trace with no rows, and on.exit() calls
+      # this after an error too, where failing would replace that error
+      no_draws <- nrow(self$traced_free_state[[1]]) == 0
+      if (no_draws) {
+        return(NULL)
+      }
       self$traced_values <- lapply(
         self$traced_free_state,
         self$model$dag$trace_values,

tests/testthat/_snaps/inference.md

@@ -229,6 +229,19 @@
       ! `Lmin` must not be larger than `Lmax`
       x `Lmin` is 10 and `Lmax` is 5
 
+# a numerical error stops mcmc() with the one_by_one advice
+
+    Code
+      mcmc(m, sampler = rwmh(epsilon = 1), warmup = 0, n_samples = 40, chains = 1,
+      initial_values = initials(r = 0.2), verbose = FALSE)
+    Condition
+      Error:
+      ! TensorFlow hit a numerical problem that caused it to error
+      greta can handle these as bad proposals if you rerun `mcmc()` with the argument `one_by_one = TRUE`.
+      This will slow down the sampler slightly.
+      The error encountered can be recovered and viewed with:
+      `greta_notes_tf_num_error()`
+
 # thin larger than n_samples is an informative error
 
     Code

tests/testthat/helpers.R

@@ -464,6 +464,30 @@ mock_mcmc <- function(n_samples = 1010) {
   iterate_progress_bar(pb, n_samples, rejects = 10, chains = 1)
 }
 
+# a model of r whose solve() errors on the singular matrix whenever round(r) is
+# 1 or -1, so mcmc() hits numerical errors. r is made by the caller, where
+# initials(r = ) looks for it
+singular_at_one_model <- function(r) {
+  singular_at_one <- diag(2) + (1 - diag(2)) * round(r)
+  y <- as_data(0.5)
+  distribution(y) <- normal(sum(solve(singular_at_one)), 1)
+  model(r)
+}
+
+# put back the samplers and TensorFlow error that an aborted mcmc() stashes,
+# when the caller exits
+local_sampler_stash <- function(.local_envir = parent.frame()) {
+  fields <- c("samplers", "tf_num_error")
+  old <- mget(intersect(fields, ls(greta_stash)), envir = greta_stash)
+  withr::defer(
+    {
+      rm(list = intersect(fields, ls(greta_stash)), envir = greta_stash)
+      list2env(old, envir = greta_stash)
+    },
+    envir = .local_envir
+  )
+}
+
 # random lkj draws, code from the rethinking package (can't load the package
 # because of stan*Travis*compiler issues)
 rlkjcorr <- function(n, eta = 1, dimension = 2) {

tests/testthat/test_inference.R

@@ -705,15 +705,54 @@ test_that("a numerical error at the end of warmup still tunes", {
   expect_equal(sampler$parameters$epsilon, exp(log_epsilon_bar))
 })
 
+test_that("a numerical error stops mcmc() with the one_by_one advice", {
+  skip_if_not(check_tf_version())
+  local_greta_seed()
+  local_sampler_stash()
+  r <- normal(0, 1)
+  m <- singular_at_one_model(r)
+
+  expect_snapshot(
+    error = TRUE,
+    mcmc(
+      m,
+      sampler = rwmh(epsilon = 1),
+      warmup = 0,
+      n_samples = 40,
+      chains = 1,
+      initial_values = initials(r = 0.2),
+      verbose = FALSE
+    )
+  )
+})
+
+test_that("stashed_samples() has no draws when several chains stop before one", {
+  skip_if_not(check_tf_version())
+  local_greta_seed()
+  local_sampler_stash()
+  r <- normal(0, 1)
+  m <- singular_at_one_model(r)
+
+  expect_error(
+    mcmc(
+      m,
+      sampler = rwmh(epsilon = 1),
+      warmup = 0,
+      n_samples = 40,
+      chains = 2,
+      initial_values = list(initials(r = 0.2), initials(r = -0.2)),
+      verbose = FALSE
+    ),
+    "one_by_one = TRUE"
+  )
+  expect_identical(stashed_samples(), list(NULL))
+})
+
 test_that("a numerical error with one_by_one repeats the draw before it", {
   skip_if_not(check_tf_version())
   local_greta_seed()
-  # solve() errors on the singular matrix whenever round(r) is 1 or -1
   r <- normal(0, 1)
-  singular_at_one <- diag(2) + (1 - diag(2)) * round(r)
-  y <- as_data(0.5)
-  distribution(y) <- normal(sum(solve(singular_at_one)), 1)
-  m <- model(r)
+  m <- singular_at_one_model(r)
 
   draws <- mcmc(
     m,
@@ -727,6 +766,7 @@ test_that("a numerical error with one_by_one repeats the draw before it", {
   )
   sampler <- get_model_info(draws)$samplers[[1]]
   expect_gt(sampler$numerical_rejections, 0)
+  expect_equal(coda::niter(draws), 40)
 
   chain <- as.numeric(draws[[1]])
   repeated <- which(duplicated(chain))
@@ -736,12 +776,8 @@ test_that("a numerical error with one_by_one repeats the draw before it", {
 test_that("numerical errors in warmup with one_by_one still tune", {
   skip_if_not(check_tf_version())
   local_greta_seed()
-  # solve() errors on the singular matrix whenever round(r) is 1 or -1
   r <- normal(0, 1)
-  singular_at_one <- diag(2) + (1 - diag(2)) * round(r)
-  y <- as_data(0.5)
-  distribution(y) <- normal(sum(solve(singular_at_one)), 1)
-  m <- model(r)
+  m <- singular_at_one_model(r)
 
   warmup <- 30
   chains <- 2