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

Summary

This PR removes formula and model-matrix work that the forward search and the projections repeat with the same result. varsel() with forward search for a Gaussian model with 50 observations and 200 predictor terms gets about 1.5 times faster. cv_varsel(validate_search = TRUE) runs the forward search once per fold, so the faster candidate selection matters most there. All results are identical to master.

This PR was assisted by Claude.

Overhead

In a profile of varsel() for a Gaussian model with N = 50 and P = 200 (forward search, 20 clusters, nterms_max = 65), with #574 applied, update.formula() took 32 % of the time and extract_terms_response() took 22 %. On master the same work happens, but the per-draw model matrices take a larger share of the time.

Changes

The PR has three commits. Each commit removes some work that master repeats with the same result.

1. Fewer update() calls and term extractions.

• validate_response_formula() splits a formula with a cbind() response into one formula per projected draw. It called update() once per draw. update() rewrites the right-hand side in the same way for each response, so the function now calls update() once and replaces the response in copies of the result.
• select_possible_terms_size() called extract_terms_response() for the already chosen terms once per candidate. The result is the same for all candidates, so the function now computes it once per search step.

2. Training-data predictions from the stored model matrices. mu_fun() (in init_refmodel()) predicts the training data in every projection and built the training-data model matrix again each time, although each subfit object stores it in x. With the default proj_predfun, mu_fun() now computes these predictions from the stored matrices with the new subprd_fitted(). The previous code runs for a custom proj_predfun, for the augmented-data projection, for fits of other classes, and if a stored matrix has fewer rows than the data (for example, after model.frame() removed rows with missing values).

3. Candidate selection from precomputed term labels. select_possible_terms_size() parsed the formula of every candidate in every step. If no search term has a smooth, group-level or offset term, the size of a candidate submodel is the number of distinct term labels plus one for the intercept. The new search_term_info() computes the term labels of the search terms once per search, and search_forward() passes them to select_possible_terms_size() (new argument term_info). A candidate with one new term label is then returned as that label without parsing; other candidates go through the previous code. For 30 search terms and 11 chosen terms, one call takes 0.9 ms instead of 18 ms.

Additional details:

In each step of the forward search, select_possible_terms_size() takes the chosen terms and all search terms and returns the candidates that make the submodel exactly one size larger. Master parses formulas for every candidate:

for each search term x:
  terms_new <- extract_terms_response(x)               # parse formula
  size_x    <- count_terms_chosen(c(chosen, x))         # parse and flatten formula
  if size_x == size:
    keep x
for each kept x:
  candidate <- flatten_formula(x + chosen - chosen)     # parse and flatten formula

For $T$ search terms, one step parses two formulas per search term and one more per kept candidate, so between $2T$ and $3T$ formulas. A forward search with $K$ steps multiplies this by $K$, and cv_varsel(validate_search = TRUE) repeats the search for every fold.

If no search term has a smooth, group-level or offset term, and every search term has an intercept, the size of a submodel is simple to compute. Let $L(x)$ be the set of term labels of search term $x$, with the variables of an interaction in sorted order, so that x2:x1 and x1:x2 are the same label. Then

$$ \mathrm{size}(\mathrm{chosen} \cup x) = \left| L(\mathrm{chosen}) \cup L(x) \right| + 1 , $$

where the $1$ counts the intercept. The new search_term_info() computes $L(x)$ once per search, and for a search term with one label also its output string. Each step then needs only set operations:

labels_chosen <- L(chosen)                              # once per step
for each search term x:
  if |labels_chosen ∪ L(x)| + 1 == size:
    if x has one label and that label is not in labels_chosen:
      candidate <- precomputed string of x              # no parsing
    else:
      candidate <- flatten_formula(x + chosen - chosen) # as in master

Compound search terms such as x1 + x2 + x1:x2 still go through the formula code for their output string, but their size check uses the labels. If a search term has a smooth, group-level or offset term, search_term_info() returns NULL, and the whole step runs as in master. test_select_terms.R compares both paths along random forward searches. For 30 search terms and 11 chosen terms, one step takes 0.9 ms instead of 18 ms.

Results

All results are identical to master (difference 0, identical rankings), also for a model with a factor, an interaction and a poly() term.

These timings were measured before commit 3, which makes the search faster still. For large n the gain is smaller, because the projections take most of the time:

Tests

The new test_select_terms.R compares select_possible_terms_size() with and without term_info along random forward searches for three sets of search terms (compound terms with interactions, a factor, poly() and I() terms; a poly() term written without a space; 15 numeric terms). It also checks that search_term_info() returns NULL for smooth and group-level terms.

The new test_formula_overhead.R checks that:

• validate_response_formula() gives the same formulas as one update() per response, for seven formula types (interactions, a factor, an offset, group-level terms, a smooth term, intercept only);
• subprd_fitted() gives the same training-data predictions as subprd() for a model with a factor and an interaction and for intercept-only submodels, and returns NULL if a stored model matrix has fewer rows than the data.

The tests of search_term_info() and subprd_fitted() fail on master (functions not found); the validate_response_formula() test also passes on master, because the change keeps the results. The full suite passes with NOT_CRAN=true: 257 tests, 56 061 expectations, 0 failures, the same 4 warnings as master. Tested with this change applied to master 73cf0ee (projpred 2.11.0.9000), R 4.6.1, brms 2.23.0, rstanarm 2.32.2, cmdstanr 0.9.0.9001, CmdStan 2.39.0, on CSC Roihu (Linux, AMD EPYC).