Skills › Scientific Agent Skills › pymc › hierarchical_model_template.py
hierarchical_model_template.py
114 lines · skills/pymc/assets/hierarchical_model_template.py at 5f2bc03
1"""Synthetic hierarchical regression with existing- and new-group predictions.
3Targets PyMC 6.3.2 / ArviZ 1.3. This model assumes independent group intercept
4and slope effects. For correlated effects use an LKJ Cholesky model instead.
5Short runs exercise APIs, not scientific validation.
6"""
8import argparse
9from pathlib import Path
11import arviz as az
12import matplotlib.pyplot as plt
13import numpy as np
14import pymc as pm
17def build_model(x, groups, y, group_names):
18 with pm.Model(coords={'groups': group_names, 'obs': np.arange(len(y))}) as model:
19 x_data = pm.Data('X_data', x, dims='obs')
20 groups_data = pm.Data('groups_data', groups, dims='obs')
21 mu_alpha = pm.Normal('mu_alpha', mu=0, sigma=5)
22 sigma_alpha = pm.HalfNormal('sigma_alpha', sigma=2)
23 mu_beta = pm.Normal('mu_beta', mu=0, sigma=2)
24 sigma_beta = pm.HalfNormal('sigma_beta', sigma=1)
25 alpha_offset = pm.Normal('alpha_offset', 0, 1, dims='groups')
26 beta_offset = pm.Normal('beta_offset', 0, 1, dims='groups')
27 alpha = pm.Deterministic('alpha', mu_alpha + sigma_alpha * alpha_offset, dims='groups')
28 beta = pm.Deterministic('beta', mu_beta + sigma_beta * beta_offset, dims='groups')
29 sigma = pm.HalfNormal('sigma', 1)
30 mu = alpha[groups_data] + beta[groups_data] * x_data
31 pm.Normal('y_obs', mu, sigma=sigma, observed=y, shape=x_data.shape[0], dims='obs')
32 return model
35def predict_new_group(idata, x_new, random_seed=42):
36 """Predict several observations sharing ONE previously unseen group.
38 Draw new group effects once per posterior draw and reuse across x_new;
39 add independent residual noise per observation. Returns (mean, outcome)
40 arrays with shape (chain, draw, new_observation).
41 """
42 x_new = np.asarray(x_new, dtype=float)
43 if x_new.ndim != 1 or not np.isfinite(x_new).all():
44 raise ValueError("x_new must be a finite vector")
45 rng = np.random.default_rng(random_seed)
46 posterior = idata.posterior
47 mu_a, mu_b = posterior['mu_alpha'].values, posterior['mu_beta'].values
48 alpha_new = mu_a + posterior['sigma_alpha'].values * rng.normal(size=mu_a.shape)
49 beta_new = mu_b + posterior['sigma_beta'].values * rng.normal(size=mu_b.shape)
50 mean = alpha_new[..., None] + beta_new[..., None] * x_new
51 outcome = mean + posterior['sigma'].values[..., None] * rng.normal(size=mean.shape)
52 return mean, outcome
55def run(draws=1000, tune=1000, chains=4, output_dir='hierarchical_results'):
56 rng = np.random.default_rng(42)
57 n_groups, n_per_group = 6, 12
58 group_names = [f'group_{i}' for i in range(n_groups)]
59 groups = np.repeat(np.arange(n_groups), n_per_group)
60 x = rng.normal(size=len(groups))
61 alpha = rng.normal(1, .6, n_groups)
62 beta = rng.normal(1.5, .3, n_groups)
63 y = alpha[groups] + beta[groups] * x + rng.normal(0, .5, len(x))
64 model = build_model(x, groups, y, group_names)
65 output = Path(output_dir)
66 output.mkdir(parents=True, exist_ok=True)
68 def save(collection, name):
69 collection.savefig(output / name, dpi=150, bbox_inches='tight')
70 plt.close(collection.viz['figure'].item())
72 with model:
73 prior = pm.sample_prior_predictive(draws=100, random_seed=40)
74 save(az.plot_ppc_dist(prior, group='prior_predictive', num_samples=30), 'prior_predictive.png')
75 # Inspect prior implications before fitting real data.
76 with model:
77 idata = pm.sample(draws=draws, tune=tune, chains=chains, cores=1,
78 nuts_sampler='pymc', target_accept=.95, random_seed=42,
79 progressbar=False)
80 pm.compute_log_likelihood(idata, progressbar=False)
81 pm.sample_posterior_predictive(idata, extend_inferencedata=True,
82 random_seed=43, progressbar=False)
83 variables = ['mu_alpha', 'sigma_alpha', 'mu_beta', 'sigma_beta', 'sigma']
84 summary = az.summary(idata, var_names=variables + ['alpha', 'beta'],
85 round_to='none', ci_kind='hdi', ci_prob=.95)
86 print(summary)
87 print('[INFO] Inspect diagnostics and prior sensitivity; six groups weakly constrain group scales.')
88 save(az.plot_trace_dist(idata, var_names=variables), 'trace.png')
89 save(az.plot_ppc_dist(idata, num_samples=30), 'posterior_predictive.png')
90 save(az.plot_forest(idata, var_names=['alpha', 'beta'], combined=True,
91 ci_kind='hdi', ci_probs=[.5, .95]), 'group_parameters.png')
92 new_x, new_groups = np.array([-1., 0., 1.]), np.array([0, 2, 4])
93 with model:
94 pm.set_data({'X_data': new_x, 'groups_data': new_groups}, coords={'obs': np.arange(3)})
95 predictions = pm.sample_posterior_predictive(idata, predictions=True,
96 random_seed=44, progressbar=False)
97 pm.set_data({'X_data': x, 'groups_data': groups}, coords={'obs': np.arange(len(y))})
98 interval = az.hdi(predictions.predictions['y_obs'], prob=.95)
99 interval.to_dataframe(name='prediction_interval').to_csv(output / 'existing_group_intervals.csv')
100 new_mean, new_outcome = predict_new_group(idata, new_x)
101 np.savez(output / 'new_group_predictions.npz', x=new_x, mean=new_mean, outcome=new_outcome)
102 idata.to_netcdf(output / 'posterior.nc')
103 predictions.to_netcdf(output / 'predictions.nc')
104 summary.to_csv(output / 'summary.csv')
105 return idata, predictions, model
108if __name__ == '__main__':
109 parser = argparse.ArgumentParser(description=__doc__)
110 parser.add_argument('--draws', type=int, default=1000)
111 parser.add_argument('--tune', type=int, default=1000)
112 parser.add_argument('--chains', type=int, default=4)
113 parser.add_argument('--output-dir', default='hierarchical_results')
114 run(**vars(parser.parse_args()))