R/plots.R

In MultiFrailty: Shared Frailty Regression Models with Inverse Gaussian, Generalized Lindley, and Gamma Frailty Distributions

#' Diagnostic and Visualization Suite for MultiFrailty Models
#'
#' Provides specialized plotting capabilities for fitted shared frailty models,
#' including frailty density overlays across all five families, baseline hazard curves,
#' Kaplan-Meier survival curves, residual diagnostic plots, coefficient forest plots,
#' and influence diagnostic plots.
#'
#' @param fit A fitted object of class \code{"multifrailty_fit"}.
#'
#' @return No return value, called for side effects (displays a Kaplan-Meier survival curve with fitted MultiFrailty model overlay).
#'
#' @references
#' Pandey, A., Hanagal, D. D., & Tyagi, S. (2022). Shared Frailty Models Based on Cancer Data. International Journal of Statistics and Reliability Engineering, 9(3), 461-474.
#'
#' Pandey, A., & Tyagi, S. (2021). Comparison of Multiplicative Frailty Models Under Weibull Baseline Distribution. Lobachevskii Journal of Mathematics, 42(13), 3184-3195.
#'
#' @export
#' @examples
#' set.seed(123)
#' dat <- r_frailty(n = 60, baseline = "weibull", bpar = c(2, 1.5), frailty = "gamma", fpar = c(0.8))
#' fit <- fit_frailty(time = dat$time, status = dat$status, baseline = "weibull", frailty = "gamma")
#' plot_survival_km(fit)
plot_survival_km <- function(fit) {
  if (!inherits(fit, "multifrailty_fit")) stop("Argument 'fit' must be of class 'multifrailty_fit'.")

  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  km_fit <- survival::survfit(survival::Surv(fit$time, fit$status) ~ 1)
  t_grid <- seq(0.01, max(fit$time), length.out = 100)
  m_surv <- predict_frailty(fit, newtime = t_grid, type = "marginal")

  graphics::plot(km_fit, xlab = "Time", ylab = "Survival Probability",
                 main = paste0("Survival Curve (", fit$baseline, " + ", fit$frailty, ")"),
                 col = "darkgray", lwd = 1.5, mark.time = TRUE)
  graphics::lines(t_grid, m_surv, col = "blue", lwd = 2)
  graphics::legend("topright", legend = c("Kaplan-Meier", "Fitted MultiFrailty"),
                   col = c("darkgray", "blue"), lty = c(1, 1), lwd = c(1.5, 2))
}

#' Baseline Hazard Plot
#' @param baseline Baseline distribution.
#' @param par Baseline parameters.
#' @param t_range Time range vector.
#' @param n_grid Grid resolution.
#' @return No return value, called for side effects (displays baseline hazard and cumulative baseline hazard plots).
#' @export
plot_baseline <- function(baseline = c("weibull", "gw"), par, t_range = c(0.01, 3), n_grid = 300) {
  baseline <- match.arg(baseline)
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  t_vec <- seq(t_range[1], t_range[2], length.out = n_grid)
  bh <- baseline_hazard(t_vec, baseline = baseline, par = par)

  graphics::par(mfrow = c(1, 2))
  graphics::plot(t_vec, bh$h0, type = "l", col = "red", lwd = 2,
                 xlab = "Time", ylab = "phi0(t)", main = paste0("Baseline Hazard (", baseline, ")"))
  graphics::plot(t_vec, bh$H0, type = "l", col = "darkgreen", lwd = 2,
                 xlab = "Time", ylab = "Phi0(t)", main = paste0("Cumulative Baseline Hazard (", baseline, ")"))
}

#' Frailty Density Overlay Plot across Families
#' @param t_max Maximum grid value.
#' @return No return value, called for side effects (displays frailty density overlay curves across distributions).
#' @export
plot_frailty_density <- function(t_max = 3) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  w_grid <- seq(0.01, t_max, length.out = 300)

  d_gamma <- stats::dgamma(w_grid, shape = 1 / 0.8, scale = 0.8)
  d_ig <- sqrt(1 / (2 * pi * 0.8 * w_grid^3)) * exp(- (w_grid - 1)^2 / (2 * 0.8 * w_grid))

  p1_gl1 <- 1.2 / (1.2 + 0.5)
  d_gl1 <- p1_gl1 * stats::dgamma(w_grid, shape = 1 / 1.2, scale = 1.2) +
    (1 - p1_gl1) * stats::dgamma(w_grid, shape = 1 / 0.5, scale = 0.5)

  p1_gl2 <- 1.5 / (1 + 1.5)
  nu_gl2 <- 1.5 * (1 + 1.5 - 0.8)
  d_gl2 <- p1_gl2 * stats::dgamma(w_grid, shape = 0.8, rate = 1.5) +
    (1 - p1_gl2) * stats::dgamma(w_grid, shape = nu_gl2, rate = 1.5)

  graphics::plot(w_grid, d_gamma, type = "l", col = "blue", lwd = 2,
                 ylim = c(0, max(d_gamma, d_ig, d_gl1, d_gl2, na.rm = TRUE) * 1.1),
                 xlab = "w (Frailty)", ylab = "g(w)", main = "Frailty Density Comparison")
  graphics::lines(w_grid, d_ig, col = "red", lwd = 2, lty = 2)
  graphics::lines(w_grid, d_gl1, col = "darkgreen", lwd = 2, lty = 3)
  graphics::lines(w_grid, d_gl2, col = "purple", lwd = 2, lty = 4)

  graphics::legend("topright", legend = c("Gamma (theta=0.8)", "IG (eta=0.8)", "GL1 (eta=1.2, eps=0.5)", "GL2 (theta=1.5, mu=0.8)"),
                   col = c("blue", "red", "darkgreen", "purple"), lty = 1:4, lwd = 2)
}

#' Q-Q Plot of Cox-Snell Residuals vs Exp(1)
#' @param fit Fitted model.
#' @return No return value, called for side effects (displays a Q-Q plot of Cox-Snell residuals versus standard exponential distribution).
#' @export
plot_qq_residuals <- function(fit) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  res <- residuals_frailty(fit)
  cs <- sort(res$cox_snell[fit$status == 1])
  n_cs <- length(cs)

  if (n_cs > 0) {
    exp_quantiles <- stats::qexp(seq(1 / (n_cs + 1), n_cs / (n_cs + 1), length.out = n_cs))
    graphics::plot(exp_quantiles, cs, xlab = "Theoretical Exp(1) Quantiles", ylab = "Cox-Snell Residuals",
                   main = "Cox-Snell Residual Q-Q Plot", pch = 19, col = "darkblue")
    graphics::abline(0, 1, col = "red", lwd = 2)
  }
}

#' Residuals vs Fitted Plot
#' @param fit Fitted model.
#' @return No return value, called for side effects (displays a plot of martingale residuals against fitted survival probabilities).
#' @export
plot_residuals_fitted <- function(fit) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  res <- residuals_frailty(fit)
  graphics::plot(res$fitted_S, res$martingale, xlab = "Fitted Survival Probabilities", ylab = "Martingale Residuals",
                 main = "Residuals vs Fitted", pch = 19, col = "darkgray")
  graphics::abline(h = 0, col = "red", lty = 2)
}

#' Scale-Location Plot
#' @param fit Fitted model.
#' @return No return value, called for side effects (displays a scale-location plot of square-root absolute standardized residuals against fitted survival probabilities).
#' @export
plot_scale_location <- function(fit) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  res <- residuals_frailty(fit)
  graphics::plot(res$fitted_S, sqrt(abs(res$standardized)), xlab = "Fitted Survival Probabilities",
                 ylab = "sqrt(|Standardized Residuals|)", main = "Scale-Location", pch = 19, col = "darkgray")
}

#' Residuals vs Leverage Plot
#' @param fit Fitted model.
#' @return No return value, called for side effects (displays a plot of standardized residuals against leverage values).
#' @export
plot_residuals_leverage <- function(fit) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  res <- residuals_frailty(fit)
  inf <- influence_frailty(fit)

  graphics::plot(inf$leverage, res$standardized, xlab = "Leverage", ylab = "Standardized Residuals",
                 main = "Residuals vs Leverage", pch = 19, col = "darkblue")
  graphics::abline(h = 0, col = "red", lty = 2)
}

#' Leverage Index Plot
#' @param fit Fitted model.
#' @return No return value, called for side effects (displays an index plot of observation leverage values).
#' @export
plot_leverage <- function(fit) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  inf <- influence_frailty(fit)
  graphics::plot(1:fit$n, inf$leverage, type = "h", xlab = "Observation Index", ylab = "Leverage",
                 main = "Leverage Values", col = "darkgreen", lwd = 2)
}

#' DFBETAS Plot
#' @param fit Fitted model.
#' @return No return value, called for side effects (displays boxplots of DFBETAS values across parameters).
#' @export
plot_dfbetas <- function(fit) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  inf <- influence_frailty(fit)
  mat <- inf$dfbetas

  if (ncol(mat) > 0) {
    graphics::boxplot(mat, main = "DFBETAS across Parameters", col = "lightblue")
    graphics::abline(h = c(-2 / sqrt(fit$n), 2 / sqrt(fit$n)), col = "red", lty = 2)
  }
}

#' Coefficient Forest Plot
#' @param fit Fitted model.
#' @return No return value, called for side effects (displays a forest plot of parameter estimates with 95\% confidence intervals).
#' @export
plot_coef_forest <- function(fit) {
  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  df <- fit$coefficients
  p <- nrow(df)

  graphics::plot(df$Estimate, 1:p, yaxt = "n", xlab = "Estimate (with 95% CI)", ylab = "",
                 main = "Coefficient Forest Plot", pch = 19, xlim = c(min(df$CI_lower) - 0.5, max(df$CI_upper) + 0.5))
  graphics::axis(2, at = 1:p, labels = rownames(df), las = 2)
  graphics::abline(v = 0, col = "red", lty = 2)

  for (i in 1:p) {
    graphics::lines(c(df$CI_lower[i], df$CI_upper[i]), c(i, i), col = "blue", lwd = 2)
  }
}

#' Full Diagnostic Plot Suite
#'
#' Displays comprehensive multi-panel diagnostic plot suite.
#'
#' @param fit Fitted model.
#' @param ask Logical; if TRUE prompts user.
#' @param save_to_file Optional file path to save plot grid.
#' @return No return value, called for side effects (displays a 2x3 grid of diagnostic plots or saves them to a file).
#' @export
plot_all <- function(fit, ask = grDevices::dev.interactive(), save_to_file = NULL) {
  if (!inherits(fit, "multifrailty_fit")) stop("Argument 'fit' must be of class 'multifrailty_fit'.")

  oldpar <- graphics::par(no.readonly = TRUE)
  on.exit(graphics::par(oldpar))

  if (!is.null(save_to_file)) {
    grDevices::pdf(save_to_file, width = 10, height = 8)
    on.exit({
      grDevices::dev.off()
      graphics::par(oldpar)
    }, add = TRUE)
  }

  graphics::par(mfrow = c(2, 3))
  plot_survival_km(fit)
  plot_qq_residuals(fit)
  plot_residuals_fitted(fit)
  plot_scale_location(fit)
  plot_residuals_leverage(fit)
  plot_coef_forest(fit)
}