R/fit.R

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

#' Low-Level Maximum Likelihood Estimator for MultiFrailty Models
#'
#' Fits shared frailty regression models using Maximum Likelihood Estimation (MLE)
#' across all 10 baseline-frailty combinations with robust optimizer fallbacks.
#'
#' @param time Primary survival/censoring time vector.
#' @param status Event indicator vector (0 = right-censored, 1 = event, 2 = left-censored, 3 = interval-censored).
#' @param x Design matrix of covariates (n x p). Default is a 0-column matrix.
#' @param baseline Character string for baseline hazard: \code{"weibull"} or \code{"gw"}.
#' @param frailty Character string for frailty family: \code{"none"}, \code{"gamma"}, \code{"ig"}, \code{"gl1"}, or \code{"gl2"}.
#' @param time2 Vector of upper interval bounds when \code{status == 3}. Default is NULL.
#' @param prog_cen Vector of progressive censoring counts. Default is NULL.
#' @param init Vector of initial values on the estimation scale. Default is NULL (automatic).
#' @param method Optimization method passed to \code{maxLik}: \code{"NR"} (Newton-Raphson) or \code{"BFGS"}.
#' @param ... Additional arguments passed to optimization algorithms.
#'
#' @return An object of class \code{"multifrailty_fit"} containing parameter estimates, standard errors,
#'   information criteria, variance-covariance matrix, and diagnostic statistics.
#'
#' @references
#' Hougaard, P. (1984). Life table methods for heterogeneous populations: distributions of frailties. Biometrika, 71(1), 75-83.
#'
#' 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 = 100, baseline = "weibull", bpar = c(2, 1.5),
#'                  frailty = "gamma", fpar = c(0.8),
#'                  x = matrix(rnorm(100), ncol = 1), beta = 0.5)
#' fit <- fit_frailty(time = dat$time, status = dat$status, x = as.matrix(dat[, "X1", drop=FALSE]),
#'                    baseline = "weibull", frailty = "gamma")
#' print(fit)
fit_frailty <- function(time, status, x = matrix(nrow = length(time), ncol = 0),
                        baseline = c("weibull", "gw"), frailty = c("none", "gamma", "ig", "gl1", "gl2"),
                        time2 = NULL, prog_cen = NULL, init = NULL, method = "NR", ...) {
  baseline <- match.arg(baseline)
  frailty <- match.arg(frailty)

  n <- length(time)
  if (n == 0) stop("Time vector cannot be empty.")
  if (length(status) != n) stop("Length of 'status' must match length of 'time'.")
  if (!is.matrix(x)) x <- as.matrix(x)
  if (nrow(x) != n) stop("Number of rows in 'x' must match length of 'time'.")

  n_cov <- ncol(x)
  n_par_base <- if (baseline == "weibull") 2L else 3L

  if (frailty == "none") {
    n_par_frailty <- 0L
  } else if (frailty %in% c("gamma", "ig")) {
    n_par_frailty <- 1L
  } else if (frailty %in% c("gl1", "gl2")) {
    n_par_frailty <- 2L
  }

  n_total <- n_par_base + n_par_frailty + n_cov

  if (is.null(init)) {
    mean_t <- mean(time, na.rm = TRUE)
    if (baseline == "weibull") {
      base_init <- c(log(pmax(mean_t, 0.1)), 0)
    } else {
      base_init <- c(log(1 / pmax(mean_t, 0.1)), 0, 0)
    }

    if (frailty == "none") {
      frail_init <- numeric(0)
    } else if (frailty %in% c("gamma", "ig")) {
      frail_init <- c(log(0.5))
    } else if (frailty == "gl1") {
      frail_init <- c(log(0.5), log(0.5))
    } else if (frailty == "gl2") {
      frail_init <- c(log(1.0), stats::qlogis(0.25))
    }

    beta_init <- rep(0.0, n_cov)
    init <- c(base_init, frail_init, beta_init)
  }

  if (length(init) != n_total) {
    stop(paste0("Length of 'init' must be ", n_total, "."))
  }

  obj_fn <- function(p) {
    loglik_frailty(p, time = time, status = status, x = x, baseline = baseline,
                   frailty = frailty, time2 = time2, prog_cen = prog_cen)
  }

  opt_res <- NULL
  converged <- FALSE

  # Method 1: maxLik Newton-Raphson
  opt_res <- tryCatch({
    maxLik::maxLik(fn = obj_fn, start = init, method = "NR",
                   control = list(iterlim = 300))
  }, error = function(e) NULL)

  if (!is.null(opt_res) && is.finite(opt_res$maximum) && opt_res$maximum > -1e10) {
    converged <- TRUE
  } else {
    # Method 2: maxLik BFGS
    opt_res <- tryCatch({
      maxLik::maxLik(fn = obj_fn, start = init, method = "BFGS",
                     control = list(iterlim = 400))
    }, error = function(e) NULL)

    if (!is.null(opt_res) && is.finite(opt_res$maximum) && opt_res$maximum > -1e10) {
      converged <- TRUE
    } else {
      # Method 3: stats::optim BFGS
      opt_optim <- tryCatch({
        stats::optim(par = init, fn = function(p) -obj_fn(p), method = "BFGS", hessian = TRUE)
      }, error = function(e) NULL)

      if (!is.null(opt_optim) && is.finite(opt_optim$value) && opt_optim$value < 1e10) {
        converged <- TRUE
        opt_res <- list(
          estimate = opt_optim$par,
          maximum = -opt_optim$value,
          hessian = -opt_optim$hessian,
          code = opt_optim$convergence
        )
      } else {
        # Method 4: stats::optim Nelder-Mead
        opt_nm <- tryCatch({
          stats::optim(par = init, fn = function(p) -obj_fn(p), method = "Nelder-Mead", hessian = TRUE)
        }, error = function(e) NULL)

        if (!is.null(opt_nm) && is.finite(opt_nm$value) && opt_nm$value < 1e10) {
          converged <- TRUE
          opt_res <- list(
            estimate = opt_nm$par,
            maximum = -opt_nm$value,
            hessian = -opt_nm$hessian,
            code = opt_nm$convergence
          )
        }
      }
    }
  }

  if (is.null(opt_res)) {
    stop("Maximum likelihood estimation failed to run.")
  }

  raw_est <- opt_res$estimate
  max_ll <- opt_res$maximum

  hess <- opt_res$hessian
  if (is.null(hess) || any(!is.finite(hess))) {
    hess <- tryCatch(numDeriv::hessian(obj_fn, raw_est), error = function(e) matrix(NA, n_total, n_total))
  }

  vcov_mat <- tryCatch({
    inv_h <- solve(-hess)
    colnames(inv_h) <- rownames(inv_h) <- NULL
    inv_h
  }, error = function(e) {
    matrix(NA_real_, nrow = n_total, ncol = n_total)
  })

  se_raw <- sqrt(pmax(diag(vcov_mat), 0))

  nat_est <- numeric(n_total)
  par_names <- character(n_total)

  if (baseline == "weibull") {
    nat_est[1:2] <- exp(raw_est[1:2])
    par_names[1:2] <- c("lambda", "gamma")
  } else {
    nat_est[1:3] <- exp(raw_est[1:3])
    par_names[1:3] <- c("delta", "zeta", "xi")
  }

  if (frailty == "gamma") {
    nat_est[n_par_base + 1] <- exp(raw_est[n_par_base + 1])
    par_names[n_par_base + 1] <- "theta"
  } else if (frailty == "ig") {
    nat_est[n_par_base + 1] <- exp(raw_est[n_par_base + 1])
    par_names[n_par_base + 1] <- "eta"
  } else if (frailty == "gl1") {
    nat_est[n_par_base + (1:2)] <- exp(raw_est[n_par_base + (1:2)])
    par_names[n_par_base + (1:2)] <- c("eta", "epsilon")
  } else if (frailty == "gl2") {
    theta_val <- exp(raw_est[n_par_base + 1])
    p_val <- pmin(pmax(stats::plogis(raw_est[n_par_base + 2]), 1e-15), 1.0 - 1e-15)
    mu_val <- (1.0 + theta_val) * p_val
    nat_est[n_par_base + 1] <- theta_val
    nat_est[n_par_base + 2] <- mu_val
    par_names[n_par_base + (1:2)] <- c("theta", "mu")
  }

  if (n_cov > 0) {
    cov_names <- colnames(x)
    if (is.null(cov_names)) cov_names <- paste0("beta_", 1:n_cov)
    nat_est[(n_par_base + n_par_frailty + 1):n_total] <- raw_est[(n_par_base + n_par_frailty + 1):n_total]
    par_names[(n_par_base + n_par_frailty + 1):n_total] <- cov_names
  }

  nat_se <- numeric(n_total)
  for (i in 1:n_total) {
    if (i <= n_par_base) {
      nat_se[i] <- nat_est[i] * se_raw[i]
    } else if (i <= n_par_base + n_par_frailty) {
      if (frailty %in% c("gamma", "ig", "gl1")) {
        nat_se[i] <- nat_est[i] * se_raw[i]
      } else if (frailty == "gl2") {
        if (i == n_par_base + 1) {
          nat_se[i] <- nat_est[i] * se_raw[i]
        } else {
          g_fn <- function(p) (1.0 + exp(p[n_par_base + 1])) * pmin(pmax(stats::plogis(p[n_par_base + 2]), 1e-15), 1.0 - 1e-15)
          grad_g <- tryCatch(numDeriv::grad(g_fn, raw_est), error = function(e) rep(0, n_total))
          nat_se[i] <- sqrt(pmax(as.numeric(t(grad_g) %*% vcov_mat %*% grad_g), 0))
        }
      }
    } else {
      nat_se[i] <- se_raw[i]
    }
  }

  frailty_par_nat <- nat_est[seq(from = n_par_base + 1, length.out = n_par_frailty)]
  names(frailty_par_nat) <- par_names[seq(from = n_par_base + 1, length.out = n_par_frailty)]

  fl_obj <- frailty_laplace(s = 0, frailty = frailty, par = frailty_par_nat)
  f_var <- fl_obj$Var

  if (frailty != "none" && !any(is.na(vcov_mat))) {
    var_fn <- function(p) {
      if (frailty %in% c("gamma", "ig")) {
        fp <- exp(p[n_par_base + 1])
      } else if (frailty == "gl1") {
        fp <- exp(p[(n_par_base + 1):(n_par_base + 2)])
      } else if (frailty == "gl2") {
        th <- exp(p[n_par_base + 1])
        m <- (1.0 + th) * pmin(pmax(stats::plogis(p[n_par_base + 2]), 1e-15), 1.0 - 1e-15)
        fp <- c(th, m)
      }
      frailty_laplace(0, frailty, fp)$Var
    }
    grad_var <- tryCatch(numDeriv::grad(var_fn, raw_est), error = function(e) rep(0, n_total))
    f_var_se <- sqrt(pmax(as.numeric(t(grad_var) %*% vcov_mat %*% grad_var), 0))
  } else {
    f_var_se <- 0.0
  }

  z_stat <- nat_est / pmax(nat_se, 1e-12)
  p_val <- 2 * (1 - stats::pnorm(abs(z_stat)))
  ci_lower <- nat_est - 1.96 * nat_se
  ci_upper <- nat_est + 1.96 * nat_se
  signif <- ifelse(p_val < 0.001, "***", ifelse(p_val < 0.01, "**", ifelse(p_val < 0.05, "*", ifelse(p_val < 0.1, ".", " "))))

  coef_df <- data.frame(
    Estimate = nat_est,
    StdErr = nat_se,
    z_stat = z_stat,
    p_value = p_val,
    CI_lower = ci_lower,
    CI_upper = ci_upper,
    Signif = signif,
    row.names = par_names,
    stringsAsFactors = FALSE
  )

  k <- n_total
  aic_val <- -2 * max_ll + 2 * k
  bic_val <- -2 * max_ll + k * log(n)
  aicc_val <- if (n - k - 1 > 0) aic_val + (2 * k * (k + 1)) / (n - k - 1) else NA_real_
  hqic_val <- -2 * max_ll + 2 * k * log(log(n))

  vif_vec <- numeric(0)
  tol_vec <- numeric(0)
  if (n_cov > 1) {
    vif_vec <- tryCatch({
      vifs <- numeric(n_cov)
      for (j in 1:n_cov) {
        fit_j <- stats::lm(x[, j] ~ x[, -j])
        r2_j <- summary(fit_j)$r.squared
        vifs[j] <- 1 / pmax(1 - r2_j, 1e-10)
      }
      names(vifs) <- colnames(x)
      vifs
    }, error = function(e) rep(NA_real_, n_cov))
    tol_vec <- 1 / vif_vec
  } else if (n_cov == 1) {
    vif_vec <- c("1" = 1.0)
    tol_vec <- c("1" = 1.0)
  }

  res <- list(
    coefficients = coef_df,
    logLik = max_ll,
    AIC = aic_val,
    BIC = bic_val,
    AICc = aicc_val,
    HQIC = hqic_val,
    vcov = vcov_mat,
    baseline = baseline,
    frailty = frailty,
    frailty_par = frailty_par_nat,
    frailty_var = f_var,
    frailty_var_se = f_var_se,
    VIF = vif_vec,
    Tolerance = tol_vec,
    F_stat = NA_real_,
    p_F_stat = NA_real_,
    n = n,
    n_cov = n_cov,
    n_par_base = n_par_base,
    n_par_frailty = n_par_frailty,
    converged = converged,
    time = time,
    status = status,
    x = x,
    time2 = time2,
    prog_cen = prog_cen,
    raw_est = raw_est
  )

  class(res) <- "multifrailty_fit"
  res
}