Skip to contents

Stratify a Learner by One or More Categorical Variables

Source: R/helpers.R
add_stratification.Rd

add_stratification() takes a learner and returns a new learner that, when trained, fits the original learner separately within each stratum of stratify_by and uses the stratum-specific fits for prediction. This is useful for enriching a super_learner() library with learners that do not smooth over pre-specified subgroups, e.g. when estimating nuisance regressions for subgroup-specific (heterogeneous treatment effect) parameters.

Usage

add_stratification(
  learner,
  stratify_by,
  min_stratum_size = 10,
  pooled_fallback = TRUE
)

Arguments

learner

A learner (see ?learners) to be fit separately within each stratum.

stratify_by

A character vector of one or more column names to stratify by. With more than one name, strata are the observed combinations of the variables (their interaction). Columns must be factor, character, logical, or integer-valued numeric; non-integer numeric columns error with a suggestion to discretize first.

min_stratum_size

Strata with fewer than this many observations in the training data are not given their own fit; they use the pooled fallback (or error if pooled_fallback = FALSE). Defaults to 10.

pooled_fallback

If TRUE (default), additionally fit the learner on the full training data and use that pooled fit for too-small strata, strata whose fits error, and strata unseen at training time. If FALSE, each of those situations errors.

Value

A new learner: a function taking (data, formula, weights = NULL, ...) and returning a prediction function of newdata. Predictions are returned in the row order of newdata. The new learner carries sl_lnr_name (e.g. "lm_stratified_by_cyl") and sl_lnr_type attributes, and inherits the wrapped learner's outcome_type_dependent_args.

Details

Formulas must not mention the stratifying variable(s). Within a stratum, the stratifying variable is constant, so including it as a predictor is at best degenerate and at worst an error for many learners. Rather than silently rewriting the user's formula, the stratified learner errors at training time if the formula mentions any of stratify_by anywhere in its right-hand side (including inside transformations like poly(), interactions, and (1 | x) terms). Note that y ~ . therefore errors, because its expansion includes the stratifying variable, which must be a column of the data; write y ~ . - stratifying_var or enumerate the predictors instead.

Pooled fallback. By default (pooled_fallback = TRUE), the original learner is also fit once on the full (unstratified) training data. This pooled fit is used:

  • at training time, for strata with fewer than min_stratum_size observations, and for strata in which the stratum-specific fit errors; and

  • at prediction time, for rows of newdata whose stratum was not present in the training data (which can happen for rare strata under cross-validation).

Each fallback signals a warning naming the affected strata; inside super_learner() these warnings are captured (not printed) and reported in the $warnings_from_* fields of the output. With pooled_fallback = FALSE, all of the above situations are errors instead, and inside super_learner() an erring learner is dropped from the ensemble by the usual machinery.

Cross-validation. When using stratified learners inside super_learner(), we strongly recommend passing the stratifying variable as strata_ids so that every cross-validation training fold contains observations from every stratum; e.g. super_learner(..., strata_ids = data$categorical_var). Note that strata_ids (which balances CV folds) and stratify_by (which stratifies a learner's fits) are related but distinct.

Supported outcome types. Currently only 'continuous' and 'binary' learners are supported (prediction functions returning one numeric value per row of newdata). Density and multiclass learners are not yet supported by add_stratification().

See also

learners add_screener super_learner

Examples

# a glm fit separately within each stratum of cyl:
lnr_lm_by_cyl <- add_stratification(lnr_lm,
  stratify_by = "cyl",
  min_stratum_size = 5
)
trained <- lnr_lm_by_cyl(mtcars, mpg ~ hp + wt)
trained(mtcars)
#>  [1] 21.62780 20.80085 25.55014 19.69825 16.78481 19.01473 15.54479 23.90637
#>  [9] 21.12359 18.67999 18.67999 15.34540 16.08533 15.97651 12.43558 11.92018
#> [17] 11.88701 28.60817 32.86786 30.56570 24.44635 16.95253 17.13752 14.95720
#> [25] 15.90342 29.96366 26.65191 27.86746 16.15551 19.69837 14.31422 21.74879

# note that formulas mentioning the stratifying variable error, including
# mpg ~ . (because its expansion includes cyl); instead write:
trained2 <- lnr_lm_by_cyl(mtcars, mpg ~ . - cyl)

# inside super_learner(), pass the stratifying variable as strata_ids so
# every training fold contains every stratum:
# \donttest{
sl <- super_learner(
  data = mtcars,
  formulas = c(.default = mpg ~ hp + wt, lm_stratified_by_cyl = mpg ~ hp + wt),
  learners = list(lnr_lm, lnr_lm_by_cyl),
  strata_ids = mtcars$cyl,
  n_folds = 3
)
# }