""" constrained_lm_gibbs.py Bayesian linear regression with linear inequality constraints. Implements Geweke (1995) "Bayesian Inference for Linear Models Subject to Linear Inequality Constraints." Model ----- y = X @ beta + eps, eps ~ N(0, sigma2 * I_n) beta ~ flat [default] or N(b0, B0) [informative prior] sigma2 ~ IG(a0/2, s0/2) [flat if a0=s0=0] Constraints: a_vec <= D @ beta <= w_vec Two key methods --------------- ghk_prob -- GHK probability simulator: estimates p_{2|1} = P(a_vec <= D @ beta <= w_vec | data) under the *unconstrained* posterior for beta. constrained_gibbs -- Gibbs sampler: draws beta from the posterior truncated to the constraint region, via component-wise truncated-Normal updates. Reference --------- Geweke, J. (1995). "Bayesian Inference for Linear Models Subject to Linear Inequality Constraints." Modelling and Prediction: Honoring Seymour Geisser, eds. W.O. Johnson and A. Zellner, 248-263. """ import time import numpy as np from scipy.linalg import cholesky from scipy.special import ndtr, ndtri # ── private helpers ─────────────────────────────────────────────────────────── def _rtn(lo, hi, mu, sigma, u): """ Draw from N(mu, sigma) truncated to (lo, hi) given u ~ Uniform(0,1). Uses the inverse-CDF method. When the truncated mass underflows to zero (e.g. the interval is 20+ conditional SDs from the mean), falls back to the Mills-ratio approximation for the nearest boundary instead of returning the midpoint, which would be inf when hi = +inf or lo = -inf. """ lo_s = (lo - mu) / sigma hi_s = (hi - mu) / sigma p_lo = float(ndtr(lo_s)) p_hi = float(ndtr(hi_s)) gap = p_hi - p_lo if gap < 1e-300: # All mass is concentrated near one boundary. # Mills-ratio approximation: E[X | X > lo] ≈ lo + sigma/lo_s (lo_s >> 0) if lo_s >= 0: # interval above mean; mass near lower bound return float(lo) + sigma / max(float(lo_s), 1e-6) elif hi_s <= 0: # interval below mean; mass near upper bound return float(hi) - sigma / max(-float(hi_s), 1e-6) else: # narrow finite interval; return midpoint lo_f = float(lo) if np.isfinite(lo) else mu - 100.0 * sigma hi_f = float(hi) if np.isfinite(hi) else mu + 100.0 * sigma return (lo_f + hi_f) * 0.5 p = float(np.clip(u * gap + p_lo, 1e-15, 1.0 - 1e-15)) return float(ndtri(p)) * sigma + mu # ── public API ──────────────────────────────────────────────────────────────── def compute_posterior(y, X, sigma2, b0=None, B0=None): """ Unconstrained Normal posterior for beta given (fixed) sigma2. Prior: beta ~ N(b0, B0) [informative, b0 and B0 in original scale] beta ~ flat [if b0 / B0 omitted] Note: B0 here is NOT scaled by sigma2; it is the prior covariance directly. For a sigma2-scaled conjugate prior, divide B0 by sigma2 before passing. Parameters ---------- y : (n,) observations X : (n, k) design matrix sigma2 : float error variance b0 : (k,) prior mean [optional] B0 : (k, k) prior covariance [optional] Returns ------- bbar : (k,) posterior mean Bbar : (k, k) posterior covariance """ y = np.asarray(y, float).ravel() X = np.asarray(X, float) XtX = X.T @ X Xty = X.T @ y if b0 is not None and B0 is not None: B0_inv = np.linalg.inv(np.asarray(B0, float)) Prec = XtX / sigma2 + B0_inv rhs = Xty / sigma2 + B0_inv @ np.asarray(b0, float) else: Prec = XtX / sigma2 rhs = Xty / sigma2 Bbar = np.linalg.inv(Prec) bbar = Bbar @ rhs return bbar, Bbar def ghk_prob(a_vec, w_vec, D, bbar, Bbar, M=20000, seed=42): """ GHK probability simulator. Estimates P(a_vec <= D @ beta <= w_vec) where beta ~ N(bbar, Bbar). Applied to the unconstrained posterior this gives the posterior probability p_{2|1} that the inequality constraints hold. The GHK simulator sequentially conditions on each element of v = D @ beta using the lower Cholesky factor of Cov(v) = D Bbar D'. Parameters ---------- a_vec, w_vec : (q,) lower and upper bounds on D @ beta Use -np.inf / np.inf for one-sided constraints. D : (q, k) linear constraint matrix bbar : (k,) unconstrained posterior mean Bbar : (k, k) unconstrained posterior covariance M : int number of GHK simulator draws seed : int Returns ------- est : float probability estimate se : float simulation standard error """ rng = np.random.default_rng(seed) a_vec = np.asarray(a_vec, float) w_vec = np.asarray(w_vec, float) D = np.asarray(D, float) bbar = np.asarray(bbar, float) Bbar = np.asarray(Bbar, float) q = len(a_vec) mu_v = D @ bbar # (q,) mean of v = D @ beta Sigma_v = D @ Bbar @ D.T # (q,q) covariance of v L = cholesky(Sigma_v, lower=True) # Sigma_v = L L' # Sequential Cholesky-factor GHK, vectorised over M paths log_weights = np.zeros(M) U = np.zeros((M, q)) # drawn standardised values u_j for j in range(q): cum = U[:, :j] @ L[j, :j] if j > 0 else np.zeros(M) # (M,) lo_j = (a_vec[j] - mu_v[j] - cum) / L[j, j] hi_j = (w_vec[j] - mu_v[j] - cum) / L[j, j] p_lo = ndtr(lo_j) p_hi = ndtr(hi_j) gap = np.clip(p_hi - p_lo, 1e-300, None) log_weights += np.log(gap) # Draw u_j ~ TN(0, 1; lo_j, hi_j) for all M paths u_unif = rng.uniform(size=M) p = np.clip(u_unif * gap + p_lo, 1e-15, 1.0 - 1e-15) U[:, j] = ndtri(p) weights = np.exp(log_weights) est = float(weights.mean()) se = float(weights.std()) / np.sqrt(M) return est, se def constrained_gibbs( y, X, D, a_vec, w_vec, sigma2=None, b0=None, B0=None, a0=0.0, s0=0.0, R=10000, burn=2000, seed=42, nprint=1000, ): """ Gibbs sampler for the constrained posterior. Samples from p(beta | y, sigma2, a_vec <= D @ beta <= w_vec) by iterating over each component beta_j, drawing from its conditional truncated-Normal distribution: beta_j | beta_{-j}, sigma2, constraints ~ TN(m_j, v_j; L_j, U_j) where m_j, v_j come from the conditional of the unconstrained Normal posterior, and L_j, U_j are the tightest bounds imposed by all rows of D given the current beta_{-j}. If sigma2 is None, also samples sigma2 | beta, y from the conjugate IG distribution each iteration. Parameters ---------- y, X : (n,), (n, k) observations and design matrix D : (q, k) linear constraint matrix a_vec, w_vec : (q,) lower / upper bounds (use ±np.inf for one-sided) sigma2 : float or None -- if None, estimate jointly with beta b0, B0 : (k,), (k, k) -- Normal prior mean, covariance [optional] a0, s0 : float -- IG(a0/2, s0/2) hyperparameters for sigma2 (flat improper prior if both are 0) R, burn : int -- total and burn-in iterations seed : int nprint : int -- print progress every nprint iters (0 = silent) Returns ------- dict with keys: 'beta' : (R_kept, k) constrained posterior draws 'sigma2' : (R_kept,) [only present when sigma2 is None on input] """ rng = np.random.default_rng(seed) y = np.asarray(y, float).ravel() X = np.asarray(X, float) D = np.asarray(D, float) a_vec = np.asarray(a_vec, float) w_vec = np.asarray(w_vec, float) n, k = X.shape q = D.shape[0] est_s2 = sigma2 is None R_kept = R - burn XtX = X.T @ X Xty = X.T @ y # Pre-parse prior if b0 is not None and B0 is not None: B0 = np.asarray(B0, float) b0 = np.asarray(b0, float) B0_inv = np.linalg.inv(B0) else: B0_inv = None b0_eff = np.zeros(k) # Helper: posterior precision matrix and mean given sig2 def _posterior_params(s2): if B0_inv is not None: Lambda = XtX / s2 + B0_inv rhs = Xty / s2 + B0_inv @ b0 else: Lambda = XtX / s2 rhs = Xty / s2 bbar = np.linalg.solve(Lambda, rhs) return Lambda, bbar # Initial values sig2 = float(np.var(y - X @ np.linalg.lstsq(X, y, rcond=None)[0]) + 1e-6) \ if est_s2 else float(sigma2) beta = np.linalg.lstsq(X, y, rcond=None)[0].copy() Lambda, bbar = _posterior_params(sig2) Dbeta = D @ beta # (q,) current D @ beta, kept in sync # Storage beta_draws = np.empty((R_kept, k)) sigma2_draws = np.empty(R_kept) if est_s2 else None kept = 0 t0 = time.time() for r in range(1, R + 1): # ── Block 1: beta_j | beta_{-j}, sigma2, constraints ───────────────── # Conditional of beta_j under N(bbar, Lambda^{-1}): # v_j = 1 / Lambda[j, j] # m_j = bbar[j] - v_j * sum_{l != j} Lambda[j, l] * (beta[l] - bbar[l]) delta = beta - bbar # (k,) centred current draw for j in range(k): v_j = 1.0 / Lambda[j, j] lam_dot = float(Lambda[j] @ delta) # Lambda[j,:] @ delta # Remove diagonal term to get sum over l != j m_j = bbar[j] - v_j * (lam_dot - Lambda[j, j] * delta[j]) # Truncation bounds on beta_j from D @ beta in [a_vec, w_vec] d_j = D[:, j] # (q,) nz = np.abs(d_j) > 1e-15 # rows with nonzero D[i,j] Lj = -np.inf Uj = np.inf if nz.any(): res = Dbeta[nz] - d_j[nz] * beta[j] # (nz_count,) lo_i = (a_vec[nz] - res) / d_j[nz] hi_i = (w_vec[nz] - res) / d_j[nz] pos = d_j[nz] > 0 neg = ~pos if pos.any(): Lj = max(Lj, float(lo_i[pos].max())) Uj = min(Uj, float(hi_i[pos].min())) if neg.any(): # Bounds flip for negative D[i,j] Lj = max(Lj, float(hi_i[neg].max())) Uj = min(Uj, float(lo_i[neg].min())) if Lj > Uj: # Numerically infeasible given current beta_{-j}: keep beta[j] continue old_j = beta[j] beta[j] = _rtn(Lj, Uj, m_j, np.sqrt(v_j), rng.uniform()) # Update cached D @ beta and delta diff = beta[j] - old_j Dbeta += d_j * diff delta[j] += diff # ── Block 2: sigma2 | beta, y ──────────────────────────────────────── if est_s2: resid = y - X @ beta ss = float(resid @ resid) a_n = a0 + n s_n = s0 + ss # sigma2 ~ IG(a_n/2, s_n/2) <=> sigma2 = s_n / chi2(a_n) sig2 = s_n / rng.chisquare(a_n) Lambda, bbar = _posterior_params(sig2) # ── Store ───────────────────────────────────────────────────────────── if r > burn: beta_draws[kept] = beta if est_s2: sigma2_draws[kept] = sig2 kept += 1 if nprint > 0 and r % nprint == 0: elapsed = time.time() - t0 print(f' Iter {r:5d}/{R} ({elapsed:.1f}s elapsed)') elapsed = time.time() - t0 print(f'Done in {elapsed / 60:.1f} min | kept draws: {R_kept}') result = {'beta': beta_draws} if est_s2: result['sigma2'] = sigma2_draws return result def posterior_summary(draws, param_names=None, true_vals=None): """ Posterior mean, SD, and 95% credible interval for each column. Parameters ---------- draws : (R, p) array of posterior draws param_names : list of p strings [optional] true_vals : (p,) array of true values; adds 'true' and 'covered' columns Returns ------- pandas.DataFrame indexed by param_names """ import pandas as pd p = draws.shape[1] if param_names is None: param_names = [f'b{j}' for j in range(p)] lo = np.quantile(draws, 0.025, axis=0) hi = np.quantile(draws, 0.975, axis=0) df = pd.DataFrame({ 'mean': draws.mean(0), 'sd': draws.std(0), 'q025': lo, 'q975': hi, }, index=param_names) if true_vals is not None: tv = np.asarray(true_vals) df['true'] = tv df['covered'] = (tv >= lo) & (tv <= hi) return df.round(4)