add property weights to KM

Author: bblodfon

R/LearnerSurvKaplan.R

@@ -19,7 +19,7 @@ LearnerSurvKaplan = R6Class("LearnerSurvKaplan",
         id = "surv.kaplan",
         predict_types = c("crank", "distr"),
         feature_types = c("logical", "integer", "numeric", "character", "factor", "ordered"),
-        properties = c("missings", "importance", "selected_features"),
+        properties = c("missings", "weights", "importance", "selected_features"),
         packages = c("survival", "distr6"),
         label = "Kaplan-Meier Estimator",
         man = "mlr3proba::mlr_learners_surv.kaplan"
@@ -59,9 +59,15 @@ LearnerSurvKaplan = R6Class("LearnerSurvKaplan",
 
   private = list(
     .train = function(task) {
-      list(model = invoke(survival::survfit, formula = task$formula(1),
-                          data = task$data(cols = task$target_names)),
-           features = task$feature_names) # keep for importance
+      fit = invoke(
+        survival::survfit,
+        formula = task$formula(1),
+        data = task$data(),
+        .args = list(weights = private$.get_weights(task))
+      )
+
+      # keep features for importance
+      list(model = fit, features = task$feature_names)
     },
 
     .predict = function(task) {

tests/testthat/test_mlr_learners_surv_kaplan.R

@@ -19,7 +19,10 @@ test_that("importance/selected", {
   learner = lrn("surv.kaplan")
   expect_error(learner$importance(), "No model stored")
   expect_error(learner$selected_features(), "No model stored")
-  learner$train(tsk("rats"))
+
+  task = tsk("rats")
+  learner$train(task)
   expect_character(learner$selected_features(), len = 0)
-  expect_named(learner$importance())
+  expect_named(learner$importance(), expected = task$feature_names)
+  expect_true(all(learner$importance() == 0))
 })