Skip to content

Linear regression

Fit a Bayesian linear regression and walk the whole workflow: check the prior, sample the posterior, then check the posterior against the data it was fitted to.

beta  ~ Normal(0, 5)
sigma ~ HalfNormal(2)
y_i   ~ Normal(beta * x_i, sigma)

Both the slope and the observation noise are inferred. The HalfNormal(2) prior on sigma suits this example's units; pick your own for your data.

Run it with python examples/simple_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

Generate data

np.random.seed(42)
N = 500
x = np.random.randn(N)
beta_true = 2.5
sigma_true = 1.5
y = beta_true * x + np.random.normal(0, sigma_true, N)
data = {"x": x, "y": y}

print(f"N={N}, beta_true={beta_true}, sigma_true={sigma_true}")

N=500, beta_true=2.5, sigma_true=1.5

Define the model

builder = rmc.ModelBuilder(data=data)
beta = builder.normal_prior("beta", mu=0.0, sigma=5.0)
sigma = builder.half_normal_prior("sigma", sigma=2.0)

# Passing `sigma` as a ParamRef makes it a parameter, inferred jointly with beta.
builder.normal_likelihood("obs", mu_expr=beta * "x", sigma=sigma, observed_key="y")
model = builder.build()

Prior predictive check

Draw from the priors alone, before seeing the data, and look at what data they imply. If the prior predictive range is absurd for your units, the priors are wrong and no amount of sampling will fix that.

prior_pred = rmc.sample_prior_predictive(model, n_samples=500, seed=0)
print(f"  Prior beta  ~ N(0, 5):  mean={prior_pred['beta'].mean():.2f}, std={prior_pred['beta'].std():.2f}")
print(f"  Prior sigma ~ HN(2):    mean={prior_pred['sigma'].mean():.2f}, std={prior_pred['sigma'].std():.2f}")
print(f"  Prior y_hat range:      [{prior_pred['obs'].min():.1f}, {prior_pred['obs'].max():.1f}]")

  Prior beta  ~ N(0, 5):  mean=-0.15, std=4.82
  Prior sigma ~ HN(2):    mean=1.65, std=1.19
  Prior y_hat range:      [-54.9, 57.6]

Sample the posterior

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
──────────────────────────────────────────────────────────────────────────────────────────────
beta           2.3866   0.0654     2.2660     2.5110       7487       5430   1.0000   0.000756
sigma          1.4663   0.0467     1.3735     1.5502       7526       5332   1.0003   0.000541
──────────────────────────────────────────────────────────────────────────────────────────────
Mean accept rate: 0.91  │  Divergences: 0

Read the diagnostics before the estimates

r_hat near 1.0 says the four chains agree with each other. ess_bulk in the thousands estimates how many independent draws these 8,000 are worth for a posterior mean. Divergences would say the sampler could not follow the posterior geometry, and any estimate below them would be suspect.

All three are estimates computed from the draws that were taken, so they can only describe where the chains went. Four chains that all miss the same region of the posterior agree with each other and report clean diagnostics. Read them as the absence of evidence of a sampling problem, not as proof of successful sampling, and certainly not as evidence that this is the right model for the data.

print(f"True beta  = {beta_true},  estimated = {fit.mean()['beta']:.4f} +/- {fit.std()['beta']:.4f}")
print(f"True sigma = {sigma_true}, estimated = {fit.mean()['sigma']:.4f} +/- {fit.std()['sigma']:.4f}")
print(f"Step sizes: {[round(s, 4) for s in fit.step_sizes()]}")

True beta  = 2.5,  estimated = 2.3866 +/- 0.0654
True sigma = 1.5, estimated = 1.4663 +/- 0.0467
Step sizes: [0.7992, 0.9171, 0.9021, 0.96]

Posterior predictive check

ppc = fit.posterior_predictive(n_samples=500, seed=42)
y_rep = ppc["obs"]  # shape: (n_samples, N)
print(f"  y_rep shape:  {y_rep.shape}")
print(f"  y_rep mean:   {y_rep.mean():.4f}  (data mean: {y.mean():.4f})")
print(f"  y_rep std:    {y_rep.std():.4f}   (data std:  {y.std():.4f})")

# Fraction of replicated datasets whose spread exceeds the observed spread. Values
# near 0 or 1 mean the model reproduces the data's spread badly; this one is fine.
ppc_p = (y_rep.std(axis=1) > y.std()).mean()
print(f"  PPC p-value (std): {ppc_p:.3f}")

  y_rep shape:  (500, 500)
  y_rep mean:   0.0180  (data mean: 0.0648)
  y_rep std:    2.7615   (data std:  2.7586)
  PPC p-value (std): 0.508

Plots

These pages carry no plots, because a committed image cannot be checked for drift the way captured text can. examples/arviz_example.py fits a linear regression with an intercept and writes trace, posterior and pair plots with ArviZ. It needs pip install "rustmc[viz]".