Adapt for upcoming flexsurv version 2.3 (#317)

* adapt for upcoming flexsurv version * add NEWS bullet * namespace `packageVersion()
tidymodels · Apr 19, 2024 · 66a9873 · 66a9873
1 parent ebda76c
commit 66a9873
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 72 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,8 @@
 # censored (development version)
 
+* Internal changes to the `predict()` methods for flexsurv models, in preparation for the upcoming flexsurv release (#317).
+
+
 # censored 0.3.0
 
 ## New features

diff --git a/R/survival_reg-data.R b/R/survival_reg-data.R
@@ -257,21 +257,7 @@ make_survival_reg_flexsurv <- function() {
     type = "hazard",
     value = list(
       pre = NULL,
-      post = function(pred, object) {
-        if (".pred" %in% names(pred)) {
-          pred %>%
-            dplyr::rowwise() %>%
-            dplyr::mutate(
-              .pred = list(dplyr::rename(.pred, .eval_time = .time))
-            ) %>%
-            dplyr::ungroup()
-        } else {
-          dplyr::rename(pred, .eval_time = .time) %>%
-            dplyr::mutate(.row = seq_len(nrow(pred))) %>%
-            tidyr::nest(.by = .row) %>%
-            dplyr::select(-.row)
-        }
-      },
+      post = flexsurv_post,
       func = c(fun = "predict"),
       args =
         list(
@@ -290,21 +276,7 @@ make_survival_reg_flexsurv <- function() {
     type = "survival",
     value = list(
       pre = NULL,
-      post = function(pred, object) {
-        if (".pred" %in% names(pred)) {
-          pred %>%
-            dplyr::rowwise() %>%
-            dplyr::mutate(
-              .pred = list(dplyr::rename(.pred, .eval_time = .time))
-            ) %>%
-            dplyr::ungroup()
-        } else {
-          dplyr::rename(pred, .eval_time = .time) %>%
-            dplyr::mutate(.row = seq_len(nrow(pred))) %>%
-            tidyr::nest(.by = .row) %>%
-            dplyr::select(-.row)
-        }
-      },
+      post = flexsurv_post,
       func = c(fun = "predict"),
       args =
         list(
@@ -442,21 +414,7 @@ make_survival_reg_flexsurvspline <- function() {
     type = "hazard",
     value = list(
       pre = NULL,
-      post = function(pred, object) {
-        if (".pred" %in% names(pred)) {
-          pred %>%
-            dplyr::rowwise() %>%
-            dplyr::mutate(
-              .pred = list(dplyr::rename(.pred, .eval_time = .time))
-            ) %>%
-            dplyr::ungroup()
-        } else {
-          dplyr::rename(pred, .eval_time = .time) %>%
-            dplyr::mutate(.row = seq_len(nrow(pred))) %>%
-            tidyr::nest(.by = .row) %>%
-            dplyr::select(-.row)
-        }
-      },
+      post = flexsurv_post,
       func = c(fun = "predict"),
       args =
         list(
@@ -475,21 +433,7 @@ make_survival_reg_flexsurvspline <- function() {
     type = "survival",
     value = list(
       pre = NULL,
-      post = function(pred, object) {
-        if (".pred" %in% names(pred)) {
-          pred %>%
-            dplyr::rowwise() %>%
-            dplyr::mutate(
-              .pred = list(dplyr::rename(.pred, .eval_time = .time))
-            ) %>%
-            dplyr::ungroup()
-        } else {
-          dplyr::rename(pred, .eval_time = .time) %>%
-            dplyr::mutate(.row = seq_len(nrow(pred))) %>%
-            tidyr::nest(.by = .row) %>%
-            dplyr::select(-.row)
-        }
-      },
+      post = flexsurv_post,
       func = c(fun = "predict"),
       args =
         list(

diff --git a/R/survival_reg-flexsurv.R b/R/survival_reg-flexsurv.R
@@ -0,0 +1,29 @@
+flexsurv_post <- function(pred, object) {
+  if (utils::packageVersion("flexsurv") < "2.3") {
+    pred <- flexsurv_rename_time(pred)
+  }
+
+  # if there's only one observation in new_data,
+  # flexsurv output isn't nested
+  if (!(".pred" %in% names(pred))) {
+    pred <- pred %>%
+      dplyr::mutate(.row = seq_len(nrow(pred))) %>%
+      tidyr::nest(.by = .row) %>%
+      dplyr::select(-.row)
+  }
+  pred 
+}
+
+flexsurv_rename_time <- function(pred){
+  if (".pred" %in% names(pred)) {
+    pred %>%
+      dplyr::rowwise() %>%
+      dplyr::mutate(
+        .pred = list(dplyr::rename(.pred, .eval_time = .time))
+      ) %>%
+      dplyr::ungroup()
+  } else {
+    pred %>%
+      dplyr::rename(.eval_time = .time)
+  }
+}
diff --git a/tests/testthat/test-survival_reg-flexsurvspline.R b/tests/testthat/test-survival_reg-flexsurvspline.R
@@ -61,12 +61,15 @@ test_that("survival probability prediction", {
     head(lung),
     type = "survival",
     times = c(0, 500, 1000)
-  ) %>%
-    dplyr::rowwise() %>%
-    dplyr::mutate(
-      .pred = list(dplyr::rename(.pred, .eval_time = .time))
-    ) %>%
-    dplyr::ungroup()
+  ) 
+  if (packageVersion("flexsurv") < "2.3") {
+    exp_pred <- exp_pred %>%
+      dplyr::rowwise() %>%
+      dplyr::mutate(
+        .pred = list(dplyr::rename(.pred, .eval_time = .time))
+      ) %>%
+      dplyr::ungroup()
+  }
 
   f_fit <- survival_reg() %>%
     set_engine("flexsurvspline", k = 1) %>%
@@ -281,12 +284,15 @@ test_that("hazard prediction", {
     head(lung),
     type = "hazard",
     times = c(0, 500, 1000)
-  ) %>%
-    dplyr::rowwise() %>%
-    dplyr::mutate(
-      .pred = list(dplyr::rename(.pred, .eval_time = .time))
-    ) %>%
-    dplyr::ungroup()
+  ) 
+  if (packageVersion("flexsurv") < "2.3") {
+    exp_pred <- exp_pred %>%
+      dplyr::rowwise() %>%
+      dplyr::mutate(
+        .pred = list(dplyr::rename(.pred, .eval_time = .time))
+      ) %>%
+      dplyr::ungroup()
+  }
 
   f_fit <- survival_reg() %>%
     set_engine("flexsurvspline", k = 1) %>%