Support list and function inputs for inits

Author: andrjohns

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: StanEstimators
 Title: Estimate Parameters for Arbitrary R Functions using 'Stan'
-Version: 0.2.0
+Version: 0.2.0.9000
 Authors@R: c(
   person(given = c("Andrew", "R."), family = "Johnson", role = c("aut", "cre"),
           email = "andrew.johnson@arjohnsonau.com",
@@ -20,7 +20,7 @@ Description: Allows for the estimation of parameters for 'R' functions using the
 License: MIT + file LICENSE
 Encoding: UTF-8
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.3.1
+RoxygenNote: 7.3.3
 NeedsCompilation: yes
 UseLTO: true
 SystemRequirements: GNU make

R/diagnose.R

@@ -3,7 +3,11 @@
 #' Check gradient estimation using Stan's 'Diagnose' method
 #'
 #' @param fn Function to estimate parameters for
-#' @param par_inits Initial values for parameters
+#' @param par_inits Initial values for parameters. This can either be a numeric vector
+#'  of initial values (which will be used for all chains), a list of numeric vectors (of length
+#'  equal to the number of chains), a function taking a single argument (the chain ID) and
+#'  returning a numeric vector of initial values, or NULL (in which case Stan will
+#'  generate initial values automatically).
 #'  (must be specified if `n_pars` is NULL)
 #' @param n_pars Number of parameters to estimate
 #'  (must be specified if `par_inits` is NULL)
@@ -51,7 +55,7 @@ stan_diagnose <- function(fn, par_inits = NULL, n_pars = NULL, additional_args =
   args <- build_stan_call(method = "diagnose",
                           method_args = "",
                           data_file = inputs$data_filepath,
-                          init = inputs$init_filepath,
+                          init = inputs$init_filepath[1],
                           seed = seed,
                           output_args = output)
   call_stan(args, inputs, quiet)

R/laplace.R

@@ -18,7 +18,11 @@ setMethod("summary", "StanLaplace", function(object, ...) {
 #' Estimate parameters using Stan's laplace algorithm
 #'
 #' @param fn Function to estimate parameters for
-#' @param par_inits Initial values for parameters
+#' @param par_inits Initial values for parameters. This can either be a numeric vector
+#'  of initial values (which will be used for all chains), a list of numeric vectors (of length
+#'  equal to the number of chains), a function taking a single argument (the chain ID) and
+#'  returning a numeric vector of initial values, or NULL (in which case Stan will
+#'  generate initial values automatically).
 #'  (must be specified if `n_pars` is NULL)
 #' @param n_pars Number of parameters to estimate
 #'  (must be specified if `par_inits` is NULL)
@@ -94,11 +98,11 @@ stan_laplace <- function(fn, par_inits = NULL, n_pars = NULL, additional_args =
     mode_vals <- opt@estimates[, -1]
   }
   mode_vals <- as.numeric(mode_vals)
-  if (length(mode_vals) != length(par_inits)) {
+  if (length(mode_vals) != length(inputs$inits[[1]])) {
     stop("The number of mode values does not match the number of parameter ",
           "inits!", .call = FALSE)
   }
-  write_inits(mode_vals, mode_file)
+  write_inits(list(mode_vals), list(mode_file))
   method_args <- list(
     mode = mode_file,
     jacobian = format_bool(jacobian),
@@ -115,7 +119,7 @@ stan_laplace <- function(fn, par_inits = NULL, n_pars = NULL, additional_args =
   args <- build_stan_call(method = "laplace",
                           method_args = method_args,
                           data_file = inputs$data_filepath,
-                          init = inputs$init_filepath,
+                          init = inputs$init_filepath[1],
                           seed = seed,
                           output_args = output)
 

R/optimize.R

@@ -18,7 +18,11 @@ setMethod("summary", "StanOptimize", function(object, ...) {
 #' Estimate parameters using Stan's optimization algorithms
 #'
 #' @param fn Function to estimate parameters for
-#' @param par_inits Initial values for parameters
+#' @param par_inits Initial values for parameters. This can either be a numeric vector
+#'  of initial values (which will be used for all chains), a list of numeric vectors (of length
+#'  equal to the number of chains), a function taking a single argument (the chain ID) and
+#'  returning a numeric vector of initial values, or NULL (in which case Stan will
+#'  generate initial values automatically).
 #'  (must be specified if `n_pars` is NULL)
 #' @param n_pars Number of parameters to estimate
 #'  (must be specified if `par_inits` is NULL)
@@ -108,7 +112,7 @@ stan_optimize <- function(fn, par_inits = NULL, n_pars = NULL, additional_args =
   args <- build_stan_call(method = "optimize",
                           method_args = method_args,
                           data_file = inputs$data_filepath,
-                          init = inputs$init_filepath,
+                          init = inputs$init_filepath[1],
                           seed = seed,
                           output_args = output)
   call_stan(args, inputs, quiet)

R/pathfinder.R

@@ -18,7 +18,11 @@ setMethod("summary", "StanPathfinder", function(object, ...) {
 #' Estimate parameters using Stan's pathfinder algorithm
 #'
 #' @param fn Function to estimate parameters for
-#' @param par_inits Initial values for parameters
+#' @param par_inits Initial values for parameters. This can either be a numeric vector
+#'  of initial values (which will be used for all chains), a list of numeric vectors (of length
+#'  equal to the number of chains), a function taking a single argument (the chain ID) and
+#'  returning a numeric vector of initial values, or NULL (in which case Stan will
+#'  generate initial values automatically).
 #'  (must be specified if `n_pars` is NULL)
 #' @param n_pars Number of parameters to estimate
 #'  (must be specified if `par_inits` is NULL)
@@ -80,7 +84,8 @@ stan_pathfinder <- function(fn, par_inits = NULL, n_pars = NULL, additional_args
                           max_lbfgs_iters = NULL, num_draws = NULL,
                           num_elbo_draws = NULL) {
   inputs <- prepare_inputs(fn, par_inits, n_pars, additional_args, grad_fun, lower, upper,
-                            globals, packages, eval_standalone, output_dir, output_basename)
+                            globals, packages, eval_standalone, output_dir, output_basename,
+                            ifelse(is.null(num_paths), 4, num_paths)) # Default to 4 pathfinders
   method_args <- list(
     init_alpha = init_alpha,
     tol_obj = tol_obj,
@@ -108,7 +113,7 @@ stan_pathfinder <- function(fn, par_inits = NULL, n_pars = NULL, additional_args
   args <- build_stan_call(method = "pathfinder",
                           method_args = method_args,
                           data_file = inputs$data_filepath,
-                          init = inputs$init_filepath,
+                          init = paste0(inputs$init_filepath, collapse = ","), # Pass all inits for multiple pathfinders
                           seed = seed,
                           output_args = output)
 

R/sample.R

@@ -18,7 +18,11 @@ setMethod("summary", "StanMCMC", function(object, ...) {
 #' Estimate parameters using Stan's sampling algorithms
 #'
 #' @param fn Function to estimate parameters for
-#' @param par_inits Initial values for parameters
+#' @param par_inits Initial values for parameters. This can either be a numeric vector
+#'  of initial values (which will be used for all chains), a list of numeric vectors (of length
+#'  equal to the number of chains), a function taking a single argument (the chain ID) and
+#'  returning a numeric vector of initial values, or NULL (in which case Stan will
+#'  generate initial values automatically).
 #'  (must be specified if `n_pars` is NULL)
 #' @param n_pars Number of parameters to estimate
 #'  (must be specified if `par_inits` is NULL)
@@ -121,7 +125,9 @@ stan_sample <- function(fn, par_inits = NULL, n_pars = NULL, additional_args = l
          call. = FALSE)
   }
   inputs <- prepare_inputs(fn, par_inits, n_pars, additional_args, grad_fun, lower, upper,
-                            globals, packages, eval_standalone, output_dir, output_basename)
+                            globals, packages, eval_standalone, output_dir, output_basename,
+                            num_chains)
+
   method_args <- list(
     algorithm = algorithm,
     algorithm_args = list(
@@ -161,7 +167,7 @@ stan_sample <- function(fn, par_inits = NULL, n_pars = NULL, additional_args = l
       args <- build_stan_call(method = "sample",
                               method_args = method_args,
                               data_file = inputs$data_filepath,
-                              init = inputs$init_filepath,
+                              init = inputs$init_filepath[chain],
                               seed = seed,
                               output_args = output,
                               id = chain)
@@ -225,7 +231,7 @@ stan_sample <- function(fn, par_inits = NULL, n_pars = NULL, additional_args = l
     args <- build_stan_call(method = "sample",
                             method_args = method_args,
                             data_file = inputs$data_filepath,
-                            init = inputs$init_filepath,
+                            init = paste0(inputs$init_filepath, collapse = ","),
                             seed = seed,
                             output_args = output)
     call_stan_impl(args, inputs)

R/util.R

@@ -19,9 +19,11 @@ inits_to_json <- function(inits) {
 }
 
 write_inits <- function(inits, init_filepath) {
-  dat_string <- inits_to_json(inits)
-  writeLines(dat_string, con = init_filepath)
-  dat_string
+  lapply(seq_len(length(inits)), function(i) {
+    dat_string <- inits_to_json(inits[[i]])
+    writeLines(dat_string, con = init_filepath[[i]])
+    dat_string
+  })
 }
 
 prepare_and_write_json <- function(what, input_list) {
@@ -44,10 +46,9 @@ with_env <- function(f, e=parent.frame()) {
   f
 }
 
-prepare_function <- function(fn, inits, extra_args_list, grad = FALSE) {
-  fn_wrapper <- function(v) { do.call(fn, c(list(v), extra_args_list)) }
+validate_function <- function(fn, inits, extra_args_list, grad = FALSE) {
   fn_type <- ifelse(isTRUE(grad), "Gradient", "Log-Likelihood")
-  test_fn <- try(invisible(fn_wrapper(inits)), silent = TRUE)
+  test_fn <- try(invisible(fn(inits)), silent = TRUE)
   correct_length <- ifelse(isTRUE(grad), length(inits), 1)
 
   if (inherits(test_fn, "try-error")) {
@@ -60,12 +61,12 @@ prepare_function <- function(fn, inits, extra_args_list, grad = FALSE) {
     stop(fn_type, " function should have return of length ", correct_length,
           ", but return was length ", length(test_fn), "instead!", call. = FALSE)
   } else {
-    fn_wrapper
+    invisible(NULL)
   }
 }
 
 prepare_inputs <- function(fn, par_inits, n_pars, extra_args_list, grad_fun, lower, upper,
-                            globals, packages, eval_standalone, output_dir, output_basename) {
+                            globals, packages, eval_standalone, output_dir, output_basename, num_chains = 1) {
   user_inits <- TRUE
   if (is.null(par_inits)) {
     if (is.null(n_pars)) {
@@ -83,7 +84,26 @@ prepare_inputs <- function(fn, par_inits, n_pars, extra_args_list, grad_fun, low
     user_inits <- FALSE
   }
 
-  fn1 <- prepare_function(fn, par_inits, extra_args_list, grad = FALSE)
+  inits <- NULL
+  if (is.list(par_inits)) {
+    if (length(par_inits) != num_chains) {
+      stop("If par_inits is a list, it must have length equal to num_chains",
+           call. = FALSE)
+    }
+    inits <- par_inits
+  } else if (is.numeric(par_inits)) {
+    inits <- lapply(seq_len(num_chains), function(i) { par_inits })
+  } else if (is.function(par_inits)) {
+    inits <- lapply(seq_len(num_chains), function(i) { par_inits(i) })
+  } else {
+    stop("par_inits must be NULL, a numeric vector, a list of numeric vectors, or a function",
+         call. = FALSE)
+  }
+
+  fn1 <- function(v) { do.call(fn, c(list(v), extra_args_list)) }
+  for (chain in seq_len(num_chains)) {
+    validate_function(fn1, inits[[chain]], extra_args_list, grad = FALSE)
+  }
   fun_globals <- NULL
   fun_packages <- NULL
   if (isTRUE(eval_standalone)) {
@@ -101,7 +121,10 @@ prepare_inputs <- function(fn, par_inits, n_pars, extra_args_list, grad_fun, low
     fun_packages <- c(gp$packages, packages)
   }
   if (!is.null(grad_fun)) {
-    gr1 <- prepare_function(grad_fun, par_inits, extra_args_list, grad = TRUE)
+    gr1 <- function(v) { do.call(grad_fun, c(list(v), extra_args_list)) }
+    for (chain in seq_len(num_chains)) {
+      validate_function(gr1, inits[[chain]], extra_args_list, grad = TRUE)
+    }
     if (isTRUE(eval_standalone)) {
       gr_gp <- future::getGlobalsAndPackages(grad_fun, globals = globals)
       fun_globals <- c(fun_globals, gr_gp$globals)
@@ -111,13 +134,13 @@ prepare_inputs <- function(fn, par_inits, n_pars, extra_args_list, grad_fun, low
     gr1 <- fn1
   }
 
-  if ((length(par_inits) > 1) && (length(lower) == 1)) {
-    lower <- rep(lower, length(par_inits))
+  if ((length(inits[[1]]) > 1) && (length(lower) == 1)) {
+    lower <- rep(lower, length(inits[[1]]))
   }
-  if ((length(par_inits) > 1) && (length(upper) == 1)) {
-    upper <- rep(upper, length(par_inits))
+  if ((length(inits[[1]]) > 1) && (length(upper) == 1)) {
+    upper <- rep(upper, length(inits[[1]]))
   }
-  bounds_types <- sapply(seq_len(length(par_inits)), function(i) {
+  bounds_types <- sapply(seq_len(length(inits[[1]])), function(i) {
     if (lower[i] != -Inf && upper[i] != Inf) {
       3
     } else if (lower[i] != -Inf) {
@@ -139,7 +162,9 @@ prepare_inputs <- function(fn, par_inits, n_pars, extra_args_list, grad_fun, low
 
   init_filepath <- NULL
   if (user_inits) {
-    init_filepath <- tempfile(fileext = ".json", tmpdir = output_dir)
+    init_filepath <- sapply(seq_len(num_chains), function(i) {
+      tempfile(fileext = ".json", tmpdir = output_dir)
+    })
   }
 
   structured <- list(
@@ -148,9 +173,9 @@ prepare_inputs <- function(fn, par_inits, n_pars, extra_args_list, grad_fun, low
     globals = fun_globals,
     packages = fun_packages,
     eval_standalone = eval_standalone,
-    inits = par_inits,
+    inits = inits,
     finite_diff = as.integer(is.null(grad_fun)),
-    Npars = length(par_inits),
+    Npars = length(inits[[1]]),
     lower_bounds = lower,
     upper_bounds = upper,
     bounds_types = bounds_types,

R/variational.R

@@ -18,7 +18,11 @@ setMethod("summary", "StanVariational", function(object, ...) {
 #' Estimate parameters using Stan's variational inference algorithms
 #'
 #' @param fn Function to estimate parameters for
-#' @param par_inits Initial values for parameters
+#' @param par_inits Initial values for parameters. This can either be a numeric vector
+#'  of initial values (which will be used for all chains), a list of numeric vectors (of length
+#'  equal to the number of chains), a function taking a single argument (the chain ID) and
+#'  returning a numeric vector of initial values, or NULL (in which case Stan will
+#'  generate initial values automatically).
 #'  (must be specified if `n_pars` is NULL)
 #' @param n_pars Number of parameters to estimate
 #'  (must be specified if `par_inits` is NULL)
@@ -104,7 +108,7 @@ stan_variational <- function(fn, par_inits = NULL, n_pars = NULL, additional_arg
   args <- build_stan_call(method = "variational",
                           method_args = method_args,
                           data_file = inputs$data_filepath,
-                          init = inputs$init_filepath,
+                          init = inputs$init_filepath[1],
                           seed = seed,
                           output_args = output)