This vignette demonstrates
advanced functionalities supported by bonsaiforest2 and the
underlying brms package, including adding an offset
to negative binomial regression models for counts,
customizing prior distributions, and applying
stratification to nuisance parameters.
For count data, the observed value often depends on the exposure time or follow-up time. Using an offset variable is the standard statistical method to properly account for this time variation by modeling the event rate (count per unit time) instead of the raw count.
In bonsaiforest2, you can include the offset directly in
the response_formula using the syntax:
count ~ trt + offset(log_fup_duration).
The example uses the shrink_data dataset to illustrate a
negative binomial model for variable count accounting for
treatment trt and the subgrouping variables
x1, x2, and x3. Each subject’s
follow-up time is provided in variable fup_duration whose
log-transformed value is included as an offset.
# Load packages and data
library(bonsaiforest2)
library(brms)
#> Loading required package: Rcpp
#> Loading 'brms' package (version 2.23.0). Useful instructions
#> can be found by typing help('brms'). A more detailed introduction
#> to the package is available through vignette('brms_overview').
#>
#> Attaching package: 'brms'
#> The following object is masked from 'package:stats':
#>
#> ar
# Use cmdstanr by default if available
if (requireNamespace("cmdstanr", quietly = TRUE)) {
prev_backend <- options(brms.backend = "cmdstanr")
}
shrink_data <- bonsaiforest2::shrink_data
# prepare offset variable
shrink_data$log_fup_duration <- log(shrink_data$fup_duration)count_model_fit <- run_brms_analysis(
data = shrink_data,
response_type = "count",
response_formula = count ~ trt + offset(log_fup_duration),
unshrunk_terms_formula = ~ 1 + x1 + x2 + x3,
shrunk_predictive_formula = ~ 0 + trt:x1 + trt:x2 + trt:x3,
shrunk_predictive_prior = "horseshoe(scale_global = 1)",
chains = 2, iter = 1000, warmup = 500, cores = 2, refresh = 0
)
#> Step 1: Preparing formula and data...
#>
#> Step 2: Fitting the brms model...
#> Using default priors for unspecified effects:
#> - intercept: brms default
#> - unshrunk terms: brms default
#> Fitting brms model...
#> Start sampling
#> Running MCMC with 2 parallel chains...
#>
#> Chain 2 finished in 2.0 seconds.
#> Chain 1 finished in 2.3 seconds.
#>
#> Both chains finished successfully.
#> Mean chain execution time: 2.1 seconds.
#> Total execution time: 2.4 seconds.
#> Warning: 3 of 1000 (0.0%) transitions ended with a divergence.
#> See https://mc-stan.org/misc/warnings for details.
#> Loading required namespace: rstan
#>
#> Analysis complete.
count_model_summary <- summary_subgroup_effects(brms_fit = count_model_fit)
#> --- Calculating specific subgroup effects... ---
#> Step 1: Identifying subgroups and creating counterfactuals...
#> ...detected subgroup variable(s): x1_onehot, x2_onehot, x3_onehot
#> Step 2: Generating posterior predictions...
#> Step 3: Calculating marginal effects...
#> Done.
plot(count_model_summary, title = "Global model for count data: All subgrouping variables")
#> Preparing data for plotting...
#> Generating plot...
#> Done.In Bayesian analysis, the choice of priors is important and can
materially affect the results. For most settings, the brms
default of flat priors for the intercept (intercept_prior)
and unshrunk terms (unshrunk_prior, includes the prior for
the main treatment effect) are sensible and do not need fine-tuning.
However, we recommend that the user carefully considers the
prior for shrunken terms, in particular those for
covariate-by-treatment interactions
(shrunk_predictive_prior).
bonsaiforest2 gives the user full flexibility to specify
priors as described via examples below.
The default priors used by
bonsaiforest2 are as follows:
| Prior Component | Parameter Name | Default Prior | |
|---|---|---|---|
| Intercept | intercept_prior |
NULL (i.e. brms default) |
|
| Unshrunk Terms | unshrunk_prior |
NULL (i.e. brms default) |
|
| Shrunk Prognostic | shrunk_prognostic_prior |
horseshoe(scale_global = 1); in case the
formula specifies random effects (pipe-pipe syntax), automatically uses
a normal prior with a half-normal hyperprior with heterogeneity
parameter 1 for the standard deviation instead |
|
| Shrunk Predictive | shrunk_predictive_prior |
horseshoe(scale_global = 1); in case the
formula specifies random effects (pipe-pipe syntax), automatically uses
a normal prior with a half-normal hyperprior with heterogeneity
parameter 1 for the standard deviation instead |
The following examples demonstrate how to customize priors in
bonsaiforest2. We use the shrink_data package
dataset with a time-to-event outcome to show different prior
specification strategies, from simple defaults to advanced hierarchical
structures.
# Prepare basic model formulas
prepared_model <- prepare_formula_model(
data = shrink_data,
response_type = "survival",
response_formula = Surv(tt_event, event_yn) ~ trt,
unshrunk_terms_formula = ~ x1 + x2,
shrunk_predictive_formula = ~ 0 + trt:x1
)
#> Response type is 'survival'. Modeling the baseline hazard explicitly using bhaz().brmsThe simplest approach: use the priors directly supported by
brms to adjust priors of different terms. Below, we assume
that the trial targeted a hazard ratio of 0.7, i.e. \(\delta_{plan}=|log(0.7)\), and set
scale_global of the regularized horseshoe prior to this
value.
fit_ex2 <- fit_brms_model(
prepared_model = prepared_model,
intercept_prior = "normal(0, 5)",
unshrunk_prior = "normal(0, 2.5)",
shrunk_predictive_prior = "horseshoe(scale_global = abs(log(0.7)), scale_slab = 2, df_slab = 4)",
chains = 2, iter = 1000, warmup = 500, cores = 2, refresh = 0
)
#> Fitting brms model...
#> Start sampling
#> Running MCMC with 2 parallel chains...
#>
#> Chain 1 finished in 1.6 seconds.
#> Chain 2 finished in 1.6 seconds.
#>
#> Both chains finished successfully.
#> Mean chain execution time: 1.6 seconds.
#> Total execution time: 1.8 seconds.
# View the priors that were automatically set
cat("\n=== Priors Used ===\n")
#>
#> === Priors Used ===
print(fit_ex2[["prior"]])
#> prior class
#> horseshoe(scale_global = abs(log(0.7)), scale_slab = 2, df_slab = 4) b
#> horseshoe(scale_global = abs(log(0.7)), scale_slab = 2, df_slab = 4) b
#> horseshoe(scale_global = abs(log(0.7)), scale_slab = 2, df_slab = 4) b
#> normal(0, 2.5) b
#> normal(0, 2.5) b
#> normal(0, 2.5) b
#> normal(0, 2.5) b
#> normal(0, 2.5) b
#> normal(0, 2.5) b
#> dirichlet(1) sbhaz
#> coef group resp dpar nlpar lb ub tag source
#> shpredeffect user
#> trt:x1_onehota shpredeffect (vectorized)
#> trt:x1_onehotb shpredeffect (vectorized)
#> unshrunktermeffect user
#> trt unshrunktermeffect (vectorized)
#> x1a unshrunktermeffect (vectorized)
#> x1b unshrunktermeffect (vectorized)
#> x2b unshrunktermeffect (vectorized)
#> x2c unshrunktermeffect (vectorized)
#> defaultAn alternative global-local shrinkage prior to the regularized
horseshoe prior is the R2-D2 prior. In the linear regression setting, it
can be motivated as specifying a prior on the model’s coefficient of
determination (\(R^2\)) first, and then
distributing the prior through to the coefficients. The R2-D2 prior is
also directly supported by brms as illustrated below.
fit_ex3 <- fit_brms_model(
prepared_model = prepared_model,
shrunk_predictive_prior = "R2D2(mean_R2 = 0.5, prec_R2 = 1)",
chains = 2, iter = 1000, warmup = 500, cores = 2, refresh = 0
)
#> Using default priors for unspecified effects:
#> - unshrunk terms: brms default
#> Fitting brms model...
#> Start sampling
#> Running MCMC with 2 parallel chains...
#>
#> Chain 2 finished in 1.2 seconds.
#> Chain 1 finished in 1.4 seconds.
#>
#> Both chains finished successfully.
#> Mean chain execution time: 1.3 seconds.
#> Total execution time: 1.5 seconds.
#> Warning: 14 of 1000 (1.0%) transitions ended with a divergence.
#> See https://mc-stan.org/misc/warnings for details.
cat("\n=== Priors Used ===\n")
#>
#> === Priors Used ===
print(fit_ex3[["prior"]])
#> prior class coef group resp dpar
#> R2D2(mean_R2 = 0.5, prec_R2 = 1) b
#> R2D2(mean_R2 = 0.5, prec_R2 = 1) b trt:x1_onehota
#> R2D2(mean_R2 = 0.5, prec_R2 = 1) b trt:x1_onehotb
#> (flat) b
#> (flat) b trt
#> (flat) b x1a
#> (flat) b x1b
#> (flat) b x2b
#> (flat) b x2c
#> dirichlet(1) sbhaz
#> nlpar lb ub tag source
#> shpredeffect user
#> shpredeffect (vectorized)
#> shpredeffect (vectorized)
#> unshrunktermeffect default
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> defaultThis example demonstrates injecting raw Stan code using
stanvars. This is necessary if you want to implement a
hierarchical structure that brms does not support natively,
such as estimating a shared variance parameter across coefficients.
# 1. Define new hyperparameters in Stan
stanvars_full_hierarchical <- brms::stanvar(
scode = " real mu_pred;\n real<lower=0> sigma_pred;\n",
block = "parameters"
) +
# Add priors for these parameters
brms::stanvar(
scode = " // Priors on the hierarchical parameters\n target += normal_lpdf(mu_pred | 0, 4); \n target += normal_lpdf(sigma_pred | 0, 1) - normal_lccdf(0 | 0, 1); \n",
block = "model"
)
# 2. Create prior that references the Stan variables
prior_full_hierarchical <- brms::set_prior("normal(mu_pred, sigma_pred)")
# 3. Pass both to fit_brms_model
fit_ex4 <- fit_brms_model(
prepared_model = prepared_model,
intercept_prior = "normal(0, 5)",
unshrunk_prior = "normal(0, 2.5)",
shrunk_prognostic_prior = "horseshoe(scale_global = 1)",
shrunk_predictive_prior = prior_full_hierarchical,
stanvars = stanvars_full_hierarchical,
chains = 2, iter = 1000, warmup = 500, cores = 2, refresh = 0
)
#> Fitting brms model...
#> Start sampling
#> Running MCMC with 2 parallel chains...
#>
#> Chain 2 finished in 3.0 seconds.
#> Chain 1 finished in 3.1 seconds.
#>
#> Both chains finished successfully.
#> Mean chain execution time: 3.1 seconds.
#> Total execution time: 3.2 seconds.
# View the used priors
cat("\n=== Priors Used ===\n")
#>
#> === Priors Used ===
print(fit_ex4[["prior"]])
#> prior class coef group resp dpar
#> normal(mu_pred, sigma_pred) b
#> normal(mu_pred, sigma_pred) b trt:x1_onehota
#> normal(mu_pred, sigma_pred) b trt:x1_onehotb
#> normal(0, 2.5) b
#> normal(0, 2.5) b trt
#> normal(0, 2.5) b x1a
#> normal(0, 2.5) b x1b
#> normal(0, 2.5) b x2b
#> normal(0, 2.5) b x2c
#> dirichlet(1) sbhaz
#> nlpar lb ub tag source
#> shpredeffect user
#> shpredeffect (vectorized)
#> shpredeffect (vectorized)
#> unshrunktermeffect user
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> default
print(fit_ex4[["stanvars"]])
#> [[1]]
#> [[1]]$name
#> [1] ""
#>
#> [[1]]$sdata
#> NULL
#>
#> [[1]]$scode
#> [1] " real mu_pred;\n real<lower=0> sigma_pred;\n"
#>
#> [[1]]$block
#> [1] "parameters"
#>
#> [[1]]$position
#> [1] "start"
#>
#> [[1]]$pll_args
#> character(0)
#>
#>
#> [[2]]
#> [[2]]$name
#> [1] ""
#>
#> [[2]]$sdata
#> NULL
#>
#> [[2]]$scode
#> [1] " // Priors on the hierarchical parameters\n target += normal_lpdf(mu_pred | 0, 4); \n target += normal_lpdf(sigma_pred | 0, 1) - normal_lccdf(0 | 0, 1); \n"
#>
#> [[2]]$block
#> [1] "model"
#>
#> [[2]]$position
#> [1] "start"
#>
#> [[2]]$pll_args
#> character(0)
#>
#>
#> attr(,"class")
#> [1] "stanvars"By default, bonsaiforest2 uses the same prior for all
coefficients assigned to the group of “unshrunken terms”, “shrunken
prognostic terms”, or “shrunken predictive terms”, respectively.
However, it is possible to fine-tune this by assigning a different prior
to one or several coefficients within a group.
The example below illustrates this. It includes the interaction
trt*x3 as an “unshrunken term” and assigns a tighter prior
to the regression coefficients for trt:x3b,
trt:x3c, and trt:x3d, respectively, compared
to other coefficients in the same group.
To implement this, we first call prepare_formula_model()
to extract the exact coefficient names that will be included in the
model. Then, we define coefficient-specific priors for these by
combining brms::set_prior calls.
# 1. Run prepare_formula_model
prepared_model_ex5 <- prepare_formula_model(
data = shrink_data,
response_type = "survival",
response_formula = Surv(tt_event, event_yn) ~ trt,
unshrunk_terms_formula = ~ x1 + trt*x3,
shrunk_predictive_formula = ~ 0 + trt:x1,
)
#> Response type is 'survival'. Modeling the baseline hazard explicitly using bhaz().
# 2. Inspect the Results
# A. The generated brms formula object
print(prepared_model_ex5$formula)
#> tt_event | cens(1 - event_yn) + bhaz(Boundary.knots = c(0, 60.500476574435), knots = c(14.0610682872923, 36.1959635027271, 46.047044984117), intercept = FALSE) ~ unshrunktermeffect + shpredeffect
#> unshrunktermeffect ~ 0 + x1 + trt * x3 + trt
#> shpredeffect ~ 0 + trt:x1_onehot
# B. The processed terms
print(prepared_model_ex5$stan_variable_names$X_unshrunktermeffect)
#> [1] "x1a" "x1b" "trt" "x3b" "x3c" "x3d" "trt:x3b"
#> [8] "trt:x3c" "trt:x3d"
cat("=== Prior Strategy ===\n")
#> === Prior Strategy ===
cat("General unshrunk prior: normal(0, 5)\n")
#> General unshrunk prior: normal(0, 5)
cat("Specific for trt:x3 interactions: normal(0, 1)\n\n")
#> Specific for trt:x3 interactions: normal(0, 1)
# Create combined prior object
# IMPORTANT: Use EXACT coefficient names from the prepared data
unshrunk_priors_combined <- c(
brms::set_prior("normal(0, 5)", class = "b"), # General
brms::set_prior("normal(0, 1)", class = "b", coef = "trt:x3b"),
brms::set_prior("normal(0, 1)", class = "b", coef = "trt:x3c"),
brms::set_prior("normal(0, 1)", class = "b", coef = "trt:x3d")
)
# Fit the model
fit_ex5 <- fit_brms_model(
prepared_model = prepared_model_ex5,
intercept_prior = "normal(0, 5)",
unshrunk_prior = unshrunk_priors_combined, # Pass the combined object
shrunk_predictive_prior = "horseshoe(scale_global = 1)",
chains = 2, iter = 1000, warmup = 500, cores = 2, refresh = 0
)
#> Fitting brms model...
#> Start sampling
#> Running MCMC with 2 parallel chains...
#>
#> Chain 1 finished in 1.8 seconds.
#> Chain 2 finished in 1.8 seconds.
#>
#> Both chains finished successfully.
#> Mean chain execution time: 1.8 seconds.
#> Total execution time: 1.9 seconds.
#> Warning: 6 of 1000 (1.0%) transitions ended with a divergence.
#> See https://mc-stan.org/misc/warnings for details.
# View the used priors
cat("\n=== Priors Used ===\n")
#>
#> === Priors Used ===
print(fit_ex5[["prior"]])
#> prior class coef group resp dpar
#> horseshoe(scale_global = 1) b
#> horseshoe(scale_global = 1) b trt:x1_onehota
#> horseshoe(scale_global = 1) b trt:x1_onehotb
#> normal(0, 5) b
#> normal(0, 5) b trt
#> normal(0, 1) b trt:x3b
#> normal(0, 1) b trt:x3c
#> normal(0, 1) b trt:x3d
#> normal(0, 5) b x1a
#> normal(0, 5) b x1b
#> normal(0, 5) b x3b
#> normal(0, 5) b x3c
#> normal(0, 5) b x3d
#> dirichlet(1) sbhaz
#> nlpar lb ub tag source
#> shpredeffect user
#> shpredeffect (vectorized)
#> shpredeffect (vectorized)
#> unshrunktermeffect user
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect user
#> unshrunktermeffect user
#> unshrunktermeffect user
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> unshrunktermeffect (vectorized)
#> defaultStratified models which assume different residual variances
(continuous endpoints), overdispersion parameters (count endpoints), or
baseline hazards (time-to-event endpoints) across strata are also
supported in bonsaiforest2 (via brms). The
argument stratification_formula defines the grouping
factor(s).
We fit a one-way model for the continuous outcome y in
shrink_data with subgroups defined by the binary covariate
x1 with adjustment for x2 and x3
as prognostic factors.
In the model specification below, we assume separate regression
coefficients for x2 and x3, respectively,
across strata defined by x1 (specified via interaction
terms x1*x2 and x1*x3) as well as different
residual variances across strata (specified via argument
stratification_formula). The only parameter that is shared
across strata is the shrinkage parameter for treatment effect
heterogeneity.
# model with x1*x2 and x1*x3 (separate regression coef per level of x1 for x2 and x3), and separate variance per x1 strata
oneway_x1_flex_strat <- run_brms_analysis(
data = shrink_data,
response_type = "continuous",
response_formula = y ~ trt,
stratification_formula = ~ x1,
unshrunk_terms_formula = ~ x1 + x1*x2 + x1*x3,
shrunk_predictive_formula = ~ (0 + trt || x1),
shrunk_predictive_prior = set_prior("normal(0, 0.3)", class = "sd"),
chains = 2, iter = 2000, warmup = 1000, cores = 2,
refresh = 0
)
#> Step 1: Preparing formula and data...
#> Applying stratification: estimating sigma by 'x1'.
#>
#> Step 2: Fitting the brms model...
#> Using default priors for unspecified effects:
#> - intercept: brms default
#> - unshrunk terms: brms default
#> Fitting brms model...
#> Start sampling
#> Running MCMC with 2 parallel chains...
#>
#> Chain 1 finished in 2.0 seconds.
#> Chain 2 finished in 2.0 seconds.
#>
#> Both chains finished successfully.
#> Mean chain execution time: 2.0 seconds.
#> Total execution time: 2.1 seconds.
#> Warning: 1 of 2000 (0.0%) transitions ended with a divergence.
#> See https://mc-stan.org/misc/warnings for details.
#>
#> Analysis complete.summary_oneway_x1_flex_strat <- summary_subgroup_effects(
brms_fit = oneway_x1_flex_strat
)
#> --- Calculating specific subgroup effects... ---
#> Step 1: Identifying subgroups and creating counterfactuals...
#> ...detected subgroup variable(s): x1
#> Step 2: Generating posterior predictions...
#> Step 3: Calculating marginal effects...
#> Done.
print(summary_oneway_x1_flex_strat$estimates)
#> # A tibble: 2 × 4
#> Subgroup Median CI_Lower CI_Upper
#> <chr> <dbl> <dbl> <dbl>
#> 1 x1: a 0.488 0.266 0.700
#> 2 x1: b 0.123 -0.193 0.436Assume that the primary analysis of the trial was stratified by
variable x2. A corresponding one-way model for
x1 which uses a Cox model stratified by x2 is
fitted below.
# Model fitting with stratified baseline hazard
fit_surv_oneway_x1_strat <- run_brms_analysis(
data = shrink_data,
response_type = "survival",
response_formula = Surv(tt_event, event_yn) ~ trt,
stratification_formula = ~ x2,
unshrunk_terms_formula = ~ x1 + x3,
shrunk_predictive_formula = ~ (0 + trt || x1),
shrunk_predictive_prior = set_prior("normal(0, abs(log(0.7)))", class = "sd"),
chains = 2, iter = 1000, warmup = 500, cores = 2, refresh = 0
)
#> Step 1: Preparing formula and data...
#> Response type is 'survival'. Modeling the baseline hazard explicitly using bhaz().
#> Applying stratification: estimating separate baseline hazards by 'x2'.
#>
#> Step 2: Fitting the brms model...
#> Using default priors for unspecified effects:
#> - unshrunk terms: brms default
#> Fitting brms model...
#> Start sampling
#> Running MCMC with 2 parallel chains...
#>
#> Chain 2 finished in 1.7 seconds.
#> Chain 1 finished in 1.9 seconds.
#>
#> Both chains finished successfully.
#> Mean chain execution time: 1.8 seconds.
#> Total execution time: 2.0 seconds.
#>
#> Analysis complete.
summary_surv_oneway_x1_strat <- summary_subgroup_effects(fit_surv_oneway_x1_strat)
#> --- Calculating specific subgroup effects... ---
#> Step 1: Identifying subgroups and creating counterfactuals...
#> ...detected subgroup variable(s): x1
#> Step 2: Generating posterior predictions...
#> Warning: Dropping 'draws_df' class as required metadata was removed.
#> Warning: Dropping 'draws_df' class as required metadata was removed.
#> Warning: Dropping 'draws_df' class as required metadata was removed.
#> Step 3: Calculating marginal effects...
#> Done.
print(summary_surv_oneway_x1_strat$estimates)
#> # A tibble: 2 × 4
#> Subgroup Median CI_Lower CI_Upper
#> <chr> <dbl> <dbl> <dbl>
#> 1 x1: a 0.384 0.265 0.540
#> 2 x1: b 0.690 0.477 0.964
plot(summary_surv_oneway_x1_strat)
#> Preparing data for plotting...
#> Generating plot...
#> Done.This vignette demonstrated several advanced functionalities supported
by bonsaiforest2:
These features provide researchers with powerful tools for complex trial analyses while maintaining the principled Bayesian shrinkage framework.