Approach 1: Two-Stage (brms + shrinkr)
Stage 1: Fit Cox Model
We begin by fitting a Cox proportional hazards model that allows the treatment effect to vary by cell type. At this stage we estimate subgroup-specific treatment effects without adding hierarchical shrinkage across cell types. That hierarchical regularization is introduced in Stage 2.
brms_uninformative <- brm(
time | cens(1 - status) ~ trt:celltype + karno + age,
data = veteran,
family = cox(),
chains = 4,
iter = 4000,
warmup = 1000,
seed = 123
)
brms_uninformative_summary <- capture.output(print(summary(brms_uninformative)))What this model does:
- Estimates a separate log hazard ratio for the test treatment in each cell type
- Adjusts for baseline performance status (
karno) and age - Leaves pooling across cell types to the second-stage hierarchical model
Results:
cat(brms_uninformative_summary, sep = "\n")
#> Family: cox
#> Links: mu = log
#> Formula: time | cens(1 - status) ~ trt:celltype + karno + age
#> Data: veteran (Number of observations: 137)
#> Draws: 4 chains, each with iter = 4000; warmup = 1000; thin = 1;
#> total post-warmup draws = 12000
#>
#> Regression Coefficients:
#> Estimate Est.Error l-95% CI
#> Intercept 4.16 0.74 2.70
#> karno -0.03 0.01 -0.04
#> age -0.01 0.01 -0.02
#> trt:celltypesquamous -0.12 0.21 -0.52
#> trt:celltypesmallcell 0.51 0.23 0.05
#> trt:celltypeadeno 0.54 0.20 0.13
#> trt:celltypelarge 0.16 0.23 -0.31
#> u-95% CI Rhat Bulk_ESS Tail_ESS
#> Intercept 5.61 1.00 9000 8957
#> karno -0.02 1.00 11647 9539
#> age 0.01 1.00 9875 8676
#> trt:celltypesquamous 0.29 1.00 5515 7429
#> trt:celltypesmallcell 0.95 1.00 5310 7315
#> trt:celltypeadeno 0.94 1.00 5398 7230
#> trt:celltypelarge 0.61 1.00 5399 7077
#>
#> Draws were sampled using sampling(NUTS). For each parameter, Bulk_ESS
#> and Tail_ESS are effective sample size measures, and Rhat is the potential
#> scale reduction factor on split chains (at convergence, Rhat = 1).Stage 2: Apply Hierarchical Shrinkage
Now we extract the cell type-specific treatment effect posteriors and apply hierarchical shrinkage.
Step 1: Extract posterior samples
brms_posteriors <- brms_uninformative %>%
spread_draws(`b_trt:celltypesquamous`, `b_trt:celltypesmallcell`,
`b_trt:celltypeadeno`, `b_trt:celltypelarge`) %>%
select(-c(.chain, .iteration, .draw)) %>%
pivot_longer(everything(), names_to = "celltype", values_to = "value") %>%
mutate(celltype = gsub("b_trt:celltype", "", celltype)) %>%
group_by(celltype) %>%
summarise(draws = list(matrix(value, ncol = 1)), .groups = "drop") %>%
deframe()brms_posteriors is a named list containing posterior
draws for each cell type.
Step 2: Fit a Gaussian mixture approximation
The fit_mixture() function approximates each subgroup
posterior with a mixture of Gaussian components. This creates a flexible
representation of the Stage 1 posterior that can be passed to
shrink().
Understanding the mixture approximation:
- Each posterior is approximated as a weighted sum of Gaussian components
- The number of components is chosen separately for each cell type
- This allows the approximation to capture skewness or heavier tails when needed
Step 3: Apply a hierarchical prior
priors_moderate <- list(
mu = dist_normal(0, 1),
tau = dist_truncated(dist_normal(0, 0.5), lower = 0)
)fit_twostage_brms <- shrink(
mixture = mix_brms,
hierarchical_priors = priors_moderate,
chains = 4,
iter = 4000,
warmup = 1000,
seed = 456
)
moderate_brms_output <- capture.output(print(fit_twostage_brms))Results:
cat(moderate_brms_output, sep = "\n")
#> # A tibble: 3 × 7
#> variable mean sd q2.5 q50 q97.5 rhat
#> <chr> <dbl> <dbl> <dbl> <dbl> <dbl> <dbl>
#> 1 mu 0.246 0.265 -0.272 0.245 0.757 1.00
#> 2 tau 0.361 0.176 0.123 0.324 0.812 1.00
#> 3 tau_squared 0.161 0.177 0.0150 0.105 0.659 1.00Interpreting the shrinkage:
- Cell type-specific estimates are pulled toward the overall mean
(
mu) - The amount of shrinkage depends on
tau, the between-cell-type heterogeneity - Smaller
tauimplies stronger pooling - Larger
tauimplies weaker pooling
