Earth offset

R/Lrnr_earth.R

@@ -68,27 +68,63 @@ Lrnr_earth <- R6Class(
   classname = "Lrnr_earth", inherit = Lrnr_base,
   portable = TRUE, class = TRUE,
   public = list(
-    initialize = function(degree = 2, penalty = 3, pmethod = "backward",
+    initialize = function(formula = NULL, degree = 2, penalty = 3, pmethod = "backward",
                           nfold = 0, ncross = 1, minspan = 0, endspan = 0,
                           ...) {
       params <- args_to_list()
       super$initialize(params = params, ...)
+    },
+    # for learners that take formula as an argument, the function
+    # process_formula that's defined in Lrnr_base needs to be redefined in
+    # the learner like below
+    process_formula = function(task) {
+      return(task)
     }
   ),
   private = list(
     .properties = c("continuous", "binomial"),
     .train = function(task) {
       args <- self$params
       outcome_type <- self$get_outcome_type(task)
-      args$x <- task$X
-      args$y <- outcome_type$format(task$Y)
+      # args$x <- task$X
+      # args$y <- outcome_type$format(task$Y)
+      args$data <- task$data
 
-      if (task$has_node("weights")) {
-        args$weights <- task$weights
+      if (!is.null(args$formula)) {
+        form <- args$formula
+        if (class(form) != "formula") form <- as.formula(form)
+
+        if (task$has_node("offset") && is.null(attr(terms(form), "offset"))) {
+          stop("Task has an offset; this needs to be specified as another term in the user-supplied formula")
+        }
+
+        # check response variable corresponds to outcome in task, if provided
+        if (attr(terms(form), "response")) {
+          if (!all.vars(form)[1] == task$nodes$outcome) {
+            stop(paste0(
+              "Outcome variable in formula ", all.vars(form)[1],
+              " does not match the task's outcome ", task$nodes$outcome
+            ))
+          }
+          formula_covars <- all.vars(form)[-1]
+        } else {
+          formula_covars <- all.vars(form)
+        }
+        # check that regressors in the formula are contained in the task
+        if (!all(formula_covars %in% task$nodes$covariates)) {
+          stop("Regressors in the formula are not covariates in task")
+        }
+      } else {
+        if (task$has_node("offset")) {
+          args$formula <- as.formula(paste(paste(task$nodes$outcome, paste("offset", "(", task$nodes$offset, ")", sep = ""), sep = "~"), paste(task$nodes$covariates, collapse = "+"), sep = "+"))
+        } else {
+          # create formula if it's not specified
+          args$formula <- as.formula(paste(task$nodes$outcome, paste(task$nodes$covariates, collapse = "+"), sep = "~"))
+        }
       }
 
-      if (task$has_node("offset")) {
-        args$offset <- task$offset
+      if (task$has_node("weights")) {
+        args$weights <- task$weights
       }
 
       if (outcome_type$type == "continuous") {
@@ -99,7 +135,7 @@ Lrnr_earth <- R6Class(
         stop("Unsupported outcome type for Lrnr_earth.")
       }
       args$glm <- glm
-      earth_fun <- utils::getS3method("earth", "default",
+      earth_fun <- utils::getS3method("earth", "formula",
         envir = getNamespace("earth")
       )
 
@@ -109,13 +145,14 @@ Lrnr_earth <- R6Class(
         utils::getS3method("earth", "fit", envir = getNamespace("earth"))
       ))
       extra_args <- extra_args[!(extra_args %in% default_args)]
+      # extra_args <- extra_args[extra_args != "offset"]
 
       fit_object <- call_with_args(earth_fun, args, other_valid = extra_args)
       return(fit_object)
     },
     .predict = function(task) {
       preds <- stats::predict(
-        object = private$.fit_object, newdata = task$X,
+        object = private$.fit_object, newdata = task$data,
         type = "response"
       )
       return(preds)

README.md

@@ -1312,6 +1312,56 @@ x
 </tr>
 <tr>
 <td style="text-align:left;">
+Lrnr\_glmtree
+</td>
+<td style="text-align:left;">
+√
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+√
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+√
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+x
+</td>
+<td style="text-align:left;">
+√
+</td>
+<td style="text-align:left;">
+x
+</td>
+</tr>
+<tr>
+<td style="text-align:left;">
 Lrnr\_grf
 </td>
 <td style="text-align:left;">

docs/pkgdown.yml

@@ -4,7 +4,7 @@ pkgdown_sha: ~
 articles:
   custom_lrnrs: custom_lrnrs.html
   intro_sl3: intro_sl3.html
-last_built: 2022-01-25T19:54Z
+last_built: 2022-02-09T01:10Z
 urls:
   reference: https://tlverse.org/sl3/reference
   article: https://tlverse.org/sl3/articles

tests/testthat/test-earth.R

@@ -31,3 +31,34 @@ test_that("Lrnr_earth produces results matching those of earth::earth", {
   # check equality of predictions
   expect_equal(preds, as.numeric(preds_classic))
 })
+
+test_that("Test Lrnr_earth earth offset matches earth::earth", {
+  covars <- c(
+    "apgar1", "apgar5", "parity", "gagebrth", "mage", "meducyrs",
+    "sexn"
+  )
+  outcome <- "haz"
+  task <- sl3_Task$new(cpp_imputed,
+    covariates = covars,
+    outcome = outcome,
+    offset = "waz"
+  )
+  # get predictions from Lrnr_* wrapper
+  set.seed(4738)
+  lrnr_earth <- make_learner(Lrnr_earth)
+  fit <- lrnr_earth$train(task)
+  preds <- fit$predict(task)
+
+  # get predictions from classic implementation
+  formula <- as.formula(paste(paste(task$nodes$outcome, paste("offset", "(", task$nodes$offset, ")", sep = ""), sep = "~"), paste(task$nodes$covariates, collapse = "+"), sep = "+"))
+  fit_classic <- earth::earth(
+    formula,
+    data = task$data, degree = 2, penalty = 3,
+    pmethod = "backward", nfold = 0, ncross = 1,
+    minspan = 0, endspan = 0
+  )
+  preds_classic <- predict(fit_classic, newdata = task$data, type = "response")
+
+  # check equality of predictions
+  expect_equal(preds, as.numeric(preds_classic))
+})