edit transformed_data

mlr-org · Jul 30, 2024 · 3bc57fd · 3bc57fd
1 parent e2b471c
commit 3bc57fd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/R/PipeOpTaskSurvClassifDiscTime.R b/R/PipeOpTaskSurvClassifDiscTime.R
@@ -203,6 +203,7 @@ PipeOpTaskSurvClassifDiscTime = R6Class("PipeOpTaskSurvClassifDiscTime",
       # map observed times back
       reps = table(long_data$id)
       long_data$obs_times = rep(time, each = rows_per_id)
+      long_data = long_data[, .(id, obs_times, tend, ped_status)]
 
       list(task_disc, long_data)
     }

diff --git a/tests/testthat/test_pipeop_trafotask_survclassif_disctime.R b/tests/testthat/test_pipeop_trafotask_survclassif_disctime.R
@@ -49,7 +49,7 @@ test_that("PipeOpTaskSurvClassifDiscTime", {
   # test rows ids are correct
   expect_equal(transformed_data$id, correct_ids)
   # check columns in the transformed data.table
-  expect_equal(sort(c("id", "ped_status", "obs_times", pred_task$feature_names)),
+  expect_equal(sort(c("id", "ped_status", "obs_times", "tend")),
                sort(colnames(transformed_data)))
 
   # `ped_status` per interval and per observation is correct