import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats
from types import SimpleNamespace
import linearmodels.iv as iv

rng = np.random.default_rng(777)
plt.rcParams['figure.figsize'] = (8, 4.5)

url = 'https://vincentarelbundock.github.io/Rdatasets/csv/sampleSelection/Mroz87.csv'
try:
    mroz = pd.read_csv(url)                       # Rdatasets Mroz87 (needs a network)
except Exception:
    # Offline fallback: the identical Mroz (1987) data ships with the `wooldridge` package.
    import wooldridge
    mroz = (wooldridge.data('mroz')
            .rename(columns={'inlf': 'lfp', 'kidslt6': 'kids5', 'kidsge6': 'kids618'})
            .copy())
mroz_w = mroz.loc[mroz['lfp'] == 1].copy()
mroz_w['lwage'] = np.log(mroz_w['wage'])
I = len(mroz_w)

y_i = mroz_w['lwage'].values
X_i_k = np.column_stack([
    np.ones(I),
    mroz_w['educ'].values,
    mroz_w['exper'].values,
    mroz_w['exper'].values ** 2,
])
Z_i_l = np.column_stack([
    np.ones(I),
    mroz_w['motheduc'].values,
    mroz_w['fatheduc'].values,
    mroz_w['exper'].values,
    mroz_w['exper'].values ** 2,
])
xnames = ['const', 'educ', 'exper', 'exper^2']
print(f'I = {I}')

I = 428

def bootstrap_estimator(estimator_fn, data, B=1_000, rng=None):
    r"""Generic nonparametric bootstrap: pass any function that takes a resampled data
    array and returns a scalar (or vector) statistic.

    All B draws are genuine resamples. The statistic on the original sample is returned
    alongside them, never *inside* them: putting $\hat\theta$ into the bootstrap array
    would contaminate every quantile computed from it.
    """
    if rng is None:
        rng = np.random.default_rng()
    I = data.shape[0] if hasattr(data, 'shape') else len(data)
    sample_stat = np.atleast_1d(estimator_fn(data))
    out = np.empty((B,) + sample_stat.shape)
    for b in range(B):
        idx = rng.integers(0, I, size=I)
        out[b] = estimator_fn(data[idx] if hasattr(data, '__getitem__') else [data[i] for i in idx])
    return SimpleNamespace(boot=out, stat=sample_stat, B=B)

# Quick sanity check: bootstrap the mean of standard-normal data
rng = np.random.default_rng(42)
data = rng.standard_normal(200)
fit_boot = bootstrap_estimator(np.mean, data, B=2000, rng=rng)
boot_means = fit_boot.boot
print(f'Sample mean                 = {data.mean():.4f}')
print(f'Bootstrap mean of means     = {boot_means.mean():.4f}')
print(f'Bootstrap SE of mean        = {boot_means.std(ddof=1):.4f}')
print(f'Theoretical SE: σ/√I        = {1/np.sqrt(200):.4f}')
print(f'\ngenuine resamples drawn   = {fit_boot.B}  (none of them the original sample)')

# Verification contract. The shape check is necessary but not sufficient: the bug this routine
# used to have produced an array of exactly the right length whose first entry was the
# original-sample statistic rather than a resample. Comparing *values* cannot detect that —
# an inserted statistic is only caught if it happens to differ numerically from every
# resample, and for a discrete statistic a resample may legitimately reproduce the sample
# value. So test at the level of what the routine actually passes to the estimator.
assert boot_means.shape[0] == 2000, f'expected 2000 resamples, got {boot_means.shape[0]}'

seen_calls = []
def spy_mean(a_i):
    """Records every array handed to it, then behaves like np.mean."""
    seen_calls.append(np.asarray(a_i).copy())
    return np.mean(a_i)

B_probe = 25
probe_i = np.random.default_rng(99).standard_normal(50)
fit_probe = bootstrap_estimator(spy_mean, probe_i, B=B_probe,
                                rng=np.random.default_rng(0))

# One evaluation on the original sample (for .stat), then exactly B on resamples.
assert len(seen_calls) == 1 + B_probe, f'estimator called {len(seen_calls)} times'
assert np.array_equal(seen_calls[0], probe_i), 'the first call is not the original sample'
assert fit_probe.boot.shape[0] == B_probe

# The decisive check: output row b must be the statistic of resample call b, so row 0 holds a
# RESAMPLE and not the original sample. Under the old implementation this failed at b = 0.
for b in range(B_probe):
    assert np.isclose(fit_probe.boot[b], np.mean(seen_calls[1 + b])), \
        f'output row {b} does not correspond to resample call {b}'
print(f'index-level check: {B_probe} resample calls, output row 0 is a resample '
      f'(not the sample statistic)')

# Cross-check the bootstrap standard error against an independent implementation.
sp_boot = stats.bootstrap((data,), np.mean, n_resamples=2000, method='percentile',
                          random_state=np.random.default_rng(4321))
se_ours, se_scipy = boot_means.std(ddof=1), sp_boot.standard_error
print(f'bootstrap SE, ours = {se_ours:.5f};  scipy.stats.bootstrap = {se_scipy:.5f};  '
      f'ratio = {se_ours/se_scipy:.4f}')
assert abs(se_ours / se_scipy - 1) < 0.10, 'bootstrap SE disagrees with scipy by >10%'
print('Verification passed: B genuine resamples, and the SE matches an independent bootstrap.')

Sample mean                 = -0.0305
Bootstrap mean of means     = -0.0317
Bootstrap SE of mean        = 0.0643
Theoretical SE: σ/√I        = 0.0707

genuine resamples drawn   = 2000  (none of them the original sample)
index-level check: 25 resample calls, output row 0 is a resample (not the sample statistic)
bootstrap SE, ours = 0.06434;  scipy.stats.bootstrap = 0.06159;  ratio = 1.0446
Verification passed: B genuine resamples, and the SE matches an independent bootstrap.

def bootstrap_cis(boot, θ_hat, α=0.05, se_hat=None, boot_se=None):
    # boot: array of B bootstrap estimates of θ
    # se_hat: SE estimate from the original sample (for studentized CI)
    # boot_se: array of B bootstrap SE estimates (for studentized CI)
    α_lo, α_hi = α/2, 1 - α/2
    q_lo = np.quantile(boot, α_lo)
    q_hi = np.quantile(boot, α_hi)
    out = {
        'percentile': (q_lo, q_hi),
        'basic':      (2*θ_hat - q_hi, 2*θ_hat - q_lo),
    }
    if se_hat is not None and boot_se is not None:
        t_boot = (boot - θ_hat) / boot_se
        t_lo = np.quantile(t_boot, α_lo)
        t_hi = np.quantile(t_boot, α_hi)
        out['studentized'] = (θ_hat - t_hi*se_hat, θ_hat - t_lo*se_hat)
    return out

def iv_2sls(X_i_k, y_i, Z_i_l):
    PZX_i_k = Z_i_l @ np.linalg.solve(Z_i_l.T @ Z_i_l, Z_i_l.T @ X_i_k)
    β_k    = np.linalg.solve(X_i_k.T @ PZX_i_k, PZX_i_k.T @ y_i)
    εhat_i = y_i - X_i_k @ β_k
    return β_k, εhat_i

def iv_2sls_with_se(X_i_k, y_i, Z_i_l):
    β_k, εhat_i = iv_2sls(X_i_k, y_i, Z_i_l)
    PZX_i_k = Z_i_l @ np.linalg.solve(Z_i_l.T @ Z_i_l, Z_i_l.T @ X_i_k)
    # Heteroskedasticity-consistent SE (HC0)
    XPZXinv_k_k = np.linalg.inv(X_i_k.T @ PZX_i_k)
    S_k_k = (PZX_i_k * (εhat_i**2)[:, None]).T @ PZX_i_k
    V_k_k = XPZXinv_k_k @ S_k_k @ XPZXinv_k_k
    return β_k, np.sqrt(np.diag(V_k_k))

β_2sls, seasy_k = iv_2sls_with_se(X_i_k, y_i, Z_i_l)
print(f'Sample 2SLS β̂_educ      = {β_2sls[1]:.4f}')
print(f'Asymptotic SE (HC0)     = {seasy_k[1]:.4f}')
print(f'Asymptotic 95% CI       = [{β_2sls[1] - 1.96*seasy_k[1]:.4f}, '
      f'{β_2sls[1] + 1.96*seasy_k[1]:.4f}]')

Sample 2SLS β̂_educ      = 0.0614
Asymptotic SE (HC0)     = 0.0332
Asymptotic 95% CI       = [-0.0036, 0.1264]

# Bootstrap the IV estimator (resample full rows of (X_i_k, Z_i_l, y_i) together)
rng = np.random.default_rng(2026)
B = 2_000
β_boot   = np.empty(B)
se_boot  = np.empty(B)
for b in range(B):
    idx = rng.integers(0, I, size=I)
    β_b, se_b = iv_2sls_with_se(X_i_k[idx], y_i[idx], Z_i_l[idx])
    β_boot[b]  = β_b[1]
    se_boot[b] = se_b[1]

cis = bootstrap_cis(β_boot, β_2sls[1], α=0.05, se_hat=seasy_k[1], boot_se=se_boot)

print(f"{'Method':<22s}{'CI lower':>10s}{'CI upper':>10s}{'width':>10s}")
ci_asy = (β_2sls[1] - 1.96*seasy_k[1], β_2sls[1] + 1.96*seasy_k[1])
print(f'{"Asymptotic Wald":<22s}{ci_asy[0]:>10.4f}{ci_asy[1]:>10.4f}{ci_asy[1]-ci_asy[0]:>10.4f}')
for name, ci in cis.items():
    print(f'{"Bootstrap " + name:<22s}{ci[0]:>10.4f}{ci[1]:>10.4f}{ci[1]-ci[0]:>10.4f}')

Method                  CI lower  CI upper     width
Asymptotic Wald          -0.0036    0.1264    0.1301
Bootstrap percentile     -0.0055    0.1264    0.1319
Bootstrap basic          -0.0036    0.1283    0.1319
Bootstrap studentized    -0.0109    0.1191    0.1300

# Visualize bootstrap distribution vs implied asymptotic normal
fig, ax = plt.subplots(figsize=(9, 4.5))
ax.hist(β_boot, bins=40, density=True, alpha=0.55, edgecolor='white',
        label=f'Bootstrap distribution (B = {B})')
xs = np.linspace(β_boot.min(), β_boot.max(), 400)
ax.plot(xs, stats.norm.pdf(xs, loc=β_2sls[1], scale=seasy_k[1]),
        color='red', lw=1.5, label=f'Asymptotic N({β_2sls[1]:.3f}, {seasy_k[1]:.3f}²)')
ax.axvline(β_2sls[1], color='k', ls=':', lw=0.8, label=r'$\hat\beta_{2SLS}$')
ax.set_xlabel(r'$\hat\beta_{educ}^*$')
ax.set_ylabel('density')
ax.set_title(r'Bootstrap distribution of 2SLS return to schooling')
ax.legend(loc='upper right')
plt.tight_layout()
plt.show()

def simulate_het(I=200, β_k=np.array([1.0, 0.5]), rng=None):
    if rng is None:
        rng = np.random.default_rng()
    X_i_k = np.column_stack([np.ones(I), rng.standard_normal(I)])
    σ_i = 0.5 * np.abs(X_i_k[:, 1])                     # strong heteroskedasticity
    ε_i = σ_i * rng.standard_normal(I)
    y_i = X_i_k @ β_k + ε_i
    return X_i_k, y_i

def ols_with_ses(X_i_k, y_i):
    β_k = np.linalg.solve(X_i_k.T @ X_i_k, X_i_k.T @ y_i)
    εhat_i = y_i - X_i_k @ β_k
    I, K = X_i_k.shape
    # Classical (homoskedastic) SE
    σ2 = (εhat_i @ εhat_i) / (I - K)
    Vhom_k_k = σ2 * np.linalg.inv(X_i_k.T @ X_i_k)
    # HC1 robust SE
    XtXinv_k_k = np.linalg.inv(X_i_k.T @ X_i_k)
    Ω_k_k = (X_i_k * (εhat_i ** 2)[:, None]).T @ X_i_k
    Vhc1_k_k = (I / (I - K)) * XtXinv_k_k @ Ω_k_k @ XtXinv_k_k
    return β_k, np.sqrt(np.diag(Vhom_k_k)), np.sqrt(np.diag(Vhc1_k_k)), εhat_i

def wild_bootstrap_ols(X_i_k, y_i, β_hat, εhat_i, B=499, rng=None):
    if rng is None:
        rng = np.random.default_rng()
    I = len(y_i)
    β_boot = np.empty((B, X_i_k.shape[1]))
    XtXinv_k_k = np.linalg.inv(X_i_k.T @ X_i_k)
    Xy_part = X_i_k.T  # so y_boot.T @ Xy_part... etc. just compute on the fly
    for b in range(B):
        w_i = rng.choice([-1.0, 1.0], size=I)            # Rademacher weights
        yboot_i = X_i_k @ β_hat + w_i * εhat_i
        β_boot[b] = XtXinv_k_k @ (X_i_k.T @ yboot_i)
    return β_boot

# Monte Carlo coverage check
rng = np.random.default_rng(2026)
M = 1_000
true_β = np.array([1.0, 0.5])
cov_hom = np.zeros(2)
cov_hc1 = np.zeros(2)
cov_wb  = np.zeros(2)
for m in range(M):
    X_i_k, y_i = simulate_het(I=200, β_k=true_β, rng=rng)
    β_k, sehom_k, sehc1_k, εhat_i = ols_with_ses(X_i_k, y_i)
    # Asymptotic CIs
    ci_hom = np.column_stack([β_k - 1.96 * sehom_k, β_k + 1.96 * sehom_k])
    ci_hc1 = np.column_stack([β_k - 1.96 * sehc1_k, β_k + 1.96 * sehc1_k])
    # Wild bootstrap CI (percentile)
    β_boot = wild_bootstrap_ols(X_i_k, y_i, β_k, εhat_i, B=399, rng=rng)
    q_lo, q_hi = np.quantile(β_boot, [0.025, 0.975], axis=0)
    cov_hom += (ci_hom[:, 0] <= true_β) & (true_β <= ci_hom[:, 1])
    cov_hc1 += (ci_hc1[:, 0] <= true_β) & (true_β <= ci_hc1[:, 1])
    cov_wb  += (q_lo        <= true_β) & (true_β <= q_hi)

cov_hom /= M; cov_hc1 /= M; cov_wb /= M

print(f'Heteroskedastic DGP, I = 200, target coverage = 95%:')
print(f"{'parameter':<10s}{'OLS (classical)':>18s}{'HC1':>10s}{'Wild boot (pct)':>18s}")
for j, name in enumerate(['β_0 (const)', 'β_1 (slope)']):
    print(f'{name:<10s}{cov_hom[j]:>18.3f}{cov_hc1[j]:>10.3f}{cov_wb[j]:>18.3f}')

Heteroskedastic DGP, I = 200, target coverage = 95%:
parameter    OLS (classical)       HC1   Wild boot (pct)
β_0 (const)             0.949     0.945             0.940
β_1 (slope)             0.726     0.942             0.929

def simulate_clustered(G=50, I_g=30, β_k=1.0, ρ=0.5, rng=None):
    if rng is None:
        rng = np.random.default_rng()
    I = G * I_g
    cluster_id = np.repeat(np.arange(G), I_g)
    # Cluster-level common shock (cluster-constant regressor)
    x_g = rng.standard_normal(G)
    X_i_k = np.column_stack([np.ones(I), np.repeat(x_g, I_g)])
    # Within-cluster equicorrelated errors: ε_i = √ρ · u_g + √(1-ρ) · v_i
    u_g = rng.standard_normal(G)
    v_i   = rng.standard_normal(I)
    ε_i   = np.sqrt(ρ) * np.repeat(u_g, I_g) + np.sqrt(1 - ρ) * v_i
    y_i   = X_i_k @ np.array([0.0, β_k]) + ε_i
    return X_i_k, y_i, cluster_id

def ols_clustered_ses(X_i_k, y_i, cluster_id):
    β_k = np.linalg.solve(X_i_k.T @ X_i_k, X_i_k.T @ y_i)
    εhat_i = y_i - X_i_k @ β_k
    I, K = X_i_k.shape
    XtXinv_k_k = np.linalg.inv(X_i_k.T @ X_i_k)
    # Classical (homoskedastic) SE
    σ2 = (εhat_i @ εhat_i) / (I - K)
    sehom_k = np.sqrt(np.diag(σ2 * XtXinv_k_k))
    # CR1 cluster-robust SE
    G = len(np.unique(cluster_id))
    S_k_k = np.zeros((K, K))
    for g in np.unique(cluster_id):
        idx = (cluster_id == g)
        Xg_i_k, εg_i = X_i_k[idx], εhat_i[idx]
        Xε_g = Xg_i_k.T @ εg_i
        S_k_k += np.outer(Xε_g, Xε_g)
    cr1_corr = (I - 1) / (I - K) * G / (G - 1)
    Vcr1_k_k = cr1_corr * XtXinv_k_k @ S_k_k @ XtXinv_k_k
    secr1_k = np.sqrt(np.diag(Vcr1_k_k))
    return β_k, sehom_k, secr1_k

def cluster_bootstrap(X_i_k, y_i, cluster_id, B=399, rng=None):
    if rng is None:
        rng = np.random.default_rng()
    clusters = np.unique(cluster_id)
    G = len(clusters)
    cluster_idx = {g: np.where(cluster_id == g)[0] for g in clusters}
    β_boot = np.empty((B, X_i_k.shape[1]))
    for b in range(B):
        sampled = rng.choice(clusters, size=G, replace=True)
        rows    = np.concatenate([cluster_idx[g] for g in sampled])
        β_boot[b] = np.linalg.solve(X_i_k[rows].T @ X_i_k[rows], X_i_k[rows].T @ y_i[rows])
    return β_boot

rng = np.random.default_rng(2026)
M = 500
true_β1 = 1.0
cov_hom = 0; cov_cr1 = 0; cov_cb = 0
for m in range(M):
    X_i_k, y_i, cid = simulate_clustered(G=50, I_g=30, β_k=true_β1, ρ=0.5, rng=rng)
    β_k, sehom_k, secr1_k = ols_clustered_ses(X_i_k, y_i, cid)
    ci_hom = (β_k[1] - 1.96*sehom_k[1], β_k[1] + 1.96*sehom_k[1])
    ci_cr1 = (β_k[1] - 1.96*secr1_k[1], β_k[1] + 1.96*secr1_k[1])
    β_b = cluster_bootstrap(X_i_k, y_i, cid, B=299, rng=rng)
    q_lo, q_hi = np.quantile(β_b[:, 1], [0.025, 0.975])
    cov_hom += (ci_hom[0] <= true_β1 <= ci_hom[1])
    cov_cr1 += (ci_cr1[0] <= true_β1 <= ci_cr1[1])
    cov_cb  += (q_lo      <= true_β1 <= q_hi)

print(f'Clustered DGP, G = 50, I_g = 30, ρ = 0.5, target coverage = 95%:')
print(f"{'method':<26s}{'coverage':>12s}")
print(f'{"OLS (homoskedastic SE)":<26s}{cov_hom/M:>12.3f}')
print(f'{"CR1 cluster-robust":<26s}{cov_cr1/M:>12.3f}')
print(f'{"Cluster bootstrap (pct)":<26s}{cov_cb /M:>12.3f}')

Clustered DGP, G = 50, I_g = 30, ρ = 0.5, target coverage = 95%:
method                        coverage
OLS (homoskedastic SE)           0.396
CR1 cluster-robust               0.930
Cluster bootstrap (pct)          0.926

# Simulate from Uniform[0, 1], compare bootstrap to true sampling distribution
rng = np.random.default_rng(2026)
I_obs = 100
M = 5_000

# True sampling distribution: I(1 - max(X_i)) ~ Exp(1) asymptotically
true_scaled = np.empty(M)
for m in range(M):
    X_i = rng.uniform(0, 1, size=I_obs)
    true_scaled[m] = I_obs * (1 - X_i.max())

# Bootstrap distribution from one sample
X_i = rng.uniform(0, 1, size=I_obs)
θ_hat = X_i.max()
boot_scaled = np.empty(M)
for m in range(M):
    Xb_i = X_i[rng.integers(0, I_obs, size=I_obs)]
    boot_scaled[m] = I_obs * (θ_hat - Xb_i.max())

p_zero = (boot_scaled == 0).mean()

fig, axes = plt.subplots(1, 2, figsize=(13, 4.2))
axes[0].hist(true_scaled,   bins=30, density=True, alpha=0.55, edgecolor='white')
xs = np.linspace(0, max(true_scaled.max(), 8), 200)
axes[0].plot(xs, stats.expon.pdf(xs), color='red', lw=1.5, label='Exp(1) limit')
axes[0].set_title(f'True sampling dist. of $I(1 - \\hat\\theta)$\n(over {M} fresh samples)')
axes[0].set_xlabel(r'$I(1 - \hat\theta_I)$'); axes[0].legend()

axes[1].hist(boot_scaled, bins=30, density=True, alpha=0.55, edgecolor='white')
axes[1].axvline(0, color='red', lw=1.5,
                label=f'P(boot = 0) = {p_zero:.3f}\n(≈ 1 - e⁻¹ = 0.632)')
axes[1].set_title(f'Bootstrap dist. of $I(\\hat\\theta - \\hat\\theta^*)$\n(from one sample, {M} resamples)')
axes[1].set_xlabel(r'$I(\hat\theta - \hat\theta_I^*)$'); axes[1].legend()
plt.tight_layout(); plt.show()

	equivariant under monotone $\phi$	needs an SE per resample	coverage error
percentile	yes	no	$O(I^{-1/2})$
basic (Hall)	no	no	$O(I^{-1/2})$
studentized	no	yes	$O(I^{-1})$
BCa	yes	no	$O(I^{-1})$

Bootstrap and simulation-based inference¶

Alfred Galichon (NYU)¶

'math+econ+code' masterclass series: econometrics¶

With python code examples¶

Learning objectives¶

References¶

Motivation: when asymptotic theory isn't enough¶

Loading our libraries and data¶

1. The nonparametric bootstrap¶

1.1 The bootstrap principle¶

1.2 The plug-in interpretation¶

1.3 DIY implementation¶

2. Bootstrap confidence intervals¶

2.1 Percentile interval¶

2.2 Basic (Hall) interval¶

2.3 Studentized (bootstrap-$t$) interval¶

2.4 Bias-corrected accelerated (BCa)¶

3. Bootstrap CI for the Mroz IV return to schooling¶

4. The wild bootstrap¶

4.1 Why a different bootstrap?¶

4.2 Application: heteroskedastic regression coverage¶

5. The cluster bootstrap¶

5.1 Algorithm¶

5.2 Application¶

6. When the bootstrap fails: nonregular estimators¶

6.1 Subsampling: a partial fix¶

7. The structure of inference so far¶

Looking ahead¶