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

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

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

# The book's listings use statistics.NormalDist wrapped in np.vectorize. We use the
# vectorized scipy equivalents, which are identical in value and far faster.
Phi = stats.norm.cdf
Phi_inv = stats.norm.ppf
print('max |scipy - statistics| on a grid :',
      np.abs(Phi(np.linspace(-4, 4, 9))
             - np.array([__import__('statistics').NormalDist().cdf(v)
                         for v in np.linspace(-4, 4, 9)])).max())

max |scipy - statistics| on a grid : 1.1102230246251565e-16

def differencing_matrix(Y, y):
    '''The (Y-1) x Y matrix M with (M eps)_{y'} = eps_{y'} - eps_y, for y' =/= y.

    M[i, j] = 1{i = j} - 1{j = y}: the identity with a column of ones subtracted at y.
    Keep the minus sign on the *column* index; putting it on the row index leaves the
    identity behind once row y is dropped, and silently computes a different integral.
    '''
    M = np.eye(Y) - np.ones((Y, 1)) @ np.eye(Y)[y][None, :]
    return M[[i for i in range(Y) if i != y], :]


Y_demo = 4
U_demo = np.array([0.0, 0.3, -0.2, 0.5])
Sigma_eps = np.eye(Y_demo)

for y in range(Y_demo):
    M = differencing_matrix(Y_demo, y)
    Sigma_eta = M @ Sigma_eps @ M.T
    z = -M @ U_demo
    exact = stats.multivariate_normal(mean=np.zeros(Y_demo - 1),
                                      cov=Sigma_eta).cdf(z)
    print(f'y = {y}:  z = {np.round(z, 3)}   pi_y = F_eta(z) = {exact:.6f}')

exact_shares = np.array([
    stats.multivariate_normal(
        mean=np.zeros(Y_demo - 1),
        cov=differencing_matrix(Y_demo, y) @ Sigma_eps @ differencing_matrix(Y_demo, y).T
    ).cdf(-differencing_matrix(Y_demo, y) @ U_demo) for y in range(Y_demo)])
print(f'\nsum over y = {exact_shares.sum():.8f}   (must be 1; nothing imposes it)')

y = 0:  z = [-0.3  0.2 -0.5]   pi_y = F_eta(z) = 0.191742
y = 1:  z = [ 0.3  0.5 -0.2]   pi_y = F_eta(z) = 0.290460
y = 2:  z = [-0.2 -0.5 -0.7]   pi_y = F_eta(z) = 0.142285
y = 3:  z = [0.5 0.2 0.7]   pi_y = F_eta(z) = 0.375519

sum over y = 0.99999046   (must be 1; nothing imposes it)

def rosenblatt_gaussian(x_i_j, L):
    '''T(x) = L Phi^{-1}(x): the Rosenblatt quantile of N(0, L L^T).'''
    return Phi_inv(x_i_j) @ L.T


# verify the lemma: pushing uniforms through T gives the target Gaussian
I_draw = 200_000 if not FAST_MODE else 100_000
J = 3
A = rng.normal(size=(J, J))
Sigma = A @ A.T + J * np.eye(J)                     # a well-conditioned covariance
L = np.linalg.cholesky(Sigma)

x_i_j = rng.uniform(size=(I_draw, J))
Z_i_j = rosenblatt_gaussian(x_i_j, L)

print('target covariance :\n', Sigma)
print('sample covariance :\n', np.cov(Z_i_j.T))
print(f'\nmax |sample - target| covariance = {np.abs(np.cov(Z_i_j.T) - Sigma).max():.4f}')
print('per-coordinate KS tests against the correct marginal normal:')
for j in range(J):
    ks = stats.kstest(Z_i_j[:, j], lambda v, s=np.sqrt(Sigma[j, j]): Phi(v / s))
    print(f'   coordinate {j}: KS = {ks.statistic:.5f}, p = {ks.pvalue:.3f}')
    assert ks.pvalue > 0.01

# triangularity: T_j depends only on x_1..x_j
x_alt = x_i_j.copy(); x_alt[:, -1] = rng.uniform(size=I_draw)   # change the last coordinate
Z_alt = rosenblatt_gaussian(x_alt, L)
print(f'\nchanging x_J leaves coordinates 1..J-1 untouched : '
      f'{np.abs(Z_alt[:, :-1] - Z_i_j[:, :-1]).max():.2e}   -> lower triangular')

target covariance :
 [[ 5.4587  -1.86895  1.29656]
 [-1.86895  4.47941 -1.04176]
 [ 1.29656 -1.04176  4.48192]]
sample covariance :
 [[ 5.47531 -1.8778   1.30065]
 [-1.8778   4.4743  -1.01336]
 [ 1.30065 -1.01336  4.4798 ]]

max |sample - target| covariance = 0.0284
per-coordinate KS tests against the correct marginal normal:
   coordinate 0: KS = 0.00219, p = 0.725
   coordinate 1: KS = 0.00203, p = 0.802
   coordinate 2: KS = 0.00217, p = 0.734

changing x_J leaves coordinates 1..J-1 untouched : 0.00e+00   -> lower triangular

# estimating a Gaussian tail probability, crude versus importance sampling
threshold = 3.5
alpha_true = 1 - Phi(threshold)
R, I_dr = 400, 4000
gen = np.random.default_rng(505)


def crude_tail(v):
    '''Draw from p = N(0,1) and count hits.'''
    return (v > threshold).mean()


def tilted_tail(v):
    '''Draw from q = N(t,1) and reweight by p/q = exp(-t v + t^2/2), with t = threshold.

    The proposal is centred on the event, so roughly half the draws land inside it.
    '''
    t = threshold
    x = v + t                                        # x ~ N(t, 1)
    return ((x > threshold) * np.exp(-t * x + t ** 2 / 2)).mean()


draws = gen.normal(size=(R, I_dr))
crude = np.array([crude_tail(d) for d in draws])
imp = np.array([tilted_tail(d) for d in gen.normal(size=(R, I_dr))])

print(f'true tail probability            : {alpha_true:.6e}')
print(f'crude Monte Carlo    mean / s.d. : {crude.mean():.6e} / {crude.std(ddof=1):.3e}')
print(f'importance sampling  mean / s.d. : {imp.mean():.6e} / {imp.std(ddof=1):.3e}')
print(f'\nvariance ratio (crude / importance) : '
      f'{(crude.std(ddof=1) / imp.std(ddof=1))**2:.1f}')
print(f'replications where crude returned exactly 0 : {(crude == 0).mean():.1%}')
assert abs(imp.mean() / alpha_true - 1) < 0.05

true tail probability            : 2.326291e-04

crude Monte Carlo    mean / s.d. : 2.243750e-04 / 2.455e-04
importance sampling  mean / s.d. : 2.325944e-04 / 7.655e-06

variance ratio (crude / importance) : 1028.6
replications where crude returned exactly 0 : 41.8%

def iar(L, z_j, x_i_j):
    '''Listing 4-02: integrated accept-reject simulator of Pr(eta <= z), eta ~ N(0, L L^T).

    Accepts on the first J-1 coordinates and integrates the last one analytically.
    '''
    cond_i_j = (Phi_inv(x_i_j)[:, None, :] * L[None, :, :]).sum(axis=2) <= z_j[None, :]
    argphi_i = (z_j[-1] - (L[-1, :-1][None, :] * Phi_inv(x_i_j[:, :-1])).sum(axis=1)) / L[-1, -1]
    return float(((cond_i_j[:, :-1]).all(axis=1) * Phi(argphi_i)).mean())


def ghk(L, z_j, x_i_j):
    '''Listing 4-02: the Geweke-Hajivassiliou-Keane simulator of Pr(eta <= z).

    Walks the Rosenblatt map one coordinate at a time, conditioning rather than
    rejecting: each step rescales the uniform onto the admissible interval and
    carries its length into the product.
    '''
    I_, J_ = x_i_j.shape
    xhat_i_j = np.zeros((I_, J_))
    xhat_i_j[:, 0] = x_i_j[:, 0] * Phi(z_j[0] / L[0, 0])
    GHK_i = Phi(z_j[0] / L[0, 0]) * np.ones(I_)
    for j in range(J_ - 1):
        arg_i = (z_j[j + 1]
                 - (L[j + 1, :(j + 1)][None, :] * Phi_inv(xhat_i_j[:, :(j + 1)])).sum(axis=1)
                 ) / L[j + 1, j + 1]
        phietc_i = Phi(arg_i)
        xhat_i_j[:, j + 1] = x_i_j[:, j + 1] * phietc_i
        GHK_i *= phietc_i
    return float(GHK_i.mean())


# listing 4-02's own experiment, with the book's legacy seed
st = np.random.RandomState(9)
I_b, B_b, J_b = (100_000 if not FAST_MODE else 50_000), 10, 2
L_b = np.array([[st.uniform() * (i <= j) for i in range(J_b)] for j in range(J_b)])
z_b = st.normal(size=J_b) / 100
x_b_i_j = st.uniform(size=(B_b, I_b, J_b))

iarSim_b = np.array([iar(L_b, z_b, x_b_i_j[b]) for b in range(B_b)])
ghkSim_b = np.array([ghk(L_b, z_b, x_b_i_j[b]) for b in range(B_b)])
exact_b = stats.multivariate_normal(mean=np.zeros(J_b), cov=L_b @ L_b.T).cdf(z_b)

print(f'exact (scipy quadrature) : {exact_b:.8f}')
print(f'IAR, mean = {iarSim_b.mean():.8f} ; std = {iarSim_b.std(ddof=1):.3e}')
print(f'GHK, mean = {ghkSim_b.mean():.8f} ; std = {ghkSim_b.std(ddof=1):.3e}')
print(f'\nboth unbiased?  |IAR - exact| = {abs(iarSim_b.mean()-exact_b):.2e}, '
      f'|GHK - exact| = {abs(ghkSim_b.mean()-exact_b):.2e}')
print(f'variance ratio  V(IAR)/V(GHK) = '
      f'{(iarSim_b.std(ddof=1)/ghkSim_b.std(ddof=1))**2:.1f}')

exact (scipy quadrature) : 0.35343745
IAR, mean = 0.35283524 ; std = 2.205e-03
GHK, mean = 0.35342694 ; std = 4.302e-05

both unbiased?  |IAR - exact| = 6.02e-04, |GHK - exact| = 1.05e-05
variance ratio  V(IAR)/V(GHK) = 2626.5

def variance_experiment(shift, I_, R, J_, gen):
    '''Compare IAR and GHK at a target probability controlled by `shift`.'''
    A_ = gen.normal(size=(J_, J_))
    Sig = A_ @ A_.T + J_ * np.eye(J_)
    Lc = np.linalg.cholesky(Sig)
    z = -shift * np.sqrt(np.diag(Sig))
    exact = stats.multivariate_normal(mean=np.zeros(J_), cov=Sig).cdf(z)
    ia, gh = [], []
    for _ in range(R):
        xs = gen.uniform(size=(I_, J_))
        ia.append(iar(Lc, z, xs)); gh.append(ghk(Lc, z, xs))
    ia, gh = np.array(ia), np.array(gh)
    return exact, ia, gh


gen = np.random.default_rng(55)
I_v = 20_000 if not FAST_MODE else 10_000
R_v = 24 if not FAST_MODE else 16

# We score by root mean squared error against the exact value rather than by the
# sample standard deviation. For a rare enough event the accept-reject simulator
# returns exactly zero in every replication: its sample s.d. is then zero, which
# would flatter it enormously, while its RMSE is the whole probability.
print(f"{'shift':>6} {'exact P':>11} {'IAR mean':>11} {'GHK mean':>11} "
      f"{'RMSE(IAR)':>11} {'RMSE(GHK)':>11} {'MSE ratio':>11} {'IAR = 0':>8}")
shifts, ratios = [], []
for shift in (0.0, 0.5, 1.0, 1.5, 2.0, 2.5):
    exact, ia, gh = variance_experiment(shift, I_v, R_v, 3, gen)
    rmse_ia = np.sqrt(np.mean((ia - exact) ** 2))
    rmse_gh = np.sqrt(np.mean((gh - exact) ** 2))
    r = (rmse_ia / rmse_gh) ** 2
    shifts.append(exact); ratios.append(r)
    print(f'{shift:>6} {exact:>11.3e} {ia.mean():>11.3e} {gh.mean():>11.3e} '
          f'{rmse_ia:>11.2e} {rmse_gh:>11.2e} {r:>11.1f} {(ia == 0).sum():>5}/{R_v}')
print('\nthe rarer the event, the larger the GHK advantage')
assert min(ratios) > 1.0

 shift     exact P    IAR mean    GHK mean   RMSE(IAR)   RMSE(GHK)   MSE ratio  IAR = 0
   0.0   2.529e-01   2.530e-01   2.528e-01    5.07e-03    9.76e-04        26.9     0/16

   0.5   4.682e-03   4.624e-03   4.690e-03    2.61e-04    2.56e-05       103.8     0/16
   1.0   6.803e-04   6.772e-04   6.800e-04    7.32e-05    2.12e-06      1187.5     0/16
   1.5   8.495e-05   8.372e-05   8.497e-05    1.22e-05    3.21e-07      1441.9     0/16

   2.0   2.622e-06   2.477e-06   2.618e-06    6.36e-07    9.68e-09      4323.2     0/16
   2.5   8.083e-07   0.000e+00   8.083e-07    8.08e-07    4.96e-10   2653862.7    16/16

the rarer the event, the larger the GHK advantage

def piprobit_y(U_y, Sigma_eps, x_i_j):
    '''Listing 4-03: the probit market share map, computed by GHK.'''
    Y_ = U_y.shape[0]
    pi_y = np.zeros(Y_)
    for y in range(Y_):
        M = differencing_matrix(Y_, y)
        L_ = np.linalg.cholesky(M @ Sigma_eps @ M.T)
        pi_y[y] = ghk(L_, -M @ U_y, x_i_j)
    return pi_y


st = np.random.RandomState(9)
I_p, Y_p = (100_000 if not FAST_MODE else 50_000), 5
U_p = np.array([1.3, 0.9, -0.7, 0.2, 0.6])
x_p = st.uniform(size=(I_p, Y_p - 1))

pi_p = piprobit_y(U_p, np.eye(Y_p), x_p)
print('listing 4-03, independent probit, Y = 5')
print('pi_y      =', pi_p)
print(f'sum(pi_y) = {pi_p.sum():.8f}')

# check 1: against quadrature, option by option
exact_p = np.array([
    stats.multivariate_normal(
        mean=np.zeros(Y_p - 1),
        cov=differencing_matrix(Y_p, y) @ np.eye(Y_p) @ differencing_matrix(Y_p, y).T
    ).cdf(-differencing_matrix(Y_p, y) @ U_p) for y in range(Y_p)])
print('\nquadrature :', exact_p)
print(f'max |GHK - quadrature| = {np.abs(pi_p - exact_p).max():.2e}')

# check 2: against a brute-force frequency simulator
eps = np.random.default_rng(551).normal(size=(400_000 if not FAST_MODE else 200_000, Y_p))
u = U_p[None, :] + eps
freq = (u == u.max(axis=1)[:, None]).mean(axis=0)
se = np.sqrt(freq * (1 - freq) / eps.shape[0])
print('\nbrute force :', freq)
print('discrepancy in brute-force standard errors :', np.abs(pi_p - freq) / se)
assert np.abs(pi_p - exact_p).max() < 5e-3, 'GHK disagrees with quadrature'
assert abs(pi_p.sum() - 1) < 5e-3, 'GHK shares do not sum to one'
assert (np.abs(pi_p - freq) / se).max() < 4, 'GHK disagrees with the brute-force count'

listing 4-03, independent probit, Y = 5
pi_y      = [0.44552 0.26824 0.01708 0.09402 0.17574]
sum(pi_y) = 1.00060431

quadrature : [0.44516 0.26805 0.01707 0.094   0.17571]
max |GHK - quadrature| = 3.64e-04

brute force : [0.44541 0.26801 0.0167  0.0937  0.17619]
discrepancy in brute-force standard errors : [0.1032  0.23571 1.3288  0.49445 0.52436]

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


travelmode = read_mec_csv('demand_travelmode/travelmodedata.csv')
travelmode['choice'] = (travelmode['choice'] == 'yes').astype(int)
options = ['car', 'air', 'train', 'bus']            # car = default, as in lecture 1
pihat_y = (travelmode.pivot(index='individual', columns='mode', values='choice')
                     .reindex(columns=options).to_numpy().mean(axis=0))
print('observed shares :', dict(zip(options, pihat_y.round(4))))

I_inv = 60_000 if not FAST_MODE else 30_000
x_inv = np.random.default_rng(56).uniform(size=(I_inv, len(options) - 1))
Sigma_id = np.eye(len(options))


def residual(U_free):
    U = np.concatenate(([0.0], U_free))             # U_car = 0
    return piprobit_y(U, Sigma_id, x_inv)[1:] - pihat_y[1:]


sol = root(residual, np.zeros(len(options) - 1), method='hybr', tol=1e-10)
U_probit = np.concatenate(([0.0], sol.x))
pi_fitted = piprobit_y(U_probit, Sigma_id, x_inv)

print(f'\nroot finder converged: {sol.success}, {sol.nfev} function evaluations')
print(f'residual at the solution : {np.abs(residual(sol.x)).max():.2e}')
print(f'fitted shares            : {pi_fitted.round(6)}   sum = {pi_fitted.sum():.6f}')

# an independent inversion, with the share map computed by quadrature instead of GHK
def piprobit_quad(U_y, Sigma_eps):
    Y_ = U_y.shape[0]
    out = np.zeros(Y_)
    for y in range(Y_):
        M = differencing_matrix(Y_, y)
        out[y] = stats.multivariate_normal(mean=np.zeros(Y_ - 1),
                                           cov=M @ Sigma_eps @ M.T).cdf(-M @ U_y)
    return out


# scipy's multivariate_normal.cdf is a randomized quadrature, so it cannot drive a
# finite-difference root finder. Instead, evaluate it AT the GHK solution: if the GHK
# inversion is right, an independent share map must reproduce the observed shares there.
pi_quad_at_sol = piprobit_quad(U_probit, Sigma_id)
print('\nquadrature share map at the GHK solution :', pi_quad_at_sol.round(6))
print('observed shares                          :', pihat_y.round(6))
print(f'max gap = {np.abs(pi_quad_at_sol - pihat_y).max():.2e}   (tolerance 2e-3)')

U_logit = np.log(pihat_y) - np.log(pihat_y[0])
print(f"\n{'option':>8} {'share':>9} {'U logit':>10} {'U probit (GHK)':>16} "
      f"{'ratio':>8}")
for y, nm in enumerate(options):
    r = f'{U_probit[y] / U_logit[y]:.3f}' if abs(U_logit[y]) > 1e-8 else '--'
    print(f'{nm:>8} {pihat_y[y]:>9.4f} {U_logit[y]:>10.4f} {U_probit[y]:>16.4f} '
          f'{r:>8}')

# what a pure scale argument predicts: sd of a logistic difference versus a Gaussian one
sd_logit_diff = np.sqrt(2 * np.pi ** 2 / 6)      # eps_y - eps_z is logistic
sd_probit_diff = np.sqrt(2.0)                    # eps_y - eps_z is N(0, 2)
print(f'\nscale factor predicted by variance matching : '
      f'{sd_probit_diff / sd_logit_diff:.4f}')
assert np.abs(residual(sol.x)).max() < 1e-6
assert np.abs(pi_quad_at_sol - pihat_y).max() < 2e-3

observed shares : {'car': np.float64(0.281), 'air': np.float64(0.2762), 'train': np.float64(0.3), 'bus': np.float64(0.1429)}

root finder converged: True, 16 function evaluations
residual at the solution : 5.55e-17
fitted shares            : [0.28163 0.27619 0.3     0.14286]   sum = 1.000682

quadrature share map at the GHK solution : [0.28146 0.27602 0.29981 0.1427 ]
observed shares                          : [0.28095 0.27619 0.3     0.14286]
max gap = 5.03e-04   (tolerance 2e-3)

  option     share    U logit   U probit (GHK)    ratio
     car    0.2810     0.0000           0.0000       --
     air    0.2762    -0.0171          -0.0146    0.854
   train    0.3000     0.0656           0.0477    0.727
     bus    0.1429    -0.6763          -0.4736    0.700

scale factor predicted by variance matching : 0.7797

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

ax[0].loglog(shifts, ratios, 'o-', color='C0', lw=1.8)
ax[0].axhline(1, color='k', lw=.8, ls=':')
ax[0].set_xlabel('target probability'); ax[0].set_ylabel('MSE(IAR) / MSE(GHK)')
ax[0].set_title('GHK gains most where the event is rare')
ax[0].invert_xaxis(); ax[0].grid(alpha=.25, which='both')

w = 0.35
idx = np.arange(len(options))
ax[1].bar(idx - w / 2, U_logit, w, color='C0', label='logit (closed form)')
ax[1].bar(idx + w / 2, U_probit, w, color='C3', label='probit (GHK inversion)')
ax[1].axhline(0, color='k', lw=.8)
ax[1].set_xticks(idx); ax[1].set_xticklabels(options)
ax[1].set_ylabel(r'$U_y$ (car normalized to 0)')
ax[1].set_title('Same ranking, different scale')
ax[1].legend(fontsize=8.5); ax[1].grid(alpha=.25, axis='y')

plt.tight_layout(); plt.show()

Simulation, importance sampling, and GHK¶

Alfred Galichon (NYU) and Antoine Jacquet (Sciences Po)¶

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

With python code examples¶

Learning objectives¶

References¶

Motivation: what to do when there is no formula¶

2. Rosenblatt quantiles: the transport map¶

3. Importance sampling, and why rejection is wasteful¶

4. From accept–reject to GHK¶

6. Inverting the travel-mode shares under a probit¶

Summing up¶

Exercises¶

Run notes¶

Simulation, importance sampling, and GHK¶

Alfred Galichon (NYU) and Antoine Jacquet (Sciences Po)¶

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

With python code examples¶

Learning objectives¶

References¶

Motivation: what to do when there is no formula¶

1. A probit market share is a Gaussian orthant probability¶

2. Rosenblatt quantiles: the transport map¶

3. Importance sampling, and why rejection is wasteful¶

4. From accept–reject to GHK¶

5. The probit market share map¶

6. Inverting the travel-mode shares under a probit¶

Summing up¶

Exercises¶

Run notes¶