[WIP] issue #139

NAMESPACE

@@ -41,6 +41,7 @@ S3method(mutate,ped)
 S3method(nest_tdc,default)
 S3method(nest_tdc,list)
 S3method(plot,pamm)
+S3method(predict,pamm)
 S3method(predictSurvProb,pamm)
 S3method(print,pamm)
 S3method(rename,nested_fdf)

R/pammfit.R

@@ -45,7 +45,9 @@ append_ped_attr <- function(pamm, ped) {
 #' pamm(
 #'  ped_status ~ s(tend) + complications,
 #'  data = tumor,
-#' trafo_args = list(formula = Surv(days, status)~complications))
+#'  trafo_args = list(
+#'    formula = Surv(days, status)~complications),
+#'    cut = seq(0, 3000, by = 50))
 #' @export
 pamm <- function(
   formula,

R/predict.R

@@ -1,3 +1,52 @@
+#' Predict hazard, cumulative hazard or survival probability
+#'
+#' @param object An object of class \code{pam_xgb}
+#' @param newdata A data set containing the same covariates as used for model
+#' fitting. If of class \code{data.frame}, the function will transform the
+#' data to the PED format using the same settings as for the data used in
+#' estimation.
+#' @param type The type of prediction desired. Either hazard (\code{type = "hazard"}),
+#' cumulative hazard (\code{type = "cumu_hazard"}) or survival probability
+#' (\code{type = "surv_prob"}).
+#' @param ... Currently not used.
+#' @importFrom stats predict
+#' @return A matrix of predictions containing one row per
+#' observation (row in newdata) and 1 column per  specified time in the
+#' \code{times} argument.
+#' @seealso pamm
+#' @export
+predict.pamm <- function(
+  object,
+  newdata,
+  type = c("hazard", "cumu_hazard", "surv_prob"),
+  ...) {
+
+  type <- match.arg(type)
+
+  if (!is.ped(newdata)) {
+    newdata <- as_ped(object, newdata)
+  }
+
+  newdata[["pred"]] <- predict(unpam(object), newdata, type = "response")
+  if (type == "cumu_hazard") {
+    newdata <- newdata %>%
+      group_by(.data$id) %>%
+      mutate(pred = cumsum(.data$pred * exp(.data$offset)))#TODO: is it correct to use offset here?
+  }
+  if (type == "surv_prob") {
+     newdata <- newdata %>%
+      group_by(.data[["id"]]) %>%
+      mutate(pred = exp(-cumsum(.data$pred * exp(.data$offset))))
+  }
+
+  newdata %>%
+    group_by(.data[["id"]]) %>%
+    filter(row_number() == n()) %>%
+    pull(.data[["pred"]]) # TODO: is the hazard/surv prob in the last available interval a useful return?
+
+}
+
+
 #' S3 method for pamm objects for compatibility with package pec
 #'
 #' @inheritParams pec::predictSurvProb

tests/testthat/test-predict-functions.R

@@ -9,8 +9,23 @@ context("Predict functions")
   pam2 <- pamm(ped_status ~ s(tend, k = 3) + complications, data = ped,
     engine = "bam", method = "fREML", discrete = TRUE)
 
+  ## predict S3 generic
+  predh <- predict(pam, newdata = tumor[21:23, ], type = "hazard")
+  predch <- predict(pam, newdata = tumor[21:23, ], type = "cumu_hazard")
+  predsp <- predict(pam, newdata = tumor[21:23, ], type = "surv_prob")
+  expect_identical(round(predh * 100, 2), c(.19, .02, .02))
+  expect_identical(round(predch, 2), c(.2, .96, .96))
+  expect_identical(round(predsp, 2), c(.81, .38, .38))
+
+  predh2 <- predict(pam2, newdata = tumor[21:23, ], type = "hazard")
+  predch2 <- predict(pam2, newdata = tumor[21:23, ], type = "cumu_hazard")
+  predsp2 <- predict(pam2, newdata = tumor[21:23, ], type = "surv_prob")
+  expect_identical(round(predh2 * 100, 2), c(.19, .02, .02))
+  expect_identical(round(predch2, 2), c(.2, .96, .96))
+  expect_identical(round(predsp2, 2), c(.81, .38, .38))
+
   ## predictSurvProb (pec) generic
-  spmat <- predictSurvProb.pamm(pam, tumor[21:23,], times = c(90, 500, 1217))
+  spmat <- predictSurvProb(pam, tumor[21:23,], times = c(90, 500, 1217))
   expect_identical(
     round(spmat, 2),
     matrix(
@@ -23,9 +38,9 @@ context("Predict functions")
       ncol = 3
     )
   )
-
-  expect_error(predictSurvProb.pamm(pam, tumor[21:23,], times = c(90, 500, 2000)))
-  spmat2 <- predictSurvProb.pamm(pam2, tumor[21:23,], times = c(90, 500, 1217))
+  expect_identical(spmat[1,1], predsp[1])
+  expect_error(predictSurvProb(pam, tumor[21:23,], times = c(90, 500, 2000)))
+  spmat2 <- predictSurvProb(pam2, tumor[21:23,], times = c(90, 500, 1217))
   expect_identical(round(spmat, 2), round(spmat2, 2))
 
   }