Add model score prediction and survival curve plotting functionality

.Rbuildignore

@@ -17,4 +17,5 @@ copilot-instructions.md
 ^doc$
 ^Meta$
 ^icon$
+^examples$
 ^CRAN-SUBMISSION$

NAMESPACE

@@ -16,6 +16,7 @@ export(br_get_x)
 export(br_get_x2)
 export(br_get_y)
 export(br_pipeline)
+export(br_predict)
 export(br_rename_models)
 export(br_run)
 export(br_set_model)
@@ -29,6 +30,7 @@ export(br_show_forest_ggstats)
 export(br_show_forest_ggstatsplot)
 export(br_show_residuals)
 export(br_show_risk_network)
+export(br_show_survival_curves)
 export(br_show_table)
 export(br_show_table_gt)
 export(breg)
@@ -67,7 +69,13 @@ importFrom(rlang,is_string)
 importFrom(rlang,is_vector)
 importFrom(rlang,set_names)
 importFrom(stats,cor)
+importFrom(stats,pchisq)
+importFrom(stats,predict)
+importFrom(stats,quantile)
+importFrom(survival,Surv)
 importFrom(survival,coxph)
+importFrom(survival,survdiff)
+importFrom(survival,survfit)
 importFrom(utils,combn)
 importFrom(utils,packageDescription)
 importFrom(utils,packageVersion)

R/00-bregr-package.R

@@ -5,11 +5,11 @@
 #' @import vctrs
 #' @import cli
 #' @importFrom mirai .progress
-#' @importFrom survival coxph
+#' @importFrom survival Surv survdiff survfit coxph
 #' @importFrom lifecycle deprecated
 #' @importFrom rlang .data .env
 #' @importFrom utils combn packageDescription packageVersion
-#' @importFrom stats cor
+#' @importFrom stats cor pchisq quantile predict
 #' @rawNamespace if (getRversion() < "4.3.0") importFrom("S7", "@")
 ## usethis namespace: end
 NULL

R/02-pipeline.R

@@ -99,18 +99,19 @@
 #' # 3. Customized model -----------
 #' dt <- data.frame(x = rnorm(100))
 #' dt$y <- rpois(100, exp(1 + dt$x))
+#' \donttest{
 #' m5 <- breg(dt) |>
 #'   br_set_y("y") |>
 #'   br_set_x("x") |>
 #'   br_set_model(method = 'quasi(variance = "mu", link = "log")') |>
 #'   br_run()
+#' }
 #'
 #' @testexamples
 #' assert_breg_obj(m)
 #' assert_breg_obj(m2)
 #' assert_breg_obj(m3)
 #' assert_breg_obj(m4)
-#' assert_breg_obj(m5)
 #' @seealso [accessors] for accessing `breg` object properties.
 NULL
 
@@ -120,10 +121,9 @@ NULL
 br_pipeline <- function(
     data, y, x, method, x2 = NULL,
     group_by = NULL,
-    n_workers = 1L, run_parallel = 1L,
+    n_workers = 1L, run_parallel = lifecycle::deprecated(),
     model_args = list(),
     run_args = list()) {
-
   if (lifecycle::is_present(run_parallel)) {
     lifecycle::deprecate_warn("1.1.0", "bregr::br_run(run_parallel = )", "bregr::br_run(n_workers = )")
     n_workers <- run_parallel

R/03-accessors.R

@@ -21,6 +21,7 @@
 #' - `br_get_results()` returns modeling result `data.frame`.
 #'
 #' @name accessors
+#' @family accessors
 #' @seealso [pipeline] for building `breg` objects.
 #' @examples
 #' m <- br_pipeline(mtcars,
@@ -188,3 +189,83 @@ br_get_results <- function(obj, tidy = FALSE, ...) {
   }
   dplyr::filter(results, ...)
 }
+
+#' Predict method for breg objects
+#'
+#' @description
+#' `r lifecycle::badge('experimental')`
+#'
+#' Generate predictions from fitted models in a `breg` object.
+#' For Cox regression models, returns linear predictors (log relative hazard).
+#' For other models, returns predicted values.
+#'
+#' @param obj A `breg` object with fitted models.
+#' @param newdata Optional data frame for predictions. If NULL, uses original data.
+#' @param idx Model index, an integer or string.
+#' @param type Type of prediction. For Cox models: "lp" (linear predictor, default)
+#' or "risk" (relative risk). For other models: "response" (default) or "link".
+#' @returns Typically, a numeric vector of predictions.
+#' @export
+#' @family accessors
+#' @examples
+#' # Cox regression example
+#' if (requireNamespace("survival", quietly = TRUE)) {
+#'   lung <- survival::lung |> dplyr::filter(ph.ecog != 3)
+#'   mds <- br_pipeline(
+#'     lung,
+#'     y = c("time", "status"),
+#'     x = c("age", "ph.ecog"),
+#'     x2 = "sex",
+#'     method = "coxph"
+#'   )
+#'   scores <- br_predict(mds)
+#'   head(scores)
+#' }
+br_predict <- function(obj, newdata = NULL, idx = NULL, type = NULL) {
+  assert_breg_obj_with_results(obj)
+
+  # Get the model to use
+  if (is.null(idx)) {
+    cli::cli_inform("{.arg idx} not set, use the first model")
+    idx <- 1
+  } else {
+    if (length(idx) != 1) {
+      cli::cli_abort("please specify one model")
+    }
+  }
+  model <- br_get_models(obj, idx)
+
+  # Get data for prediction
+  if (is.null(newdata)) {
+    newdata <- br_get_data(obj)
+  }
+
+  # Determine prediction type
+  model_class <- class(model)[1]
+  if (is.null(type)) {
+    if (model_class == "coxph") {
+      # https://www.rdocumentation.org/packages/survival/versions/3.8-3/topics/predict.coxph
+      type <- "lp" # linear predictor (log relative hazard)
+    } else {
+      type <- "response" # predicted response
+    }
+    cli::cli_inform("{.arg type} is not specified, use {type} for the model")
+  }
+
+  # Generate predictions
+  tryCatch(
+    {
+      predictions <- predict(model, newdata = newdata, type = type)
+
+      # Handle missing values
+      if (any(is.na(predictions))) {
+        cli::cli_warn("Some predictions are NA, consider checking your data for missing values")
+      }
+
+      predictions
+    },
+    error = function(e) {
+      cli::cli_abort("Failed to generate predictions: {e$message}")
+    }
+  )
+}

R/04-show.R

@@ -527,6 +527,174 @@ br_show_table_gt <- function(
   t
 }
 
+#' Show survival curves based on model scores
+#'
+#' @description
+#' `r lifecycle::badge('experimental')`
+#'
+#' Generate survival curves by grouping observations based on model prediction scores.
+#' This function is specifically designed for Cox regression models and creates
+#' survival curves comparing different risk groups.
+#'
+#' @param breg A `breg` object with fitted Cox regression models.
+#' @param idx Index or name of the model to use for prediction.
+#' If NULL, uses the first model.
+#' @param n_groups Number of groups to create based on score quantiles. Default is 3.
+#' @param group_labels Custom labels for the groups. If NULL, uses "Low", "Medium", "High"
+#' for 3 groups or "Q1", "Q2", etc. for other numbers.
+#' @param title Plot title. If NULL, generates automatic title.
+#' @param subtitle Plot subtitle.
+#' @returns A ggplot2 object showing survival curves.
+#' @export
+#' @family br_show
+#' @examples
+#' \donttest{
+#' # Cox regression example with survival curves
+#' if (requireNamespace("survival", quietly = TRUE)) {
+#'   lung <- survival::lung |> dplyr::filter(ph.ecog != 3)
+#'   mds <- br_pipeline(
+#'     lung,
+#'     y = c("time", "status"),
+#'     x = c("age", "ph.ecog"),
+#'     x2 = "sex",
+#'     method = "coxph"
+#'   )
+#'   p <- br_show_survival_curves(mds)
+#'   print(p)
+#' }
+#' }
+br_show_survival_curves <- function(breg,
+                                    idx = NULL,
+                                    n_groups = 3,
+                                    group_labels = NULL,
+                                    title = NULL,
+                                    subtitle = NULL) {
+  assert_breg_obj_with_results(breg)
+
+  # Get the model to use
+  if (is.null(idx)) {
+    cli::cli_inform("{.arg idx} not set, use the first model")
+    idx <- 1
+  } else {
+    if (length(idx) != 1) {
+      cli::cli_abort("please specify one model")
+    }
+  }
+  model <- br_get_models(breg, idx)
+
+  # Check if we have survival models
+  if (!inherits(model, "coxph")) {
+    cli::cli_abort("survival curves are only supported for Cox regression models (coxph)")
+  }
+
+  # Get predictions (risk scores)
+  scores <- br_predict(breg, idx = idx, type = "lp")
+  data <- br_get_data(breg)
+
+  # Get survival variables
+  y_vars <- br_get_y(breg)
+  if (length(y_vars) != 2) {
+    cli::cli_abort("Expected 2 survival variables (time, status), got {length(y_vars)}")
+  }
+
+  time_var <- y_vars[1]
+  status_var <- y_vars[2]
+
+  # Create score groups
+  score_quantiles <- quantile(scores, probs = seq(0, 1, length.out = n_groups + 1), na.rm = TRUE)
+
+  # Handle case where scores have little variation
+  if (length(unique(score_quantiles)) <= 2) {
+    cli::cli_warn("Score values have limited variation, grouping may not be meaningful")
+  }
+
+  score_groups <- cut(scores,
+    breaks = score_quantiles,
+    include.lowest = TRUE,
+    labels = FALSE
+  )
+
+  # Create group labels
+  if (is.null(group_labels)) {
+    if (n_groups == 3) {
+      group_labels <- c("Low Risk", "Medium Risk", "High Risk")
+    } else if (n_groups == 2) {
+      group_labels <- c("Low Risk", "High Risk")
+    } else {
+      group_labels <- paste0("Q", 1:n_groups)
+    }
+  } else if (length(group_labels) != n_groups) {
+    cli::cli_abort("Length of group_labels ({length(group_labels)}) must match n_groups ({n_groups})")
+  }
+
+  # Prepare data for survival analysis
+  surv_data <- data |>
+    dplyr::mutate(
+      .score = scores,
+      .score_group = factor(score_groups, labels = group_labels),
+      .time = .data[[time_var]],
+      .status = .data[[status_var]]
+    ) |>
+    dplyr::filter(!is.na(.data$.score_group))
+
+  # Create survival object
+  surv_obj <- survival::Surv(surv_data$.time, surv_data$.status)
+
+  # Fit Kaplan-Meier curves for each group
+  km_fit <- survival::survfit(surv_obj ~ .score_group, data = surv_data)
+
+  # Create survival curve data for plotting
+  surv_summary <- summary(km_fit)
+
+  plot_data <- data.frame(
+    time = surv_summary$time,
+    surv = surv_summary$surv,
+    group = surv_summary$strata,
+    upper = surv_summary$upper,
+    lower = surv_summary$lower
+  )
+
+  # Clean group names
+  plot_data$group <- gsub(".*=", "", plot_data$group)
+
+  # Create the plot
+  if (is.null(title)) {
+    focal_var <- names(br_get_models(breg))[idx %||% 1]
+    title <- paste("Survival Curves by", focal_var, "Model Score")
+  }
+
+  p <- ggplot2::ggplot(plot_data, ggplot2::aes(x = .data$time, y = .data$surv, color = .data$group)) +
+    ggplot2::geom_step(linewidth = 1) +
+    ggplot2::geom_ribbon(ggplot2::aes(ymin = .data$lower, ymax = .data$upper, fill = .data$group),
+      alpha = 0.3, color = NA
+    ) +
+    ggplot2::scale_y_continuous(limits = c(0, 1), labels = function(x) paste0(x * 100, "%")) +
+    ggplot2::labs(
+      title = title,
+      subtitle = subtitle,
+      x = "Time",
+      y = "Survival Probability",
+      color = "Risk Group",
+      fill = "Risk Group"
+    ) +
+    ggplot2::theme_minimal() +
+    ggplot2::theme(
+      legend.position = "bottom",
+      plot.title = ggplot2::element_text(hjust = 0.5),
+      plot.subtitle = ggplot2::element_text(hjust = 0.5)
+    )
+
+  # Add log-rank test p-value if requested
+  if (n_groups > 1) {
+    logrank_test <- survival::survdiff(surv_obj ~ .score_group, data = surv_data)
+    p_value <- pchisq(logrank_test$chisq, df = length(logrank_test$n) - 1, lower.tail = FALSE)
+
+    p <- p + ggplot2::labs(caption = paste("Log-rank test p-value:", format.pval(p_value, digits = 3)))
+  }
+
+  return(p)
+}
+
 #' Show residuals vs fitted plot for regression models
 #'
 #' @description

README.Rmd

@@ -113,7 +113,7 @@ mds <- br_pipeline(
 
 Run in parallel:
 
-```{r}
+```{r warning=FALSE}
 mds_p <- br_pipeline(
   lung,
   y = c("time", "status"),
@@ -183,6 +183,25 @@ mds2 <- br_pipeline(
 br_show_risk_network(mds2)
 ```
 
+#### Model Score Prediction and Survival Curves
+
+For Cox-PH models, you can generate model predictions (risk scores) and create survival curves grouped by these scores:
+
+```{r}
+# Generate model predictions
+scores <- br_predict(mds2, idx = "ph.ecog")
+head(scores)
+```
+
+```{r dpi=150, fig.height=6}
+# Create survival curves based on model scores
+br_show_survival_curves(
+  mds2,
+  idx = "ph.ecog",
+  n_groups = 3,
+  title = "Survival Curves by 'ph.ecog' Model Risk Score"
+)
+```
 
 ### Table