R/fit_spm_jags.R

In BDIB: Bayesian Dynamic Information Borrowing for Longitudinal Data with Informative Dropout

#' Fit SPM with MNAR monotone dropout in JAGS (continuous or binary outcomes)
#'
#' Unified JAGS interface for the shared-parameter model (SPM) with MNAR dropout:
#' - \code{family="gaussian"}: Continuous outcome.
#' - \code{family="bernoulli"}: Binary outcome.
#'
#' @param df Long-format \code{data.frame} with columns \code{id}, \code{time}, \code{R}, \code{Y}.
#'   \code{R=1} means observed, \code{R=0} means missing.
#' @param family Either \code{"gaussian"} or \code{"bernoulli"} (probit).
#' @param prior_theta Optional prior for \eqn{\theta}.
#'   \itemize{
#'     \item \code{family="gaussian"}: \code{RBesT::normMix} object (mixture Normal on \eqn{\theta}).
#'     \item \code{family="bernoulli"}: a \code{"betaMix"} object (mixture Beta on \eqn{\theta}).
#'   }
#'   If \code{NULL}, defaults to a very weak one-component prior.
#' @param params_to_save Character vector of parameter names to monitor.
#'   If \code{NULL}, sensible defaults are used for each family.
#' @param n.chains,n.iter,n.burnin,n.thin MCMC settings passed to \code{jagsUI::jags}.
#' @param parallel Logical; whether to run chains in parallel (via \code{jagsUI}).
#' @return Numeric vector of posterior draws of \eqn{\theta}.
#' @export
#' @examples
#' # ---- Example 1: Continuous outcome with weak prior ----
#' set.seed(1000)
#'
#' # 1) Generate longitudinal data with MNAR monotone dropout
#' t_vec <- seq(4, 48, 4)
#' gen <- generate_spm(
#'   n = 80, t_vec = t_vec, family = "gaussian",
#'   alpha0 = -4, alpha1 = 0.05, rho = 1, sigma_b = 1,
#'   theta = -3.2,   # total change at the last visit (= slope * max(t))
#'   sigma_e = 1
#' )
#' df <- gen$dat  # long-format data: id, time, R, Y
#'
#' # 2) Define a weakly informative prior for theta (one Normal component with large variance)
#' prior <- RBesT::mixnorm(c(1, 0, 1000))  # weight=1, mean=0, sd=1000
#'
#' # 3) Fit the SPM (MNAR) and obtain posterior draws of theta
#' \donttest{
#' # Full MCMC fitting takes longer than a brief example check.
#' draws_theta <- fit_spm_jags(
#'   df = df, family = "gaussian",
#'   prior_theta = prior,
#'   params_to_save = c("theta","slope","alpha0","alpha1","rho","sigma_e","sigma_b"),
#'   n.chains = 3, n.iter = 10000, n.burnin = 5000, n.thin = 2, parallel = FALSE
#' )
#'
#' # Summarize posterior results
#' summary(draws_theta)    # posterior mean, median, quantiles
#' mean(draws_theta)       # posterior mean of theta
#' }
#'
#' # ---- Example 2: Binary outcome (probit link) with weak prior ----
#' set.seed(1000)
#' # 1) Generate longitudinal binary data with MNAR monotone dropout
#' t_vec <- seq(4, 48, 4)
#' gen <- generate_spm(
#'   n       = 200,
#'   t_vec   = t_vec,
#'   family  = "bernoulli",
#'   alpha0  = -4,       # dropout logit-intercept
#'   alpha1  =  0.05,    # dropout logit-slope in time
#'   rho     = -1,       # association between random effect and dropout
#'   sigma_b =  1,       # SD of random intercept b_i
#'   theta   =  0.4,    # target marginal Pr(Y=1) at final visit
#'   beta0   = -2        # outcome probit-intercept
#' )
#' df <- gen$dat  # long-format data: id, time, R, Y
#'
#' # 2) Define a weakly informative prior for theta:
#' #    Beta(1,1) (Uniform on (0,1)) as a one-component mixture
#' prior <- RBesT::mixbeta(c(1, 1, 1))  # weight=1, a=1, b=1
#'
#' # 3) Fit the SPM (MNAR) with probit link for the binary outcome
#' #    and obtain posterior draws of theta (final-visit response rate)
#' \donttest{
#' # Full MCMC fitting takes longer than a brief example check.
#' draws_theta <- fit_spm_jags(
#'   df             = df,
#'   family         = "bernoulli",
#'   prior_theta    = prior,
#'   params_to_save = c("theta","slope","beta0","alpha0","alpha1","rho","sigma_b"),
#'   n.chains       = 3,
#'   n.iter         = 5000,
#'   n.burnin       = 3000,
#'   n.thin         = 2,
#'   parallel       = FALSE
#' )
#'
#' # 4) Summarize posterior results
#' summary(draws_theta)   # posterior mean, median, quantiles of theta
#' mean(draws_theta)      # posterior mean of theta
#' }
fit_spm_jags <- function(
    df,
    family = c("gaussian","bernoulli"),
    prior_theta = NULL,
    params_to_save = NULL,
    n.chains = 3, n.iter = 10000, n.burnin = 5000, n.thin = 2,
    parallel = FALSE
) {
  family <- match.arg(family)
  stopifnot(is.data.frame(df),
            all(c("id","time","R","Y") %in% names(df)))

  # ---- Build prior for theta ----
  if (family == "gaussian") {
    # Expect RBesT::normMix: 3 x K matrix (row1=w, row2=mu, row3=sd)
    if (is.null(prior_theta)) {
      mix_list <- list(K = 1L, w = 1, mu = 0, tau = 1/1e3) # very flat
    } else {
      stopifnot(inherits(prior_theta, "normMix"))
      mix_list <- list(
        K   = ncol(prior_theta),
        w   = as.numeric(prior_theta[1, ]),
        mu  = as.numeric(prior_theta[2, ]),
        tau = 1 / as.numeric(prior_theta[3, ])^2  # precision = 1/sd^2
      )
    }
  } else {
    # Binary outcome: expect "betaMix": 3 x K matrix (row1=w, row2=a, row3=b)
    if (is.null(prior_theta)) {
      mix_list <- list(K = 1L, w = 1, a = 1, b = 1)  # Uniform(0,1)
    } else {
      stopifnot(inherits(prior_theta, "betaMix"))
      mix_list <- list(
        K = ncol(prior_theta),
        w = as.numeric(prior_theta[1, ]),
        a = as.numeric(prior_theta[2, ]),
        b = as.numeric(prior_theta[3, ])
      )
    }
  }

  # ---- Data assembly for JAGS ----
  J         <- length(unique(df$id))
  id        <- df$id
  time_vec  <- as.numeric(df$time)
  K         <- length(unique(time_vec))
  t_max     <- max(time_vec)
  R_matrix  <- t(matrix(df$R, K, J))
  obs_idx   <- which(df$R == 1L)
  y_obs     <- df$Y[obs_idx]

  # Defaults for monitored params
  if (is.null(params_to_save)) {
    if (family == "gaussian") {
      params_to_save <- c("theta","slope","alpha0","alpha1","rho","sigma_e","sigma_b")
    } else {
      params_to_save <- c("theta","slope","beta0","alpha0","alpha1","rho","sigma_b")
    }
  }

  # Common data list
  data_jags <- list(
    N_obs   = length(obs_idx),
    J       = J,
    K       = K,
    id      = id,
    t       = time_vec,
    t_max   = t_max,
    R       = R_matrix,
    obs_idx = obs_idx,
    y       = y_obs
  )

  # Prior parameters into data list
  if (family == "gaussian") {
    data_jags$K_prior  <- mix_list$K
    data_jags$w_prior  <- mix_list$w
    data_jags$mu_prior <- mix_list$mu
    data_jags$tau_prior<- mix_list$tau

    # ---- JAGS model (Continuous outcome) ----
    modelString <- "
    model {
      # Mixture prior on theta
      z_prior ~ dcat(w_prior[1:K_prior])
      theta   ~ dnorm(mu_prior[z_prior], tau_prior[z_prior])
      slope   <- theta / t_max

      # Priors
      alpha0 ~ dnorm(0, 0.001)
      alpha1 ~ dnorm(0, 0.001)
      rho    ~ dnorm(0, 0.001)

      tau_e ~ dgamma(0.01, 0.01);  sigma_e <- pow(tau_e, -0.5)
      tau_b ~ dgamma(0.01, 0.01);  sigma_b <- pow(tau_b, -0.5)

      for (j in 1:J) { b[j] ~ dnorm(0, tau_b) }

      # Outcome model for observed Y
      for (n in 1:N_obs) {
        mu_y[n] <- slope * t[obs_idx[n]] + b[id[obs_idx[n]]]
        y[n] ~ dnorm(mu_y[n], tau_e)
      }

      # MNAR dropout model (monotone)
      for (j in 1:J) {
        logit(p_miss[j,1]) <- alpha0 + alpha1 * t[1] + rho * b[j]
        R[j,1] ~ dbern(1 - p_miss[j,1])
        for (k in 2:K) {
          logit(p_miss[j,k]) <- alpha0 + alpha1 * t[k] + rho * b[j]
          R[j,k] ~ dbern(R[j,k-1] * (1 - p_miss[j,k]))
        }
      }
    }"

    inits <- function() list(
      z_prior = sample(seq_len(mix_list$K), 1),
      theta   = rnorm(1, 0, 1),
      alpha0  = 0, alpha1 = 0, rho = 0,
      tau_e   = 1, tau_b = 1
    )

  } else {
    data_jags$K_prior <- mix_list$K
    data_jags$w_prior <- mix_list$w
    data_jags$a_prior <- mix_list$a
    data_jags$b_prior <- mix_list$b

    # ---- JAGS model (Binary outcome) ----
    # theta is bounded away from {0,1} by a tiny epsilon to avoid numerical issues
    modelString <- "
    model {
      # Mixture Beta prior on theta (final-visit marginal Pr(Y=1))
      z_prior   ~ dcat(w_prior[1:K_prior])
      theta_raw ~ dbeta(a_prior[z_prior], b_prior[z_prior])
      theta     <- (1 - 2.0E-6) * theta_raw + 1.0E-6
      slope <- (sqrt(1 + pow(sigma_b, 2)) *  probit(theta) - beta0) / t_max

      # Priors for missingness model & random effects
      alpha0 ~ dnorm(0, 0.001)
      alpha1 ~ dnorm(0, 0.001)
      rho    ~ dnorm(0, 0.001)
      beta0 ~ dnorm(0, 0.001)

      tau_b ~ dgamma(0.01, 0.01);  sigma_b <- pow(tau_b, -0.5)

      # Random effects
      for (j in 1:J) { b[j] ~ dnorm(0, tau_b) }

      # Outcome model: Bernoulli with probit link
      for (n in 1:N_obs) {
        linpred[n] <- beta0 + slope * t[obs_idx[n]] + b[id[obs_idx[n]]]
        p[n] <- phi(linpred[n])
        y[n] ~ dbern(p[n])
      }

      # MNAR dropout model (monotone)
      for (j in 1:J) {
        logit(p_miss[j,1]) <- alpha0 + alpha1 * t[1] + rho * b[j]
        R[j,1] ~ dbern(1 - p_miss[j,1])
        for (k in 2:K) {
          logit(p_miss[j,k]) <- alpha0 + alpha1 * t[k] + rho * b[j]
          R[j,k] ~ dbern(R[j,k-1] * (1 - p_miss[j,k]))
        }
      }
    }"

    inits <- function() list(
      z_prior   = sample(seq_len(mix_list$K), 1),
      theta_raw = 0.5,  # inside (0,1)
      beta0     = 0,
      alpha0    = 0, alpha1 = 0, rho = 0,
      tau_b     = 1
    )
  }

  # ---- Run JAGS ----
  fit <- jagsUI::jags(
    data               = data_jags,
    inits              = inits,
    parameters.to.save = params_to_save,
    model.file         = textConnection(modelString),
    n.chains           = n.chains,
    n.iter             = n.iter,
    n.burnin           = n.burnin,
    n.thin             = n.thin,
    parallel           = parallel
  )

  as.vector(fit$sims.list$theta)
}