Tobit — PyMC cross-check (pm.Censored)¶

Companion to tobit_python.ipynb. PyMC's pm.Censored wraps a base distribution with censoring bounds: an observation equal to the bound is treated as censored (CDF), otherwise exact. For left-censored hours that's pm.Censored(pm.Normal.dist(mu, sigma), lower=0, upper=inf, observed=hours).

In [1]:
import numpy as np, pandas as pd, pymc as pm, pytensor.tensor as pt, arviz as az
from tobit_gibbs import tobit_gibbs
d = pd.read_csv('mroz_hours.csv')
Xm = np.column_stack([np.ones(len(d)), d.nwifeinc, d.educ, d.exper, d.expersq, d.age, d.kidslt6, d.kidsge6])
ym = d.hours.to_numpy(float); k = Xm.shape[1]
with pm.Model() as m:
    beta = pm.Normal('beta', 0, 1e4, shape=k); sigma = pm.HalfNormal('sigma', 3000.0)
    dist = pm.Normal.dist(mu=pt.dot(Xm, beta), sigma=sigma)
    pm.Censored('hours', dist, lower=0.0, upper=np.inf, observed=ym)
    idata = pm.sample(2000, tune=2000, chains=4, target_accept=0.9, random_seed=1, progressbar=False)
po = idata.posterior; nm=['intercept','nwifeinc','educ','exper','expersq','age','kidslt6','kidsge6']
b = po['beta'].values.reshape(-1, k)
print('%-10s %14s' % ('', 'PyMC Tobit'))
for j in range(k): print('%-10s %10.2f [%.0f, %.0f]' % (nm[j], b[:,j].mean(), np.percentile(b[:,j],2.5), np.percentile(b[:,j],97.5)))
print('sigma = %.0f   max r_hat %.3f' % (float(po['sigma'].mean()), float(az.summary(idata,var_names=['beta','sigma'])['r_hat'].max())))
g = tobit_gibbs(ym, Xm, np.where(ym>0,ym,-np.inf), np.where(ym>0,ym,0.0), ym>0, R=12000, burn=3000, seed=1)
print('\nfrom-scratch vs PyMC (key coefs): educ %.1f/%.1f  kidslt6 %.0f/%.0f  sigma %.0f/%.0f'
      % (g['beta'].mean(0)[2], b[:,2].mean(), g['beta'].mean(0)[6], b[:,6].mean(), g['sigma'].mean(), float(po['sigma'].mean())))

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [beta, sigma]

Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 19 seconds.

               PyMC Tobit
intercept      951.70 [77, 1854]
nwifeinc        -8.93 [-18, -0]
educ            81.40 [38, 124]
exper          133.09 [100, 167]
expersq         -1.90 [-3, -1]
age            -54.69 [-69, -40]
kidslt6       -902.79 [-1135, -683]
kidsge6        -15.34 [-92, 62]
sigma = 1138   max r_hat 1.000

from-scratch vs PyMC (key coefs): educ 81.7/81.4  kidslt6 -900/-903  sigma 1136/1138

Mixing: data-augmentation Gibbs vs marginalised NUTS¶

The two engines sample different objects. The from-scratch Gibbs imputes the ~43%-censored latents $y^*$ each sweep (data augmentation); those latents are correlated with $(\beta,\sigma^2)$, so the chain's autocorrelation grows with the censored fraction — worst for $\sigma$, which feeds on the imputed tail. PyMC uses pm.Censored, which marginalises the latents out (the analytic normal-CDF likelihood) and lets NUTS sample $(\beta,\sigma)$ directly. We compare effective sample size (ESS) per draw and per second, plus the $\sigma$ trace and autocorrelation.

In [2]:
import time
# fresh, timed runs for a fair comparison
t0 = time.time()
gg = tobit_gibbs(ym, Xm, np.where(ym>0,ym,-np.inf), np.where(ym>0,ym,0.0), ym>0, R=12000, burn=3000, seed=7)
t_g = time.time() - t0
with m:
    t0 = time.time()
    idn = pm.sample(1500, tune=1500, chains=4, target_accept=0.9, random_seed=7, progressbar=False)
    t_n = time.time() - t0
nd_g = gg['beta'].shape[0]; nd_n = 1500 * 4
def ess_arr(x): return float(az.ess(np.asarray(x).reshape(1, -1)))
ess_n = az.ess(idn, var_names=['beta','sigma'])
items = [('sigma', gg['sigma'], float(ess_n['sigma'].values)),
         ('intercept', gg['beta'][:,0], float(ess_n['beta'].values[0])),
         ('educ', gg['beta'][:,2], float(ess_n['beta'].values[2]))]
print('Gibbs: %d draws in %.1fs   |   NUTS: %d draws in %.1fs\n' % (nd_g, t_g, nd_n, t_n))
print('%-10s | %8s %7s %8s | %8s %7s %8s' % ('param','Gibbs ESS','/draw','/sec','NUTS ESS','/draw','/sec'))
for name, gdraws, en in items:
    eg = ess_arr(gdraws)
    print('%-10s | %8.0f %7.2f %8.0f | %8.0f %7.2f %8.0f' %
          (name, eg, eg/nd_g, eg/t_g, en, en/nd_n, en/t_n))

C:\Users\user\anaconda3\envs\pymc-env\Lib\site-packages\threadpoolctl.py:1226: RuntimeWarning: 
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
the same time. Both libraries are known to be incompatible and this
can cause random crashes or deadlocks on Linux when loaded in the
same Python program.
Using threadpoolctl may cause crashes or deadlocks. For more
information and possible workarounds, please see
    https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md

  warnings.warn(msg, RuntimeWarning)
Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [beta, sigma]

Sampling 4 chains for 1_500 tune and 1_500 draw iterations (6_000 + 6_000 draws total) took 16 seconds.

Gibbs: 9000 draws in 3.0s   |   NUTS: 6000 draws in 17.5s

param      | Gibbs ESS   /draw     /sec | NUTS ESS   /draw     /sec
sigma      |     2130    0.24      704 |     5703    0.95      327
intercept  |     6188    0.69     2045 |     2793    0.47      160
educ       |     5607    0.62     1854 |     3703    0.62      212
In [3]:
import matplotlib.pyplot as plt
sig_g = gg['sigma']; sig_n = idn.posterior['sigma'].values
def acf(x, L=50):
    x = np.asarray(x) - np.mean(x); nn = len(x); v = float(np.dot(x, x))
    return np.array([np.dot(x[:nn-l], x[l:])/v for l in range(L)])
fig, ax = plt.subplots(1, 2, figsize=(13, 4.3))
ax[0].plot(sig_g[:600], color='orange', lw=.5, label='Gibbs (data augmentation)')
ax[0].plot(sig_n[0,:600], color='steelblue', lw=.5, label='NUTS (pm.Censored)')
ax[0].set_xlabel('iteration'); ax[0].set_ylabel('σ'); ax[0].set_title('σ trace (first 600 draws)'); ax[0].legend(fontsize=8)
ac_g = acf(sig_g, 50); ac_n = acf(sig_n[0], 50)
ax[1].plot(ac_g, color='orange', lw=1.5, label='Gibbs'); ax[1].plot(ac_n, color='steelblue', lw=1.5, label='NUTS')
ax[1].axhline(0, color='k', lw=.5); ax[1].set_xlabel('lag'); ax[1].set_ylabel('autocorrelation')
ax[1].set_title('σ autocorrelation — DA-Gibbs decays slower'); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.savefig('tobit_mixing.png', dpi=120, bbox_inches='tight'); plt.show()
No description has been provided for this image

Mixing — what to read¶

  • Per draw, NUTS mixes better: its $\sigma$ ESS/draw is far higher and its autocorrelation collapses within a few lags, while the data-augmentation Gibbs $\sigma$ shows visibly slower ACF decay — the cost of imputing the 43% censored latents (β's are much less affected; the augmentation hits the variance hardest).
  • Per second, it's closer (and can flip): each Gibbs sweep is cheap (closed-form draws, no gradients), so despite the worse per-draw mixing the Gibbs can be competitive or better on ESS/second — the honest, full-picture verdict.
  • Takeaway: augmentation trades mixing efficiency for cheap conjugate steps; marginalise-then-HMC trades expensive steps for low autocorrelation. Under heavy censoring the gap in per-draw mixing is real — exactly the caveat flagged in the from-scratch sampler section.

Results¶

pm.Censored reproduces the from-scratch Tobit: educ ≈ +81, a child under 6 ≈ −900 hours, σ ≈ 1130, r̂ ≈ 1.0 — and both match AER::tobit. The data-augmentation Gibbs (impute) and PyMC's CDF-based censored likelihood are two routes to the same posterior.