""" Bayesian Hierarchical Linear Model -- Gibbs sampler. Replicates bayesm::rhierLinearModel (Rossi, Allenby & McCulloch 2005, Ch. 3). Model ----- y_i = X_i @ beta_i + eps_i, eps_i ~ N(0, sigma2_i * I_ni) beta_i = Delta.T @ z_i + u_i, u_i ~ N(0, Vbeta) Priors ------ vec(Delta) | Vbeta ~ N(vec(Deltabar), Vbeta (x) A^{-1}) Vbeta ~ IW(nu, V) sigma2_i ~ IG(nu_e/2, nu_e*ssq_i/2) iid, ssq_i may be unit-specific """ import numpy as np from scipy.stats import invwishart # ── private helpers ─────────────────────────────────────────────────────────── def _chol_draw(Prec, rhs, rng): """x ~ N(Prec^{-1} rhs, Prec^{-1}) via Cholesky of the precision matrix.""" L = np.linalg.cholesky(Prec) # L L' = Prec L_inv_rhs = np.linalg.solve(L, rhs) mu_post = np.linalg.solve(L.T, L_inv_rhs) # = Prec^{-1} rhs return mu_post + np.linalg.solve(L.T, rng.standard_normal(len(rhs))) def _draw_Delta(betas, Z, Vbeta, Deltabar, A, rng): """Delta | Vbeta, betas ~ MN(Delta_tilde, Vbeta, A_tilde^{-1}).""" A_tilde = A + Z.T @ Z A_tilde_inv = np.linalg.inv(A_tilde) Delta_tilde = A_tilde_inv @ (A @ Deltabar + Z.T @ betas) L_A = np.linalg.cholesky(A_tilde_inv) L_V = np.linalg.cholesky(Vbeta) return Delta_tilde + L_A @ rng.standard_normal(Delta_tilde.shape) @ L_V.T def _draw_Vbeta(betas, Z, Delta, nu, V, rng): """Vbeta | betas, Delta ~ IW(nu + nreg, V + resid' resid).""" resid = betas - Z @ Delta return invwishart.rvs(df=nu + len(betas), scale=V + resid.T @ resid, random_state=rng) # ── public API ──────────────────────────────────────────────────────────────── def rhier_linear_model(regdata, Z, Prior=None, Mcmc=None, seed=42): """ Gibbs sampler for the Hierarchical Linear Model. Parameters ---------- regdata : list of dicts {'y': (n_i,), 'X': (n_i, k)} Z : (nreg, nz) unit-level covariates (include intercept column) Prior : dict with optional keys Deltabar : (nz, k) prior mean of Delta, default 0 A : (nz, nz) prior precision for Delta, default 0.01*I nu : int IW df for Vbeta, default k+3 V : (k, k) IW scale for Vbeta, default (k+2)*I nu_e : float IG df for sigma2, default 3 ssq : float or (nreg,) IG scale, default 1 Mcmc : dict with keys R, burn (default 0.2*R), keep (default 1), nprint (default R//10) Returns ------- dict with betadraw : (nreg, k, R_kept) Deltadraw : (R_kept, nz*k) column-major (matches R vec()) Vbetadraw : (R_kept, k*k) column-major sigmasqdraw : (R_kept, nreg) """ if Prior is None: Prior = {} if Mcmc is None: Mcmc = {'R': 10000} rng = np.random.default_rng(seed) nreg = len(regdata) k = regdata[0]['X'].shape[1] nz = Z.shape[1] Deltabar = np.asarray(Prior.get('Deltabar', np.zeros((nz, k))), float) A = np.asarray(Prior.get('A', 0.01 * np.eye(nz)), float) nu = int(Prior.get('nu', k + 3)) V = np.asarray(Prior.get('V', float(k + 2) * np.eye(k)), float) nu_e = float(Prior.get('nu_e', 3.0)) ssq_raw = Prior.get('ssq', 1.0) ssq = (np.full(nreg, ssq_raw, float) if np.isscalar(ssq_raw) else np.asarray(ssq_raw, float)) R = int(Mcmc['R']) burn = int(Mcmc.get('burn', 0.2 * R)) keep = int(Mcmc.get('keep', 1)) nprint = int(Mcmc.get('nprint', max(R // 10, 1))) R_kept = (R - burn) // keep # precompute sufficient statistics once ys = [d['y'] for d in regdata] Xs = [d['X'] for d in regdata] XtXs = [X.T @ X for X in Xs] Xtys = [X.T @ y for X, y in zip(Xs, ys)] ns = [len(y) for y in ys] # storage betadraw = np.zeros((nreg, k, R_kept)) Deltadraw = np.zeros((R_kept, nz * k)) Vbetadraw = np.zeros((R_kept, k * k)) sigmasqdraw = np.zeros((R_kept, nreg)) # initialise betas = np.zeros((nreg, k)) sigma2 = ssq.copy() Delta = Deltabar.copy() Vbeta = V / max(nu - k - 1, 1) i_kept = 0 for r in range(R): Vbeta_inv = np.linalg.inv(Vbeta) mu_prior = Z @ Delta # nreg × k # 1. beta_i | Delta, Vbeta, sigma2_i for i in range(nreg): Prec = Vbeta_inv + XtXs[i] / sigma2[i] rhs = Vbeta_inv @ mu_prior[i] + Xtys[i] / sigma2[i] betas[i] = _chol_draw(Prec, rhs, rng) # 2. sigma2_i | beta_i (IG draw via Gamma) for i in range(nreg): e = ys[i] - Xs[i] @ betas[i] scale_i = (nu_e * ssq[i] + float(e @ e)) / 2.0 sigma2[i] = scale_i / rng.gamma((nu_e + ns[i]) / 2.0) # 3. Delta | betas, Vbeta Delta = _draw_Delta(betas, Z, Vbeta, Deltabar, A, rng) # 4. Vbeta | betas, Delta Vbeta = _draw_Vbeta(betas, Z, Delta, nu, V, rng) if nprint and (r + 1) % nprint == 0: print(f' Iter {r+1:6d}/{R}') if r >= burn and (r - burn) % keep == 0: betadraw[:, :, i_kept] = betas Deltadraw[i_kept] = Delta.ravel(order='F') Vbetadraw[i_kept] = Vbeta.ravel(order='F') sigmasqdraw[i_kept] = sigma2 i_kept += 1 return { 'betadraw': betadraw, 'Deltadraw': Deltadraw, 'Vbetadraw': Vbetadraw, 'sigmasqdraw': sigmasqdraw, } def simulate_hlm(nreg, k, Z, true_Delta, true_Vbeta, true_sigma2=1.0, ni=20, seed=0): """Simulate data from the HLM with known parameters.""" rng = np.random.default_rng(seed) L_V = np.linalg.cholesky(true_Vbeta) betas = Z @ true_Delta + rng.standard_normal((nreg, k)) @ L_V.T if np.isscalar(ni): ni = [ni] * nreg sig = (np.full(nreg, true_sigma2, float) if np.isscalar(true_sigma2) else np.asarray(true_sigma2, float)) regdata = [] for i in range(nreg): X = np.column_stack([np.ones(ni[i]), rng.standard_normal((ni[i], k - 1))]) y = X @ betas[i] + rng.normal(0.0, np.sqrt(sig[i]), ni[i]) regdata.append({'y': y, 'X': X}) return regdata, betas def posterior_summary(draws, param_names=None, true_vals=None): """Posterior mean, SD, 95% CrI for each column of a (R, p) draws array.""" import pandas as pd p = draws.shape[1] if param_names is None: param_names = [f'p{j}' for j in range(p)] rows = [] for j, name in enumerate(param_names): d = draws[:, j] lo, hi = np.percentile(d, [2.5, 97.5]) row = {'param': name, 'mean': d.mean(), 'sd': d.std(), 'q025': lo, 'q975': hi} if true_vals is not None: row['true'] = float(true_vals[j]) row['covered'] = bool(lo <= true_vals[j] <= hi) rows.append(row) return pd.DataFrame(rows).set_index('param')