Add RStan backend with multivariate observation model (Model B)

.Rbuildignore

@@ -29,3 +29,8 @@
 ^_quarto\.yml$
 ^\.quarto$
 ^\.devcontainer$
+^inst/stan/.*\.o$
+^inst/stan/.*\.so$
+^inst/stan/.*\.d$
+^inst/stan/.*_rng\.hpp$
+^inst/stan/.*\.rds$

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: serodynamics
 Title: What the Package Does (One Line, Title Case)
-Version: 0.0.0.9047
+Version: 0.0.0.9048
 Authors@R: c(
     person("Peter", "Teunis", , "p.teunis@emory.edu", role = c("aut", "cph"),
            comment = "Author of the method and original code."),
@@ -37,6 +37,7 @@ Suggests:
     knitr,
     readr,
     rlist,
+    rstan,
     sessioninfo,
     spelling,
     stringr,

NAMESPACE

@@ -15,9 +15,12 @@ export(plot_jags_trace)
 export(plot_predicted_curve)
 export(post_summ)
 export(postprocess_jags_output)
+export(postprocess_stan_output)
 export(prep_data)
+export(prep_data_stan)
 export(prep_priors)
 export(run_mod)
+export(run_mod_stan)
 export(serodynamics_example)
 export(sim_case_data)
 export(sim_n_obs)

NEWS.md

@@ -1,5 +1,15 @@
 # serodynamics (development version)
 
+* Added RStan backend with multivariate observation model (Model B):
+  `run_mod_stan()`, `prep_data_stan()`, and `postprocess_stan_output()`.
+  Model B replaces K independent univariate normal likelihoods with a
+  single K-variate normal likelihood per time point, capturing residual
+  correlations across antigen-isotype pairs.  Output is an `sr_model`
+  tibble identical in structure to `run_mod()`, so all existing plot and
+  summary functions work without modification.  Requires the optional
+  dependency `rstan` (`install.packages("rstan")`).  The JAGS backend is
+  fully backward-compatible and unchanged.
+
 * Added dev container configuration for persistent, cached development environment
   that includes R, JAGS, and all dependencies preinstalled, making Copilot
   Workspace sessions much faster.

R/postprocess_stan_output.R

@@ -0,0 +1,78 @@
+#' Postprocess Stan model output
+#'
+#' @description
+#' Converts a `stanfit` object returned by [rstan::sampling()] into an
+#' `sr_model` tibble with the same column layout as [run_mod()].
+#'
+#' Extracted parameters are the natural-scale quantities computed in the
+#' Stan `generated quantities` block:
+#' `y0_nat`, `y1_nat`, `t1_nat`, `alpha_nat`, `shape_nat`.
+#'
+#' @param stan_fit A `stanfit` object from [rstan::sampling()].
+#' @param ids Character vector of subject IDs (length `N`).
+#' @param antigen_isos Character vector of antigen-isotype names (length `K`).
+#' @param stratification Character scalar: stratification label used to
+#'   populate the `Stratification` column.  Defaults to `"None"`.
+#'
+#' @returns A [tibble::tbl_df] with columns:
+#'   `Iteration`, `Chain`, `Parameter`, `Iso_type`, `Stratification`,
+#'   `Subject`, `value`.
+#'
+#' @seealso [run_mod_stan()]
+#' @export
+postprocess_stan_output <- function(
+    stan_fit,
+    ids,
+    antigen_isos,
+    stratification = "None") {
+
+  # Extract natural-scale parameters from generated quantities
+  param_names <- c("y0_nat", "y1_nat", "t1_nat", "alpha_nat", "shape_nat")
+  param_labels <- c("y0", "y1", "t1", "alpha", "shape")
+
+  N <- length(ids)
+  K <- length(antigen_isos)
+
+  # rstan::extract with permuted = FALSE returns [iter, chains, parameters]
+  all_samples <- rstan::extract(stan_fit, pars = param_names, permuted = FALSE)
+
+  n_iter   <- dim(all_samples)[1]
+  n_chains <- dim(all_samples)[2]
+
+  # Build long-format tibble for all subjects, biomarkers, and parameters
+  rows <- vector("list", N * K * length(param_labels))
+  idx  <- 1L
+
+  for (k in seq_len(K)) {
+    for (subj in seq_len(N)) {
+      for (p in seq_along(param_labels)) {
+        # Parameter name in Stan: e.g. "y0_nat[subj,k]"
+        stan_par_name <- paste0(param_names[p], "[", subj, ",", k, "]")
+        col_pos <- which(dimnames(all_samples)[[3]] == stan_par_name)
+
+        if (length(col_pos) == 0L) {
+          cli::cli_abort(
+            "Stan parameter {.val {stan_par_name}} not found in stanfit object."
+          )
+        }
+
+        # Flatten across chains: Iteration index increases across chains
+        for (chain in seq_len(n_chains)) {
+          iter_vals <- all_samples[, chain, col_pos]
+          rows[[idx]] <- tibble::tibble(
+            Iteration      = seq_len(n_iter),
+            Chain          = chain,
+            Parameter      = param_labels[p],
+            Iso_type       = antigen_isos[k],
+            Stratification = stratification,
+            Subject        = ids[subj],
+            value          = as.numeric(iter_vals)
+          )
+          idx <- idx + 1L
+        }
+      }
+    }
+  }
+
+  dplyr::bind_rows(rows)
+}

R/prep_data_stan.R

@@ -0,0 +1,74 @@
+#' Prepare data for Stan Model B
+#'
+#' @description
+#' Converts the output of [prep_data()] into a named list suitable for
+#' passing to `rstan::sampling()`. Missing values in the padded arrays are
+#' replaced with `0`; Stan only accesses entries up to `n_obs[i]` per
+#' subject, so the padding values are never used in the likelihood.
+#'
+#' The `mu_hyp` and `sigma_hyp` arguments correspond to the hyperprior
+#' parameters from [prep_priors()]:
+#' - `mu_hyp` maps to `mu_hyp_param` (prior means on the five log-scale
+#'   curve parameters)
+#' - `sigma_hyp` maps to `1 / sqrt(prec_hyp_param)` (prior standard
+#'   deviations derived from the JAGS precision hyperparameter)
+#'
+#' @param prepped_jags_data A `prepped_jags_data` list returned by
+#'   [prep_data()] (with `add_newperson = FALSE`).
+#' @param mu_hyp A numeric vector of length 5 giving the prior means for
+#'   the five log-scale curve parameters
+#'   (`log_y0`, `log_delta`, `log_t1`, `log_alpha`, `log_shape_minus_1`).
+#'   Defaults match the `mu_hyp_param` argument of [prep_priors()].
+#' @param sigma_hyp A numeric vector of length 5 giving the prior standard
+#'   deviations for the five log-scale curve parameters.
+#'   Defaults are derived from the `prec_hyp_param` argument of
+#'   [prep_priors()] as `1 / sqrt(prec_hyp_param)`.
+#'
+#' @returns A named [list] suitable for passing to `rstan::sampling()`.
+#'   Contains the following elements:
+#'   - `N` — number of subjects
+#'   - `K` — number of antigen-isotype pairs
+#'   - `max_obs` — maximum number of observations per subject
+#'   - `n_obs` — integer vector of length `N` giving actual observations
+#'   - `time_obs` — `N × max_obs` matrix of observation times
+#'     (padded with `0`)
+#'   - `log_y_obs` — `N × max_obs × K` array of log antibody levels
+#'     (padded with `0`)
+#'   - `mu_hyp` — vector of length 5 (hyperprior means)
+#'   - `sigma_hyp` — vector of length 5 (hyperprior SDs)
+#'
+#' @seealso [prep_data()], [run_mod_stan()]
+#' @export
+#'
+#' @examples
+#' set.seed(1)
+#' raw_data <-
+#'   serocalculator::typhoid_curves_nostrat_100 |>
+#'   sim_case_data(n = 5)
+#' prepped <- prep_data(raw_data, add_newperson = FALSE)
+#' stan_data <- prep_data_stan(prepped)
+#' str(stan_data)
+prep_data_stan <- function(
+    prepped_jags_data,
+    mu_hyp = c(1.0, 7.0, 1.0, -4.0, -1.0),
+    sigma_hyp = 1.0 / sqrt(c(1.0, 0.00001, 1.0, 0.001, 1.0))) {
+
+  smpl_t <- prepped_jags_data$smpl.t
+  logy   <- prepped_jags_data$logy
+  nsmpl  <- prepped_jags_data$nsmpl
+
+  # Replace NA padding values with 0; Stan only reads positions <= n_obs[i]
+  smpl_t[is.na(smpl_t)] <- 0.0
+  logy[is.na(logy)]     <- 0.0
+
+  list(
+    N        = unname(prepped_jags_data$nsubj),
+    K        = unname(prepped_jags_data$n_antigen_isos),
+    max_obs  = unname(ncol(smpl_t)),
+    n_obs    = as.integer(nsmpl),
+    time_obs = unname(smpl_t),
+    log_y_obs = unname(logy),
+    mu_hyp   = mu_hyp,
+    sigma_hyp = sigma_hyp
+  )
+}