import numpy as np
import pandas as pd
from scipy.optimize import minimize
from scipy import stats
import matplotlib.pyplot as plt

rng = np.random.default_rng(8)          # lecture number
FAST_MODE = True

np.set_printoptions(precision=4, suppress=True)

from pathlib import Path


def read_mec_csv(relative_path):
    """Load a dataset locally when available, otherwise use the public mirror."""
    for parent in (Path.cwd(), *Path.cwd().parents):
        local = parent / 'mec_datasets' / relative_path
        if local.is_file():
            return pd.read_csv(local)
    url = ('https://raw.githubusercontent.com/math-econ-code/'
           'mec_datasets/main/' + relative_path)
    try:
        return pd.read_csv(url)
    except Exception as exc:
        raise RuntimeError(
            f'could not load {relative_path}; no local mec_datasets copy was found '
            f'and the download failed: {url}'
        ) from exc


travel = read_mec_csv('demand_travelmode/travelmodedata.csv')
pihat_iy = np.where(travel['choice'] == 'yes', 1, 0).astype(float)
Y = travel['mode'].nunique()
I = travel.shape[0] // Y
cov = travel[['travel', 'income', 'gcost']].values
Phi_iy_k = np.column_stack([cov[:, 0], -(cov[:, 0] * cov[:, 1]), -cov[:, 2]])
Phi_iy_k = (Phi_iy_k - Phi_iy_k.mean(0)) / Phi_iy_k.std(0, ddof=1)
K = Phi_iy_k.shape[1]
regressors = ['travel time', '-(travel time x income)', '-generalized cost']


def loglike(lam):
    A = (Phi_iy_k @ lam).reshape(I, Y)
    return float(pihat_iy @ (Phi_iy_k @ lam) - np.logaddexp.reduce(A, axis=1).sum())


def grad_loglike(lam):
    A = (Phi_iy_k @ lam).reshape(I, Y)
    p = np.exp(A - np.logaddexp.reduce(A, axis=1)[:, None]).reshape(-1)
    return (pihat_iy - p) @ Phi_iy_k


def hess_loglike(lam):
    A = (Phi_iy_k @ lam).reshape(I, Y)
    p = np.exp(A - np.logaddexp.reduce(A, axis=1)[:, None]).reshape(I, Y)
    P3 = Phi_iy_k.reshape(I, Y, K)
    Am = np.einsum('iy,iyk,iyl->kl', p, P3, P3)
    m = np.einsum('iy,iyk->ik', p, P3)
    return -(Am - np.einsum('ik,il->kl', m, m))


TAU = 5.0                                        # prior standard deviation
logpost = lambda lam: loglike(lam) - 0.5 * np.sum(lam ** 2) / TAU ** 2
grad_logpost = lambda lam: grad_loglike(lam) - lam / TAU ** 2

fit = lambda obj, jac: minimize(obj, np.zeros(K), jac=jac, method='BFGS',
                                options={'gtol': 1e-12}).x

lam_map = fit(lambda l: -logpost(l), lambda l: -grad_logpost(l))
lam_mle = fit(lambda l: -loglike(l), lambda l: -grad_loglike(l))

gamma = 1 / (2 * TAU ** 2)
lam_ridge = fit(lambda l: -(loglike(l) - gamma * np.sum(l ** 2)),
                lambda l: -(grad_loglike(l) - 2 * gamma * l))

print(f'prior standard deviation tau = {TAU}, so gamma = 1/(2 tau^2) = {gamma}')
print(f"\n{'':>26} {'lambda_1':>10} {'lambda_2':>10} {'lambda_3':>10}")
for nm, v in [('MLE (lecture 3)', lam_mle), ('MAP, Gaussian prior', lam_map),
              ('ridge (lecture 4)', lam_ridge)]:
    print(f'{nm:>26} ' + ' '.join(f'{x:>10.6f}' for x in v))
print(f'\nmax |MAP - ridge| = {np.abs(lam_map - lam_ridge).max():.2e}   (must be zero)')
assert np.abs(lam_map - lam_ridge).max() < 1e-8

prior standard deviation tau = 5.0, so gamma = 1/(2 tau^2) = 0.02

                             lambda_1   lambda_2   lambda_3
           MLE (lecture 3)   0.186243   0.468979   0.550577
       MAP, Gaussian prior   0.185145   0.467735   0.549878
         ridge (lecture 4)   0.185145   0.467735   0.549878

max |MAP - ridge| = 0.00e+00   (must be zero)

H_post = -(hess_loglike(lam_map) - np.eye(K) / TAU ** 2)
Sigma_lap = np.linalg.inv(H_post)
sd_lap = np.sqrt(np.diag(Sigma_lap))
sd_mle = np.sqrt(np.diag(np.linalg.inv(-hess_loglike(lam_mle))))

print(f"{'':>26} {'sd_1':>9} {'sd_2':>9} {'sd_3':>9}")
print(f"{'Laplace, at the MAP':>26} " + ' '.join(f'{v:>9.4f}' for v in sd_lap))
print(f"{'asymptotic, lecture 3':>26} " + ' '.join(f'{v:>9.4f}' for v in sd_mle))
print(f'\nmax difference = {np.abs(sd_lap - sd_mle).max():.4f}')

                                sd_1      sd_2      sd_3
       Laplace, at the MAP    0.1885    0.2356    0.1819
     asymptotic, lecture 3    0.1889    0.2361    0.1820

max difference = 0.0004

chol_lap = np.linalg.cholesky(Sigma_lap)
scale = 2.38 / np.sqrt(K)


def rwm(nsteps, x0, gen):
    '''Random-walk Metropolis on the log-posterior, preconditioned by Sigma_Lap.'''
    x = x0.copy()
    lp = logpost(x)
    out = np.empty((nsteps, K))
    naccept = 0
    for t in range(nsteps):
        prop = x + scale * (chol_lap @ gen.normal(size=K))
        lpp = logpost(prop)
        if np.log(gen.uniform()) < lpp - lp:
            x, lp, naccept = prop, lpp, naccept + 1
        out[t] = x
    return out, naccept / nsteps


nsteps = 40_000 if not FAST_MODE else 20_000
chain, acc_rate = rwm(nsteps, lam_map, np.random.default_rng(81))
burn_naive = 2_000                                # the usual guess; section 3 revisits it
post = chain[burn_naive:]

print(f'{nsteps:,} steps, acceptance rate {acc_rate:.3f}')
print(f"\n{'':>26} {'lambda_1':>10} {'lambda_2':>10} {'lambda_3':>10}")
print(f"{'posterior mean':>26} " + ' '.join(f'{v:>10.4f}' for v in post.mean(0)))
print(f"{'MAP':>26} " + ' '.join(f'{v:>10.4f}' for v in lam_map))
print(f"{'posterior sd':>26} " + ' '.join(f'{v:>10.4f}' for v in post.std(0, ddof=1)))
print(f"{'Laplace sd':>26} " + ' '.join(f'{v:>10.4f}' for v in sd_lap))


def ess(x):
    '''Effective sample size from the initial-positive-sequence autocorrelations.'''
    n = len(x)
    z = x - x.mean()
    f = np.fft.rfft(z, 2 * n)
    ac = np.fft.irfft(f * np.conj(f))[:n].real
    ac /= ac[0]
    s = 0.0
    for k in range(1, n):
        if ac[k] < 0.05:
            break
        s += ac[k]
    return n / (1 + 2 * s)


es = [ess(post[:, k]) for k in range(K)]
print(f'\neffective sample size : ' + ', '.join(f'{e:,.0f}' for e in es)
      + f'   ({100*min(es)/len(post):.1f}% of {len(post):,} draws)')
mc_se = post.std(0, ddof=1) / np.sqrt(np.array(es))
print('Monte Carlo standard error of the posterior mean :', mc_se.round(4))
print('|posterior mean - MAP| in Monte Carlo standard errors :',
      (np.abs(post.mean(0) - lam_map) / mc_se).round(2))

20,000 steps, acceptance rate 0.322

                             lambda_1   lambda_2   lambda_3
            posterior mean     0.2081     0.4904     0.5539
                       MAP     0.1851     0.4677     0.5499
              posterior sd     0.1929     0.2406     0.1817
                Laplace sd     0.1885     0.2356     0.1819

effective sample size : 1,796, 1,747, 1,930   (9.7% of 18,000 draws)
Monte Carlo standard error of the posterior mean : [0.0046 0.0058 0.0041]
|posterior mean - MAP| in Monte Carlo standard errors : [5.04 3.94 0.96]

Sigma_inv = np.linalg.inv(Sigma_lap)


def log_prop_dens(z, mean):
    '''Log density of the proposal N(mean, scale^2 Sigma_Lap), up to a common constant.'''
    dz = (z - mean) / scale
    return -0.5 * dz @ Sigma_inv @ dz


def maximal_coupling(mx, my, gen):
    '''Maximal coupling of the two proposal distributions, centred at mx and my.

    Returns a pair with the correct marginals that is equal with the largest
    probability any coupling of the two can achieve.
    '''
    X = mx + scale * (chol_lap @ gen.normal(size=K))
    if np.log(gen.uniform()) + log_prop_dens(X, mx) <= log_prop_dens(X, my):
        return X, X                                    # the proposals coincide
    while True:                                        # otherwise draw Y from the residual
        Y = my + scale * (chol_lap @ gen.normal(size=K))
        if np.log(gen.uniform()) + log_prop_dens(Y, my) > log_prop_dens(Y, mx):
            return X, Y


def meeting_time(lag, gen, maxit=20_000):
    '''Run an L-lagged coupled random-walk Metropolis and return the meeting time.'''
    x = lam_map + scale * (chol_lap @ gen.normal(size=K))
    y = lam_map + scale * (chol_lap @ gen.normal(size=K))
    lpx, lpy = logpost(x), logpost(y)
    for _ in range(lag):                               # advance X alone by L steps
        prop = x + scale * (chol_lap @ gen.normal(size=K))
        lpp = logpost(prop)
        if np.log(gen.uniform()) < lpp - lpx:
            x, lpx = prop, lpp
    for n in range(1, maxit + 1):
        px, py = maximal_coupling(x, y, gen)
        lppx, lppy = logpost(px), logpost(py)
        u = np.log(gen.uniform())                      # one uniform for both chains
        if u < lppx - lpx:
            x, lpx = px, lppx
        if u < lppy - lpy:
            y, lpy = py, lppy
        if np.array_equal(x, y):
            return n + lag
    return np.inf


LAG = 30
nrep = 400 if not FAST_MODE else 250
gen = np.random.default_rng(83)
taus = np.array([meeting_time(LAG, gen) for _ in range(nrep)])
print(f'lag L = {LAG}, {nrep} replicates')
print(f'  all chains met      : {bool(np.isfinite(taus).all())}')
print(f'  meeting time tau    : mean {taus.mean():.1f}, median {np.median(taus):.0f}, '
      f'max {taus.max():.0f}')
assert np.isfinite(taus).all()


def tv_bound(t, taus, lag):
    return min(1.0, np.mean(np.maximum(0, np.ceil((taus - lag - t) / lag))))


grid = np.arange(0, 260, 5)
bounds = np.array([tv_bound(t, taus, LAG) for t in grid])
t_star = int(grid[np.argmax(bounds < 0.01)])
print(f'\n  upper bound on ||L(X_t) - pi||_TV falls below 0.01 at t = {t_star}')
print(f'  the burn-in guessed in section 2 was {burn_naive:,}, '
      f'{burn_naive // max(t_star,1)}x larger than needed')

lag L = 30, 250 replicates
  all chains met      : True
  meeting time tau    : mean 38.5, median 36, max 72

  upper bound on ||L(X_t) - pi||_TV falls below 0.01 at t = 35
  the burn-in guessed in section 2 was 2,000, 57x larger than needed

wide = travel.pivot(index='individual', columns='mode', values='choice')
y_bin = (wide['car'] == 'yes').to_numpy().astype(float)
car = travel[travel['mode'] == 'car'].sort_values('individual')
X_bin = np.column_stack([np.ones(I),
                         car['travel'].to_numpy(),
                         car['gcost'].to_numpy(),
                         car['income'].to_numpy()])
X_bin[:, 1:] = (X_bin[:, 1:] - X_bin[:, 1:].mean(0)) / X_bin[:, 1:].std(0, ddof=1)
names_bin = ['intercept', 'travel time', 'generalized cost', 'income']
print(f'binary probit: {int(y_bin.sum())} of {I} travellers drove')


def albert_chib(nsteps, gen, tau=5.0):
    '''Albert and Chib (1993) data augmentation for the binary probit.'''
    I_b, K_b = X_bin.shape
    V = np.linalg.inv(X_bin.T @ X_bin + np.eye(K_b) / tau ** 2)
    cholV = np.linalg.cholesky(V)
    beta = np.zeros(K_b)
    out = np.empty((nsteps, K_b))
    for t in range(nsteps):
        m = X_bin @ beta
        lo = np.where(y_bin > 0, stats.norm.cdf(-m), 0.0)
        hi = np.where(y_bin > 0, 1.0, stats.norm.cdf(-m))
        u = lo + (hi - lo) * gen.uniform(size=I_b)
        z = m + stats.norm.ppf(np.clip(u, 1e-12, 1 - 1e-12))     # truncated normal
        beta = V @ (X_bin.T @ z) + cholV @ gen.normal(size=K_b)
        out[t] = beta
    return out


nsteps_ac = 20_000 if not FAST_MODE else 10_000
draws = albert_chib(nsteps_ac, np.random.default_rng(84))[500:]

# maximum likelihood for comparison
def nll_probit(b):
    m = X_bin @ b
    return -float(np.sum(np.log(np.clip(np.where(y_bin > 0, stats.norm.cdf(m),
                                                 stats.norm.cdf(-m)), 1e-300, None))))


b_mle = minimize(nll_probit, np.zeros(X_bin.shape[1]), method='BFGS',
                 options={'gtol': 1e-10})
Hnum = np.zeros((X_bin.shape[1],) * 2)
eps = 1e-5
for a in range(X_bin.shape[1]):
    for b_ in range(X_bin.shape[1]):
        e1, e2 = np.zeros(X_bin.shape[1]), np.zeros(X_bin.shape[1])
        e1[a] = eps; e2[b_] = eps
        Hnum[a, b_] = (nll_probit(b_mle.x + e1 + e2) - nll_probit(b_mle.x + e1 - e2)
                       - nll_probit(b_mle.x - e1 + e2)
                       + nll_probit(b_mle.x - e1 - e2)) / (4 * eps ** 2)
se_mle = np.sqrt(np.diag(np.linalg.inv(Hnum)))

print(f"\n{'':>20} {'post. mean':>11} {'post. sd':>10} {'MLE':>11} {'MLE s.e.':>10}"
      f" {'gap / s.e.':>11}")
for k, nm in enumerate(names_bin):
    gap = abs(draws[:, k].mean() - b_mle.x[k]) / se_mle[k]
    print(f'{nm:>20} {draws[:,k].mean():>11.4f} {draws[:,k].std(ddof=1):>10.4f} '
          f'{b_mle.x[k]:>11.4f} {se_mle[k]:>10.4f} {gap:>11.2f}')
assert (np.abs(draws.mean(0) - b_mle.x) / se_mle).max() < 0.5

binary probit: 59 of 210 travellers drove

                      post. mean   post. sd         MLE   MLE s.e.  gap / s.e.
           intercept     -0.6294     0.0977     -0.6223     0.0962        0.07
         travel time     -0.5162     0.4698     -0.5309     0.4678        0.03
    generalized cost      0.3834     0.4680      0.4010     0.4649        0.04
              income      0.3390     0.0997      0.3327     0.0982        0.06

fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))

ax[0].plot(grid, bounds, 'o-', color='C0', ms=4, label='upper bound on TV distance')
ax[0].axhline(0.01, color='0.4', ls=':', lw=1.2, label='0.01')
ax[0].axvline(t_star, color='C3', ls='--', lw=1.3, label=f'bound < 0.01 at t = {t_star}')
ax[0].set_yscale('log'); ax[0].set_ylim(1e-3, 2)
ax[0].set_xlabel('iteration $t$'); ax[0].set_ylabel(r'$\|\mathcal{L}(X_t)-\pi\|_{TV}$ bound')
ax[0].set_title('A computed burn-in, from coupled chains')
ax[0].legend(fontsize=8.5); ax[0].grid(alpha=.25, which='both')

k_show = 2
ax[1].hist(post[:, k_show], bins=70, density=True, color='C0', alpha=.55,
           edgecolor='none', label='MCMC posterior')
xs = np.linspace(*ax[1].get_xlim(), 400)
ax[1].plot(xs, stats.norm.pdf(xs, lam_map[k_show], sd_lap[k_show]), color='C3', lw=2,
           label='Laplace approximation')
ax[1].axvline(lam_mle[k_show], color='k', ls='--', lw=1.2, label='MLE (lecture 3)')
ax[1].set_xlabel(rf'$\lambda_{k_show+1}$: {regressors[k_show]}')
ax[1].set_ylabel('density')
ax[1].set_title('Posterior versus the asymptotic approximation')
ax[1].legend(fontsize=8.5); ax[1].grid(alpha=.25)

plt.tight_layout(); plt.show()

Bayesian and MCMC methods for discrete choice¶

Alfred Galichon (NYU)¶

'math+econ+code' masterclass series: discrete choice¶

With python code examples¶

Learning objectives¶

References¶

Motivation: the same object, approached from the other side¶

1. The posterior is a Gibbs measure, and the prior is the penalty¶

2. Random-walk Metropolis¶

3. How long is burn-in? Coupled chains give an answer¶

4. Data augmentation: the Albert–Chib sampler for the probit¶

5. What is missing¶

Summing up¶

Exercises¶

Run notes¶