The ICC of posteriors

rmu_reliability.R

#' Calculate Relative Measurement Uncertainty (RMU) from an rvar
#'
#' This function takes a vector of `rvar`s (one for each subject) and estimates
#' the reliability of the posteriors using the Relative Measurement Uncertainty (RMU)
#' method, as described by Bignardi et al. (2024). It is designed to be robust
#' to the common case where the underlying array of a vector rvar from a model
#' has a trailing dimension of 1 (e.g., [draws, subjects, 1]).
#'
#' @param intercepts_rvar A vector `rvar` object from the `posterior` package.
#'   Each element of the vector represents the posterior distribution for a
#'   different subject or group.
#' @param n_rmu_draws The number of RMU samples to generate for the reliability
#'   posterior distribution. A larger number provides a smoother posterior estimate.
#' @param verbose A logical value indicating whether to print operational messages.
#'
#' @return A scalar `rvar` object. This object contains the full posterior
#'   distribution of the reliability estimate (RMU).
#'
#' @examples
#' if (requireNamespace("posterior", quietly = TRUE)) {
#'
#'   # Example 1: A simple rvar with a 2D underlying array
#'   rvar_2d <- posterior::rvar(matrix(rnorm(1000 * 20), nrow = 1000, ncol = 20))
#'   rmu_2d <- reliability(rvar_2d)
#'   print(rmu_2d)
#'
#'   # Example 2: A more complex rvar with a 3D [draws, subjects, 1] array
#'   draws_array_3d <- array(rnorm(1000 * 20 * 1), dim = c(1000, 20, 1))
#'   rvar_3d <- posterior::rvar(draws_array_3d)
#'   rmu_3d <- reliability(rvar_3d)
#'   print(rmu_3d)
#'   posterior::summarise_draws(rmu_3d, "mean", "hdci")
#' }

reliability <- function(
    intercepts_rvar,
    n_rmu_draws = 4000,
    verbose = TRUE
) {
  # --- Input Validation ---
  if (!posterior::is_rvar(intercepts_rvar)) {
    stop("Input must be an 'rvar' object from the posterior package.")
  }
  
  # --- Data Preparation & Array Squeezing ---
  
  # Extract the underlying array from the rvar
  draws_array <- posterior::draws_of(intercepts_rvar)
  d <- dim(draws_array)

  if (length(d) == 2) {
    # Case 1: The array is already a 2D [draws, subjects] matrix.
    if (verbose) print("Input rvar has a 2D underlying array [draws, subjects].")
    draws_matrix_from_rvar <- draws_array
  } else if (length(d) == 3 && d[3] == 1) {
    # Case 2: The array is 3D [draws, subjects, 1]. This is the common case.
    if (verbose) print("Input rvar has a 3D array [draws, subjects, 1]. Squeezing to 2D.")
    # We use your exact solution to slice and drop the trailing dimension.
    draws_matrix_from_rvar <- draws_array[,,1, drop = TRUE]
  } else {
    # Case 3: The array is some other shape we don't support.
    stop(paste0(
      "The input 'rvar' is not a vector of scalar posteriors.\n",
      "Its underlying array has unsupported dimensions: [", paste(d, collapse = ", "), "].\n",
      "The function expects a [draws, subjects] or [draws, subjects, 1] array."
    ))
  }

  N_subjects <- ncol(draws_matrix_from_rvar)
  N_posterior_draws <- nrow(draws_matrix_from_rvar)

  if (verbose) {
    base::print(paste0("Number of subjects (N): ", N_subjects))
    base::print(paste0("Number of posterior draws per subject (K): ", N_posterior_draws))
  }
  
  # The paper's N x K logic requires subjects in rows and draws in columns.
  subject_draws_matrix <- t(draws_matrix_from_rvar)

  # Check for NAs after potential squeezing
  if (anyNA(subject_draws_matrix)) {
    stop("NA values detected in the draws of the rvar. Please handle missing values.")
  }

  # --- RMU Calculation Loop ---
  
  reliability_posterior_draws <- numeric(n_rmu_draws)

  for (i in 1:n_rmu_draws) {
    # Sample WITH REPLACEMENT to generate two vectors of draws
    cols1 <- sample(1:N_posterior_draws, size = N_subjects, replace = TRUE)
    cols2 <- sample(1:N_posterior_draws, size = N_subjects, replace = TRUE)
    
    vec1 <- subject_draws_matrix[cbind(1:N_subjects, cols1)]
    vec2 <- subject_draws_matrix[cbind(1:N_subjects, cols2)]

    reliability_posterior_draws[i] <- stats::cor(vec1, vec2, method = "pearson")
  }
  
  # --- Convert to rvar and Return ---
  rmu_rvar <- posterior::rvar(reliability_posterior_draws)

  return(rmu_rvar)
}