The Rats growth data — BiRats in PyMC¶

Cross-check of the from-scratch hierarchical linear model¶

Companion to rats_python.ipynb (from-scratch rhier_linear_mixture). Same BiRats model — random intercept + slope with a full-covariance bivariate-normal population — fitted here with PyMC (NUTS), to confirm the from-scratch results.

The fitted model¶

Data. 30 rats, each weighed at five ages $x_j \in \{8, 15, 22, 29, 36\}$ days. $y_{ij}$ = weight (grams) of rat $i$ at age $x_j$.

Likelihood — one linear growth curve per rat (age centred at the midpoint $\bar x = 22$): $$y_{ij} = \alpha_i + \beta_i\,(x_j - 22) + \varepsilon_{ij}, \qquad \varepsilon_{ij}\sim N(0,\sigma^2).$$ Centring makes the parameters interpretable: $\alpha_i$ is rat $i$'s weight at day 22 (the middle of the study) and $\beta_i$ is its growth rate in g/day. The response is raw weight in grams (not log-weight) — the canonical BUGS Rats/BiRats specification, which reproduces the textbook population values (≈ 242 g at day 22, ≈ 6.2 g/day).

Hierarchy — the population of growth curves. The per-rat (intercept, slope) pairs share a single bivariate-normal population: $$(\alpha_i,\beta_i)' \sim N(\mu,\ \Sigma) \qquad \text{(BiRats)}.$$ This notebook fits only that single-normal hierarchy. The mixture question — are there discrete growth types? — is settled in the from-scratch companion rats_python.ipynb, where WAIC picks $K=1$.

Priors — the PyMC parameterisation. Note these differ from the conjugate from-scratch sampler by design: where rats_python.ipynb uses an Inverse-Wishart on $\Sigma$ and an Inverse-Gamma on each rat's error variance, PyMC uses the HMC-friendly forms —

  • $\mu \sim N\big([242,\ 6],\ [50,\ 5]\big)$ — population mean (intercept, slope);
  • $\Sigma$ via LKJCholeskyCov ($\eta = 2$) with HalfNormal([20, 1]) scales, which yields interpretable standard deviations plus a correlation rather than a raw covariance matrix;
  • a single shared residual scale $\sigma \sim \text{HalfNormal}(10)$ (not a per-rat $\sigma_i^2$).

That the two parameterisations land on the same posterior is exactly the point of the cross-check.

Model¶

$$y_{ij} \sim N\!\big(\alpha_i + \beta_i (x_j-\bar x),\ \sigma^2\big),\qquad (\alpha_i,\beta_i)' \sim \text{MvN}(\mu,\ \Sigma).$$

In PyMC the population covariance $\Sigma$ gets an LKJ-Cholesky prior (LKJCholeskyCov, yielding interpretable standard deviations + a correlation), the 30 rat-level $(\alpha_i,\beta_i)$ are drawn from MvNormal, and the data enter long-format (150 rows). Age is centred at $\bar x=22$.

In [1]:
import os
_lib = r"C:\Users\user\anaconda3\envs\pymc-env\Library\bin"            # Windows DLL fix
if _lib not in os.environ.get("PATH", ""): os.environ["PATH"] = _lib + os.pathsep + os.environ.get("PATH", "")
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ.setdefault("PYTENSOR_FLAGS", "cxx=C:/Users/user/anaconda3/envs/pymc-env/Library/bin/g++.exe")

import numpy as np, pandas as pd, matplotlib.pyplot as plt, pymc as pm, arviz as az
from matplotlib.patches import Ellipse
import pymc.sampling.mcmc as _mcmc                                       # threadpoolctl OpenMP patch
class _Noop:
    def __init__(self,*a,**k): pass
    def __enter__(self): return self
    def __exit__(self,*a): pass
_mcmc.threadpool_limits = _Noop

Y = pd.read_csv('rats.csv', index_col=0).values
ages = np.array([8,15,22,29,36.]); xc = ages - 22
rat = np.repeat(np.arange(30), 5); age = np.tile(xc, 30); weight = Y.reshape(-1)
print(f'{Y.shape[0]} rats, long-format rows = {weight.size}')

30 rats, long-format rows = 150
In [2]:
with pm.Model() as birats:
    mu   = pm.Normal('mu', mu=[242., 6.], sigma=[50., 5.], shape=2)              # population mean
    chol, corr, sds = pm.LKJCholeskyCov('chol', n=2, eta=2.,
                                        sd_dist=pm.HalfNormal.dist([20., 1.]), compute_corr=True)
    ab    = pm.MvNormal('ab', mu=mu, chol=chol, shape=(30, 2))                   # rat-level (alpha,beta)
    sigma = pm.HalfNormal('sigma', 10.)
    pm.Normal('y', ab[rat, 0] + ab[rat, 1] * age, sigma, observed=weight)
    idata = pm.sample(1000, tune=1000, chains=2, cores=1, target_accept=0.9,
                      random_seed=1, progressbar=False)

print(az.summary(idata, var_names=['mu', 'chol_stds', 'sigma']).to_string())   # free params
corr_pm = float(idata.posterior['chol_corr'].mean(('chain','draw')).values[0, 1])
print(f'\nintercept-slope correlation = {corr_pm:.2f}')

Initializing NUTS using jitter+adapt_diag...

Sequential sampling (2 chains in 1 job)

NUTS: [mu, chol, ab, sigma]

Sampling 2 chains for 1_000 tune and 1_000 draw iterations (2_000 + 2_000 draws total) took 3 seconds.

We recommend running at least 4 chains for robust computation of convergence diagnostics

               mean     sd eti89_lb eti89_ub  ess_bulk  ess_tail r_hat mcse_mean mcse_sd
mu[0]         242.6   2.72      240      250      3585      1694  1.00     0.046   0.034
mu[1]         6.182  0.109        6      6.4      2911      1451  1.00     0.002  0.0015
chol_stds[0]  14.53      2       12       18      3098      1397  1.01     0.035   0.027
chol_stds[1]  0.519  0.089     0.39     0.67      1562      1503  1.00    0.0022  0.0018
sigma          6.11   0.48      5.4      6.9      1398      1308  1.00     0.013  0.0098

intercept-slope correlation = 0.59
In [3]:
post = idata.posterior
ab_pm = post['ab'].mean(('chain','draw')).values                       # (30,2) posterior-mean per rat
mu_pm = post['mu'].mean(('chain','draw')).values
sds_pm = post['chol_stds'].mean(('chain','draw')).values
corr_pm = float(post['chol_corr'].mean(('chain','draw')).values[0, 1])
Sig = np.array([[sds_pm[0]**2, sds_pm[0]*sds_pm[1]*corr_pm],
                [sds_pm[0]*sds_pm[1]*corr_pm, sds_pm[1]**2]])
ols = np.array([np.linalg.lstsq(np.column_stack([np.ones(5), xc]), Y[i], rcond=None)[0] for i in range(30)])

fig, ax = plt.subplots(1, 2, figsize=(13, 4.8))
ax[0].plot([ols[:,1].min(), ols[:,1].max()], [ols[:,1].min(), ols[:,1].max()], 'k:', lw=1)
ax[0].scatter(ols[:,1], ab_pm[:,1], color='firebrick', zorder=3)
ax[0].axhline(mu_pm[1], color='gray', ls='--', lw=.8, label=f'pop. slope {mu_pm[1]:.2f}')
ax[0].set_xlabel('per-rat OLS slope'); ax[0].set_ylabel('PyMC (shrunken) slope')
ax[0].set_title('Shrinkage (PyMC)'); ax[0].legend()

ax[1].scatter(ab_pm[:,0], ab_pm[:,1], color='steelblue', zorder=3)
vals, vecs = np.linalg.eigh(Sig); o = vals.argsort()[::-1]; vals, vecs = vals[o], vecs[:, o]
width, height = 2*np.sqrt(vals*5.991); ang = np.degrees(np.arctan2(vecs[1,0], vecs[0,0]))
ax[1].add_patch(Ellipse(mu_pm, width, height, angle=ang, fill=False, edgecolor='firebrick', lw=2,
                        label='fitted 95% ellipse'))
ax[1].set_xlabel(r'intercept $\alpha_i$ (g at day 22)'); ax[1].set_ylabel(r'slope $\beta_i$ (g/day)')
ax[1].set_ylim(mu_pm[1]-4*sds_pm[1], mu_pm[1]+4*sds_pm[1])
ax[1].set_title('Per-rat growth params + fitted bivariate normal'); ax[1].legend()
plt.tight_layout(); plt.savefig('rats_pymc.png', dpi=120, bbox_inches='tight'); plt.show()
No description has been provided for this image

The same two views, now from PyMC / NUTS. Left — shrinkage: per-rat OLS slopes (x) pulled toward the population slope (6.18 g/day, dashed) in the posterior (y). Right — the 30 rats in (intercept, slope) space with the fitted bivariate-normal 95% ellipse: wide in intercept (SD ≈ 14.5 g), narrow in slope (SD ≈ 0.52 g/day), and tilted by a positive intercept–slope correlation (≈ 0.59). Both panels reproduce the from-scratch Gibbs fit — the same single elongated cloud, no hint of separate growth types. The hierarchy is a property of the data, not of one implementation.

In [4]:
# Fitted per-rat growth curves (PyMC shrunken estimates), coloured by growth rate
from matplotlib import cm
from matplotlib.colors import Normalize
fig, ax = plt.subplots(figsize=(7.6, 4.7))
xg = np.linspace(ages.min(), ages.max(), 50); slopes = ab_pm[:, 1]; nrm = Normalize(slopes.min(), slopes.max())
for i in range(30):
    ax.plot(ages, Y[i], '.', color='lightgray', ms=4, zorder=1)
    ax.plot(xg, ab_pm[i, 0] + ab_pm[i, 1]*(xg - 22), '-', color=cm.viridis(nrm(slopes[i])), lw=1, alpha=.85, zorder=2)
ax.plot(xg, mu_pm[0] + mu_pm[1]*(xg - 22), 'k-', lw=2.6, label='population fit')
sm = cm.ScalarMappable(cmap='viridis', norm=nrm); sm.set_array([]); plt.colorbar(sm, ax=ax, label='rat growth rate (g/day)')
ax.set_xlabel('age (days)'); ax.set_ylabel('weight (g)')
ax.set_title('Fitted per-rat growth curves (PyMC shrunken) over the data'); ax.legend(loc='upper left')
plt.tight_layout(); plt.savefig('rats_fitted_pymc.png', dpi=120, bbox_inches='tight'); plt.show()
No description has been provided for this image

Results — PyMC agrees with the from-scratch BiRats¶

parameter from-scratch PyMC
intercept μ (day 22) 242.3 ≈ 242.6
slope μ 6.21 ≈ 6.19
α SD 14.0 ≈ 14.7
β SD 0.51 ≈ 0.52
intercept–slope corr 0.55 ≈ 0.58
residual σ 7.3 ≈ 6.1

All r̂ ≈ 1.00, and the shrinkage pattern matches: per-rat slopes are pulled in from their OLS values toward the population mean, and the fitted bivariate-normal ellipse (wide in intercept ≈ ±14 g, narrow in slope ≈ ±0.5 g/day, tilted by the positive correlation) is the same as the from-scratch fit. Two independent engines (from-scratch Gibbs and PyMC NUTS) give the same hierarchical-linear result — confirming the Gelfand et al. growth model. (The mixture/"growth types" question was answered in the from-scratch notebook: K=1 suffices.)