tidymodels / censored

Parsnip wrappers for survival models
https://censored.tidymodels.org/
Other
123 stars 12 forks source link

Fix output format #284

Closed hfrick closed 9 months ago

hfrick commented 9 months ago

Closes #274 and closes #275

I introduce a wrapper survival_prob_pecRpart() because we want to move away from relying on eval_time added to the spec, see #236.

library(censored)
#> Loading required package: parsnip
#> Loading required package: survival

mod <- decision_tree() %>%
  set_mode("censored regression") %>%
  set_engine("rpart") %>%
  fit(Surv(time, status) ~ ., data = lung)

predict(mod, lung[2,], type = "survival", eval_time = c(100, 500)) %>% str()
#> tibble [1 × 1] (S3: tbl_df/tbl/data.frame)
#>  $ .pred:List of 1
#>   ..$ : tibble [2 × 2] (S3: tbl_df/tbl/data.frame)
#>   .. ..$ .eval_time    : num [1:2] 100 500
#>   .. ..$ .pred_survival: num [1:2] 0.75 0.0739

lung_orsf <- na.omit(lung)
mod <- rand_forest() %>%
  set_mode("censored regression") %>%
  set_engine("aorsf") %>%
  fit(Surv(time, status) ~ ., data = lung_orsf)

predict(mod, lung[2,], type = "survival", eval_time = c(100, 500)) %>% str()
#> tibble [1 × 1] (S3: tbl_df/tbl/data.frame)
#>  $ .pred:List of 1
#>   ..$ : tibble [2 × 2] (S3: tbl_df/tbl/data.frame)
#>   .. ..$ .eval_time    : num [1:2] 100 500
#>   .. ..$ .pred_survival: num [1:2] 0.932 0.35

Created on 2023-12-20 with reprex v2.0.2

github-actions[bot] commented 9 months ago

This pull request has been automatically locked. If you believe you have found a related problem, please file a new issue (with a reprex: https://reprex.tidyverse.org) and link to this issue.