R/mixedbiastest.R

In mixedbiastest: Bias Diagnostic for Linear Mixed Models

# ===== mixedbiastest.R =====
#' Bias Diagnostic for Linear Mixed Models
#'
#' Performs a permutation test to assess the bias of fixed effects in a linear mixed model fitted with `lmer`.
#' This function computes the test statistic and performs the permutation test, returning an object of class `"mixedbiastest"`.
#'
#' @param model An object of class `lmerMod` (or more generally `merMod`) fitted using `lmer` from the `lme4` package.
#' @param n_permutations Integer. Number of permutations to perform (default is 10000).
#' @param k_list Optional list of numeric vectors. Each vector specifies a linear combination of fixed effects to test. If `NULL`, each fixed effect is tested individually.
#' @param verbose Logical. If `TRUE`, prints detailed messages during execution.
#' @details
#' **Note:** This function currently supports only models with diagonal random effects covariance matrices (i.e., the G matrix is diagonal). The methodology for non-diagonal G matrices is described in Karl and Zimmerman (2021), but is not implemented in this version of the package.
#'
#' See the \code{\link{list_fixed_effects}} function if you would like to estimate the bias of a contrast of fixed effects.
#'
#' @section Acknowledgments:
#' Development of this package was assisted by GPT o1-preview, which helped in constructing the structure of much of the code and the roxygen documentation. The code is based on the R code provided by Karl and Zimmerman (2020).
#'
#' @return An object of class \code{"mixedbiastest"} containing:
#' \describe{
#'   \item{\code{results_table}}{A data frame with the test results for each fixed effect or contrast, including bias estimates and p-values.}
#'   \item{\code{permutation_values}}{A list of numeric vectors containing permutation values for each fixed effect or contrast.}
#'   \item{\code{model}}{The original \code{lmerMod} model object provided as input.}
#' }
#' @references
#' Karl, A. T., & Zimmerman, D. L. (2021). A diagnostic for bias in linear mixed model estimators induced by dependence between the random effects and the corresponding model matrix. \emph{Journal of Statistical Planning and Inference}, \emph{212}, 70–80. \doi{10.1016/j.jspi.2020.06.004}
#'
#' Karl, A., & Zimmerman, D. (2020). Data and Code Supplement for 'A Diagnostic for Bias in Linear Mixed Model Estimators Induced by Dependence Between the Random Effects and the Corresponding Model Matrix'. Mendeley Data, V1. \doi{10.17632/tmynggddfm.1}
#'
#' @examples
#' if (requireNamespace("plm", quietly = TRUE) && requireNamespace("lme4", quietly = TRUE)) {
#' library(lme4)
#' data("Gasoline", package = "plm")
#' mixed_model <- lmer(lgaspcar ~ lincomep + lrpmg + lcarpcap + (1 | country),
#'                     data = Gasoline)
#' result <- mixedbiastest(mixed_model)
#' print(result); plot(result)
#' }
#' if (requireNamespace("lme4", quietly = TRUE)) {
#' library(lme4)
#' example_model <- lmer(y ~ x + (1 | group), data = example_data)
#' result2 <- mixedbiastest(example_model)
#' print(result2); plot(result2)
#'
#' # Simulate data
#' set.seed(123)
#' n_groups <- 30
#' n_obs_per_group <- 10
#' group <- rep(1:n_groups, each = n_obs_per_group)
#' x <- runif(n_groups * n_obs_per_group)
#' beta0 <- 2; beta1 <- 5
#' sigma_u <- 1; sigma_e <- 0.5
#' u <- rnorm(n_groups, 0, sigma_u)
#' e <- rnorm(n_groups * n_obs_per_group, 0, sigma_e)
#' y <- beta0 + beta1 * x + u[group] + e
#'
#' data_sim <- data.frame(y = y, x = x, group = factor(group))
#' model3 <- lmer(y ~ x + (1 | group), data = data_sim)
#' result3 <- mixedbiastest(model3, verbose = TRUE)
#' plot(result3)
#' }
#' @importFrom lme4 getME fixef VarCorr
#' @importFrom stats sigma
#' @import Matrix
#' @export
mixedbiastest <- function(model, n_permutations = 10000, k_list = NULL, verbose = FALSE) {

  if (!inherits(model, "merMod")) {
    stop("The 'model' must be an object of class 'merMod' from the 'lme4' package.")
  }

  if (lme4::isSingular(model, tol = 1e-4)) {
    warning("The model fit is singular. Some variance components are zero.")
  }

  # Extract model components
  X <- lme4::getME(model, "X")       # n x p
  Z <- lme4::getME(model, "Z")       # n x q
  y <- lme4::getME(model, "y")       # n
  beta_hat <- lme4::fixef(model)     # p
  b_hat <- lme4::getME(model, "b")   # q

  sigma2_hat <- stats::sigma(model)^2

  VarCorr_list <- lme4::VarCorr(model)
  cnms  <- lme4::getME(model, "cnms")
  flist <- lme4::getME(model, "flist")
  Gp    <- lme4::getME(model, "Gp")

  if (verbose) {
    message("Dimensions of X: ", paste(dim(X), collapse = " x "))
    message("Dimensions of Z: ", paste(dim(Z), collapse = " x "))
    message("Class of Z: ", paste(class(Z), collapse = ", "))
    message("Length of y: ", length(y))
    message("Length of beta_hat: ", length(beta_hat))
    message("Length of b_hat: ", length(b_hat))
    message("Estimated residual variance (sigma^2): ", sigma2_hat)
  }

  n <- length(y)
  p <- length(beta_hat)

  G_blocks <- list()
  G_inv_blocks <- list()
  b_hat_list <- list()

  isDiagonal <- function(mat) {
    offL <- abs(mat[lower.tri(mat)]) < .Machine$double.eps
    offU <- abs(mat[upper.tri(mat)]) < .Machine$double.eps
    all(offL) && all(offU)
  }

  # Build G and G^{-1} blocks (never drop a term; clamp tiny variances)
  for (i in seq_along(cnms)) {
    grouping_factor_name <- names(cnms)[i]
    n_re_per_group <- length(cnms[[i]])
    n_groups <- nlevels(flist[[grouping_factor_name]])

    start_idx <- Gp[i] + 1
    end_idx   <- Gp[i + 1]

    b_hat_term <- b_hat[start_idx:end_idx]
    # Column-major: each column = RE coefficient, rows = groups
    b_hat_matrix <- matrix(b_hat_term, nrow = n_groups, ncol = n_re_per_group, byrow = FALSE)
    b_hat_list[[i]] <- b_hat_matrix

    cov_mat <- as.matrix(VarCorr_list[[grouping_factor_name]])

    if (verbose) {
      message("Term ", i, " (", grouping_factor_name, ")")
      message("  Covariance dims: ", paste(dim(cov_mat), collapse = " x "))
      message("  RE per group: ", n_re_per_group, " ; groups: ", n_groups)
    }

    if (!isDiagonal(cov_mat)) {
      stop(paste("The covariance matrix for random effect term", i, "is not diagonal.",
                 "Only diagonal random effects covariance matrices are supported."))
    }

    diag_elements <- diag(cov_mat)
    eps <- 1e-10
    diag_elements[diag_elements < eps] <- eps

    n_b <- n_re_per_group * n_groups
    G_block     <- Matrix::Diagonal(n = n_b, x = rep(diag_elements, times = n_groups))
    G_inv_block <- Matrix::Diagonal(n = n_b, x = rep(1 / diag_elements, times = n_groups))
    G_blocks[[i]]     <- G_block
    G_inv_blocks[[i]] <- G_inv_block

    if (verbose) {
      message("  G block dims: ", paste(dim(G_block), collapse = " x "))
    }
  }

  G     <- Matrix::bdiag(G_blocks)
  G_inv <- Matrix::bdiag(G_inv_blocks)

  if (verbose) {
    message("Combined G dims: ", paste(dim(G), collapse = " x "))
  }

  # Helpers: Cholesky-based inverse with safe fallback
  chol2inv_safe <- function(M_dense, name = "matrix") {
    M_dense <- 0.5 * (M_dense + t(M_dense))
    tryCatch({
      C <- chol(M_dense)
      chol2inv(C)
    }, error = function(e) {
      if (verbose) message("Cholesky failed for ", name, " (", conditionMessage(e), "). Falling back to solve().")
      solve(M_dense)
    })
  }

  # Avoid explicit V^{-1}
  X_Rinv  <- (1 / sigma2_hat) * X
  Z_R_inv <- (1 / sigma2_hat) * Z

  # T_inv = (Z' R^{-1} Z + G^{-1})^{-1}
  T_mat  <- Matrix::crossprod(Z_R_inv, Z) + G_inv
  T_mat  <- Matrix::forceSymmetric(T_mat)
  T_inv_dense <- chol2inv_safe(as.matrix(T_mat), name = "Z'R^{-1}Z + G^{-1}")
  T_inv <- Matrix::Matrix(T_inv_dense, sparse = FALSE)

  # A = Z' R^{-1} X
  A <- Matrix::crossprod(Z, X_Rinv)   # q x p

  # XVX = X'R^{-1}X - t(A) %*% T_inv %*% A
  XVX <- Matrix::crossprod(X_Rinv, X) - Matrix::t(A) %*% T_inv %*% A
  XVX <- Matrix::forceSymmetric(XVX)
  XVX_inv <- chol2inv_safe(as.matrix(XVX), name = "X'V^{-1}X") # dense p x p

  # XVZ = X'V^{-1}Z = t(A) - t(A) %*% T_inv %*% (Z'R^{-1}Z)
  B <- Matrix::crossprod(Z_R_inv, Z)  # q x q
  XVZ_dense <- as.matrix(Matrix::t(A) - Matrix::t(A) %*% T_inv %*% B)  # p x q

  if (verbose) {
    message("Dims of XVX: ", paste(dim(XVX), collapse = " x "))
  }

  # Default k vectors if none supplied
  if (is.null(k_list)) {
    k_list <- lapply(seq_len(p), function(j) { k <- rep(0, p); k[j] <- 1; k })
    names(k_list) <- names(beta_hat)
  } else {
    if (!is.list(k_list)) stop("k_list must be a list of numeric vectors.")
    for (i in seq_along(k_list)) {
      k <- k_list[[i]]
      if (!is.numeric(k) || length(k) != p) {
        stop("Each k vector in k_list must be numeric and of length equal to the number of fixed effects.")
      }
    }
  }

  if (is.null(names(k_list))) {
    names(k_list) <- paste0("Contrast_", seq_along(k_list))
  }

  results <- data.frame(
    Fixed_Effect = names(k_list),
    Mean_Permuted_Value = NA_real_,
    P_Value = NA_real_,
    Bias_Estimate = NA_real_,
    stringsAsFactors = FALSE
  )
  permutation_values_list <- list()

  # --- NEW: precompute segment indices for each random-effect term ---
  seg_idx <- lapply(seq_along(cnms), function(k_term) {
    seq.int(Gp[k_term] + 1L, Gp[k_term + 1L])
  })

  # Main loop over contrasts
  for (j in seq_along(k_list)) {
    k <- k_list[[j]]
    fe_name <- names(k_list)[j]

    # nu_hat = k' (X'V^{-1}X)^{-1} X'V^{-1}Z  -> (1 x q)
    nu_hat <- as.numeric(k %*% XVX_inv %*% XVZ_dense)

    observed_value <- sum(nu_hat * b_hat)

    if (verbose) {
      message("\nBias diagnostic for: ", fe_name)
      message("  Observed (nu' * b_hat): ", observed_value)
    }

    # Permutation test
    permutation_values <- numeric(n_permutations)

    for (i_perm in seq_len(n_permutations)) {
      permuted_b_hat <- numeric(length(b_hat))
      for (k_term in seq_along(cnms)) {
        b_matrix <- b_hat_list[[k_term]]
        # Independent column-wise permutations (fast + equivalent)
        permuted_b_matrix <- apply(b_matrix, 2, sample)
        # Write back into original segment (column-major)
        permuted_b_hat[ seg_idx[[k_term]] ] <- as.vector(permuted_b_matrix)
      }
      permutation_values[i_perm] <- sum(nu_hat * permuted_b_hat)
    }

    permutation_values <- permutation_values[is.finite(permutation_values)]

    if (length(permutation_values) == 0L) {
      p_value <- NA_real_
      mean_perm <- NA_real_
    } else {
      # two-sided with +1 correction
      p_value <- (sum(abs(permutation_values) >= abs(observed_value)) + 1) /
        (length(permutation_values) + 1)
      mean_perm <- mean(permutation_values)
    }

    results$Mean_Permuted_Value[j] <- round(mean_perm, 5)
    results$P_Value[j]             <- round(as.numeric(p_value), 5)
    results$Bias_Estimate[j]       <- round(observed_value, 5)

    permutation_values_list[[fe_name]] <- permutation_values
  }

  result_object <- list(
    results_table = results,
    permutation_values = permutation_values_list,
    model = model
  )
  class(result_object) <- "mixedbiastest"
  return(result_object)
}