""" Bayesian Hierarchical SUR (hSUR) -- Gibbs sampler. Model ----- Store i (i = 1,...,nreg), M equations: y_{ij} = X_{ij} @ beta_j^(i) + eps_{ij}, j = 1,...,M eps_{it} ~ N_M(0, Omega) cross-equation correlated errors per period Each store has its own coefficient vector: beta^(i) = (beta_1^(i)', ..., beta_M^(i)') of length K = sum(k_j) Hierarchy --------- beta^(i) = Delta.T @ z_i + u_i, u_i ~ N(0, Vbeta) z_i : (nz,) store-level covariates (include an intercept column) Delta: (nz, K) population regression of beta on z Priors ------ vec(Delta) | Vbeta ~ N(vec(Deltabar), Vbeta (x) A^{-1}) (matrix normal) Vbeta ~ IW(nu_beta, V_beta) Omega ~ IW(nu, V) 4-block Gibbs ------------- (1) beta^(i) | Omega, Delta, Vbeta, y_i, X_i -- N (one draw per store) Q_i = Vbeta^{-1} + [[omega^{pq} X_{ip}' X_{iq}]] b_i = Vbeta^{-1} (Delta.T z_i) + [[sum_q omega^{pq} X_{ip}' y_{iq}]] (2) Omega | {beta^(i)}, Y -- IW(nu + N_total, V + sum_i E_i'E_i) N_total = sum_i T_i, E_i[:,j] = y_{ij} - X_{ij} beta_j^(i) (3) Delta | {beta^(i)}, Vbeta -- MN (same as rhierLinearModel) A_tilde = A + Z'Z, Delta_tilde = A_tilde^{-1}(A Deltabar + Z' B) B = [beta^(1)', ..., beta^(nreg)'] (nreg x K) (4) Vbeta | {beta^(i)}, Delta -- IW(nu_beta + nreg, V_beta + D'D) D_i = beta^(i) - Delta.T z_i """ import numpy as np from scipy.stats import invwishart from scipy.linalg import cholesky, solve_triangular # ── private helpers ─────────────────────────────────────────────────────────── def _chol_draw(Q, b, rng): """Draw x ~ N(Q^{-1} b, Q^{-1}) via Cholesky of the precision matrix.""" L = cholesky(Q, lower=True) v = solve_triangular(L, b, lower=True) mu = solve_triangular(L.T, v, lower=False) return mu + solve_triangular(L.T, rng.standard_normal(len(b)), lower=False) def _draw_Delta(betas, Z, Vbeta, Deltabar, A, rng): """Delta | Vbeta, betas ~ MN(Delta_tilde, Vbeta, A_tilde^{-1}). betas : (nreg, K); Z : (nreg, nz); Delta : (nz, K) """ 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 + D'D).""" resid = betas - Z @ Delta # (nreg, K) return invwishart.rvs(df=nu + len(betas), scale=V + resid.T @ resid, random_state=rng) def _build_store_stats(y_list, X_list, M): """Precompute S[p,q] = X_p'X_q and XtY[p,q] = X_p'y_q for one store.""" S = [[X_list[p].T @ X_list[q] for q in range(M)] for p in range(M)] XtY = [[X_list[p].T @ y_list[q] for q in range(M)] for p in range(M)] return S, XtY def _draw_beta_i(S_i, XtY_i, starts, Omega_inv, Vbeta_inv, mu_i, rng): """Draw beta^(i) ~ N(Q_i^{-1} b_i, Q_i^{-1}). S_i : M x M list of (k_p, k_q) matrices X_p' X_q XtY_i : M x M list of (k_p,) vectors X_p' y_q mu_i : (K,) prior mean Delta.T @ z_i """ K = starts[-1] M = Omega_inv.shape[0] Q = Vbeta_inv.copy() b = Vbeta_inv @ mu_i for p in range(M): sp = starts[p]; ep = starts[p + 1] for q in range(M): sq = starts[q]; eq = starts[q + 1] Q[sp:ep, sq:eq] += Omega_inv[p, q] * S_i[p][q] b[sp:ep] += Omega_inv[p, q] * XtY_i[p][q] return _chol_draw(Q, b, rng) # ── public API ──────────────────────────────────────────────────────────────── def rhsur_gibbs(regdata, Z=None, Prior=None, Mcmc=None, seed=42): """ 4-block Gibbs sampler for the Hierarchical SUR model. Parameters ---------- regdata : list of nreg dicts, each with keys 'y_list' : list of M arrays, each (T_i,) 'X_list' : list of M arrays, each (T_i, k_j) Alternatively pass {'Y': (T_i, M), 'X': (T_i, k)} for the common-X case (same regressors across all equations in store i). Z : (nreg, nz) store-level covariates including intercept. Default: column of ones (intercept-only hierarchy). Prior : dict with optional keys Deltabar : (nz, K) prior mean of Delta, default 0 A : (nz, nz) prior precision of Delta, default 0.01*I nu : int IW df for Omega, default M+3 V : (M, M) IW scale for Omega, default (M+2)*I nu_beta : int IW df for Vbeta, default K+3 V_beta : (K, K) IW scale for Vbeta, default (K+2)*I 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, K, R_kept) store-level coefficient draws Deltadraw : (R_kept, nz*K) column-major vec(Delta) Vbetadraw : (R_kept, K*K) column-major vec(Vbeta) Omegadraw : (R_kept, M*M) column-major vec(Omega) """ if Prior is None: Prior = {} if Mcmc is None: Mcmc = {'R': 10000} rng = np.random.default_rng(seed) nreg = len(regdata) # ── normalise regdata to y_list / X_list ───────────────────────────────── def _normalise(d): if 'y_list' in d: return d['y_list'], d['X_list'] Y = d['Y']; X = d['X'] M_ = Y.shape[1] return [Y[:, j] for j in range(M_)], [X] * M_ y_data = [_normalise(d) for d in regdata] # list of (y_list, X_list) M = len(y_data[0][0]) ks = [y_data[0][1][j].shape[1] for j in range(M)] K = sum(ks) starts = np.concatenate([[0], np.cumsum(ks)]).astype(int) # ── Z ──────────────────────────────────────────────────────────────────── if Z is None: Z = np.ones((nreg, 1)) Z = np.asarray(Z, float) nz = Z.shape[1] # ── priors ─────────────────────────────────────────────────────────────── 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', M + 3)) V = np.asarray(Prior.get('V', float(M + 2) * np.eye(M)), float) nu_beta = int(Prior.get('nu_beta', K + 3)) V_beta = np.asarray(Prior.get('V_beta', float(K + 2) * np.eye(K)), float) # ── MCMC settings ──────────────────────────────────────────────────────── 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 stats per store ───────────────────────────────── Ts = [len(y_data[i][0][0]) for i in range(nreg)] N_tot = sum(Ts) stats = [_build_store_stats(y_data[i][0], y_data[i][1], M) for i in range(nreg)] # stats[i] = (S_i, XtY_i) # ── storage ─────────────────────────────────────────────────────────────── betadraw = np.zeros((nreg, K, R_kept)) Deltadraw = np.zeros((R_kept, nz * K)) Vbetadraw = np.zeros((R_kept, K * K)) Omegadraw = np.zeros((R_kept, M * M)) # ── initialise ──────────────────────────────────────────────────────────── betas = np.zeros((nreg, K)) Delta = Deltabar.copy() Vbeta = V_beta / max(nu_beta - K - 1, 1) Omega = V / max(nu - M - 1, 1) i_kept = 0 for r in range(R): Omega_inv = np.linalg.inv(Omega) Vbeta_inv = np.linalg.inv(Vbeta) mus = Z @ Delta # (nreg, K) prior means # ── block 1: beta^(i) | Omega, Delta, Vbeta ────────────────────────── for i in range(nreg): S_i, XtY_i = stats[i] betas[i] = _draw_beta_i(S_i, XtY_i, starts, Omega_inv, Vbeta_inv, mus[i], rng) # ── block 2: Omega | {beta^(i)} ─────────────────────────────────────── S_eps = np.zeros((M, M)) for i in range(nreg): ys_i, Xs_i = y_data[i] E_i = np.column_stack( [ys_i[j] - Xs_i[j] @ betas[i, starts[j]:starts[j + 1]] for j in range(M)]) # (T_i, M) S_eps += E_i.T @ E_i Omega = invwishart.rvs(df=nu + N_tot, scale=V + S_eps, random_state=rng) # ── block 3: Delta | {beta^(i)}, Vbeta ─────────────────────────────── Delta = _draw_Delta(betas, Z, Vbeta, Deltabar, A, rng) # ── block 4: Vbeta | {beta^(i)}, Delta ─────────────────────────────── Vbeta = _draw_Vbeta(betas, Z, Delta, nu_beta, V_beta, 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') Omegadraw[i_kept] = Omega.ravel(order='F') i_kept += 1 return {'betadraw': betadraw, 'Deltadraw': Deltadraw, 'Vbetadraw': Vbetadraw, 'Omegadraw': Omegadraw} def simulate_hsur(nreg, M, k, T, true_Delta, true_Vbeta, true_Omega, Z=None, seed=0): """ Simulate data from the hierarchical SUR model. Parameters ---------- nreg : number of stores M : number of equations k : regressors per equation (int, same for all) T : observations per store (int or list of ints) true_Delta : (nz, K) population regression matrix true_Vbeta : (K, K) store-level coefficient covariance true_Omega : (M, M) cross-equation error covariance Z : (nreg, nz) store covariates; default ones (intercept only) seed : random seed Returns ------- regdata : list of nreg dicts with 'y_list' and 'X_list' betas : (nreg, K) true store-level coefficients Z : (nreg, nz) store covariates used """ rng = np.random.default_rng(seed) K = M * k Ts = [T] * nreg if np.isscalar(T) else list(T) if Z is None: Z = np.ones((nreg, 1)) Z = np.asarray(Z, float) # draw store-level coefficients from hierarchy L_V = np.linalg.cholesky(true_Vbeta) mus = Z @ true_Delta # (nreg, K) betas = mus + rng.standard_normal((nreg, K)) @ L_V.T L_O = np.linalg.cholesky(true_Omega) starts = np.arange(0, K + 1, k) regdata = [] for i in range(nreg): Ti = Ts[i] E_i = rng.standard_normal((Ti, M)) @ L_O.T # correlated errors y_list = []; X_list = [] for j in range(M): X_ij = np.column_stack([np.ones(Ti), rng.standard_normal((Ti, k - 1))]) y_ij = X_ij @ betas[i, starts[j]:starts[j + 1]] + E_i[:, j] y_list.append(y_ij) X_list.append(X_ij) regdata.append({'y_list': y_list, 'X_list': X_list}) return regdata, betas, Z 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')