R/gsurv.R

In R4VN: Health Data Analysis and Publication-Ready Reporting

#' Survival and Cumulative Incidence Curves
#'
#' Draws Kaplan-Meier survival, cumulative risk, or competing-risk cumulative
#' incidence curves from a `tabsurv()` result. Uses the existing R4VN base-graphics
#' engine and can add confidence intervals, censor marks, log-rank p-values,
#' median lines, and number-at-risk tables. Examples below are self-contained.
#'
#' @param x An object returned by `tabsurv()`.
#' @param type NULL, "survival", "risk", or "cif". NULL chooses automatically.
#' @param xlab,ylab Axis titles retained for convenience.
#' @param xtitle,ytitle R4VN-style aliases for the x- and y-axis titles; when supplied they override `xlab`/`ylab`.
#' @param xlim,ylim Optional axis limits.
#' @param breaks Numeric spacing between x ticks, or explicit x tick positions.
#' @param percent Show y values as percentages.
#' @param ci Show confidence limits.
#' @param censor Show censor marks for ordinary Kaplan-Meier curves.
#' @param ci_color,ci_alpha,ci_linetype,ci_line_width Confidence-limit colour, transparency, line type, and line width. `ci_color = NULL` follows each curve colour.
#' @param censor_color,censor_pch,censor_size Censor-mark colour, point symbol, and size. `censor_color = NULL` follows each curve colour.
#' @param median_color,median_lty,median_lwd Colour, line type, and width for median-survival reference lines.
#' @param ref_color,ref_lty,ref_lwd Colour, line type, and width for user-specified `xline`/`yline` reference lines.
#' @param color,palette Curve colors; compatible with other R4VN graphs.
#' @param linetype Line types, recycled across groups.
#' @param line_width Curve line width.
#' @param size Base text size, consistent with other R4VN graph functions.
#' @param labels Optional replacement labels for curve groups.
#' @param legend TRUE/FALSE or a character legend title.
#' @param legend_position Base graphics legend position.
#' @param pvalue Annotate the log-rank p-value when available.
#' @param median Draw median-survival reference lines.
#' @param risk_table Add number-at-risk table below the plot.
#' @param risk_at Time points for the risk table. Defaults to plot ticks or `tabsurv(at=)`.
#' @param xline,yline Optional reference lines at x- and y-axis values.
#' @param vline,hline Deprecated aliases for `xline` and `yline`.
#' @param title,subtitle,note Plot annotations.
#' @param theme,file,width,height,dpi,show,bg Graph-output controls.
#'
#' @return An `r4vn_graph` object.
#' @examples
#' if (requireNamespace("survival", quietly = TRUE)) {
#'   d <- data.frame(
#'     time = c(5, 8, 10, 12, 15, 18, 20, 22),
#'     event = c(1, 0, 1, 1, 0, 1, 0, 1),
#'     group = factor(rep(c("A", "B"), 4))
#'   )
#'   s <- tabsurv(time, event, by = group, data = d, km = TRUE, show = FALSE)
#'   gsurv(s, show = FALSE)
#'   gsurv(s, type = "risk", ci = TRUE, pvalue = TRUE, show = FALSE)
#'   gsurv(s, risk_table = TRUE, risk_at = c(0, 10, 20), show = FALSE)
#' }
#' @export
gsurv <- function(x, type = NULL,
                  xlab = NULL, ylab = NULL, xtitle = NULL, ytitle = NULL,
                  xlim = NULL, ylim = NULL,
                  breaks = NULL, percent = TRUE,
                  ci = FALSE, censor = TRUE,
                  color = NULL, palette = "default", linetype = NULL,
                  line_width = 1.5,
                  ci_color = NULL, ci_alpha = 0.45, ci_linetype = 3, ci_line_width = NULL,
                  censor_color = NULL, censor_pch = 3, censor_size = 0.7,
                  median_color = "gray50", median_lty = 2, median_lwd = 1,
                  ref_color = "gray55", ref_lty = 3, ref_lwd = 1,
                  size = 11,
                  labels = NULL, legend = TRUE, legend_position = "topright",
                  pvalue = FALSE, median = FALSE,
                  risk_table = FALSE, risk_at = NULL,
                  xline = NULL, yline = NULL,
                  title = NULL, subtitle = NULL, note = NULL,
                  theme = "journal",
                  file = NULL, width = 7, height = NULL, dpi = 300,
                  show = TRUE, bg = "white",
                  vline = NULL, hline = NULL) {
  if (!inherits(x, "r4vn_surv")) stop("`x` must be an object returned by `tabsurv()`.", call. = FALSE)
  if (is.null(x$curve) || !nrow(x$curve)) stop("No survival/CIF curve is stored in this object. Run `tabsurv(..., km = TRUE)`.", call. = FALSE)
  call <- match.call()
  refs <- .r4vn_reference_lines(xline, yline, vline, hline)

  competing <- isTRUE(x$metadata$competing)
  if (is.null(type)) type <- if (competing) "cif" else "survival"
  type <- match.arg(tolower(type), c("survival", "risk", "cif"))
  if (type == "cif" && !competing) warning("`type = 'cif'` requested for ordinary survival; drawing cumulative risk 1-S(t).", call. = FALSE)

  dat <- x$curve
  if (type == "survival") {
    if (competing) stop("A single overall survival curve is not defined from the stored competing-risk CIF object. Use `type = 'cif'`.", call. = FALSE)
    dat$y <- dat$estimate; dat$lo <- dat$lower; dat$hi <- dat$upper
    default_y <- "Survival probability"
  } else if (type == "risk" && !competing) {
    dat$y <- 1 - dat$estimate; dat$lo <- 1 - dat$upper; dat$hi <- 1 - dat$lower
    default_y <- "Cumulative incidence"
  } else {
    dat$y <- dat$estimate; dat$lo <- dat$lower; dat$hi <- dat$upper
    default_y <- "Cumulative incidence"
  }

  groups <- unique(dat$group)
  ng <- length(groups)
  if (is.null(labels)) labels2 <- groups else {
    labels2 <- as.character(labels)
    if (!is.null(names(labels2)) && any(nzchar(names(labels2)))) {
      tmp <- groups; hit <- match(groups, names(labels2)); tmp[!is.na(hit)] <- labels2[hit[!is.na(hit)]]; labels2 <- tmp
    } else if (length(labels2) != ng) stop(sprintf("`labels` must have %d entries.", ng), call. = FALSE)
  }

  cols <- if (exists(".r4vn_colors", mode = "function", inherits = TRUE)) {
    .r4vn_colors(ng, color, palette, 1)
  } else {
    if (is.null(color)) grDevices::hcl.colors(ng, "Dark 3") else rep(as.character(color), length.out = ng)
  }
  lty <- if (is.null(linetype)) rep(1L, ng) else rep(linetype, length.out = ng)
  if (!is.numeric(line_width) || length(line_width) != 1L || is.na(line_width) || line_width <= 0) stop("`line_width` must be positive.", call. = FALSE)
  if (is.null(ci_line_width)) ci_line_width <- max(0.8, line_width * 0.7)
  if (!is.numeric(ci_line_width) || length(ci_line_width) != 1L || is.na(ci_line_width) || ci_line_width <= 0) stop("`ci_line_width` must be positive.", call. = FALSE)
  if (!is.numeric(ci_alpha) || length(ci_alpha) != 1L || is.na(ci_alpha) || ci_alpha < 0 || ci_alpha > 1) stop("`ci_alpha` must be between 0 and 1.", call. = FALSE)
  if (!is.numeric(censor_size) || length(censor_size) != 1L || is.na(censor_size) || censor_size <= 0) stop("`censor_size` must be positive.", call. = FALSE)
  if (!is.numeric(median_lwd) || length(median_lwd) != 1L || is.na(median_lwd) || median_lwd <= 0) stop("`median_lwd` must be positive.", call. = FALSE)
  if (!is.numeric(ref_lwd) || length(ref_lwd) != 1L || is.na(ref_lwd) || ref_lwd <= 0) stop("`ref_lwd` must be positive.", call. = FALSE)
  ci_cols <- if (is.null(ci_color)) vapply(cols, grDevices::adjustcolor, character(1), alpha.f = ci_alpha) else .r4vn_colors(ng, ci_color, palette, ci_alpha)
  censor_cols <- if (is.null(censor_color)) cols else .r4vn_colors(ng, censor_color, palette, 1)
  censor_pch <- rep(censor_pch, length.out = ng)
  ci_linetype <- rep(ci_linetype, length.out = ng)
  if (!is.numeric(size) || length(size) != 1L || is.na(size) || size <= 0) stop("`size` must be positive.", call. = FALSE)
  if (!is.null(breaks) && length(breaks) == 1L && (!is.numeric(breaks) || !is.finite(breaks) || breaks <= 0)) stop("A single `breaks` value must be a positive number.", call. = FALSE)

  xmax_data <- max(dat$time, na.rm = TRUE)
  if (is.null(xlim)) xlim <- c(0, xmax_data)
  if (length(breaks) == 1L && is.numeric(breaks)) {
    xb <- seq(ceiling(xlim[1L] / breaks) * breaks, floor(xlim[2L] / breaks) * breaks, by = breaks)
    if (!length(xb) || xb[1L] > xlim[1L]) xb <- c(xlim[1L], xb)
  } else if (is.numeric(breaks) && length(breaks) > 1L) xb <- breaks else xb <- pretty(xlim, n = 5)
  xb <- xb[xb >= xlim[1L] & xb <= xlim[2L]]

  if (is.null(ylim)) ylim <- c(0, 1)
  if (percent) {
    ylab_default <- default_y
    yticks <- pretty(ylim, n = 5)
    yticklabs <- paste0(round(yticks * 100), "%")
  } else {
    ylab_default <- default_y
    yticks <- pretty(ylim, n = 5); yticklabs <- format(yticks, trim = TRUE)
  }
  if (!is.null(xtitle)) xlab <- xtitle
  if (!is.null(ytitle)) ylab <- ytitle
  if (is.null(xlab)) {
    u <- x$metadata$unit
    xlab <- if (is.null(u)) "Follow-up time" else paste0("Follow-up time (", .r4vn_surv_unit_label(u, TRUE), ")")
  }
  if (is.null(ylab)) ylab <- ylab_default
  if (is.null(title)) title <- x$title
  legend_title <- if (is.character(legend) && length(legend) == 1L) legend else NULL
  show_legend <- !identical(legend, FALSE) && ng > 1L

  if (isTRUE(risk_table)) {
    if (is.null(risk_at)) {
      risk_at <- if (!is.null(x$metadata$at) && length(x$metadata$at)) sort(unique(c(0, x$metadata$at))) else xb
    }
    risk_at <- risk_at[risk_at >= xlim[1L] & risk_at <= xlim[2L]]
    ad <- x$analysis_data
    gp <- if (".by" %in% names(ad)) ad$.by else NULL
    rt <- .r4vn_surv_at_risk(ad, risk_at, gp)
  } else rt <- NULL

  if (is.null(height)) height <- if (risk_table) 6.4 else 5

  draw_main <- function() {
    graphics::plot(NA, NA, xlim = xlim, ylim = ylim, axes = FALSE,
                   xlab = xlab, ylab = ylab, main = title, xaxs = "i", yaxs = "i")
    graphics::axis(1, at = xb, labels = format(xb, trim = TRUE))
    graphics::axis(2, at = yticks, labels = yticklabs, las = 1)
    graphics::box(bty = if (theme %in% c("clean", "minimal")) "l" else "o")
    .r4vn_add_reference_lines(
      xline = refs$xline, yline = refs$yline,
      color = ref_color, lty = ref_lty, lwd = ref_lwd
    )

    for (i in seq_along(groups)) {
      z <- dat[dat$group == groups[i], , drop = FALSE]
      if (isTRUE(ci)) {
        graphics::lines(z$time, z$lo, type = "s", col = ci_cols[i], lty = ci_linetype[i], lwd = ci_line_width)
        graphics::lines(z$time, z$hi, type = "s", col = ci_cols[i], lty = ci_linetype[i], lwd = ci_line_width)
      }
      graphics::lines(z$time, z$y, type = "s", col = cols[i], lty = lty[i], lwd = line_width)
      if (isTRUE(censor) && !competing && type %in% c("survival", "risk")) {
        cc <- z$n.censor > 0 & z$time > 0
        if (any(cc, na.rm = TRUE)) graphics::points(z$time[cc], z$y[cc], pch = censor_pch[i], cex = censor_size, col = censor_cols[i])
      }
    }

    if (isTRUE(median) && !competing && type == "survival" && !is.null(x$median)) {
      graphics::abline(h = .5, lty = median_lty, lwd = median_lwd, col = median_color)
      med <- x$median$median[is.finite(x$median$median)]
      if (length(med)) graphics::abline(v = med, lty = median_lty, lwd = median_lwd, col = median_color)
    }

    if (show_legend) {
      leg <- labels2
      if (!is.null(legend_title)) leg <- paste0(labels2)
      graphics::legend(legend_position, legend = leg, title = legend_title,
                       col = cols, lty = lty, lwd = line_width, bty = "n", cex = size / 12)
    }
    if (isTRUE(pvalue) && !is.null(x$logrank) && nrow(x$logrank)) {
      txt <- paste0("Log-rank p ", .r4vn_surv_fmt_p(x$logrank$p[1L], 3))
      graphics::legend("bottomleft", legend = txt, bty = "n", cex = size / 12)
    }
    if (!is.null(subtitle)) graphics::mtext(subtitle, side = 3, line = .25, adj = 0, cex = size / 12)
    if (!is.null(note) && !risk_table) graphics::mtext(note, side = 1, line = 3.5, adj = 0, cex = size / 13)
  }

  draw_risk <- function() {
    graphics::plot.new()
    graphics::plot.window(xlim = xlim, ylim = c(0.5, ng + 1.6), xaxs = "i", yaxs = "i")
    graphics::text(xlim[1L], ng + 1.3, labels = "Number at risk", adj = c(0, .5), font = 2, cex = size / 12)
    for (i in seq_along(groups)) {
      y <- ng - i + 1
      graphics::text(xlim[1L], y, labels = labels2[i], adj = c(0, .5), cex = size / 13)
      z <- rt[rt$group == groups[i], , drop = FALSE]
      z <- z[match(risk_at, z$time), , drop = FALSE]
      graphics::text(risk_at, rep(y, length(risk_at)), labels = z$n.risk, cex = size / 13)
    }
    graphics::axis(1, at = risk_at, labels = format(risk_at, trim = TRUE), tick = FALSE, line = -0.2, cex.axis = size / 13)
    if (!is.null(note)) graphics::mtext(note, side = 1, line = 2.0, adj = 0, cex = size / 13)
  }

  draw <- function() {
    old <- graphics::par(no.readonly = TRUE); on.exit(graphics::par(old), add = TRUE)
    if (risk_table) {
      graphics::layout(matrix(c(1, 2), nrow = 2), heights = c(4.1, max(1.1, .55 * ng + .8)))
      graphics::par(mar = c(2.0, 4.7, 3.2, 1.2), mgp = c(2.6, .75, 0), las = 1,
                    cex.axis = size / 11, cex.lab = size / 11, cex.main = size / 10)
      draw_main()
      graphics::par(mar = c(2.5, 4.7, .2, 1.2))
      draw_risk()
    } else {
      if (exists(".r4vn_theme", mode = "function", inherits = TRUE)) {
        invisible(.r4vn_theme(theme, size, note))
      } else graphics::par(mar = c(if (is.null(note)) 4.2 else 5.2, 4.7, 3.2, 1.2), las = 1)
      draw_main()
    }
  }

  if (exists(".r4vn_render", mode = "function", inherits = TRUE)) {
    .r4vn_render(draw, file, show, width, height, dpi, bg)
  } else {
    if (!is.null(file)) {
      ext <- tolower(tools::file_ext(file))
      if (ext == "png") grDevices::png(file, width = width, height = height, units = "in", res = dpi, bg = bg)
      else if (ext == "pdf") grDevices::pdf(file, width = width, height = height, bg = bg)
      else stop("Without the R4VN graph engine, fallback export supports png or pdf only.", call. = FALSE)
      draw(); grDevices::dev.off()
    }
    if (show) draw()
  }

  outdata <- dat[, c("group", "time", "y", "lo", "hi", "n.risk", "n.event", "n.censor"), drop = FALSE]
  names(outdata)[names(outdata) == "y"] <- "estimate"
  names(outdata)[names(outdata) == "lo"] <- "lower"
  names(outdata)[names(outdata) == "hi"] <- "upper"
  if (exists(".r4vn_graph_result", mode = "function", inherits = TRUE)) {
    out <- .r4vn_graph_result(paste0("survival-", type), outdata, call, file, draw = draw)
    out$risk_table <- rt
    return(invisible(out))
  }
  invisible(structure(list(type = paste0("survival-", type), data = outdata, risk_table = rt,
                           call = call, file = file), class = "r4vn_graph"))
}