R/score_tirt_stan.R

In autoFC: Automatic Toolkit for Construction, Optimization, Scoring and Simulation of Forced-Choice Tests

#' Run Stan TIRT Model and Extract Formatted Scores
#'
#' @param stan_data The data list generated by prepare_tirt_stan_data().
#' @param chains Integer. Number of MCMC chains. Defaults to 4.
#' @param parallel_chains Integer. How many chains to run simultaneously. Defaults to 4.
#' @param threads_per_chain Integer. CPU threads allocated INSIDE each chain. Defaults to 2.
#' @param iter_warmup Integer. Warmup iterations. Defaults to 1000.
#' @param iter_sampling Integer. Sampling iterations. Defaults to 1000.
#' @param init Same as the `init` parameter in rstan.
#'
#' @return A list containing `scores` (Data frame of traits) and `fit` (The cmdstanr fit object).
#' @export
score_tirt_stan <- function(stan_data,
                            chains = 4,
                            parallel_chains = 4,
                            threads_per_chain = 2,
                            iter_warmup = 1000,
                            iter_sampling = 1000,
                            init = 0) {

  if (!requireNamespace("cmdstanr", quietly = TRUE)) {
    stop(
      "This function requires the 'cmdstanr' package.\n",
      "Install with: install.packages(\"cmdstanr\", ",
      "repos = c(\"https://stan-dev.r-universe.dev\", getOption(\"repos\")))\n",
      "Then: cmdstanr::install_cmdstan()"
    )
  }

  # 1. Compile the Model (cmdstanr is smart: it only compiles if it hasn't already)
  cat("Checking/Compiling Stan model...compilation required the first time you run the model\n")
  stan_file = system.file("stan", "tirt_model.stan", package = "autoFC")
  mod <- cmdstanr::cmdstan_model(stan_file, cpp_options = list(stan_threads = TRUE))

  # 2. Run the Sampler
  cat("Starting MCMC Sampling...\n")
  fit <- mod$sample(
    data = stan_data,
    chains = chains,
    parallel_chains = parallel_chains,
    threads_per_chain = threads_per_chain,
    iter_warmup = iter_warmup,
    iter_sampling = iter_sampling,
    init = init,               # Crucial for avoiding -Inf errors!
    refresh = 100
  )

  # 3. Extract the Scores using the helper we wrote earlier
  cat("Extracting Trait Scores and SEs...\n")

  trait_names <- attr(stan_data, "trait_names")
  N <- stan_data$N
  D <- stan_data$D

  theta_summary <- fit$summary(variables = "theta", c("mean", "sd"))

  clean_scores <- data.frame(matrix(NA, nrow = N, ncol = D * 2))
  col_names <- c()
  for (t in trait_names) col_names <- c(col_names, t, paste0(t, "_SE"))
  colnames(clean_scores) <- col_names

  for (d in 1:D) {
    pattern <- sprintf("^theta\\[\\d+,%d\\]$", d)
    trait_rows <- theta_summary[grepl(pattern, theta_summary$variable), ]
    clean_scores[, (d - 1) * 2 + 1] <- trait_rows$mean
    clean_scores[, (d - 1) * 2 + 2] <- trait_rows$sd
  }

  cat("Done!\n")

  # Return both the scores and the fit object
  return(list(
    scores = clean_scores,
    fit = fit
  ))
}