R/predict.R

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

#' Predictions for MultiFrailty Regression Models
#'
#' Computes survival probabilities, hazard rates, median survival times, expected survival times,
#' risk scores, marginal survival curves, or future survival forecasts for a fitted \code{multifrailty} model.
#'
#' @param fit A fitted object of class \code{"multifrailty_fit"}.
#' @param newdata Optional data frame of new covariate values. If NULL, uses training data.
#' @param newtime Optional vector of evaluation time points. If NULL, uses default grid.
#' @param type Type of prediction: \code{"survival"}, \code{"hazard"}, \code{"median"}, \code{"expected"},
#'   \code{"risk"}, \code{"marginal"}, or \code{"forecast"}.
#' @param window Optional forecast window or horizon parameter.
#'
#' @return Vector or matrix of predictions depending on \code{type}.
#'
#' @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")
#' pred_surv <- predict_frailty(fit, type = "survival", newtime = c(1, 2, 3))
predict_frailty <- function(fit, newdata = NULL, newtime = NULL,
                            type = c("survival", "hazard", "median", "expected", "risk", "marginal", "forecast"),
                            window = NULL) {
  if (!inherits(fit, "multifrailty_fit")) {
    stop("Argument 'fit' must be of class 'multifrailty_fit'.")
  }
  type <- match.arg(type)

  b_type <- fit$baseline
  f_type <- fit$frailty
  coef_est <- fit$coefficients$Estimate

  n_base <- fit$n_par_base
  n_frail <- fit$n_par_frailty
  n_cov <- fit$n_cov

  bpar <- coef_est[1:n_base]
  fpar <- if (n_frail > 0) coef_est[(n_base + 1):(n_base + n_frail)] else numeric(0)
  beta <- if (n_cov > 0) coef_est[(n_base + n_frail + 1):(n_base + n_frail + n_cov)] else numeric(0)

  if (is.null(newdata)) {
    x_mat <- fit$x
  } else {
    if (n_cov == 0) {
      x_mat <- matrix(nrow = nrow(newdata), ncol = 0)
    } else {
      cov_names <- names(beta)
      if (all(cov_names %in% colnames(newdata))) {
        x_mat <- as.matrix(newdata[, cov_names, drop = FALSE])
      } else {
        x_mat <- as.matrix(newdata[, 1:n_cov, drop = FALSE])
      }
    }
  }

  n_sub <- max(1, nrow(x_mat))
  rho <- if (n_cov > 0) as.vector(exp(x_mat %*% beta)) else rep(1.0, n_sub)

  if (is.null(newtime)) {
    newtime <- if (type %in% c("survival", "hazard", "marginal")) seq(0.1, max(fit$time, 10), length.out = 100) else fit$time
  }

  if (type == "risk") {
    return(rho)
  }

  if (type %in% c("survival", "hazard")) {
    res_mat <- matrix(NA_real_, nrow = n_sub, ncol = length(newtime))
    colnames(res_mat) <- paste0("t=", round(newtime, 2))
    for (i in 1:n_sub) {
      ff <- frailty_functions(t = newtime, eta = rho[i], frailty = f_type, fpar = fpar,
                              baseline = b_type, bpar = bpar)
      res_mat[i, ] <- if (type == "survival") ff$S else ff$h
    }
    return(res_mat)
  }

  if (type == "marginal") {
    surv_m <- predict_frailty(fit, newdata = newdata, newtime = newtime, type = "survival")
    return(colMeans(surv_m))
  }

  if (type == "median") {
    med_vec <- numeric(n_sub)
    for (i in 1:n_sub) {
      fn_root <- function(t_val) {
        frailty_functions(t_val, eta = rho[i], frailty = f_type, fpar = fpar,
                          baseline = b_type, bpar = bpar)$S - 0.5
      }
      med_vec[i] <- tryCatch({
        stats::uniroot(fn_root, lower = 1e-4, upper = 1000)$root
      }, error = function(e) NA_real_)
    }
    return(med_vec)
  }

  if (type == "expected") {
    exp_vec <- numeric(n_sub)
    for (i in 1:n_sub) {
      exp_vec[i] <- tryCatch({
        stats::integrate(function(t_val) {
          frailty_functions(t_val, eta = rho[i], frailty = f_type, fpar = fpar,
                            baseline = b_type, bpar = bpar)$S
        }, lower = 1e-5, upper = 500)$value
      }, error = function(e) NA_real_)
    }
    return(exp_vec)
  }

  if (type == "forecast") {
    horizon <- if (is.null(window)) max(fit$time) + 5 else window
    grid_t <- seq(max(fit$time), max(fit$time) + horizon, length.out = 50)
    return(predict_frailty(fit, newdata = newdata, newtime = grid_t, type = "survival"))
  }
}

#' Risk Score Predictions
#' @param fit Fitted \code{multifrailty_fit} object.
#' @param newdata Data frame of new observations.
#' @param times Optional time grid.
#' @return Numeric vector of risk scores \code{exp(X \%*\% beta)}.
#' @export
risk_predict <- function(fit, newdata, times = NULL) {
  predict_frailty(fit, newdata = newdata, newtime = times, type = "risk")
}

#' Survival Probability at Specific Time Points
#' @param fit Fitted \code{multifrailty_fit} object.
#' @param times Vector of target time points.
#' @param newdata Optional new data frame.
#' @return Matrix of survival probabilities.
#' @export
survival_at <- function(fit, times, newdata = NULL) {
  predict_frailty(fit, newdata = newdata, newtime = times, type = "survival")
}

#' Future Survival Forecast
#' @param fit Fitted \code{multifrailty_fit} object.
#' @param horizon Forecast horizon time.
#' @param n_grid Number of evaluation points. Default 200.
#' @param newdata Optional new data frame.
#' @return Matrix of forecasted survival probabilities.
#' @export
forecast_frailty <- function(fit, horizon, n_grid = 200, newdata = NULL) {
  predict_frailty(fit, newdata = newdata, newtime = seq(max(fit$time), max(fit$time) + horizon, length.out = n_grid), type = "survival")
}