Add arguments cvfun and cvrefbuilder to get_refmodel.stanreg()

NEWS.md

@@ -27,6 +27,7 @@ If you read this from a place other than <https://mc-stan.org/projpred/news/inde
 
 ## Minor changes
 
+* `get_refmodel.stanreg()` has new arguments `cvfun` and `cvrefbuilder` for K-fold CV. For example, a custom `cvrefbuilder` lets custom latent-projection functions use the draws of the fold-wise refits. (GitHub: #604)
 * New function `custom_stat()` defines a user-defined performance statistic for argument `stats` of `summary.vsel()`, `plot.vsel()` and `performances()`, and argument `stat` of `suggest_size()`. (GitHub: #566, #590)
 * Faster `cl_agg()` and conversion of arrays to augmented-rows matrices. For example, LOO CV with validated search for the latent projection with 4000 projected draws is about 3-4 times faster. (GitHub: #600)
 

R/extend_family.R

@@ -245,20 +245,23 @@
 #'
 #' In [cv_varsel()] with `cv_method = "kfold"`, the reference model objects of
 #' the \eqn{K} folds are created by the `cvrefbuilder` function (see
-#' [init_refmodel()]). The `cvrefbuilder` functions of [get_refmodel.stanreg()]
-#' and [brms::get_refmodel.brmsfit()] pass user-supplied functions
-#' `latent_ilink`, `latent_ll_oscale`, and `latent_ppd_oscale` (and a
-#' user-supplied argument `dis`, see [init_refmodel()]) unchanged to the
-#' fold-wise reference model objects. If these functions take reference model
-#' draws from their enclosing environment (e.g., draws of a shape parameter),
-#' or if `dis` contains such draws, then the folds use the draws from the
-#' full-data fit. To use the draws from the fold-wise refits, define the
-#' functions (and compute `dis`) inside a function that takes the fitted model
-#' object as its argument and calls [get_refmodel()]. Then set element
-#' `cvrefbuilder` of the reference model object to a function that calls this
-#' outer function with the fold-wise refit. (Methods [get_refmodel.stanreg()]
-#' and [brms::get_refmodel.brmsfit()] do not accept a `cvrefbuilder` argument.)
-#' The latent-projection vignette shows an example.
+#' [init_refmodel()]). The default `cvrefbuilder` functions of
+#' [get_refmodel.stanreg()] and [brms::get_refmodel.brmsfit()] pass
+#' user-supplied functions `latent_ilink`, `latent_ll_oscale`, and
+#' `latent_ppd_oscale` (and a user-supplied argument `dis`, see
+#' [init_refmodel()]) unchanged to the fold-wise reference model objects. If
+#' these functions take reference model draws from their enclosing environment
+#' (e.g., draws of a shape parameter), or if `dis` contains such draws, then the
+#' folds use the draws from the full-data fit. To use the draws from the
+#' fold-wise refits, define the functions (and compute `dis`) inside a function
+#' that takes the fitted model object as its argument and calls
+#' [get_refmodel()]. Then define a function that calls this outer function with
+#' the fold-wise refit. For an \pkg{rstanarm} reference model, pass this
+#' function as argument `cvrefbuilder` of [get_refmodel.stanreg()]. Method
+#' [brms::get_refmodel.brmsfit()] does not accept a `cvrefbuilder` argument;
+#' for a \pkg{brms} reference model, set element `cvrefbuilder` of the
+#' reference model object to this function instead. The latent-projection
+#' vignette shows an example.
 #'
 #' If `family$cats` is not `NULL` (after taking `latent_y_unqs` into account),
 #' then the response vector resulting from `extract_model_data` (see

R/refmodel.R

@@ -95,6 +95,8 @@
 #'   `NULL` for using an internal default. Only one of `cvfits` and `cvfun`
 #'   needs to be provided (for \eqn{K}-fold CV). Note that `cvfits` takes
 #'   precedence over `cvfun`, i.e., if both are provided, `cvfits` is used.
+#'   For [get_refmodel.stanreg()], `cvfun = NULL` uses an internal default
+#'   based on [rstanarm::kfold.stanreg()].
 #' @param cvrefbuilder For \eqn{K}-fold CV only. A function that, given a
 #'   reference model fit for fold \eqn{k \in \{1, ..., K\}}{k = 1, ..., K},
 #'   returns an object of the same type as [init_refmodel()] does. The reference
@@ -107,7 +109,12 @@
 #'   `projpred_k` are appended as attributes). Argument `cvrefbuilder` may be
 #'   `NULL` for using an internal default: [get_refmodel()] if `object` is not
 #'   `NULL` and a function calling [init_refmodel()] appropriately (with the
-#'   assumption `dis = 0`) if `object` is `NULL`.
+#'   assumption `dis = 0`) if `object` is `NULL`. For [get_refmodel.stanreg()],
+#'   `cvrefbuilder = NULL` uses an internal default that calls [get_refmodel()]
+#'   on the fold-wise fit with the arguments `latent`, `dis`, and `...` of the
+#'   original call and with `called_from_cvrefbuilder = TRUE`. A custom
+#'   `cvrefbuilder` is needed, for example, if the reference model uses custom
+#'   functions that depend on the posterior draws (see [extend_family()]).
 #' @param called_from_cvrefbuilder A single logical value indicating whether
 #'   [init_refmodel()] is called from a `cvrefbuilder` function (`TRUE`) or not
 #'   (`FALSE`). Currently, `TRUE` only causes some warnings to be suppressed
@@ -792,10 +799,17 @@ get_refmodel.default <- function(object, family = NULL, ...) {
 
 #' @rdname refmodel-init-get
 #' @export
-get_refmodel.stanreg <- function(object, latent = FALSE, dis = NULL, ...) {
+get_refmodel.stanreg <- function(object, latent = FALSE, dis = NULL,
+                                 cvfun = NULL, cvrefbuilder = NULL, ...) {
   if (!requireNamespace("rstanarm", quietly = TRUE)) {
     stop("Please install the 'rstanarm' package.")
   }
+  if (!is.null(cvfun) && !is.function(cvfun)) {
+    stop("Argument `cvfun` must be a function or `NULL`.")
+  }
+  if (!is.null(cvrefbuilder) && !is.function(cvrefbuilder)) {
+    stop("Argument `cvrefbuilder` must be a function or `NULL`.")
+  }
 
   # Family ------------------------------------------------------------------
 
@@ -988,24 +1002,28 @@ get_refmodel.stanreg <- function(object, latent = FALSE, dis = NULL, ...) {
     return(linpred_out)
   }
 
-  cvfun <- function(folds) {
-    # Use `cores = 1` because of rstanarm issue #551. In fact, this issue only
-    # affects Windows systems, but since `cores = 1` leads to an *inner*
-    # parallelization (i.e., across chains, not across CV folds) with
-    # `stan_cores <- getOption("mc.cores", 1)` cores, this should also be
-    # suitable for other systems:
-    kfold(
-      object, K = max(folds), save_fits = TRUE, folds = folds, cores = 1
-    )$fits[, "fit"]
-  }
-
-  # Keep the user's `dis` (possibly `NULL`) for the folds. The default set below
-  # uses the full-data `sigma`; with `NULL`, each fold takes `sigma` from its
-  # own refit.
-  dis_user <- dis
-  cvrefbuilder <- function(cvfit) {
-    get_refmodel(cvfit, latent = latent, dis = dis_user,
-                 called_from_cvrefbuilder = TRUE, ...)
+  if (is.null(cvfun)) {
+    cvfun <- function(folds) {
+      # Use `cores = 1` because of rstanarm issue #551. In fact, this issue
+      # only affects Windows systems, but since `cores = 1` leads to an *inner*
+      # parallelization (i.e., across chains, not across CV folds) with
+      # `stan_cores <- getOption("mc.cores", 1)` cores, this should also be
+      # suitable for other systems:
+      kfold(
+        object, K = max(folds), save_fits = TRUE, folds = folds, cores = 1
+      )$fits[, "fit"]
+    }
+  }
+
+  if (is.null(cvrefbuilder)) {
+    # Keep the user's `dis` (possibly `NULL`) for the folds. The default set
+    # below uses the full-data `sigma`; with `NULL`, each fold takes `sigma`
+    # from its own refit.
+    dis_user <- dis
+    cvrefbuilder <- function(cvfit) {
+      get_refmodel(cvfit, latent = latent, dis = dis_user,
+                   called_from_cvrefbuilder = TRUE, ...)
+    }
   }
 
   # Miscellaneous -----------------------------------------------------------

tests/testthat/test_refmodel.R

@@ -228,6 +228,86 @@ test_that(paste(
   expect_identical(refmod_user$cvrefbuilder(fit_fold)$dis, dis_user)
 })
 
+test_that(paste(
+  "get_refmodel.stanreg() uses user-supplied `cvfun` and `cvrefbuilder`",
+  "functions in K-fold CV"
+), {
+  tstsetups <- grep("^rstanarm\\..*\\.gauss\\.", names(fits), value = TRUE)
+  skip_if_not(length(tstsetups) >= 1)
+  fit <- fits[[tstsetups[1]]]
+
+  # The user-supplied `cvfun` returns the full-data fit for each fold, so no
+  # refit is needed:
+  cvfun_folds <- list()
+  cvfun_user <- function(folds) {
+    cvfun_folds[[length(cvfun_folds) + 1L]] <<- folds
+    rep(list(fit), max(folds))
+  }
+  cvrefbuilder_ks <- integer()
+  cvrefbuilder_user <- function(cvfit) {
+    cvrefbuilder_ks <<- c(cvrefbuilder_ks, cvfit$projpred_k)
+    get_refmodel(cvfit, called_from_cvrefbuilder = TRUE)
+  }
+  refmod <- get_refmodel(fit, cvfun = cvfun_user,
+                         cvrefbuilder = cvrefbuilder_user)
+  expect_identical(refmod$cvfun, cvfun_user)
+  expect_identical(refmod$cvrefbuilder, cvrefbuilder_user)
+
+  folds <- rep_len(seq_len(K_tst), refmod$nobs)
+  cvfits <- run_cvfun(refmod, folds = folds)
+  expect_length(cvfun_folds, 1)
+  expect_identical(cvfun_folds[[1]], folds)
+  expect_identical(attr(cvfits, "folds"), folds)
+  expect_length(cvfits, K_tst)
+
+  cvvs <- cv_varsel(refmod, cv_method = "kfold", K = K_tst, cvfits = cvfits,
+                    nclusters = nclusters_tst,
+                    nclusters_pred = nclusters_pred_tst, nterms_max = 1,
+                    verbose = FALSE, seed = seed_tst)
+  expect_s3_class(cvvs, "vsel")
+  expect_identical(cvrefbuilder_ks, seq_len(K_tst))
+})
+
+test_that(paste(
+  "get_refmodel.stanreg() with only one of `cvfun` and `cvrefbuilder` keeps",
+  "the default of the other, and the default `cvrefbuilder` does not pass a",
+  "user-supplied `cvfun` to the folds"
+), {
+  tstsetups <- grep("^rstanarm\\..*\\.gauss\\.", names(fits), value = TRUE)
+  skip_if_not(length(tstsetups) >= 2)
+  # The second fit stands in for a fold refit:
+  fit_full <- fits[[tstsetups[1]]]
+  fit_fold <- fits[[tstsetups[2]]]
+  sigma_fold <- data.frame(fit_fold)[, "sigma"]
+  refmod_dflt <- get_refmodel(fit_full)
+
+  cvfun_user <- function(folds) stop("Not used.")
+  refmod <- get_refmodel(fit_full, cvfun = cvfun_user)
+  expect_identical(refmod$cvfun, cvfun_user)
+  refmod_fold <- refmod$cvrefbuilder(fit_fold)
+  expect_false(identical(refmod_fold$cvfun, cvfun_user))
+  expect_identical(refmod_fold$dis, sigma_fold)
+
+  cvrefbuilder_user <- function(cvfit) stop("Not used.")
+  refmod <- get_refmodel(fit_full, cvrefbuilder = cvrefbuilder_user)
+  expect_identical(refmod$cvrefbuilder, cvrefbuilder_user)
+  # The default `cvfun` is kept:
+  expect_identical(body(refmod$cvfun), body(refmod_dflt$cvfun))
+})
+
+test_that(paste(
+  "get_refmodel.stanreg() rejects a `cvfun` or `cvrefbuilder` that is not a",
+  "function"
+), {
+  tstsetups <- grep("^rstanarm\\.", names(fits), value = TRUE)
+  skip_if_not(length(tstsetups) >= 1)
+  fit <- fits[[tstsetups[1]]]
+  expect_error(get_refmodel(fit, cvfun = "a"),
+               "^Argument `cvfun` must be a function or `NULL`\\.$")
+  expect_error(get_refmodel(fit, cvrefbuilder = 1),
+               "^Argument `cvrefbuilder` must be a function or `NULL`\\.$")
+})
+
 # predict.refmodel() ------------------------------------------------------
 
 test_that("invalid `type` fails", {