Skills › Scientific Agent Skills › pymc › linear_regression_template.py
linear_regression_template.py
99 lines · skills/pymc/assets/linear_regression_template.py at 5f2bc03
1"""Runnable synthetic PyMC 6.3.2 / ArviZ 1.3 linear regression template.
3Replace the synthetic data with your measurements, preserve training scaling,
4and choose priors in outcome units. Short runs only exercise the workflow.
5"""
7import argparse
8from pathlib import Path
10import arviz as az
11import matplotlib.pyplot as plt
12import numpy as np
13import pymc as pm
16def standardize_predictors(x):
17 """Fit a scaling transform to finite, nonconstant training columns."""
18 x = np.asarray(x, dtype=float)
19 if x.ndim != 2 or not np.isfinite(x).all():
20 raise ValueError("Predictors must be a finite 2D array")
21 mean, scale = x.mean(axis=0), x.std(axis=0)
22 if (scale == 0).any():
23 raise ValueError("Remove constant columns before standardizing")
24 return (x - mean) / scale, mean, scale
27def build_model(x_scaled, y, predictor_names):
28 coords = {'predictors': predictor_names, 'obs_id': np.arange(len(y))}
29 with pm.Model(coords=coords) as model:
30 x_data = pm.Data('X_scaled', x_scaled, dims=('obs_id', 'predictors'))
31 alpha = pm.Normal('alpha', mu=0, sigma=1)
32 beta = pm.Normal('beta', mu=0, sigma=1, dims='predictors')
33 sigma = pm.HalfNormal('sigma', sigma=1)
34 mu = pm.Deterministic('mu', alpha + pm.math.dot(x_data, beta), dims='obs_id')
35 pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y,
36 shape=x_data.shape[0], dims='obs_id')
37 return model
40def run(draws=1000, tune=1000, chains=4, output_dir='linear_results'):
41 rng = np.random.default_rng(42)
42 x = rng.normal(size=(60, 3))
43 y = 0.5 + x @ np.array([1.5, -0.8, 2.1]) + rng.normal(0, 0.5, len(x))
44 x_scaled, x_mean, x_scale = standardize_predictors(x)
45 names = ['predictor1', 'predictor2', 'predictor3']
46 model = build_model(x_scaled, y, names)
47 output = Path(output_dir)
48 output.mkdir(parents=True, exist_ok=True)
50 def save(collection, name):
51 collection.savefig(output / name, dpi=150, bbox_inches='tight')
52 plt.close(collection.viz['figure'].item())
54 with model:
55 prior = pm.sample_prior_predictive(draws=100, random_seed=41)
56 save(az.plot_ppc_dist(prior, group='prior_predictive', num_samples=30), 'prior_predictive.png')
57 # Inspect the prior predictive plot before fitting real data.
58 with model:
59 idata = pm.sample(draws=draws, tune=tune, chains=chains, cores=1,
60 nuts_sampler='pymc', target_accept=0.9, random_seed=42,
61 progressbar=False)
62 pm.compute_log_likelihood(idata, progressbar=False)
63 pm.sample_posterior_predictive(idata, extend_inferencedata=True,
64 random_seed=43, progressbar=False)
65 summary = az.summary(idata, var_names=['alpha', 'beta', 'sigma'], round_to='none',
66 ci_kind='hdi', ci_prob=0.95)
67 print(summary)
68 print('[INFO] Check R-hat, bulk/tail ESS, MCSE, divergences and BFMI before interpretation.')
69 save(az.plot_trace_dist(idata, var_names=['alpha', 'beta', 'sigma']), 'trace.png')
70 save(az.plot_ppc_dist(idata, num_samples=30), 'posterior_predictive.png')
71 save(az.plot_dist(idata, var_names=['alpha', 'beta', 'sigma'], ci_kind='hdi', ci_prob=.95), 'posterior.png')
72 save(az.plot_forest(idata, var_names=['beta'], combined=True, ci_kind='hdi', ci_probs=[.5, .95]), 'coefficients.png')
74 x_new = rng.normal(size=(5, 3))
75 with model:
76 pm.set_data({'X_scaled': (x_new - x_mean) / x_scale},
77 coords={'obs_id': np.arange(len(x_new))})
78 predictions = pm.sample_posterior_predictive(
79 idata, var_names=['mu', 'y_obs'], predictions=True, random_seed=44,
80 progressbar=False)
81 # Restore the fitting state before any further training log-likelihood work.
82 pm.set_data({'X_scaled': x_scaled}, coords={'obs_id': np.arange(len(y))})
83 # mu is conditional mean uncertainty; y_obs includes observation noise.
84 interval = az.hdi(predictions.predictions['y_obs'], prob=.95)
85 interval.to_dataframe(name='prediction_interval').to_csv(output / 'prediction_intervals.csv')
86 idata.to_netcdf(output / 'posterior.nc')
87 predictions.to_netcdf(output / 'predictions.nc')
88 summary.to_csv(output / 'summary.csv')
89 np.savez(output / 'training_scaling.npz', mean=x_mean, scale=x_scale)
90 return idata, predictions, model
93if __name__ == '__main__':
94 parser = argparse.ArgumentParser(description=__doc__)
95 parser.add_argument('--draws', type=int, default=1000)
96 parser.add_argument('--tune', type=int, default=1000)
97 parser.add_argument('--chains', type=int, default=4)
98 parser.add_argument('--output-dir', default='linear_results')
99 run(**vars(parser.parse_args()))