Hierarchical models
Partial pooling across J groups, in the shape of the "eight schools" model.
mu_global ~ Normal(0, 10) global mean hyperprior
sigma_group ~ HalfNormal(5) between-group scale hyperprior
mu_j ~ Normal(mu_global, sigma_group) group mean, j = 0 .. J-1
y_ij ~ Normal(mu_j, sigma_obs) observations within a group
mu_global and sigma_group are hyperparameters: the prior on each group mean is
itself estimated. That is what ties the groups together. A group with few
observations is pulled toward the global mean; a group with many stays near its own
sample mean.
The model is written in the conditional, "centered" form above, which is the form
that reads like the mathematics. rustmc compiles eligible scalar hierarchies to
noncentered sampling coordinates internally, so the awkward geometry of the centered
form does not reach the sampler. mu_j is still what you see in summaries,
diagnostics and posterior draws.
This is the end-to-end workflow for that model: prior predictive check, fit, recovery
of the simulated values, shrinkage and posterior predictive check. Two other examples
fit an eight-group Gaussian hierarchy for different reasons. partial_pooling_template.py
maps the builder surface: which priors accept a parameter as a hyperparameter, how the
compiled coordinates are named, and what is not supported yet. site_effects.py pools
a vector of site effects over unequal sample counts, written noncentered by hand.
Run it with python examples/hierarchical_example.py from the repository root. The blocks below are
cells of that one file and share its state, so they assume the blocks above
them have run.
import numpy as np
import rustmc as rmc
from hierarchical_templates import build_centered_normal_partial_pooling
Simulate data
rng = np.random.default_rng(42)
J = 8 # number of groups
sigma_obs = 2.0 # known within-group noise
N_per_group = 30 # observations per group
mu_global_true = 2.5
sigma_group_true = 3.0
mu_true = rng.normal(mu_global_true, sigma_group_true, J)
ys = [rng.normal(mu_true[j], sigma_obs, N_per_group) for j in range(J)]
data = {f"y_{j}": ys[j] for j in range(J)}
print("Simulated data")
print(f" True mu_global = {mu_global_true:.2f}")
print(f" True sigma_group = {sigma_group_true:.2f}")
print(f" True mu_j = {np.round(mu_true, 3).tolist()}")
Simulated data
True mu_global = 2.50
True sigma_group = 3.00
True mu_j = [3.414, -0.62, 4.751, 5.322, -3.353, -1.407, 2.884, 1.551]
Build the hierarchy
# `build_centered_normal_partial_pooling` lives in examples/hierarchical_templates.py
# and does nothing you could not write inline: one `normal_prior` for the global
# mean, one `half_normal_prior` for the between-group scale, then one
# `normal_prior(mu=mu_global, sigma=sigma_group)` and one likelihood per group.
builder = rmc.ModelBuilder(data=data)
template = build_centered_normal_partial_pooling(
builder,
observed_keys=[f"y_{j}" for j in range(J)],
sigma_obs=sigma_obs,
)
mu_global = template.mu_global
sigma_group = template.sigma_group
mu_j = template.group_params
model = builder.build()
Prior predictive check
prior_pred = rmc.sample_prior_predictive(model, n_samples=200, seed=0)
print(f" mu_global prior: mean={prior_pred['mu_global'].mean():.2f}, std={prior_pred['mu_global'].std():.2f}")
print(f" sigma_group prior: mean={prior_pred['sigma_group'].mean():.2f}, std={prior_pred['sigma_group'].std():.2f}")
for j in range(J):
key = f"obs_{j}"
if key in prior_pred:
prior_y = prior_pred[key]
print(f" Group {j} prior y range: [{prior_y.min():.1f}, {prior_y.max():.1f}]")
mu_global prior: mean=0.20, std=9.17
sigma_group prior: mean=4.05, std=2.91
Group 0 prior y range: [-35.2, 30.4]
Group 1 prior y range: [-35.9, 38.7]
Group 2 prior y range: [-27.4, 42.6]
Group 3 prior y range: [-44.9, 38.0]
Group 4 prior y range: [-46.8, 37.4]
Group 5 prior y range: [-33.2, 34.6]
Group 6 prior y range: [-36.9, 31.6]
Group 7 prior y range: [-35.7, 33.6]
Sample
fit = rmc.sample(
model_spec=model,
chains=4,
draws=2000,
warmup=1000,
seed=42,
)
print(fit.summary())
4 chains × 2000 draws per chain
Parameter mean std hdi_3% hdi_97% ess_bulk ess_tail r_hat mcse_mean
──────────────────────────────────────────────────────────────────────────────────────────────
mu_global 1.3507 1.3124 -0.8568 4.0559 896 1124 1.0044 0.044351
sigma_group 3.5300 1.0647 1.8405 5.4506 696 1149 1.0051 0.040945
mu_0 3.6672 0.3569 3.0189 4.3508 10167 6401 1.0005 0.003539
mu_1 -0.4870 0.3610 -1.1705 0.1866 10421 6715 1.0003 0.003538
mu_2 4.1951 0.3649 3.5170 4.8808 9743 6749 1.0002 0.003699
mu_3 5.2365 0.3651 4.5668 5.9224 11182 7228 1.0005 0.003450
mu_4 -3.4856 0.3605 -4.1611 -2.8083 10306 7084 0.9998 0.003552
mu_5 -1.3029 0.3613 -1.9706 -0.6305 10077 5793 1.0005 0.003597
mu_6 2.8564 0.3634 2.1738 3.5527 10481 5961 1.0007 0.003551
mu_7 1.1118 0.3627 0.4668 1.8277 9655 6049 1.0002 0.003690
──────────────────────────────────────────────────────────────────────────────────────────────
Mean accept rate: 0.91 │ Divergences: 0
Recover the parameters
means = fit.mean()
stds = fit.std()
print(f"{'Parameter':<15} {'True':>8} {'Estimate':>10} {'Std':>8}")
print("-" * 45)
print(f"{'mu_global':<15} {mu_global_true:>8.2f} {means['mu_global']:>10.4f} {stds['mu_global']:>8.4f}")
print(f"{'sigma_group':<15} {sigma_group_true:>8.2f} {means['sigma_group']:>10.4f} {stds['sigma_group']:>8.4f}")
for j in range(J):
key = f"mu_{j}"
print(f" {key:<13} {mu_true[j]:>8.2f} {means[key]:>10.4f} {stds[key]:>8.4f}")
print()
print("Step sizes:", [round(s, 5) for s in fit.step_sizes()])
print("Divergences:", fit.divergences())
Parameter True Estimate Std
---------------------------------------------
mu_global 2.50 1.3507 1.3124
sigma_group 3.00 3.5300 1.0647
mu_0 3.41 3.6672 0.3569
mu_1 -0.62 -0.4870 0.3610
mu_2 4.75 4.1951 0.3649
mu_3 5.32 5.2365 0.3651
mu_4 -3.35 -3.4856 0.3605
mu_5 -1.41 -1.3029 0.3613
mu_6 2.88 2.8564 0.3634
mu_7 1.55 1.1118 0.3627
Step sizes: [0.1072, 0.10161, 0.07235, 0.10609]
Divergences: [0, 0, 0, 0]
What the hyperparameters can and cannot say
mu_global has a wide posterior and it should. Eight group means drawn from a
distribution carry about as much information about that distribution's mean as
eight observations do, so the interval stays broad however many draws you take.
sigma_group is estimated from the same eight numbers and on this dataset comes
out above the value that generated it, by about half a posterior standard
deviation. That is one draw of eight groups, not evidence about the estimator;
it is what a weakly determined scale looks like.
The centered form of this model is the textbook case of Neal's funnel, where the
sampler stalls in the neck and reports divergent transitions. rustmc compiles
this hierarchy to noncentered coordinates, which is the standard remedy, and
this run reports none. examples/partial_pooling_template.py shows the rewrite
in CompiledModel.param_names. If you write a hierarchy rustmc cannot recognise
and see divergences, reparameterise it by hand before trusting the
hyperparameter estimates.
Partial pooling
print("Partial pooling effect (shrinkage toward the global mean):")
print(f" Global mean estimate: {means['mu_global']:.2f}")
sample_means = [ys[j].mean() for j in range(J)]
for j in range(J):
est = means[f"mu_{j}"]
raw = sample_means[j]
print(f" Group {j}: raw={raw:+.2f} pooled={est:+.2f} true={mu_true[j]:+.2f}")
print()
print("Each group has 30 observations and known noise, so the data pins mu_j down")
print("and the pull toward the global mean is small. Shrinkage grows as a group's")
print("sample size falls; examples/site_effects.py shows it with unequal counts.")
Partial pooling effect (shrinkage toward the global mean):
Global mean estimate: 1.35
Group 0: raw=+3.69 pooled=+3.67 true=+3.41
Group 1: raw=-0.51 pooled=-0.49 true=-0.62
Group 2: raw=+4.23 pooled=+4.20 true=+4.75
Group 3: raw=+5.29 pooled=+5.24 true=+5.32
Group 4: raw=-3.56 pooled=-3.49 true=-3.35
Group 5: raw=-1.34 pooled=-1.30 true=-1.41
Group 6: raw=+2.87 pooled=+2.86 true=+2.88
Group 7: raw=+1.11 pooled=+1.11 true=+1.55
Each group has 30 observations and known noise, so the data pins mu_j down
and the pull toward the global mean is small. Shrinkage grows as a group's
sample size falls; examples/site_effects.py shows it with unequal counts.
Posterior predictive check
ppc = fit.posterior_predictive(n_samples=500, seed=42)
print(f" Likelihood keys in PPC: {sorted(ppc.keys())}")
for j in range(J):
key = f"obs_{j}"
if key in ppc:
y_rep = ppc[key] # (n_samples, N_per_group)
y_obs = ys[j]
lower = y_rep.mean(axis=0) - 2 * y_rep.std(axis=0)
upper = y_rep.mean(axis=0) + 2 * y_rep.std(axis=0)
inside = ((lower < y_obs) & (y_obs < upper)).mean()
print(f" Group {j}: obs mean={y_obs.mean():.2f} ppc mean={y_rep.mean():.2f} within +/-2 sd={inside:.2%}")
Likelihood keys in PPC: ['obs_0', 'obs_1', 'obs_2', 'obs_3', 'obs_4', 'obs_5', 'obs_6', 'obs_7']
Group 0: obs mean=3.69 ppc mean=3.67 within +/-2 sd=96.67%
Group 1: obs mean=-0.51 ppc mean=-0.47 within +/-2 sd=100.00%
Group 2: obs mean=4.23 ppc mean=4.14 within +/-2 sd=100.00%
Group 3: obs mean=5.29 ppc mean=5.27 within +/-2 sd=100.00%
Group 4: obs mean=-3.56 ppc mean=-3.50 within +/-2 sd=90.00%
Group 5: obs mean=-1.34 ppc mean=-1.30 within +/-2 sd=100.00%
Group 6: obs mean=2.87 ppc mean=2.84 within +/-2 sd=96.67%
Group 7: obs mean=1.11 ppc mean=1.12 within +/-2 sd=93.33%