I’ve written two functions which might be of use to others. They are surely not fully optimized but quick enough for my use cases.
- Predictions from blrm() fit which can include the invidual random effects. I needed this for a project where we’re really interested in the ATE, because an intervention applies to a whole population, not for 1:1 decision making. And from my current understanding based on this, we need to include the RE in the standardization for that.
predict_blrm()
predict_blrm_manual <- function(fit, newdata, include_re = FALSE) {
# --- 0) Prerequisite Checks ---
if (!is.data.frame(newdata)) {
stop("`newdata` must be a data frame.")
}
if (include_re && !"id" %in% colnames(newdata)) {
stop("`newdata` must contain a 'id' column when include_re=TRUE")
}
N_new <- nrow(newdata)
# --- 1) Get the design matrix for the new data ---
X <- rms::predictrms(fit, newdata = newdata, type = 'x')
# --- 2) Extract posterior draws and model information ---
draws <- fit$draws
ndraws <- nrow(draws)
ns <- fit$non.slopes
K <- ns + 1
intercept_draws <- draws[, 1:ns, drop = FALSE]
beta_draws <- draws[, (ns + 1):ncol(draws), drop = FALSE]
# --- 3) Calculate the fixed-effect linear predictor (eta_fixed) ---
eta_fixed <- beta_draws %*% t(X)
# --- 4) Get the random effect draws (u_draws) if requested ---
u_draws <- matrix(0, nrow = ndraws, ncol = N_new)
if (include_re) {
if (length(unique(newdata$id)) > 1) {
stop("`newdata` contains multiple unique patient IDs. This version is restricted to one at a time.")
}
if (!requireNamespace("posterior", quietly = TRUE)) {
stop("The 'posterior' package is required to extract random effect draws.")
}
sf <- rmsb::stanGet(fit)
all_gamma_draws <- posterior::as_draws_matrix(sf$draws(variables = "gamma"))
patient_ids <- newdata$id
u_draws <- all_gamma_draws[, patient_ids, drop = FALSE]
}
# --- 5) Calculate the total linear predictor (eta) ---
stopifnot(all(dim(eta_fixed) == dim(u_draws)))
eta <- eta_fixed + u_draws
# --- 6) Calculate probabilities for each observation and each draw ---
p_cat_draws_list <- vector("list", N_new)
for (i in 1:N_new) {
eta_i <- eta[, i]
s_ge <- plogis(sweep(intercept_draws, 1, eta_i, "+"))
p_cat_draws <- matrix(NA_real_, nrow = ndraws, ncol = K)
p_cat_draws[, 1] <- 1 - s_ge[, 1]
if (K > 2) {
for (kk in 2:(K - 1)) {
p_cat_draws[, kk] <- s_ge[, kk - 1] - s_ge[, kk]
}
}
p_cat_draws[, K] <- s_ge[, K - 1]
colnames(p_cat_draws) <- paste0("y=", 1:K)
p_cat_draws_list[[i]] <- p_cat_draws
}
# --- 7) Reshape draws into a single 3D array ---
draws_array <- array(
data = NA_real_,
dim = c(ndraws, N_new, K)
)
for (i in 1:N_new) {
draws_array[, i, ] <- p_cat_draws_list[[i]]
}
dimnames(draws_array) <- list(
NULL,
rownames(newdata),
colnames(p_cat_draws_list[[1]])
)
return(draws_array)
}
- Function which calculates uncondtional (conditional on baseline-states only) SOPs for a second order markov model. I’m relatively confident that I’m correctly unconditionalizing the SOPs but a review is warranted.
Second Order Markov SOPs
#' Calculate Second-Order State Occupation Probabilities
#'
#' This function calculates the unconditional state occupation probabilities for a
#' second-order Markov model.
#'
#' @param model_fit The fitted model object (rmsb or brms)
#' @param baseline_data A single-row data frame containing the initial states and
#' all baseline predictor values required by the model.
#' @param col_yprev A string with the column name for the state at t-1 (e.g., "yprev").
#' @param col_ypprev A string with the column name for the state at t-2 (e.g., "ypprev").
#' @param col_time A string with the column name that represents time in the model (e.g., "week").
#' @param times A numeric vector specifying the time points for calculation (e.g., `1:26`).
#' @param n_states An integer for the number of levels in the ordinal outcome.
#' @param absorbing_states A numeric vector listing the absorbing states.
#'
#' @return A list containing two arrays:
#' \item{unconditional}{A 3D array `(draws, states, time)` with the unconditional state probabilities.}
#' @export
sop_markov_second_order <- function(model_fit,
baseline_data,
col_yprev,
col_ypprev,
col_time,
times,
n_states,
absorbing_states,
include_re = FALSE) {
# --- Detect model type and get number of draws ---
if (inherits(model_fit, "brmsfit")) {
detected_type <- "brms"
if (!requireNamespace("marginaleffects", quietly = TRUE) || !requireNamespace("posterior", quietly = TRUE)) {
stop("Packages 'marginaleffects' and 'posterior' are required for brms models.")
}
n_draws <- posterior::ndraws(model_fit)
} else if (inherits(model_fit, "blrm")) {
detected_type <- "rmsb"
n_draws <- nrow(model_fit$draws)
} else {
stop("Unsupported model object class. Please provide a model from 'brms' (class 'brmsfit') or 'rmsb' (class 'blrm').")
}
# --- Parameters ---
M <- n_states
T_horizon <- max(times)
# --- Initial Conditions ---
initial_state_t_minus_1 <- as.numeric(baseline_data[[col_ypprev]])
initial_state_t_0 <- as.numeric(baseline_data[[col_yprev]])
# --- Data Structures to Store Results ---
P_joint_bayes <- array(0, dim = c(n_draws, M, M, T_horizon + 1),
dimnames = list(paste0("draw_", 1:n_draws),
paste0("t-1=", 1:M),
paste0("t=", 1:M),
paste0("time=", 0:T_horizon)))
P_uncond_bayes <- array(0, dim = c(n_draws, M, T_horizon + 1),
dimnames = list(paste0("draw_", 1:n_draws),
paste0("State_", 1:M),
paste0("t=", 0:T_horizon)))
# Set the known probability at t=0
P_joint_bayes[, initial_state_t_minus_1, initial_state_t_0, 1] <- 1.0
P_uncond_bayes[, initial_state_t_0, 1] <- 1.0
# --- Create the static grid for batch predictions ---
history_grid <- expand.grid(
h = factor(1:M, levels = 1:M), # X_{t-2}
j = factor(1:M, levels = 1:M) # X_{t-1}
)
names(history_grid) <- c(col_ypprev, col_yprev)
# --- Dynamically create the full prediction dataset ---
# Identify static covariates by excluding time-varying/history columns
cols_to_exclude <- c(col_yprev, col_ypprev, col_time)
static_covariate_names <- setdiff(names(baseline_data), cols_to_exclude)
static_covariates <- baseline_data[, static_covariate_names, drop = FALSE]
# Combine static covariates with the grid of all possible histories
prediction_data <- cbind(static_covariates, history_grid)
# --- Main Iteration Loop ---
for (t in times) {
# Identify predictable histories
predictable_idx <- which(
!(prediction_data[[col_ypprev]] %in% absorbing_states) &
!(prediction_data[[col_yprev]] %in% absorbing_states)
)
n_predictable <- length(predictable_idx)
# Subset the data for prediction and set the current time
prediction_data_subset <- prediction_data[predictable_idx, ]
prediction_data_subset[[col_time]] <- t
pred_array_3d <- NULL
if (detected_type == "brms") {
# refactor previous states to remove absorbing states
prediction_data_subset[[col_yprev]] <- factor(prediction_data_subset[[col_yprev]], levels = c(1:M)[-absorbing_states])
prediction_data_subset[[col_ypprev]] <- factor(prediction_data_subset[[col_ypprev]], levels = c(1:M)[-absorbing_states])
re_arg <- if (include_re) NULL else NA
brms_pred_long <- marginaleffects::predictions(model_fit,
newdata = prediction_data_subset,
type = "response",
re_formula = re_arg) |>
marginaleffects::posterior_draws()
pred_array_3d <- array(0, dim = c(n_draws, n_predictable, M))
indices <- cbind(brms_pred_long$drawid,
brms_pred_long$rowid,
as.numeric(brms_pred_long$group))
pred_array_3d[indices] <- brms_pred_long$draw
} else { # detected_type == "rmsb"
# dimension 1 = 4000 draws
# dimension 2 = X_{t-1} and X_{t-2} combinations (rows from prediction_data_subset)
# dimension 3 = cumulative probabilities (e.g., 4 columns for 5 states)
pred_array_3d <- predict_blrm_manual(model_fit, prediction_data_subset, include_re = include_re)
}
state_probs_all_draws <- do.call(rbind, lapply(1:dim(pred_array_3d)[1], function(d) pred_array_3d[d, , ]))
# --- Assemble transition probabilities for each draw ---
# Create a 4D array to hold all transition probabilities: P(X_t=l | X_{t-2}=h, X_{t-1}=j)
# Dims: (draw, h, j, l)
# So [, 1, 1, 1] give all draws for P(l = 1 | h = 1, j = 1)
# [, 1, 1, ] for all draws for P(l | h = 1, j = 1) (rows must sum to 1)
cond_probs_4d_array <- array(0, dim = c(n_draws, M, M, M))
for (d in 1:n_draws) {
# Create a placeholder matrix for this specific draw
# each row is combination of (h, j), each column is l (rows must sum to 1)
# this is for one draw only
cond_probs_matrix_d <- matrix(0, nrow = M * M, ncol = M)
# Handle absorbing states (P(j->j) = 1 if j is absorbing)
absorbing_idx <- which(prediction_data[[col_yprev]] %in% absorbing_states)
if (length(absorbing_idx) > 0) {
absorbing_state <- as.numeric(as.character(prediction_data[[col_yprev]][absorbing_idx]))
rows_and_cols <- cbind(absorbing_idx, absorbing_state)
cond_probs_matrix_d[rows_and_cols] <- 1.0
}
# Get state probabilities for the current draw
start_row <- (d - 1) * n_predictable + 1
end_row <- d * n_predictable
draw_d_probs <- state_probs_all_draws[start_row:end_row, ]
# Place them into the correct rows
cond_probs_matrix_d[predictable_idx, ] <- draw_d_probs
# Reshape and store in the final 4D array
cond_probs_4d_array[d, , , ] <- array(cond_probs_matrix_d, dim = c(M, M, M))
}
# --- Update Probabilities ---
# Get joint probabilities from the previous step for all draws: P(X_{t-2}, X_{t-1})
prev_P_joint_all_draws <- P_joint_bayes[, , , t]
current_P_joint_all_draws <- array(0, dim = c(n_draws, M, M))
for (d in 1:n_draws) {
# This is P(X_{t-2}=h, X_{t-1}=j) for draw d
prev_P_joint_d <- prev_P_joint_all_draws[d, , ]
# P(X_t = l | X_{t-1} = j, X_{t-2} = h) for this draw
cond_probs_d <- cond_probs_4d_array[d, , , ]
# Law of Total Probability: P(X_t=l|h,j) * P(h,j)
weighted_probs <- cond_probs_d * array(prev_P_joint_d, dim = dim(cond_probs_d))
# Marginalize over the t-2 state (h) to get P(X_{t-1}, X_t)
# = Summation over h (the state at t-2), resulting in P(X_{t-1}=j, X_t=l)
current_P_joint_d <- apply(weighted_probs, c(2, 3), sum)
current_P_joint_all_draws[d, , ] <- current_P_joint_d
# Marginalize over the t-1 state (j) to get P(X_t)
uncond_probs_t <- colSums(current_P_joint_d)
# Consider normalizing to prevent floating point error accumulation (not doing this for now)
P_uncond_bayes[d, , t + 1] <- uncond_probs_t
# P_uncond_bayes[d, , t + 1] <- uncond_probs_t / sum(uncond_probs_t)
}
P_joint_bayes[, , , t + 1] <- current_P_joint_all_draws
}
return(P_uncond_bayes)
}
M is the set of possible states, l is state at time t, j is state t-1, h is state t-2.
-
We start with a known history at t=0, so the joint probability of the observed j_{obs} and h_{obs} and states for X_{t-1} and X_{t-2} is 1.0
-
We use this known probability as the “weight” to calculate the joint probabilities for t=1, t=0.
- P(X_{t=1} = l, X_{t-1} = j\_{obs}) = P(X_{t=1} | X_{t-1}=j_{obs}, X_{t-2}=h_{obs}) \times 1
- So for t=1 we have M non-zero joint probabilities, one for each possible state M at time t=1.
-
We then use the joint probabilities from t = 1, t = 0 as the weights to calculate the joint probabilities for t=2, t = 1.
- P(X_{t=2} = l, X_{t-1} = j) = \sum_{h=1}^{M} P(X_{t=2} | X_{t-1}=j, X_{t-2}=h) \times P(X_{t-1}=j, X_{t-2}=h)
- Now for t=2 we have M \times M joint probabilities, one for each combination of possible states at times t=1 and t=2.
-
…and so on.
-
Final step for any time t: Unconditional probability P(X_t=l), we sum the joint probabilities from the previous step over all possible previous states j.
- P(X_t=l)=\sum_{j=1}^M P(X_t=l,X_{t−1}=j)
plus handling of absorbing states I haven’t written in plain language yet.