From d3b455f50cdb7eb54ef57ae3b9e7704ec4b8eec1 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Fri, 26 Jan 2024 14:51:27 -0500 Subject: [PATCH] update pathfinder args for psis_resample and lp_calculate (#903) * update pathfinder args for psis_resample and lp_calculate --------- Co-authored-by: Andrew Johnson --- R/args.R | 20 ++++++++++++++++++-- R/model.R | 17 ++++++++++++++++- man/model-method-pathfinder.Rd | 13 +++++++++++++ tests/testthat/test-model-pathfinder.R | 4 +++- 4 files changed, 50 insertions(+), 4 deletions(-) diff --git a/R/args.R b/R/args.R index 123022a5..089f6f20 100644 --- a/R/args.R +++ b/R/args.R @@ -566,7 +566,9 @@ PathfinderArgs <- R6::R6Class( num_paths = NULL, max_lbfgs_iters = NULL, num_elbo_draws = NULL, - save_single_paths = NULL) { + save_single_paths = NULL, + psis_resample = NULL, + calculate_lp = NULL) { self$init_alpha <- init_alpha self$tol_obj <- tol_obj self$tol_rel_obj <- tol_rel_obj @@ -580,6 +582,8 @@ PathfinderArgs <- R6::R6Class( self$max_lbfgs_iters <- max_lbfgs_iters self$num_elbo_draws <- num_elbo_draws self$save_single_paths <- save_single_paths + self$psis_resample <- psis_resample + self$calculate_lp <- calculate_lp invisible(self) }, @@ -608,7 +612,9 @@ PathfinderArgs <- R6::R6Class( .make_arg("num_paths"), .make_arg("max_lbfgs_iters"), .make_arg("num_elbo_draws"), - .make_arg("save_single_paths") + .make_arg("save_single_paths"), + .make_arg("psis_resample"), + .make_arg("calculate_lp") ) new_args <- do.call(c, new_args) c(args, new_args) @@ -966,6 +972,16 @@ validate_pathfinder_args <- function(self) { if (!is.null(self$save_single_paths)) { self$save_single_paths <- 0 } + if (!is.null(self$psis_resample) && is.logical(self$psis_resample)) { + self$psis_resample = as.integer(self$psis_resample) + } + checkmate::assert_integerish(self$psis_resample, null.ok = TRUE, + lower = 0, upper = 1, len = 1) + if (!is.null(self$calculate_lp) && is.logical(self$calculate_lp)) { + self$calculate_lp = as.integer(self$calculate_lp) + } + checkmate::assert_integerish(self$calculate_lp, null.ok = TRUE, + lower = 0, upper = 1, len = 1) # check args only available for lbfgs and bfgs diff --git a/R/model.R b/R/model.R index e11d5c41..289558f0 100644 --- a/R/model.R +++ b/R/model.R @@ -1887,6 +1887,17 @@ CmdStanModel$set("public", name = "variational", value = variational) #' calculating the ELBO of the approximation at each iteration of LBFGS. #' @param save_single_paths (logical) Whether to save the results of single #' pathfinder runs in multi-pathfinder. +#' @param psis_resample (logical) Whether to perform pareto smoothed importance sampling. +#' If `TRUE`, the number of draws returned will be equal to `draws`. +#' If `FALSE`, the number of draws returned will be equal to `single_path_draws * num_paths`. +#' @param calculate_lp (logical) Whether to calculate the log probability of the draws. +#' If `TRUE`, the log probability will be calculated and given in the output. +#' If `FALSE`, the log probability will only be returned for draws used to determine the +#' ELBO in the pathfinder steps. All other draws will have a log probability of `NA`. +#' A value of `FALSE` will also turn off pareto smoothed importance sampling as the +#' lp calculation is needed for PSIS. +#' @param save_single_paths (logical) Whether to save the results of single +#' pathfinder runs in multi-pathfinder. #' @return A [`CmdStanPathfinder`] object. #' #' @template seealso-docs @@ -1915,6 +1926,8 @@ pathfinder <- function(data = NULL, max_lbfgs_iters = NULL, num_elbo_draws = NULL, save_single_paths = NULL, + psis_resample = NULL, + calculate_lp = NULL, show_messages = TRUE, show_exceptions = TRUE) { procs <- CmdStanProcs$new( @@ -1940,7 +1953,9 @@ pathfinder <- function(data = NULL, num_paths = num_paths, max_lbfgs_iters = max_lbfgs_iters, num_elbo_draws = num_elbo_draws, - save_single_paths = save_single_paths + save_single_paths = save_single_paths, + psis_resample = psis_resample, + calculate_lp = calculate_lp ) args <- CmdStanArgs$new( method_args = pathfinder_args, diff --git a/man/model-method-pathfinder.Rd b/man/model-method-pathfinder.Rd index aacf1fbc..db3185a3 100644 --- a/man/model-method-pathfinder.Rd +++ b/man/model-method-pathfinder.Rd @@ -29,6 +29,8 @@ pathfinder( max_lbfgs_iters = NULL, num_elbo_draws = NULL, save_single_paths = NULL, + psis_resample = NULL, + calculate_lp = NULL, show_messages = TRUE, show_exceptions = TRUE ) @@ -167,6 +169,17 @@ calculating the ELBO of the approximation at each iteration of LBFGS.} \item{save_single_paths}{(logical) Whether to save the results of single pathfinder runs in multi-pathfinder.} +\item{psis_resample}{(logical) Whether to perform pareto smoothed importance sampling. +If \code{TRUE}, the number of draws returned will be equal to \code{draws}. +If \code{FALSE}, the number of draws returned will be equal to \code{single_path_draws * num_paths}.} + +\item{calculate_lp}{(logical) Whether to calculate the log probability of the draws. +If \code{TRUE}, the log probability will be calculated and given in the output. +If \code{FALSE}, the log probability will only be returned for draws used to determine the +ELBO in the pathfinder steps. All other draws will have a log probability of \code{NA}. +A value of \code{FALSE} will also turn off pareto smoothed importance sampling as the +lp calculation is needed for PSIS.} + \item{show_messages}{(logical) When \code{TRUE} (the default), prints all output during the execution process, such as iteration numbers and elapsed times. If the output is silenced then the \code{\link[=fit-method-output]{$output()}} method of diff --git a/tests/testthat/test-model-pathfinder.R b/tests/testthat/test-model-pathfinder.R index 2f23cc9b..0860ba27 100644 --- a/tests/testthat/test-model-pathfinder.R +++ b/tests/testthat/test-model-pathfinder.R @@ -30,7 +30,9 @@ ok_arg_values <- list( draws = 100, num_paths = 4, max_lbfgs_iters = 100, - save_single_paths = FALSE) + save_single_paths = FALSE, + calculate_lp = TRUE, + psis_resample=TRUE) # using any one of these should cause sample() to error bad_arg_values <- list(