Stratify learner fits by a single variable

Format

R6Class object.

Value

Learner object with methods for training and prediction. See Lrnr_base for documentation on learners.

Parameters

learner="learner"

An initialized Lrnr_* object.

variable_stratify="variable_stratify"

character giving the variable in the covariates on which to stratify. Supports only variables with discrete levels coded as numeric.

...

Other parameters passed directly to learner$train. See its documentation for details.

Examples

library(data.table)

# load example data set
data(cpp_imputed)
setDT(cpp_imputed)

# use covariates of intest and the outcome to build a task object
covars <- c("apgar1", "apgar5", "sexn")
task <- sl3_Task$new(cpp_imputed, covariates = covars, outcome = "haz")

hal_lrnr <- Lrnr_hal9001$new(fit_control = list(n_folds = 3))
stratified_hal <- Lrnr_stratified$new(
  learner = hal_lrnr,
  variable_stratify = "sexn"
)

# stratified learner
set.seed(123)
stratified_hal_fit <- stratified_hal$train(task)
stratified_prediction <- stratified_hal_fit$predict(task = task)