Bayesian Linear Model with Inequality Constraints

Python · R  ·  Download Gibbs sampler

Model

Standard Bayesian linear regression leaves the posterior for β\beta unconstrained. When economic theory or physical plausibility demands that certain linear combinations of coefficients satisfy inequality restrictions, the unconstrained posterior is replaced by its truncation to the feasible region.

y=Xβ+ε,εN(0,σ2In),aDβwy = X\beta + \varepsilon, \quad \varepsilon \sim N(0,\,\sigma^2 I_n), \quad a \le D\beta \le w

GHK Probability Simulator

Evaluating the posterior probability that all constraints hold — p21=P(aDβwdata)p_{2|1} = P(a \le D\beta \le w \mid \text{data}) — requires integrating a multivariate Normal over a polyhedral region. The GHK simulator (Geweke 1991; Hajivassiliou & McFadden 1998) decomposes this integral into a product of univariate truncated-Normal probabilities using the lower Cholesky factor of Cov(Dβ)=DBˉD\text{Cov}(D\beta) = D\bar{B}D^\top, making high-dimensional constraint probabilities tractable.

p21=P(aDβwdata)=j=1qP ⁣(a~jujw~j)p_{2|1} = P(a \le D\beta \le w \mid \text{data}) = \prod_{j=1}^{q} P\!\left(\tilde{a}_j \le u_j \le \tilde{w}_j\right)

where uju_j are sequentially standardised draws from the Cholesky decomposition of DBˉDD\bar{B}D^\top. Each factor is a univariate truncated-Normal probability evaluated via the Normal CDF.

2-Block Constrained Gibbs — Geweke (1995)

When p21p_{2|1} is small, rejection sampling from the unconstrained posterior is impractical — it requires 1/p21\sim 1/p_{2|1} draws per accepted sample. The constrained Gibbs sampler avoids this by cycling through each βj\beta_j component-wise, drawing from its conditional truncated-Normal distribution. Each draw always satisfies the constraints, so efficiency is independent of p21p_{2|1}.

BlockDrawDistribution
(1) βjβj,σ2,constraints\beta_j \mid \beta_{-j},\, \sigma^2,\, \text{constraints}Truncated normal TN(mj,vj;  Lj,Uj)\text{TN}(m_j, v_j;\; L_j, U_j)
(2) σ2β,y\sigma^2 \mid \beta,\, yInverse-Gamma conjugate IG(an/2,  sn/2)\text{IG}(a_n/2,\; s_n/2)

The truncation bounds Lj,UjL_j, U_j are derived row-by-row from Dβ[a,w]D\beta \in [a, w], solving for βj\beta_j given the current βj\beta_{-j}. Block (2) is skipped when σ2\sigma^2 is treated as known.

Notebooks

The main notebook applies all three algorithms to the UCI Automobile Imports dataset (195 cars, 1985). Log-price is regressed on log engine size, log horsepower, log curb weight, log city MPG, and three additional regressors. Seven economic sign constraints are imposed: larger engine, more power, and heavier weight should raise price; better fuel economy should lower it; bore, stroke, and compression ratio should enter positively. Two constraints (bore and stroke) are violated at OLS due to multicollinearity — making p214×104p_{2|1} \approx 4 \times 10^{-4} and rejection sampling completely impractical. Two companion notebooks cross-check the machinery against general-purpose tools. The HMC notebook (NumPyro NUTS) shows the key trade-off: for simple sign constraints, bounded-prior HMC agrees with the Gibbs sampler, but on a tight polytope (β1+β2<0.3\beta_1+\beta_2<0.3, a thin feasible triangle) HMC hits the hard walls — divergences rise, ESS collapses — while Geweke's truncated-Gibbs samples it exactly and efficiently. That contrast is precisely why the specialised sampler exists: sign constraints are easy, polytopes are not. The R notebook maps each example piece onto a standard package — TruncatedNormal::pmvnorm (the GHK probability via Botev minimax tilting), tmvtnorm::rtmvnorm (the truncated-MVN Gibbs posterior), and bain (the constraint as a Bayes factor) — all agreeing on identical data.

Downloads

Gibbs Sampler — Source Code

"""
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)

References