add exclude argument to subset_draws

R/subset_draws.R

@@ -3,19 +3,25 @@
 #' Subset [`draws`] objects by variables, iterations, chains, and draws indices.
 #'
 #' @template args-methods-x
-#' @param variable (character vector) The variables to select. All elements of
-#'   non-scalar variables can be selected at once.
+#' @param variable (character vector) The variables to select. All
+#'   elements of non-scalar variables can be selected at once.
 #' @param iteration (integer vector) The iteration indices to select.
 #' @param chain (integer vector) The chain indices to select.
-#' @param draw (integer vector) The draw indices to be select. Subsetting draw
-#'   indices will lead to an automatic merging of chains via [`merge_chains`].
+#' @param draw (integer vector) The draw indices to be
+#'   select. Subsetting draw indices will lead to an automatic merging
+#'   of chains via [`merge_chains`].
 #' @param regex (logical) Should `variable` should be treated as a
-#'   (vector of) regular expressions? Any variable in `x` matching at least one
-#'   of the regular expressions will be selected. Defaults to `FALSE`.
-#' @param unique (logical) Should duplicated selection of chains, iterations, or
-#'   draws be allowed? If `TRUE` (the default) only unique chains, iterations,
-#'   and draws are selected regardless of how often they appear in the
-#'   respective selecting arguments.
+#'   (vector of) regular expressions? Any variable in `x` matching at
+#'   least one of the regular expressions will be selected. Defaults
+#'   to `FALSE`.
+#' @param unique (logical) Should duplicated selection of chains,
+#'   iterations, or draws be allowed? If `TRUE` (the default) only
+#'   unique chains, iterations, and draws are selected regardless of
+#'   how often they appear in the respective selecting arguments.
+#' @param exclude (logical) Should the matched selection be excluded?
+#'   If `FALSE` (the default) the matched subset of draws will be
+#'   returned.  If `TRUE` the draws excluding the matched subset will
+#'   be returned.
 #'
 #' @template args-methods-dots
 #' @template return-draws
@@ -46,7 +52,7 @@ subset_draws <- function(x, ...) {
 #' @export
 subset_draws.draws_matrix <- function(x, variable = NULL, iteration = NULL,
                                       chain = NULL, draw = NULL, regex = FALSE,
-                                      unique = TRUE, ...) {
+                                      unique = TRUE, exclude = FALSE, ...) {
   if (all_null(variable, iteration, chain, draw)) {
     return(x)
   }
@@ -55,6 +61,14 @@ subset_draws.draws_matrix <- function(x, variable = NULL, iteration = NULL,
   iteration <- check_iteration_ids(iteration, x, unique = unique)
   chain <- check_chain_ids(chain, x, unique = unique)
   draw <- check_draw_ids(draw, x, unique = unique)
+
+  if (exclude) {
+    variable <- setdiff(variables(x), variable)
+    iteration <- setdiff(iteration_ids(x), iteration)
+    chain <- setdiff(chain_ids(x), chain)
+    draw <- setdiff(draw_ids(x), draw)
+  }
+
   x <- prepare_subsetting(x, iteration, chain, draw)
   x <- .subset_draws(x, iteration, chain, draw, variable, reserved = TRUE)
   if (!is.null(chain) || !is.null(iteration)) {
@@ -67,15 +81,32 @@ subset_draws.draws_matrix <- function(x, variable = NULL, iteration = NULL,
 #' @export
 subset_draws.draws_array <- function(x, variable = NULL, iteration = NULL,
                                      chain = NULL, draw = NULL, regex = FALSE,
-                                     unique = TRUE, ...) {
+                                     unique = TRUE, exclude = FALSE, ...) {
   if (all_null(variable, iteration, chain, draw)) {
     return(x)
   }
+
   x <- repair_draws(x)
   variable <- check_existing_variables(variable, x, regex = regex)
   iteration <- check_iteration_ids(iteration, x, unique = unique)
   chain <- check_chain_ids(chain, x, unique = unique)
   draw <- check_draw_ids(draw, x, unique = unique)
+
+  if (exclude) {
+    if (!is.null(variable)) {
+      variable <- setdiff(variables(x), variable)
+    }
+    if (!is.null(iteration)) {
+      iteration <- setdiff(iteration_ids(x), iteration)
+    }
+    if (!is.null(chain)) {
+      chain <- setdiff(chain_ids(x), chain)
+    }
+    if (!is.null(draw)) {
+      draw <- setdiff(draw_ids(x), draw)
+    }
+  }
+
   x <- prepare_subsetting(x, iteration, chain, draw)
   if (!is.null(draw)) {
     iteration <- draw
@@ -91,16 +122,33 @@ subset_draws.draws_array <- function(x, variable = NULL, iteration = NULL,
 #' @export
 subset_draws.draws_df <- function(x, variable = NULL, iteration = NULL,
                                   chain = NULL, draw = NULL, regex = FALSE,
-                                  unique = TRUE, ...) {
+                                  unique = TRUE, exclude = FALSE, ...) {
   if (all_null(variable, iteration, chain, draw)) {
     return(x)
   }
+
   x <- repair_draws(x)
   unique <- as_one_logical(unique)
   variable <- check_existing_variables(variable, x, regex = regex)
   iteration <- check_iteration_ids(iteration, x, unique = unique)
   chain <- check_chain_ids(chain, x, unique = unique)
   draw <- check_draw_ids(draw, x, unique = unique)
+
+  if (exclude) {
+    if (!is.null(variable)) {
+      variable <- setdiff(variables(x), variable)
+    }
+    if (!is.null(iteration)) {
+      iteration <- setdiff(iteration_ids(x), iteration)
+    }
+    if (!is.null(chain)) {
+      chain <- setdiff(chain_ids(x), chain)
+    }
+    if (!is.null(draw)) {
+      draw <- setdiff(draw_ids(x), draw)
+    }
+  }
+
   x <- prepare_subsetting(x, iteration, chain, draw)
   x <- .subset_draws(
     x, iteration, chain, draw, variable,
@@ -113,15 +161,32 @@ subset_draws.draws_df <- function(x, variable = NULL, iteration = NULL,
 #' @export
 subset_draws.draws_list <- function(x, variable = NULL, iteration = NULL,
                                     chain = NULL, draw = NULL, regex = FALSE,
-                                    unique = TRUE, ...) {
+                                    unique = TRUE, exclude = FALSE, ...) {
   if (all_null(variable, iteration, chain, draw)) {
     return(x)
   }
+
   x <- repair_draws(x)
   variable <- check_existing_variables(variable, x, regex = regex)
   iteration <- check_iteration_ids(iteration, x, unique = unique)
   chain <- check_chain_ids(chain, x, unique = unique)
   draw <- check_draw_ids(draw, x, unique = unique)
+
+  if (exclude) {
+    if (!is.null(variable)) {
+      variable <- setdiff(variables(x), variable)
+    }
+    if (!is.null(iteration)) {
+      iteration <- setdiff(iteration_ids(x), iteration)
+    }
+    if (!is.null(chain)) {
+      chain <- setdiff(chain_ids(x), chain)
+    }
+    if (!is.null(draw)) {
+      draw <- setdiff(draw_ids(x), draw)
+    }
+  }
+
   x <- prepare_subsetting(x, iteration, chain, draw)
   if (!is.null(draw)) {
     iteration <- draw
@@ -137,15 +202,32 @@ subset_draws.draws_list <- function(x, variable = NULL, iteration = NULL,
 #' @export
 subset_draws.draws_rvars <- function(x, variable = NULL, iteration = NULL,
                                      chain = NULL, draw = NULL, regex = FALSE,
-                                     unique = TRUE, ...) {
+                                     unique = TRUE, exclude = FALSE, ...) {
   if (all_null(variable, iteration, chain, draw)) {
     return(x)
   }
+
   x <- repair_draws(x)
   variable <- check_existing_variables(variable, x, regex = regex)
   iteration <- check_iteration_ids(iteration, x, unique = unique)
   chain <- check_chain_ids(chain, x, unique = unique)
   draw <- check_draw_ids(draw, x, unique = unique)
+
+  if (exclude) {
+    if (!is.null(variable)) {
+      variable <- setdiff(variables(x), variable)
+    }
+    if (!is.null(iteration)) {
+      iteration <- setdiff(iteration_ids(x), iteration)
+    }
+    if (!is.null(chain)) {
+      chain <- setdiff(chain_ids(x), chain)
+    }
+    if (!is.null(draw)) {
+      draw <- setdiff(draw_ids(x), draw)
+    }
+  }
+
   x <- prepare_subsetting(x, iteration, chain, draw)
   if (!is.null(draw)) {
     iteration <- draw