R/polyreg-methods.R

In cifmodeling: Visualization and Polytomous Modeling of Survival and Competing Risks

#' @export
print.cifplot <- function(x, ...) {
  p <- x$plot %||% x$patchwork
  if (!is.null(p)) print(p)
  invisible(x)
}

#' @export
print.cifpanel <- function(x, ...) {
  p <- x$patchwork
  if (!is.null(p)) print(p)
  invisible(x)
}


#' Methods for polyreg objects
#'
#' S3 methods to extract coefficients, variance-covariance matrix,
#' sample size, formatted summaries, and tidy/glance/augment
#' from objects returned by `polyreg()`.
#'
#' @name polyreg-methods
#'
#' @param object A polyreg object returned by `polyreg()`.
#' @param type Character string; one of `"default"`, `"sandwich"`,
#'   or `"bootstrap"`. When `"default"`, the function chooses between
#'   sandwich and bootstrap variance based on the original `polyreg()`
#'   settings, using `outcome.type`, `report.sandwich.conf`, and
#'   `report.boot.conf`. (Used only by `vcov.polyreg()`.)
#' @param variance Explicit covariance choice, `"if"` or `"bootstrap"`.
#'   `NULL` preserves legacy selection. Explicit choices fail if unavailable.
#' @param conf.int Numeric interval level for `tidy.polyreg()`, defaulting to
#'   the level stored in the fit. A changed level uses normal Wald intervals;
#'   stored BCa intervals require refitting to change their level.
#'   A logical value is also accepted: `TRUE` includes confidence limits and
#'   `FALSE` omits them. Use `conf.level` to select the level in this form.
#' @param conf.level Optional confidence level for the broom-compatible
#'   logical `conf.int` interface. Defaults to the fitted level.
#' @param exponentiate Logical; exponentiate effect estimates and confidence
#'   limits for reporting. Standard errors and statistics stay on the log scale.
#' @param vcov Compatibility argument for `modelsummary`. Only `NULL` is
#'   supported by `tidy.polyreg()`. Use `variance` to choose the covariance estimator.
#' @param coef_rename Compatibility argument for `modelsummary`. Only `FALSE`
#'   is supported by `tidy.polyreg()`.
#' @param gof_map Table formatting argument forwarded by `modelsummary` and
#'   ignored by `glance.polyreg()`.
#' @param x Object to be printed or summarised. Typically a
#'   `"summary.polyreg"` object for `print.summary.polyreg()`, or
#'   a `"polyreg"` object for `tidy.polyreg()`, `glance.polyreg()`,
#'   `augment.polyreg()`, and `effect_label.polyreg()`.
#' @param digits Number of digits to print for parameter estimates
#'   or effect measures. Used by `print.summary.polyreg()` and
#'   `effect_label.polyreg()`.
#' @param event Character string indicating which event to extract.
#'   For `effect_label.polyreg()` and `glance.polyreg()` this is
#'   one of `"event1"` or `"event2"`. For `tidy.polyreg()` it can
#'   also be `"both"` to return rows for all events.
#' @param add.time.point Logical; if `TRUE`, `effect_label.polyreg()`
#'   appends the time point to the label (e.g., “at 5 years”).
#' @param add.outcome Logical; if `TRUE`, `effect_label.polyreg()`
#'   appends the outcome/event description (e.g., “of event 1”).
#' @param add.exposure.levels Logical; if `TRUE`, `effect_label.polyreg()`
#'   includes the exposure level in the label (e.g., treatment group).
#' @param add.conf Logical; if `TRUE`, `effect_label.polyreg()`
#'   includes a confidence interval in the label.
#' @param add.p Logical; if `TRUE`, `effect_label.polyreg()`
#'   includes a p-value or thresholded p-value (e.g. p < 0.05).
#' @param value.time Optional numeric value overriding the time point
#'   stored in the `"polyreg"` object when constructing labels in
#'   `effect_label.polyreg()`.
#' @param unit.time Optional character string giving the time unit
#'   to display in labels constructed by `effect_label.polyreg()`,
#'   such as `"years"` or `"months"`.
#' @param p_digits Integer; number of digits used to format p-values
#'   in `effect_label.polyreg()`.
#' @param p_cut Numeric threshold used by `effect_label.polyreg()`
#'   to decide between printing `p < p_cut` and an exact p-value.
#' @param ... Further arguments passed to or from methods.
#'
#' @return
#' \itemize{
#'   \item `coef.polyreg()` returns a numeric vector of regression
#'     coefficients.
#'   \item `vcov.polyreg()` returns a variance-covariance matrix.
#'   \item `nobs.polyreg()` returns the number of observations.
#'   \item `summary.polyreg()` returns a list of tidy and glance
#'     summaries by event.
#'   \item `print.summary.polyreg()` is called for its side effect
#'     of printing a formatted, modelsummary-like table to the
#'     console and returns `x` invisibly.
#'   \item `tidy.polyreg()` returns a data frame of tidy coefficients
#'     by event.
#'   \item `glance.polyreg()` returns a data frame of model-level
#'     summaries by event.
#'   \item `augment.polyreg()` returns an augmented data frame
#'     containing diagnostics, weights, and predicted CIFs.
#' }
#'
#' @seealso [polyreg()] for log odds product modeling of CIFs
#' @export
#' @rdname polyreg-methods
coef.polyreg <- function(object, ...) {
  object$coef
}

#' @export
#' @rdname polyreg-methods
vcov.polyreg <- function(object,
                         type = c("default", "sandwich", "bootstrap"),
                         variance = NULL,
                         ...) {
  output_check_dots(...)
  if (!is.null(variance)) {
    variance <- match.arg(variance, c("if", "bootstrap"))
    requested <- if (variance == "if") "sandwich" else "bootstrap"
    if (!missing(type) && !identical(match.arg(type), "default") && match.arg(type) != requested)
      stop("type and variance specify different covariance estimators")
    type <- requested
    value <- if (variance == "if") object[["vcov"]] else object$vcov_bootstrap
    if (is.null(value) || !is.matrix(value) || !length(value) || any(!is.finite(value)))
      stop("Requested variance is unavailable: ", variance)
    return(value)
  }
  type <- match.arg(type)

  outcome.type         <- object$outcome.type
  report.sandwich.conf <- object$boot.method$report.sandwich.conf
  report.boot.conf     <- object$boot.method$report.boot.conf

  V_sand <- object[["vcov"]]
  V_boot <- object$vcov_bootstrap

  ok_mat <- function(V) {
    !is.null(V) && is.matrix(V) && length(V) > 0 && any(is.finite(V))
  }

  if (type == "default") {
    if (outcome.type %in% c("competing-risk", "survival", "binomial")) {
      type <- if (isFALSE(report.sandwich.conf)) "bootstrap" else "sandwich"
    } else if (outcome.type %in% c("proportional-survival", "proportional-competing-risk")) {
      type <- if (isFALSE(report.boot.conf)) "sandwich" else "bootstrap"
    } else {
      type <- "sandwich"
    }
  }

  if (type == "sandwich") {
    if (ok_mat(V_sand)) return(V_sand)
    if (ok_mat(V_boot)) {
      warning(
        "Sandwich variance is not available; ",
        "returning bootstrap variance instead.",
        call. = FALSE
      )
      return(V_boot)
    }
    stop("No valid sandwich or bootstrap variance-covariance matrix.")
  }

  if (type == "bootstrap") {
    if (ok_mat(V_boot)) return(V_boot)
    if (ok_mat(V_sand)) {
      warning(
        "Bootstrap variance is not available; ",
        "returning sandwich variance instead.",
        call. = FALSE
      )
      return(V_sand)
    }
    stop("No valid bootstrap or sandwich variance-covariance matrix.")
  }

  stop("Internal error in vcov.polyreg: unsupported 'type'.")
}

#' @export
#' @rdname polyreg-methods
nobs.polyreg <- function(object, ...) {
  nrow(object$diagnostics)
}

#' @export
#' @rdname polyreg-methods
summary.polyreg <- function(object, ...) {
  s <- object$summary
  class(s) <- c("summary.polyreg", class(s))
  s
}

#' @export
#' @rdname polyreg-methods
print.summary.polyreg <- function(x,
                                  digits = 3,
                                  ...) {
  event_names <- names(x)

  cat("\n")
  cat(sprintf("%-20s", ""))
  for (ev in event_names) {
    cat(sprintf("  %-12s", ev))
  }
  cat("\n")

  cat(strrep("-", 20 + 2 + 12 * length(event_names)), "\n")

  terms_all <- unique(unlist(lapply(x, function(ev) ev$tidy$term)))

  for (term in terms_all) {

    cat(sprintf("%-20s", term), "\n")

    est_row  <- character(length(event_names))
    ci_row   <- character(length(event_names))
    p_row    <- character(length(event_names))

    for (j in seq_along(event_names)) {
      ev <- x[[event_names[j]]]
      td <- ev$tidy[ev$tidy$term == term, , drop = FALSE]
      if (nrow(td) == 0) {
        est_row[j] <- ""
        ci_row[j]  <- ""
        p_row[j]   <- ""
      } else {
        a <- exp(td$estimate[1])
        l <- exp(td$conf.low[1])
        h <- exp(td$conf.high[1])
        est_row[j] <- sprintf("%.*f", digits, a)
        ci_row[j]  <- sprintf("[%.*f, %.*f]",
                              digits, l,
                              digits, h)
        p_row[j]   <- sprintf("(p=%0.3f)", td$p.value[1])
      }
    }

    cat(sprintf("%-20s", ""))
    for (val in est_row) cat(sprintf("  %-12s", val))
    cat("\n")

    cat(sprintf("%-20s", ""))
    for (val in ci_row) cat(sprintf("  %-12s", val))
    cat("\n")

    cat(sprintf("%-20s", ""))
    for (val in p_row) cat(sprintf("  %-12s", val))
    cat("\n\n")
  }

  cat(strrep("-", 20 + 2 + 12 * length(event_names)), "\n\n")

  glance_rows <- c(
    "effect.measure",
    "n.events",
    "median.follow.up",
    "range.follow.up",
    "n.parameters",
    "converged.by",
    "nleqslv.message"
  )

  for (rowname in glance_rows) {
    cat(sprintf("%-20s", rowname))

    for (evnm in event_names) {
      gl  <- x[[evnm]]$glance
      val <- if (rowname %in% names(gl)) gl[[rowname]][1] else ""
      cat(sprintf("  %-12s", as.character(val)))
    }
    cat("\n")
  }

  cat("\n")
  invisible(x)
}

#' @export
#' @rdname polyreg-methods
effect_label.polyreg <- function(x,
                                 event               = c("event1", "event2"),
                                 add.time.point      = TRUE,
                                 add.outcome         = TRUE,
                                 add.exposure.levels = TRUE,
                                 add.conf            = TRUE,
                                 add.p               = TRUE,
                                 value.time          = NULL,
                                 unit.time           = NULL,
                                 digits              = 2,
                                 p_digits            = 2,
                                 p_cut               = 0.05,
                                 ...) {

  event <- match.arg(event)

  if (event == "event1") {
    td <- generics::tidy(x, event = "event1")
    eff <- x$estimand$effect.measure1
    ev_id <- 1L
  } else {
    td <- generics::tidy(x, event = if (x$outcome.type %in% c("competing-risk",
                                                              "proportional-competing-risk")) {
      "event2"
    } else {
      "event1"
    })
    eff  <- x$estimand$effect.measure2
    ev_id <- 2L
  }

  exposure <- x$exposure
  is_exp <- grepl(paste0("^", exposure, ","), td$term)
  td_exp <- td[is_exp, , drop = FALSE]

  if (nrow(td_exp) == 0L) {
    stop("No exposure effect terms were found in tidy(x) for this polyreg object.",
         call. = FALSE)
  }

  if (is.null(value.time)) {
    tp <- x$estimand$time.point
    if (length(tp) == 1L && is.finite(tp)) {
      value.time <- tp
    }
  }

  if (is.null(unit.time)) {
    unit.time <- if (!is.null(value.time)) "years" else ""
  }

  labels <- character(nrow(td_exp))

  for (i in seq_len(nrow(td_exp))) {
    est <- exp(td_exp$estimate[i])
    lcl <- exp(td_exp$conf.low[i])
    ucl <- exp(td_exp$conf.high[i])
    p   <- td_exp$p.value[i]

    level_str <- ""
    if (add.exposure.levels) {
      s <- td_exp$term[i]
      s <- sub(paste0("^", exposure, ",\\s*"), "", s)
      level_str <- s
    }

    eff_part <- eff
    if (add.exposure.levels && nzchar(level_str)) {
      eff_part <- sprintf("%s (%s)", eff, level_str)
    }

    prefix <- eff_part
    if (add.outcome) {
      prefix <- paste(prefix, sprintf("of event %d", ev_id))
    }

    if (add.time.point && !is.null(value.time)) {
      if (nzchar(unit.time)) {
        tp_txt <- sprintf("at %g %s", value.time, unit.time)
      } else {
        tp_txt <- sprintf("at %g", value.time)
      }
      prefix <- paste(prefix, tp_txt)
    }

    est_txt <- sprintf("%.*f", digits, est)
    parts   <- character(0)

    if (add.conf) {
      ci_txt <- sprintf("%.*f to %.*f", digits, lcl, digits, ucl)
      parts  <- c(parts, ci_txt)
    }

    if (add.p && !is.na(p)) {
      if (p < p_cut) {
        p_txt <- sprintf("p < %.*f", p_digits, p_cut)
      } else {
        p_txt <- sprintf("p = %.*f", p_digits, p)
      }
      parts <- c(parts, p_txt)
    }

    if (length(parts) > 0L) {
      inner <- sprintf("%s (%s)", est_txt, paste(parts, collapse = ", "))
    } else {
      inner <- est_txt
    }
    labels[i] <- sprintf("%s = %s", prefix, inner)
  }
  paste(labels, collapse = "\n")
}

#' @export
#' @rdname polyreg-methods
tidy.polyreg <- function(x,
                         event = c("event1", "event2", "both"),
                         conf.int = x$conf.int, variance = NULL,
                         exponentiate = FALSE, conf.level = NULL,
                         ..., vcov = NULL, coef_rename = FALSE) {
  output_check_dots(...)
  if (!is.null(vcov))
    stop("Custom vcov is not supported; use variance instead.", call. = FALSE)
  if (!identical(coef_rename, FALSE))
    stop("Only coef_rename = FALSE is supported by tidy.polyreg().", call. = FALSE)
  if (!is.logical(exponentiate) || length(exponentiate) != 1L || is.na(exponentiate))
    stop("exponentiate must be TRUE or FALSE")
  if (!is.null(conf.level)) output_check_level(conf.level, "conf.level")
  include_ci <- TRUE
  if (is.logical(conf.int)) {
    if (length(conf.int) != 1L || is.na(conf.int)) stop("conf.int must be TRUE or FALSE")
    include_ci <- conf.int
    conf.int <- conf.level %||% x$conf.int
  } else {
    output_check_level(conf.int)
    if (!missing(conf.int) && !is.null(conf.level) &&
        !isTRUE(all.equal(conf.int, conf.level)))
      stop("conf.int and conf.level specify different confidence levels")
    conf.int <- conf.level %||% conf.int
  }
  output_check_level(conf.int)
  if (!is.null(variance)) variance <- match.arg(variance, c("if", "bootstrap"))
  event <- match.arg(event)
  s     <- x$summary
  ot    <- x$outcome.type

  if (ot %in% c("competing-risk", "proportional-competing-risk")) {
    s1 <- s$event1
    s2 <- s$event2
  } else {
    s1 <- s[[1L]]
    s2 <- NULL
  }

  make_df <- function(part, label) {
    if (is.null(part)) return(NULL)
    df <- part$tidy
    recompute_ci <- include_ci &&
      (!is.null(variance) || !isTRUE(all.equal(conf.int, x$conf.int)))
    recompute <- !is.null(variance) || recompute_ci
    if (recompute) {
      index <- attr(df, "coefficient.index")
      if (is.null(index)) stop("Coefficient indices are unavailable; refit this legacy polyreg model")
      selected <- variance
      if (is.null(selected)) selected <- if (isTRUE(x$boot.method$report.sandwich.conf)) "if" else "bootstrap"
      if (recompute_ci && selected == "bootstrap" && isTRUE(x$boot.method$boot.bca) &&
          x$outcome.type %in% c("proportional-survival", "proportional-competing-risk"))
        stop("BCa intervals cannot be reconstructed from covariance; use the stored intervals or refit")
      V <- vcov.polyreg(x, variance = selected)
      beta <- if (selected == "bootstrap" && !is.null(x$bootstrap$boot.coef)) {
        x$bootstrap$boot.coef[index]
      } else x$coef[index]
      se <- sqrt(pmax(0, diag(V)[index]))
      df$estimate <- unname(beta)
      df$std.error <- unname(se)
      df$p.value <- 2 * stats::pnorm(abs(beta / se), lower.tail = FALSE)
      if (recompute_ci) {
        z <- stats::qnorm((1 + conf.int) / 2)
        df$conf.low <- unname(beta - z * se)
        df$conf.high <- unname(beta + z * se)
      }
    }
    df$event <- label
    df$statistic <- df$estimate / df$std.error
    if (exponentiate) {
      df$estimate <- exp(df$estimate)
      df$conf.low <- exp(df$conf.low)
      df$conf.high <- exp(df$conf.high)
    }
    if (!include_ci) df[c("conf.low", "conf.high")] <- NULL
    attr(df, "conf.int") <- conf.int
    attr(df, "interval") <- if (include_ci && any(is.finite(df$conf.low) & is.finite(df$conf.high)))
      "confidence" else "none"
    selected <- if (!is.null(variance)) variance else if
      (isTRUE(x$boot.method$report.sandwich.conf)) "if" else if
      (isTRUE(x$boot.method$report.boot.conf)) "bootstrap" else "none"
    attr(df, "variance.info") <- list(type = selected, source = if (selected == "if")
      "polyreg estimating-equation covariance" else selected,
      smoothing.uncertainty = "not applicable",
      interval.approximation = if (selected == "bootstrap" && isTRUE(x$boot.method$boot.bca) &&
          x$outcome.type %in% c("proportional-survival", "proportional-competing-risk"))
        "BCa" else "normal")
    df
  }

  if (event == "event1") {
    return(make_df(s1, "event1"))
  }

  if (event == "event2") {
    if (is.null(s2)) {
      stop("No competing event in this polyreg model.", call. = FALSE)
    }
    return(make_df(s2, "event2"))
  }

  out <- rbind(
    make_df(s1, "event1"),
    make_df(s2, "event2")
  )
  rownames(out) <- NULL
  return(out)
}

#' @export
#' @rdname polyreg-methods
glance.polyreg <- function(x,
                           event = c("event1", "event2"),
                           ..., gof_map = NULL) {
  output_check_dots(...)
  event <- match.arg(event)
  s     <- x$summary
  ot    <- x$outcome.type

  if (ot %in% c("competing-risk", "proportional-competing-risk")) {
    s1 <- s$event1
    s2 <- s$event2
  } else {
    s1 <- s[[1L]]
    s2 <- NULL
  }

  make_df <- function(part, label) {
    if (is.null(part)) return(NULL)
    df <- part$glance
    df$nobs <- if (is.data.frame(x$diagnostics)) nrow(x$diagnostics) else NA_integer_
    df$event <- label
    df
  }

  if (event == "event1") {
    return(make_df(s1, "event1"))
  }

  if (event == "event2") {
    if (is.null(s2)) {
      stop("No competing event in this polyreg model.", call. = FALSE)
    }
    return(make_df(s2, "event2"))
  }

  out <- rbind(
    make_df(s1, "event1"),
    make_df(s2, "event2")
  )
  rownames(out) <- NULL
  out
}

#' @export
#' @rdname polyreg-methods
augment.polyreg <- function(x, ...) {
  output_check_dots(...)
  df <- x$diagnostics

  if ("ip.weight" %in% names(df)) {
    names(df)[names(df) == "ip.weight"] <- ".weights"
  }

  if ("potential.CIFs" %in% names(df)) {
    cif <- df$potential.CIFs
    df$potential.CIFs <- NULL

    if (is.matrix(cif) || is.data.frame(cif)) {
      cif_df <- as.data.frame(cif)
      if (is.null(colnames(cif_df))) {
        colnames(cif_df) <- paste0("event", seq_len(ncol(cif_df)))
      }
      colnames(cif_df) <- paste0(".fitted_", colnames(cif_df))
      df <- cbind(df, cif_df)
    } else {
      df$.fitted <- cif
    }
  }

  if ("influence.function" %in% names(df)) {
    names(df)[names(df) == "influence.function"] <- ".influence"
  }
  df
}