""" Bayesian Hierarchical Linear Model Mixture -- Gibbs sampler. Replicates bayesm::rhierLinearMixture (Rossi, Allenby & McCulloch 2005, Ch. 5). Model ----- y_i = X_i @ beta_i + eps_i, eps_i ~ N(0, sigma2_i I_Ti) beta_i | s_i=k ~ MvN(mu_k + Delta @ z_i, Sigma_k) s_i | pi ~ Categorical(pi) pi ~ Dirichlet(a * 1_K) mu_k ~ MvN(mubar, (Amu * I_p)^{-1}) [independent of Sigma_k] Sigma_k ~ IW(nu, V) sigma2_i ~ IG(nu_e/2, nu_e * ssq_i / 2) If Z=None: Delta = 0 (pure mixture, no store covariates) If Z given: vec(Delta) ~ N(deltabar, Ad^{-1}) Z should be centred (no intercept column), matching bayesm convention. Identification -------------- Components are sorted by mu_k[0] (intercept) ascending after each (mu_k, Sigma_k) draw to resolve label switching. 6-block Gibbs ------------- (1) beta_i | s_i, sigma2_i, mu_{s_i}, Sigma_{s_i} -- Normal conjugate (2) s_i | beta_i, pi, {mu_k, Sigma_k} -- Categorical (3a) mu_k | {beta_i: s_i=k}, Sigma_k -- Normal conjugate (3b) Sigma_k | {beta_i: s_i=k}, mu_k -- IW conjugate (4) pi | {s_i} -- Dirichlet (5) sigma2_i | beta_i, y_i, X_i -- IG conjugate (6) Delta | {beta_i, s_i, mu_{s_i}, {Sigma_k}} -- GLS Normal [if Z] Delta convention ---------------- Delta is stored as (p, nz). delta_vec = Delta.ravel() [C/row-major]: delta_vec[j*nz + l] = Delta[j, l] = effect of z[:,l] on beta[:,j]. Deltadraw columns follow this row-major ordering (differs from bayesm's column-major vec(); convert with .reshape(p,nz).ravel('F') if needed). """ import time import numpy as np from scipy.stats import invwishart from scipy.linalg import cholesky, solve_triangular # ── private helpers ─────────────────────────────────────────────────────────── def _chol_draw(Prec, rhs, rng): """x ~ N(Prec^{-1} rhs, Prec^{-1}) via Cholesky of the precision matrix.""" L = cholesky(Prec, lower=True) v = solve_triangular(L, rhs, lower=True) mu = solve_triangular(L.T, v, lower=False) z = solve_triangular(L.T, rng.standard_normal(len(rhs)), lower=False) return mu + z def _log_mvn(x, mu, L): """Log density of MvN(mu, L L') at x. L is lower Cholesky of Sigma.""" p = len(x) dev = x - mu v = solve_triangular(L, dev, lower=True) return (-0.5 * p * np.log(2.0 * np.pi) - np.sum(np.log(np.diag(L))) - 0.5 * float(v @ v)) # ── public API ──────────────────────────────────────────────────────────────── def rhier_linear_mixture(regdata, K=2, Z=None, Prior=None, Mcmc=None, seed=42): """ Gibbs sampler for the Hierarchical Linear Model Mixture. Parameters ---------- regdata : list of nreg dicts, each with 'y': (T_i,) and 'X': (T_i, p) K : int, number of mixture components (default 2) Z : (nreg, nz) store-level covariates (centred, no intercept). If None, pure mixture with component means mu_k only. Prior : dict with optional keys a : float or (K,) Dirichlet concentration, default 5. mubar : (p,) prior mean for mu_k, default 0 Amu : float precision for mu_k prior, default 0.1 nu : int IW df for Sigma_k, default p+3 V : (p,p) IW scale for Sigma_k, default (p+2)*I nu_e : float IG df for sigma2_i, default 3. ssq : float|(nreg,) IG scale for sigma2_i, default 1. Ad : (p,p) precision for Delta, default 0.01*I_p [only if Z provided; for nz>1 use (p*nz, p*nz)] deltabar : (p,)|(p*nz,) prior mean for delta_vec, default 0 Mcmc : dict with optional keys R : int total draws, default 10000 burn : int burn-in, default 0.2*R keep : int thinning, default 1 nprint : int print frequency, default R//10 Returns ------- dict with betadraw : (nreg, p, R_kept) store-level coefficient draws sdraw : (R_kept, nreg) component assignment, 0-indexed, sorted mudraw : (R_kept, K, p) component mean draws Sigmadraw : (R_kept, K, p, p) component covariance draws pidraw : (R_kept, K) mixing weight draws sigmasqdraw : (R_kept, nreg) error variance draws Deltadraw : (R_kept, p*nz) [only if Z provided, row-major vec] """ if Prior is None: Prior = {} if Mcmc is None: Mcmc = {'R': 10000} rng = np.random.default_rng(seed) nreg = len(regdata) p = regdata[0]['X'].shape[1] # ── priors ─────────────────────────────────────────────────────────────── a_raw = Prior.get('a', 5.) a_dir = (np.full(K, float(a_raw)) if np.isscalar(a_raw) else np.asarray(a_raw, float)) mubar = np.asarray(Prior.get('mubar', np.zeros(p)), float) Amu = float(Prior.get('Amu', 0.1)) A_mu = Amu * np.eye(p) # (p, p) precision nu0 = int(Prior.get('nu', p + 3)) V0 = np.asarray(Prior.get('V', float(p + 2) * np.eye(p)), 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)) # ── MCMC settings ──────────────────────────────────────────────────────── R = int(Mcmc['R']) burn = int(Mcmc.get('burn', int(0.2 * R))) keep = int(Mcmc.get('keep', 1)) nprint = int(Mcmc.get('nprint', max(R // 10, 1))) R_kept = (R - burn) // keep # ── data ───────────────────────────────────────────────────────────────── ys = [np.asarray(d['y'], float) for d in regdata] Xs = [np.asarray(d['X'], float) for d in regdata] Ts = np.array([len(y) for y in ys]) XtXs = [X.T @ X for X in Xs] Xtys = [X.T @ y for X, y in zip(Xs, ys)] # ── Z covariate (optional) ──────────────────────────────────────────────── use_Z = Z is not None if use_Z: Z_arr = np.asarray(Z, float) # (nreg, nz) nz = Z_arr.shape[1] dim_d = p * nz # dimension of delta_vec Ad = np.asarray(Prior.get('Ad', 0.01 * np.eye(dim_d)), float) d0 = np.asarray(Prior.get('deltabar', np.zeros(dim_d)), float) Delta = np.zeros((p, nz)) # current Delta matrix else: nz = 0 # ── initialise ──────────────────────────────────────────────────────────── sigma2 = ssq.copy() mu_k = np.tile(mubar, (K, 1)).copy() # (K, p) Sigma_k = np.stack([V0 / max(nu0 - p - 1, 1)] * K) # (K, p, p) Sig_inv = np.stack([np.linalg.inv(Sigma_k[k]) for k in range(K)]) # (K,p,p) pi_k = np.full(K, 1.0 / K) s = np.arange(nreg, dtype=int) % K # round-robin init betas = np.zeros((nreg, p)) # ── storage ─────────────────────────────────────────────────────────────── betadraw = np.zeros((nreg, p, R_kept)) sdraw = np.zeros((R_kept, nreg), dtype=int) mudraw = np.zeros((R_kept, K, p)) Sigmadraw = np.zeros((R_kept, K, p, p)) pidraw = np.zeros((R_kept, K)) sigmasqdraw = np.zeros((R_kept, nreg)) if use_Z: Deltadraw = np.zeros((R_kept, dim_d)) i_kept = 0 t0 = time.time() for r in range(R): # ── block 1: beta_i | s_i, sigma2_i, mu_{s_i}, Sigma_{s_i} ────────── for i in range(nreg): k_i = s[i] mu_eff = (mu_k[k_i] + Delta @ Z_arr[i]) if use_Z else mu_k[k_i] Prec_i = Sig_inv[k_i] + XtXs[i] / sigma2[i] rhs_i = Sig_inv[k_i] @ mu_eff + Xtys[i] / sigma2[i] betas[i] = _chol_draw(Prec_i, rhs_i, rng) # ── block 2: s_i | beta_i, pi, {mu_k, Sigma_k} ─────────────────────── L_ks = [cholesky(Sigma_k[k], lower=True) for k in range(K)] log_pi = np.log(np.clip(pi_k, 1e-300, None)) for i in range(nreg): log_w = np.empty(K) for k in range(K): mu_eff = (mu_k[k] + Delta @ Z_arr[i]) if use_Z else mu_k[k] log_w[k] = log_pi[k] + _log_mvn(betas[i], mu_eff, L_ks[k]) log_w -= log_w.max() w = np.exp(log_w) w /= w.sum() s[i] = rng.choice(K, p=w) # ── block 3: (mu_k, Sigma_k) | {beta_i: s_i=k}, Delta ─────────────── for k in range(K): mask = s == k n_k = int(mask.sum()) beff_k = betas[mask].copy() # (n_k, p) if use_Z: beff_k -= (Delta @ Z_arr[mask].T).T # subtract Delta z_i shift # 3a: mu_k | beff_k, Sigma_k (Normal conjugate, independent prior) if n_k > 0: Prec_mu = A_mu + n_k * Sig_inv[k] rhs_mu = A_mu @ mubar + Sig_inv[k] @ beff_k.sum(0) else: Prec_mu = A_mu.copy() rhs_mu = A_mu @ mubar mu_k[k] = _chol_draw(Prec_mu, rhs_mu, rng) # 3b: Sigma_k | beff_k, mu_k (IW conjugate) if n_k > 0: dev = beff_k - mu_k[k] scatter = dev.T @ dev Sigma_k[k] = invwishart.rvs(df=nu0 + n_k, scale=V0 + scatter, random_state=rng) else: Sigma_k[k] = invwishart.rvs(df=nu0, scale=V0, random_state=rng) Sig_inv[k] = np.linalg.inv(Sigma_k[k]) # label-switching: sort components by intercept ascending order = np.argsort(mu_k[:, 0]) mu_k = mu_k[order] Sigma_k = Sigma_k[order] Sig_inv = Sig_inv[order] s_new = np.empty_like(s) for new_k, old_k in enumerate(order): s_new[s == old_k] = new_k s = s_new # ── block 4: pi | {s_i} ─────────────────────────────────────────────── counts = np.bincount(s, minlength=K).astype(float) pi_k = rng.dirichlet(a_dir + counts) # ── block 5: sigma2_i | beta_i, y_i, X_i ───────────────────────────── 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 + Ts[i]) / 2.0) # ── block 6: Delta | {beta_i, s_i, mu_{s_i}, Sigma_k} [if Z] ───────── if use_Z: # Full GLS update for delta_vec = Delta.ravel() [row-major, shape p*nz] # Model: beta_i - mu_{s_i} = Delta @ z_i + u_i, u_i ~ N(0, Sigma_{s_i}) # Precision: P = Ad + sum_i kron(Sig_inv[s_i], outer(z_i, z_i)) # rhs: r = Ad @ d0 + sum_i kron(Sig_inv[s_i] @ bi_eff, z_i) P_D = Ad.copy() rhs_D = Ad @ d0 for i in range(nreg): k_i = s[i] zi = Z_arr[i] # (nz,) bi_eff = betas[i] - mu_k[k_i] # (p,) P_D += np.kron(Sig_inv[k_i], np.outer(zi, zi)) rhs_D += np.kron(Sig_inv[k_i] @ bi_eff, zi) delta_vec = _chol_draw(P_D, rhs_D, rng) # (p*nz,) row-major Delta = delta_vec.reshape(p, nz) # (p, nz) if nprint and (r + 1) % nprint == 0: elapsed = time.time() - t0 print(f' Iter {r+1:6d}/{R} ({elapsed:.0f}s)') if r >= burn and (r - burn) % keep == 0: betadraw[:, :, i_kept] = betas sdraw[i_kept] = s mudraw[i_kept] = mu_k Sigmadraw[i_kept] = Sigma_k pidraw[i_kept] = pi_k sigmasqdraw[i_kept] = sigma2 if use_Z: Deltadraw[i_kept] = delta_vec i_kept += 1 elapsed = time.time() - t0 print(f'Done in {elapsed / 60:.1f} min | kept: {R_kept}') out = { 'betadraw': betadraw, 'sdraw': sdraw, 'mudraw': mudraw, 'Sigmadraw': Sigmadraw, 'pidraw': pidraw, 'sigmasqdraw': sigmasqdraw, } if use_Z: out['Deltadraw'] = Deltadraw return out def simulate_hlm_mixture(nreg, p, K, true_pi, true_mu, true_Sigma, true_sigma2=1.0, ni=20, seed=0): """ Simulate data from the HLM Mixture model with known parameters. Parameters ---------- nreg : int, number of units (stores) p : int, number of regression coefficients K : int, number of mixture components true_pi : (K,) mixing weights (must sum to 1) true_mu : (K, p) component means true_Sigma : (K, p, p) component covariances true_sigma2 : float or (nreg,), error variance per unit ni : int or list, observations per unit seed : int, random seed Returns ------- regdata : list of nreg dicts with 'y', 'X' betas : (nreg, p) true store-level coefficients comps : (nreg,) true component assignments (0-indexed) """ rng = np.random.default_rng(seed) L_sigs = [np.linalg.cholesky(true_Sigma[k]) for k in range(K)] sig2 = (np.full(nreg, float(true_sigma2)) if np.isscalar(true_sigma2) else np.asarray(true_sigma2, float)) ns = [ni] * nreg if np.isscalar(ni) else list(ni) comps = rng.choice(K, size=nreg, p=true_pi) betas = np.array([ true_mu[comps[i]] + L_sigs[comps[i]] @ rng.standard_normal(p) for i in range(nreg) ]) regdata = [] for i in range(nreg): X = np.column_stack([np.ones(ns[i]), rng.standard_normal((ns[i], p - 1))]) y = X @ betas[i] + rng.normal(0., np.sqrt(sig2[i]), ns[i]) regdata.append({'y': y, 'X': X}) return regdata, betas, comps def posterior_summary(draws, param_names=None, true_vals=None): """Posterior mean, SD, 95% CrI for each column of a (R, q) draws array.""" import pandas as pd q = draws.shape[1] if param_names is None: param_names = [f'p{j}' for j in range(q)] 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')