Avoid repeated formula updates, term extraction and training-data model matrices

NEWS.md

@@ -7,6 +7,8 @@ If you read this from a place other than <https://mc-stan.org/projpred/news/inde
 
 ## Major changes
 
+* Less formula and model-matrix work in the forward search and in the projections: the candidate submodels are selected from precomputed term labels, and the predictions for the training data use the stored submodel model matrices. For example, `varsel()` with forward search for a Gaussian model with 50 observations and 200 predictor terms is about 1.5 times faster. (GitHub: #578)
+
 ## Minor changes
 
 ## Bug fixes and maintenance

R/divergence_minimizers.R

@@ -1158,6 +1158,27 @@ check_conv_s <- function(fit_s) {
 
 # Prediction functions for submodels --------------------------------------
 
+# Predictions of `subfit` objects for the data they were fitted to, computed
+# from the stored model matrices `x` (so no new model matrix is built). `nobs`
+# is the number of observations in the training data. Returns `NULL` if a fit
+# is not of class `subfit` or its `x` does not have `nobs` rows (for example,
+# after model.frame() removed rows with missing values).
+subprd_fitted <- function(fits, nobs) {
+  if (!all(vapply(fits, function(fit) {
+    inherits(fit, "subfit") && NROW(fit$x) == nobs
+  }, logical(1)))) {
+    return(NULL)
+  }
+  prd_list <- lapply(fits, function(fit) {
+    # As in predict.subfit():
+    if (is.null(fit$beta)) {
+      return(as.matrix(rep(fit$alpha, nobs)))
+    }
+    return(cbind(1, fit$x) %*% rbind(fit$alpha, fit$beta))
+  })
+  return(do.call(cbind, prd_list))
+}
+
 subprd <- function(fits, newdata) {
   prd_list <- lapply(fits, function(fit) {
     is_glmmPQL <- inherits(fit, "glmmPQL")

R/formula.R

@@ -158,9 +158,15 @@ validate_response_formula <- function(formula) {
   response <- terms_$response
 
   if (length(response) > 1) {
-    return(lapply(response, function(r) {
-      update(formula, paste0(r, " ~ ."))
-    }))
+    # update() rewrites the right-hand side in the same way for each response,
+    # so call it only once and then replace the response (for many draws, the
+    # update() calls are costly):
+    formula_1 <- update(formula, paste0(response[1], " ~ ."))
+    return(c(list(formula_1), lapply(response[-1], function(r) {
+      formula_r <- formula_1
+      formula_r[[2]] <- str2lang(r)
+      formula_r
+    })))
   } else {
     return(list(formula))
   }
@@ -645,26 +651,108 @@ count_terms_in_group_term <- function(term) {
   return(length(terms_))
 }
 
+# Term labels with the variables of an interaction in sorted order, so that, for
+# example, "x2:x1" and "x1:x2" are equal (as for terms()).
+norm_term_labels <- function(labels) {
+  vapply(strsplit(labels, ":", fixed = TRUE), function(parts) {
+    paste(sort(trimws(parts)), collapse = ":")
+  }, character(1))
+}
+
+# For each search term: its term labels (see norm_term_labels()) and, for a
+# search term with one label, the string that select_possible_terms_size()
+# returns for it. Computed once per search, so that
+# select_possible_terms_size() does not parse the same formulas again for each
+# candidate and step. Returns `NULL` if a search term has a smooth, group-level
+# or offset term or no intercept (the fast path does not cover these cases).
+search_term_info <- function(search_terms) {
+  info <- lapply(search_terms, function(x) {
+    if (identical(x, "1")) {
+      return(list(labels = character(0), label_out = "1"))
+    }
+    tt <- extract_terms_response(make_formula(x))
+    if (length(tt$additive_terms) > 0 || length(tt$group_terms) > 0 ||
+        length(tt$offset_terms) > 0 || !isTRUE(tt$global_intercept)) {
+      return(NULL)
+    }
+    lbls <- c(tt$individual_terms, tt$interaction_terms)
+    label_out <- NA_character_
+    if (length(lbls) == 1) {
+      label_out <- to_character_rhs(flatten_formula(make_formula(x)))
+    }
+    return(list(labels = norm_term_labels(lbls), label_out = label_out))
+  })
+  if (any(vapply(info, is.null, logical(1)))) {
+    return(NULL)
+  }
+  return(setNames(info, search_terms))
+}
+
+# select_possible_terms_size() for search terms without smooth, group-level
+# and offset terms (see search_term_info()). Without these terms, the size of a
+# candidate submodel is the number of its distinct term labels plus one for the
+# intercept, and a candidate with one new label is returned as that label.
+# Other candidates go through the formula code of select_possible_terms_size().
+select_possible_terms_size_simple <- function(chosen, terms, size, term_info,
+                                              terms_chosen) {
+  labels_chosen <- norm_term_labels(c(terms_chosen$individual_terms,
+                                      terms_chosen$interaction_terms))
+  is_valid <- vapply(terms, function(x) {
+    length(union(labels_chosen, term_info[[x]]$labels)) + 1L == size
+  }, logical(1))
+  valid <- terms[is_valid]
+  if (length(valid) == 0) {
+    return(NULL)
+  }
+  if (length(chosen) > 0) {
+    add_chosen <- paste0(" + ", paste(chosen, collapse = "+"))
+    remove_chosen <- paste0(" - ",
+                            paste(gsub("\\+", "-", chosen), collapse = "-"))
+  } else {
+    add_chosen <- ""
+    remove_chosen <- ""
+  }
+  out <- vapply(valid, function(x) {
+    info_x <- term_info[[x]]
+    if (!is.na(info_x$label_out) && !any(info_x$labels %in% labels_chosen)) {
+      return(info_x$label_out)
+    }
+    return(to_character_rhs(flatten_formula(make_formula(
+      paste(x, add_chosen, remove_chosen)
+    ))))
+  }, character(1))
+  return(unique(unname(out)))
+}
+
 ## Select next possible terms without surpassing a specific size
 ## @param chosen A list of currently chosen terms
 ## @param terms A list of all possible terms
 ## @param size Maximum allowed size
-select_possible_terms_size <- function(chosen, terms, size) {
+## @param term_info Optional output of search_term_info() for `terms`
+select_possible_terms_size <- function(chosen, terms, size, term_info = NULL) {
   if (size < 1) {
     stop("size must be at least 1")
   }
 
+  # The terms of the already chosen submodel (the same for all candidates):
+  terms_chosen <- extract_terms_response(make_formula(c(chosen)))
+  if (!is.null(term_info) && all(terms %in% names(term_info)) &&
+      length(terms_chosen$additive_terms) == 0 &&
+      length(terms_chosen$group_terms) == 0 &&
+      isTRUE(terms_chosen$global_intercept)) {
+    return(select_possible_terms_size_simple(chosen, terms, size, term_info,
+                                             terms_chosen))
+  }
+  additive <- unlist(regmatches(
+    terms_chosen$additive_terms,
+    gregexpr("(?<=\\().*?(?=\\))",
+             terms_chosen$additive_terms,
+             perl = TRUE)
+  ))
   valid_sub_trm_combs <- lapply(terms, function(x) {
     ## if we are adding a linear term whose smooth is already
     ## included, we reject it
-    terms <- extract_terms_response(make_formula(c(chosen)))
     terms_new <- extract_terms_response(make_formula(x))
-    additive <- unlist(regmatches(
-      terms$additive_terms,
-      gregexpr("(?<=\\().*?(?=\\))",
-               terms$additive_terms,
-               perl = TRUE)
-    ))
     linear <- terms_new$individual_terms
     dups <- setdiff(linear[!is.na(match(linear, additive))], chosen)
 

R/refmodel.R

@@ -1379,6 +1379,10 @@ init_refmodel <- function(object, data, formula, family, ref_predfun = NULL,
     }
   }
 
+  # With the default `proj_predfun` (and no augmented-data projection), the
+  # predictions for the training data can use the stored submodel model
+  # matrices (see mu_fun() below):
+  proj_predfun_is_subprd <- is.null(proj_predfun) && !family$for_augdat
   if (is.null(proj_predfun)) {
     if (!family$for_augdat) {
       proj_predfun <- subprd
@@ -1398,13 +1402,21 @@ init_refmodel <- function(object, data, formula, family, ref_predfun = NULL,
 
   mu_fun <- function(fits, obs = NULL, newdata = NULL, offset = NULL,
                      transform = TRUE) {
+    pred_sub <- NULL
+    if (proj_predfun_is_subprd && is.null(obs) && is.null(newdata)) {
+      # Training data: use the model matrices stored in the submodel fits
+      # instead of building them again for every projection:
+      pred_sub <- subprd_fitted(fits, nobs = nrow(data))
+    }
     newdata <- fetch_data(data, obs = obs, newdata = newdata)
     if (is.null(offset)) {
       offset <- rep(0, nrow(newdata))
     } else {
       stopifnot(length(offset) %in% c(1L, nrow(newdata)))
     }
-    pred_sub <- proj_predfun(fits, newdata = newdata)
+    if (is.null(pred_sub)) {
+      pred_sub <- proj_predfun(fits, newdata = newdata)
+    }
     if (family$family %in% fams_neg_linpred()) {
       pred_sub <- pred_sub - offset
     } else {

R/search.R

@@ -9,10 +9,12 @@ search_forward <- function(p_ref, refmodel, nterms_max, verbose = TRUE,
 
   chosen <- character()
   outdmins <- c()
+  term_info <- search_term_info(search_terms)
 
   for (size in seq_len(nterms_max_with_icpt)) {
     # Determine candidate predictors for the current size:
-    cands <- select_possible_terms_size(chosen, search_terms, size = size)
+    cands <- select_possible_terms_size(chosen, search_terms, size = size,
+                                        term_info = term_info)
     if (is.null(cands))
       next
     full_cands <- lapply(cands, function(cand) c(chosen, cand))

tests/testthat/test_formula_overhead.R

@@ -0,0 +1,49 @@
+context("formula and model-matrix shortcuts of the projections")
+
+test_that(paste(
+  "validate_response_formula() gives the same formulas as one update() per",
+  "response"
+), {
+  rhs_list <- c(
+    "x1 + x2", "x1 * f", "x1 + offset(o)", "x1 + (1 | g)",
+    "x1 + (x2 | g) + (1 | h)", "s(x1) + x2", "1"
+  )
+  resps <- paste0(".y_draw.", 1:4)
+  for (rhs in rhs_list) {
+    fml <- as.formula(paste0("cbind(", paste(resps, collapse = ", "), ") ~ ",
+                             rhs))
+    fmls_expected <- lapply(resps, function(r) update(fml, paste0(r, " ~ .")))
+    expect_identical(validate_response_formula(fml), fmls_expected,
+                     info = rhs)
+  }
+})
+
+test_that(paste(
+  "subprd_fitted() gives the training-data predictions of subprd() and",
+  "returns NULL if a model matrix has fewer rows than the data"
+), {
+  set.seed(8413)
+  n <- 30
+  dat <- data.frame(x1 = rnorm(n), x2 = rnorm(n),
+                    f = factor(sample(c("a", "b", "c"), n, replace = TRUE)))
+  dat$y1 <- dat$x1 - 0.5 * dat$x2 + (dat$f == "b") + rnorm(n, sd = 0.3)
+  dat$y2 <- dat$y1 + rnorm(n, sd = 0.1)
+  fmls <- list(y1 ~ x1 * f + x2, y2 ~ x1 * f + x2)
+  fits <- lapply(fmls, fit_glm_ridge_callback, data = dat,
+                 family = extend_family(gaussian()), weights = rep(1, n))
+  expect_equal(subprd_fitted(fits, nobs = n), subprd(fits, newdata = dat),
+               tolerance = 1e-12)
+  # Intercept-only submodels:
+  fits_icpt <- lapply(list(y1 ~ 1, y2 ~ 1), fit_glm_ridge_callback,
+                      data = dat, family = extend_family(gaussian()),
+                      weights = rep(1, n))
+  expect_equal(subprd_fitted(fits_icpt, nobs = n),
+               subprd(fits_icpt, newdata = dat), tolerance = 1e-12)
+  # If model.frame() removed rows with missing values, the stored model
+  # matrices have fewer rows than the data:
+  fits_short <- lapply(fits, function(fit) {
+    fit$x <- fit$x[-5, , drop = FALSE]
+    fit
+  })
+  expect_null(subprd_fitted(fits_short, nobs = n))
+})

tests/testthat/test_select_terms.R

@@ -0,0 +1,36 @@
+context("select_possible_terms_size() with precomputed term information")
+
+test_that(paste(
+  "select_possible_terms_size() gives the same candidates with and without",
+  "`term_info`"
+), {
+  term_sets <- list(
+    union("1", split_formula(
+      y ~ x1 + x2 + x3 + f + x1:x2 + poly(x4, 2) + I(x5^2) + x3:x6,
+      data = NULL
+    )),
+    c("1", "x1", "x2", "x1:x2", "x3", "f", "poly(x4,2)"),
+    c("1", paste0("x", 1:15))
+  )
+  for (st in term_sets) {
+    info <- search_term_info(st)
+    expect_false(is.null(info))
+    for (rep_i in 1:10) {
+      set.seed(rep_i)
+      chosen <- character()
+      for (size in 1:8) {
+        cands <- select_possible_terms_size(chosen, st, size)
+        cands_info <- select_possible_terms_size(chosen, st, size,
+                                                 term_info = info)
+        expect_identical(cands_info, cands)
+        if (is.null(cands)) next
+        chosen <- c(chosen, cands[sample.int(length(cands), 1)])
+      }
+    }
+  }
+})
+
+test_that("search_term_info() returns `NULL` for smooth or group-level terms", {
+  expect_null(search_term_info(c("1", "x1", "s(x2)")))
+  expect_null(search_term_info(c("1", "x1", "(1 | g)")))
+})