Keep the data-dependent bases

NEWS.md

@@ -11,6 +11,8 @@ If you read this from a place other than <https://mc-stan.org/projpred/news/inde
 
 ## Bug fixes and maintenance
 
+* Fixed bug in the predictions of submodels for new data: Terms with a data-dependent basis, such as `poly()` or `splines::ns()`, used a basis computed from the new data instead of the basis of the training data. This affected `proj_linpred()` and `proj_predict()` with `newdata`, `varsel()` with `d_test`, and K-fold CV, and caused `cv_varsel()` with `validate_search = TRUE` and LOO CV to fail with `poly()` terms. (GitHub: #573)
+
 
 # projpred 2.11.0
 

R/divergence_minimizers.R

@@ -215,8 +215,10 @@ fit_glm_ridge_callback <- function(formula, data,
   }
   # Post-processing:
   rownames(fit$beta) <- colnames(x)
+  # `terms` keeps the data-dependent bases of the training data (attribute
+  # "predvars"), for predictions for new data:
   sub <- nlist(alpha = fit$beta0, beta = fit$beta, w = fit$w, formula, x, y,
-               xlvls)
+               xlvls, terms = attr(fr, "terms"))
   class(sub) <- "subfit"
   return(sub)
 }
@@ -1274,6 +1276,29 @@ subprd_augdat_binom <- function(fits, newdata) {
   return(array(augprd, dim = c(nrow(augprd), 1L, ncol(augprd))))
 }
 
+# The terms of `object$formula` (without the response) for new data, with the
+# data-dependent bases of the training data (attribute "predvars" of
+# `object$terms`, for example the coefficients of poly()). Without them,
+# model.matrix() would compute such bases again from the new data. Returns the
+# plain terms if `object$terms` is missing or does not contain all variables.
+subfit_terms_newdata <- function(object) {
+  tt <- delete.response(terms(object$formula))
+  tt_train <- object$terms
+  if (is.null(tt_train) || is.null(attr(tt_train, "predvars"))) {
+    return(tt)
+  }
+  vars <- as.list(attr(tt, "variables"))[-1]
+  vars_train <- as.list(attr(tt_train, "variables"))[-1]
+  predvars_train <- as.list(attr(tt_train, "predvars"))[-1]
+  idx <- match(vapply(vars, deparse_combine, character(1)),
+               vapply(vars_train, deparse_combine, character(1)))
+  if (anyNA(idx)) {
+    return(tt)
+  }
+  attr(tt, "predvars") <- as.call(c(quote(list), predvars_train[idx]))
+  return(tt)
+}
+
 ## FIXME: find a way that allows us to remove this
 #' @noRd
 #' @export
@@ -1296,7 +1321,7 @@ predict.subfit <- function(object, newdata = NULL, ...) {
       # more convenient to let users specify contrasts directly. At that
       # occasion, contrasts should also be tested thoroughly (not done until
       # now).
-      x <- model.matrix(delete.response(terms(object$formula)), data = newdata,
+      x <- model.matrix(subfit_terms_newdata(object), data = newdata,
                         xlev = object$xlvls)
       # Need to ensure that the columns of `x` match the coefficients (note that
       # `rownames(beta)` is the same as `colnames(object$x)`, see the tests):

R/search.R

@@ -413,7 +413,8 @@ search_L1 <- function(p_ref, refmodel, nterms_max, penalty, search_control) {
     }
 
     sub <- nlist(alpha = search_path$alpha[nterms + 1], beta,
-                 w = search_path$w[, nterms + 1], formula, x, xlvls)
+                 w = search_path$w[, nterms + 1], formula, x, xlvls,
+                 terms = attr(fr, "terms"))
     class(sub) <- "subfit"
     return(list(sub))
   })

tests/testthat/helpers/testers.R

@@ -796,7 +796,8 @@ outdmin_tester_trad <- function(
         ]
       }
     }
-    subfit_nms <- c("alpha", "beta", "w", "formula", "x", "y", "xlvls")
+    subfit_nms <- c("alpha", "beta", "w", "formula", "x", "y", "xlvls",
+                    "terms")
     if (from_vsel_L1_search) {
       subfit_nms <- setdiff(subfit_nms, "y")
     }

tests/testthat/test_subfit_predvars.R

@@ -0,0 +1,74 @@
+context("submodel predictions for new data with data-dependent bases")
+
+# A Gaussian reference model with a term that has a data-dependent basis (default
+# poly()) and simulated draws of the coefficients. Its predictions for new data
+# use the terms of the training model frame (with "predvars"), so the reference
+# model itself has no problem with such terms.
+refm_predvars <- function(fml = y ~ x1 + x2 + poly(x3, 2), n = 80,
+                          ndraws = 200, seed = 6283) {
+  set.seed(seed)
+  dat <- data.frame(x1 = rnorm(n), x2 = rnorm(n), x3 = rnorm(n))
+  dat$y <- 0.8 * dat$x1 + 0.5 * dat$x3^2 + rnorm(n)
+  mf <- model.frame(fml, dat)
+  tt <- delete.response(attr(mf, "terms"))
+  mm <- model.matrix(tt, mf)
+  beta <- drop(solve(crossprod(mm), crossprod(mm, dat$y))) +
+    matrix(rnorm(ncol(mm) * ndraws, sd = 0.05), ncol(mm))
+  fit <- list(beta = beta, tt = tt, data = dat)
+  ref_predfun <- function(fit, newdata = NULL) {
+    nd <- if (is.null(newdata)) fit$data else newdata
+    model.matrix(fit$tt, model.frame(fit$tt, nd)) %*% fit$beta
+  }
+  suppressWarnings(init_refmodel(fit, data = dat, formula = fml,
+                                 family = gaussian(), ref_predfun = ref_predfun,
+                                 dis = rep(1, ndraws)))
+}
+
+test_that(paste(
+  "projected predictions for new data with a poly() term equal the",
+  "predictions for the same training observations"
+), {
+  refm <- refm_predvars()
+  prj <- project(refm, predictor_terms = c("x1", "poly(x3, 2)"), ndraws = 20,
+                 seed = 3906, verbose = FALSE)
+  pred_train <- proj_linpred(prj)$pred
+  pred_new <- proj_linpred(prj, newdata = refm$fetch_data()[1:10, ])$pred
+  expect_equal(unname(pred_new), unname(pred_train[, 1:10]), tolerance = 1e-10)
+})
+
+test_that(paste(
+  "projected predictions for new data with an ns() term equal the",
+  "predictions for the same training observations"
+), {
+  skip_if_not_installed("splines")
+  refm <- refm_predvars(y ~ x1 + x2 + splines::ns(x3, df = 3))
+  prj <- project(refm, predictor_terms = c("x1", "splines::ns(x3, df = 3)"),
+                 ndraws = 20, seed = 3906, verbose = FALSE)
+  pred_train <- proj_linpred(prj)$pred
+  pred_new <- proj_linpred(prj, newdata = refm$fetch_data()[1:10, ])$pred
+  expect_equal(unname(pred_new), unname(pred_train[, 1:10]), tolerance = 1e-10)
+})
+
+test_that(paste(
+  "predictions of L1-search submodels for new data with a poly() term equal",
+  "the predictions for the same training observations"
+), {
+  refm <- refm_predvars()
+  vs <- suppressWarnings(varsel(refm, method = "L1", refit_prj = FALSE,
+                                nterms_max = 3, seed = 3906, verbose = FALSE))
+  prj <- suppressWarnings(project(vs, nterms = 3, refit_prj = FALSE,
+                                  verbose = FALSE))
+  pred_train <- proj_linpred(prj)$pred
+  pred_new <- proj_linpred(prj, newdata = refm$fetch_data()[1:10, ])$pred
+  expect_equal(unname(drop(pred_new)), unname(drop(pred_train)[1:10]),
+               tolerance = 1e-10)
+})
+
+test_that("the LOO CV with validated search works with a poly() term", {
+  refm <- refm_predvars()
+  cv <- suppressWarnings(cv_varsel(
+    refm, cv_method = "LOO", validate_search = TRUE, nloo = 5, nclusters = 3,
+    ndraws_pred = 20, nterms_max = 3, seed = 3906, verbose = FALSE
+  ))
+  expect_length(ranking(cv)$fulldata, 3)
+})