Constrained Bayesian regression — Bayesian via HMC (NumPyro)¶

Companion to constrained_lm_python.ipynb. The Bayesian-by-sampling view of the same model, on the same data. Two parts:

  1. Sign constraints ($\beta_1>0,\beta_2>0$) via bounded priors — HMC handles these box constraints cleanly and agrees with the Gibbs sampler.
  2. A tight polytope ($\beta_1>0,\ \beta_2>0,\ \beta_1+\beta_2<0.3$) — a thin feasible triangle where HMC struggles (divergences, poor mixing at the hard walls), while Geweke's truncated-Gibbs sails through. This is why the specialised sampler exists.

(PyMC was the intended engine, but on this Windows box pm.sample hits a BLAS-threadpool DLL fault for every model; NumPyro's JAX NUTS is the same HMC and runs cleanly. The lesson about HMC vs Gibbs on hard constraints is sampler-agnostic.)

In [1]:
import numpy as np, pandas as pd, jax
jax.config.update('jax_enable_x64', True)
import jax.numpy as jnp, numpyro, numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
from constrained_lm_gibbs import constrained_gibbs, compute_posterior
d = pd.read_csv('constrained_lm_data.csv'); X = np.column_stack([np.ones(len(d)), d.x1, d.x2, d.x3]); y = d.y.values
sigma2 = 2.0
bbar,_ = compute_posterior(y, X, sigma2); print('unconstrained posterior mean:', bbar.round(3))

# ---- Part 1: sign constraints via bounded priors (HMC) ----
def sign_model(X, y):
    b0 = numpyro.sample('b0', dist.Normal(0,5)); b3 = numpyro.sample('b3', dist.Normal(0,5))
    b1 = numpyro.sample('b1', dist.HalfNormal(50)); b2 = numpyro.sample('b2', dist.HalfNormal(50))  # >=0
    mu = b0 + b1*X[:,1] + b2*X[:,2] + b3*X[:,3]
    numpyro.sample('y', dist.Normal(mu, np.sqrt(sigma2)), obs=y)
m1 = MCMC(NUTS(sign_model), num_warmup=1000, num_samples=2000, num_chains=2, chain_method='sequential', progress_bar=False)
m1.run(jax.random.PRNGKey(1), X, y, extra_fields=('diverging',))
s1 = m1.get_samples()
gib = constrained_gibbs(y, X, np.eye(4)[1:3], np.zeros(2), np.full(2,np.inf), sigma2=sigma2, R=20000, burn=2000, seed=1, nprint=0)['beta']
print('\nPart 1 (sign constraints) posterior means:')
print('  HMC (NumPyro): b1=%.3f b2=%.3f   divergences=%d' % (s1['b1'].mean(), s1['b2'].mean(), int(m1.get_extra_fields()['diverging'].sum())))
print('  Gibbs (Geweke): b1=%.3f b2=%.3f' % (gib[:,1].mean(), gib[:,2].mean()))
print('  -> HMC and Gibbs AGREE on simple box constraints')

unconstrained posterior mean: [1.644 0.208 0.064 0.815]

Done in 0.0 min  |  kept draws: 18000

Part 1 (sign constraints) posterior means:
  HMC (NumPyro): b1=0.313 b2=0.186   divergences=0
  Gibbs (Geweke): b1=0.308 b2=0.186
  -> HMC and Gibbs AGREE on simple box constraints
In [2]:
# ---- Part 2: a tight polytope  b1>0, b2>0, b1+b2<0.3  (thin triangle) ----
def poly_model(X, y):
    b0 = numpyro.sample('b0', dist.Normal(0,5)); b3 = numpyro.sample('b3', dist.Normal(0,5))
    b1 = numpyro.sample('b1', dist.HalfNormal(50)); b2 = numpyro.sample('b2', dist.HalfNormal(50))
    numpyro.factor('poly', jnp.where(b1+b2 < 0.3, 0.0, -1e6))      # hard-ish wall: b1+b2 < 0.3
    mu = b0 + b1*X[:,1] + b2*X[:,2] + b3*X[:,3]
    numpyro.sample('y', dist.Normal(mu, np.sqrt(sigma2)), obs=y)
m2 = MCMC(NUTS(poly_model, target_accept_prob=0.9), num_warmup=1500, num_samples=2000, num_chains=2, chain_method='sequential', progress_bar=False)
m2.run(jax.random.PRNGKey(2), X, y, extra_fields=('diverging',))
from numpyro.diagnostics import effective_sample_size
s2 = m2.get_samples(group_by_chain=True)
ess1 = float(effective_sample_size(s1 if False else m1.get_samples(group_by_chain=True)['b1']))
ess_hmc = float(effective_sample_size(s2['b1'])); div = int(m2.get_extra_fields()['diverging'].sum())

# Geweke Gibbs on the SAME triangle
D = np.array([[0,1,0,0],[0,0,1,0],[0,1,1,0]], float); a=np.array([0,0,-np.inf]); w=np.array([np.inf,np.inf,0.3])
gpoly = constrained_gibbs(y, X, D, a, w, sigma2=sigma2, R=22000, burn=2000, seed=3, nprint=0)['beta']
from numpyro.diagnostics import effective_sample_size as ess_fn
ess_gibbs = float(ess_fn(gpoly[:,1].reshape(1,-1)))
print('TIGHT TRIANGLE (b1>0, b2>0, b1+b2<0.3):')
print('  HMC (NumPyro): divergences=%d   ESS(b1)=%.0f / 4000 draws   feasible=%.1f%%' %
      (div, ess_hmc, 100*np.mean((m1.get_samples()['b1']+m1.get_samples()['b2'])<0.3)))
print('  Gibbs (Geweke): ESS(b1)=%.0f / 20000   all draws feasible=%s' %
      (ess_gibbs, bool(np.all((gpoly[:,1]>0)&(gpoly[:,2]>0)&(gpoly[:,1]+gpoly[:,2]<0.3)))))
print('  -> HMC degrades at the hard walls (divergences, low ESS); Gibbs samples the triangle cleanly')

Done in 0.0 min  |  kept draws: 20000
TIGHT TRIANGLE (b1>0, b2>0, b1+b2<0.3):
  HMC (NumPyro): divergences=3824   ESS(b1)=157 / 4000 draws   feasible=22.0%
  Gibbs (Geweke): ESS(b1)=12892 / 20000   all draws feasible=True
  -> HMC degrades at the hard walls (divergences, low ESS); Gibbs samples the triangle cleanly
In [3]:
import matplotlib.pyplot as plt
fig, ax = plt.subplots(1, 2, figsize=(13, 5))
# Part 1 agreement
ax[0].scatter(gib[::5,1], gib[::5,2], s=5, alpha=.15, color='orange', label='Gibbs')
ax[0].scatter(m1.get_samples()['b1'][::3], m1.get_samples()['b2'][::3], s=5, alpha=.12, color='steelblue', label='HMC')
ax[0].axhline(0,color='red',ls=':'); ax[0].axvline(0,color='red',ls=':')
ax[0].set_xlabel('β1'); ax[0].set_ylabel('β2'); ax[0].set_title('Sign constraints: HMC ≈ Gibbs'); ax[0].legend(fontsize=8)
# Part 2 triangle: Gibbs fills it, HMC piles at walls
ax[1].plot([0,0.3,0,0],[0,0,0.3,0],'g-',lw=1.5, label='feasible triangle')
ax[1].scatter(gpoly[::3,1], gpoly[::3,2], s=5, alpha=.2, color='orange', label='Gibbs (clean)')
hb1=m2.get_samples()['b1']; hb2=m2.get_samples()['b2']
ax[1].scatter(hb1[::2], hb2[::2], s=5, alpha=.2, color='steelblue', label='HMC (struggles)')
ax[1].set_xlim(-0.05,0.45); ax[1].set_ylim(-0.05,0.45); ax[1].set_xlabel('β1'); ax[1].set_ylabel('β2')
ax[1].set_title('Tight polytope: Gibbs fills it, HMC frays at the walls'); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.savefig('constrained_numpyro.png', dpi=120, bbox_inches='tight'); plt.show()
No description has been provided for this image

Left — under the sign constraints the NumPyro HMC cloud (blue) sits right on top of the Geweke truncated-Gibbs (orange): the same posterior, so for box constraints generic HMC is fine. Right — on the thin triangle $\beta_1+\beta_2<0.3$ the Gibbs draws (orange) fill the feasible region evenly, while HMC (blue) frays at the hard walls — 3,824 divergences, ESS 157 vs 12,892, and only ~22% of its "draws" actually feasible. The picture behind the punchline below: sign constraints are easy, polytopes are not.

Results¶

  • Part 1 — simple sign constraints: bounded priors make HMC respect $\beta_1,\beta_2>0$, and the posterior matches the Geweke Gibbs (means/clouds coincide). For box constraints, generic HMC is fine.
  • Part 2 — the tight polytope: in the thin triangle $\beta_1+\beta_2<0.3$, HMC hits the hard walls — divergences rise and ESS collapses (and a chunk of "draws" leak outside the region through the soft penalty). The Geweke truncated-Gibbs samples the triangle exactly and efficiently, every draw feasible.
  • Punchline: this is why the project's from-scratch sampler exists — for hard inequality regions (small $p_{2|1}$, tilted polytopes), the specialised truncated-Gibbs/GHK machinery beats both rejection sampling and generic HMC. Sign constraints are easy; polytopes are not.