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'])

# We need motheduc and fatheduc for the instruments
needed = ['lwage', 'educ', 'exper', 'motheduc', 'fatheduc']
print(f'I = {len(mroz_w)} working women')
print(f'Missing values:\n{mroz_w[needed].isna().sum()}')

I = 428 working women
Missing values:
lwage       0
educ        0
exper       0
motheduc    0
fatheduc    0
dtype: int64

def simulate_iv(I=2_000, β=0.5, γ=1.0, π=0.5, ρ_xw=0.5, rng=None):
    if rng is None:
        rng = np.random.default_rng()
    z_i = rng.standard_normal(I)
    w_i = rng.standard_normal(I)                     # the omitted variable
    # x is generated with a contribution from z (instrument relevance)
    # and from w (the source of endogeneity, since w will go into ε)
    v_i = rng.standard_normal(I)                     # idiosyncratic shock to x
    x_i = π * z_i + ρ_xw * w_i + np.sqrt(1 - ρ_xw**2) * v_i
    u_i = rng.standard_normal(I)
    ε_i = γ * w_i + u_i                              # w enters ε via γ
    y_i = β * x_i + ε_i
    return x_i, y_i, z_i

rng_mc = np.random.default_rng(2026)
β_true = 0.5
R = 500
βols_r = np.empty(R)
βiv_r  = np.empty(R)

for r in range(R):
    x_i, y_i, z_i = simulate_iv(I=2_000, β=β_true, rng=rng_mc)
    βols_r[r] = (x_i @ y_i) / (x_i @ x_i)
    βiv_r[r]  = (z_i @ y_i) / (z_i @ x_i)

print(f'Truth:                  β = {β_true:.3f}')
print(f'OLS  mean(β̂) = {βols_r.mean():.4f}    std(β̂) = {βols_r.std():.4f}')
print(f'IV   mean(β̂) = {βiv_r.mean():.4f}    std(β̂) = {βiv_r.std():.4f}')

Truth:                  β = 0.500
OLS  mean(β̂) = 0.8998    std(β̂) = 0.0267
IV   mean(β̂) = 0.5014    std(β̂) = 0.0655

def tsls(X_i_k, y_i, Z_i_l):
    # Two-stage least squares with classical and HC0 standard errors.
    #   X_i_k : (I, K) regressors (endogenous + exogenous)
    #   y_i   : (I,)   outcome
    #   Z_i_l : (I, L) instruments (must include the exogenous part of X)
    I, K = X_i_k.shape
    ZtZinv_l_l = np.linalg.inv(Z_i_l.T @ Z_i_l)
    PZX_i_k    = Z_i_l @ (ZtZinv_l_l @ (Z_i_l.T @ X_i_k))       # = P_Z X
    XPZX_k_k   = X_i_k.T @ PZX_i_k
    XPZy_k     = X_i_k.T @ (Z_i_l @ (ZtZinv_l_l @ (Z_i_l.T @ y_i)))
    β_k        = np.linalg.solve(XPZX_k_k, XPZy_k)
    εhat_i     = y_i - X_i_k @ β_k
    XPZXinv_k_k = np.linalg.inv(XPZX_k_k)

    # Classical (homoskedastic) variance
    σ2 = (εhat_i @ εhat_i) / (I - K)
    Vcl_k_k = σ2 * XPZXinv_k_k

    # HC0 robust variance (sandwich form)
    Ω_k_k = (PZX_i_k * (εhat_i ** 2)[:, None]).T @ PZX_i_k
    Vrob_k_k = XPZXinv_k_k @ Ω_k_k @ XPZXinv_k_k

    return SimpleNamespace(
        β_k=β_k, εhat_i=εhat_i,
        Vcl_k_k=Vcl_k_k, secl_k=np.sqrt(np.diag(Vcl_k_k)),
        Vrob_k_k=Vrob_k_k, serob_k=np.sqrt(np.diag(Vrob_k_k)),
        I=I, K=K, L=Z_i_l.shape[1],
    )

# Just-identified IV reduces to 2SLS when L = K
def iv_just(X_i_k, y_i, Z_i_l):
    I, K = X_i_k.shape
    β_k = np.linalg.solve(Z_i_l.T @ X_i_k, Z_i_l.T @ y_i)
    εhat_i = y_i - X_i_k @ β_k
    σ2 = (εhat_i @ εhat_i) / (I - K)
    ZtXinv_k_l = np.linalg.inv(Z_i_l.T @ X_i_k)
    Vcl_k_k = σ2 * ZtXinv_k_l @ (Z_i_l.T @ Z_i_l) @ ZtXinv_k_l.T

    # HC0 sandwich, the just-identified analogue of the 2SLS one below:
    #   V = (Z'X)^{-1} [ Σ_i ε̂_i² z_i z_i' ] (X'Z)^{-1}
    Ω_l_l = (Z_i_l * (εhat_i ** 2)[:, None]).T @ Z_i_l
    Vrob_k_k = ZtXinv_k_l @ Ω_l_l @ ZtXinv_k_l.T

    return SimpleNamespace(β_k=β_k, εhat_i=εhat_i,
                           Vcl_k_k=Vcl_k_k, secl_k=np.sqrt(np.diag(Vcl_k_k)),
                           Vrob_k_k=Vrob_k_k, serob_k=np.sqrt(np.diag(Vrob_k_k)),
                           I=I, K=K)

y_i = mroz_w['lwage'].values
I = len(y_i)

# Endogenous regressor + exogenous controls
X_i_k = np.column_stack([
    np.ones(I),
    mroz_w['educ'].values,                       # ENDOGENOUS
    mroz_w['exper'].values,
    mroz_w['exper'].values ** 2,
])
xnames = ['const', 'educ', 'exper', 'exper^2']

# Just-identified instruments: replace educ by motheduc
Zjust_i_l = np.column_stack([
    np.ones(I),
    mroz_w['motheduc'].values,
    mroz_w['exper'].values,
    mroz_w['exper'].values ** 2,
])

# Over-identified instruments: replace educ by motheduc + fatheduc
Zover_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,
])
print(f'X_i_k has shape {X_i_k.shape}; Zjust_i_l has shape {Zjust_i_l.shape}; Zover_i_l has shape {Zover_i_l.shape}')

X_i_k has shape (428, 4); Zjust_i_l has shape (428, 4); Zover_i_l has shape (428, 5)

# OLS
βols_k = np.linalg.solve(X_i_k.T @ X_i_k, X_i_k.T @ y_i)
εhatols_i = y_i - X_i_k @ βols_k
σ2_ols = (εhatols_i @ εhatols_i) / (I - X_i_k.shape[1])
seols_k = np.sqrt(np.diag(σ2_ols * np.linalg.inv(X_i_k.T @ X_i_k)))

# Just-identified IV (Zjust_i_l is square at L = K = 4)
fit_iv = iv_just(X_i_k, y_i, Zjust_i_l)

# Over-identified 2SLS
fit_2sls = tsls(X_i_k, y_i, Zover_i_l)

# Robust standard errors are the ones to lead with: nothing in §1-§3 assumed homoskedasticity,
# and §5.1 will show the Mroz residuals are not obviously homoskedastic either. The classical
# ones are kept alongside so the two can be compared rather than silently swapped.
seolsrob_k = np.sqrt(np.diag(
    np.linalg.inv(X_i_k.T @ X_i_k)
    @ ((X_i_k * (εhatols_i ** 2)[:, None]).T @ X_i_k)
    @ np.linalg.inv(X_i_k.T @ X_i_k)))

print('standard errors: HC0 robust first, classical homoskedastic in brackets\n')
print(f"{'name':<10s}{'OLS':>26s}{'IV (motheduc)':>26s}{'2SLS (mother+father)':>28s}")
for k, name in enumerate(xnames):
    o  = f'{βols_k[k]:>8.4f} {seolsrob_k[k]:.4f} [{seols_k[k]:.4f}]'
    i_ = f'{fit_iv.β_k[k]:>8.4f} {fit_iv.serob_k[k]:.4f} [{fit_iv.secl_k[k]:.4f}]'
    s  = f'{fit_2sls.β_k[k]:>8.4f} {fit_2sls.serob_k[k]:.4f} [{fit_2sls.secl_k[k]:.4f}]'
    print(f'{name:<10s}{o:>26s}{i_:>26s}{s:>28s}')

print(f'\nratio robust/classical on educ:  OLS {seolsrob_k[1]/seols_k[1]:.3f}   '
      f'IV {fit_iv.serob_k[1]/fit_iv.secl_k[1]:.3f}   '
      f'2SLS {fit_2sls.serob_k[1]/fit_2sls.secl_k[1]:.3f}')

standard errors: HC0 robust first, classical homoskedastic in brackets

name                             OLS             IV (motheduc)        2SLS (mother+father)
const        -0.5220 0.2007 [0.1986]    0.1982 0.4869 [0.4729]      0.0481 0.4278 [0.4003]
educ          0.1075 0.0132 [0.0141]    0.0493 0.0379 [0.0374]      0.0614 0.0332 [0.0314]
exper         0.0416 0.0152 [0.0132]    0.0449 0.0155 [0.0136]      0.0442 0.0155 [0.0134]
exper^2      -0.0008 0.0004 [0.0004]   -0.0009 0.0004 [0.0004]     -0.0009 0.0004 [0.0004]

ratio robust/classical on educ:  OLS 0.930   IV 1.011   2SLS 1.056

# linearmodels expects a specific input layout; the API distinguishes
# 'exogenous' regressors, 'endogenous' regressors, and 'instruments'.
data = mroz_w.copy()
data['expersq'] = data['exper'] ** 2

mod_2sls = iv.IV2SLS.from_formula(
    'lwage ~ 1 + exper + expersq + [educ ~ motheduc + fatheduc]',
    data=data,
)
res_2sls = mod_2sls.fit(cov_type='unadjusted', debiased=True)   # I-K denominator, matching the DIY s.e.

# Order coefficients to match our DIY xnames (linearmodels uses 'Intercept')
lm_index = ['Intercept', 'educ', 'exper', 'expersq']
sm_β_k  = res_2sls.params.reindex(lm_index).values
sm_se_k = res_2sls.std_errors.reindex(lm_index).values

print(f"{'name':<10s}{'DIY 2SLS β̂':>14s}{'lm β̂':>12s}{'DIY se':>12s}{'lm se':>12s}")
for name, b1, b2, s1, s2 in zip(xnames, fit_2sls.β_k, sm_β_k, fit_2sls.secl_k, sm_se_k):
    print(f'{name:<10s}{b1:>14.6f}{b2:>12.6f}{s1:>12.6f}{s2:>12.6f}')
print(f'\nmax |β̂_DIY - β̂_lm|  = {np.max(np.abs(fit_2sls.β_k - sm_β_k)):.2e}')
print(f'max |se_DIY - se_lm| = {np.max(np.abs(fit_2sls.secl_k - sm_se_k)):.2e}')

name         DIY 2SLS β̂       lm β̂      DIY se       lm se
const           0.048100    0.048100    0.400328    0.400328
educ            0.061397    0.061397    0.031437    0.031437
exper           0.044170    0.044170    0.013432    0.013432
exper^2        -0.000899   -0.000899    0.000402    0.000402

max |β̂_DIY - β̂_lm|  = 1.61e-13
max |se_DIY - se_lm| = 4.95e-14

# Sargan J-statistic
εhat_i = fit_2sls.εhat_i
σ2_hat = (εhat_i @ εhat_i) / I
PZεhat_i = Zover_i_l @ np.linalg.solve(Zover_i_l.T @ Zover_i_l, Zover_i_l.T @ εhat_i)
J = (εhat_i @ PZεhat_i) / σ2_hat
df_J = Zover_i_l.shape[1] - X_i_k.shape[1]    # L - K = 5 - 4 = 1
p_J = 1 - stats.chi2.cdf(J, df=df_J)
print(f'Sargan J-statistic    = {J:.4f}')
print(f'χ²_{df_J} 95% critical = {stats.chi2.ppf(0.95, df_J):.4f}')
print(f'p-value               = {p_J:.4f}')

# Hausman test on the educ coefficient (the only suspected endogenous regressor)
k_educ = 1
β_diff = fit_2sls.β_k[k_educ] - βols_k[k_educ]
Vols_k_k = σ2_ols * np.linalg.inv(X_i_k.T @ X_i_k)
V_diff = fit_2sls.Vcl_k_k[k_educ, k_educ] - Vols_k_k[k_educ, k_educ]
if V_diff > 0:
    H = β_diff ** 2 / V_diff
    p_H = 1 - stats.chi2.cdf(H, df=1)
    print(f'\nHausman statistic     = {H:.4f}  (df = 1)')
    print(f'p-value               = {p_H:.4f}')
else:
    print(f'\nV_diff = {V_diff:.2e} not positive — Hausman test inconclusive')

Sargan J-statistic    = 0.3781
χ²_1 95% critical = 3.8415
p-value               = 0.5386

Hausman statistic     = 2.6957  (df = 1)
p-value               = 0.1006

# First-stage regression: educ on Zover_i_l (incl. constant, exper, exper^2)
educ_i = mroz_w['educ'].values
βfs_l = np.linalg.solve(Zover_i_l.T @ Zover_i_l, Zover_i_l.T @ educ_i)
educhat_i = Zover_i_l @ βfs_l
ssr_full = ((educ_i - educhat_i) ** 2).sum()

# Restricted: drop motheduc and fatheduc (positions 1 and 2 in Zover_i_l)
Zrestr_i_l = Zover_i_l[:, [0, 3, 4]]
βfsr_l = np.linalg.solve(Zrestr_i_l.T @ Zrestr_i_l, Zrestr_i_l.T @ educ_i)
ssr_restricted = ((educ_i - Zrestr_i_l @ βfsr_l) ** 2).sum()

Q = 2  # number of excluded instruments
F_first_stage = ((ssr_restricted - ssr_full) / Q) / (ssr_full / (I - Zover_i_l.shape[1]))
p_F = 1 - stats.f.cdf(F_first_stage, Q, I - Zover_i_l.shape[1])
print(f'Conventional (homoskedastic) first-stage F               = {F_first_stage:.3f}')
print(f'p-value                                                  = {p_F:.4e}')

# The same joint exclusion restriction, tested with an HC0 covariance. This does not assume
# homoskedastic first-stage errors, and it is the statistic the F > 10 rule was NOT calibrated on.
εfs_i = educ_i - educhat_i
ZtZinv = np.linalg.inv(Zover_i_l.T @ Zover_i_l)
Vfs_l_l = ZtZinv @ ((Zover_i_l * (εfs_i ** 2)[:, None]).T @ Zover_i_l) @ ZtZinv
R_q_l = np.zeros((Q, Zover_i_l.shape[1]))
R_q_l[0, 1] = R_q_l[1, 2] = 1.0                     # motheduc, fatheduc
dev_q = R_q_l @ βfs_l
W_rob = dev_q @ np.linalg.solve(R_q_l @ Vfs_l_l @ R_q_l.T, dev_q)
print(f'Heteroskedasticity-robust Wald / Q (same restriction)    = {W_rob / Q:.3f}')
print(f'   robust p-value                                        = '
      f'{1 - stats.chi2.cdf(W_rob, Q):.4e}')

print(f'\nStaiger–Stock rule of thumb (F > 10)?                    = '
      f'{"Yes" if F_first_stage > 10 else "No"}')
print('   — but see the caveats in the text: the rule is calibrated for the conventional')
print('     homoskedastic F with a single endogenous regressor, not for the robust statistic.')

Conventional (homoskedastic) first-stage F               = 55.400
p-value                                                  = 1.1102e-16
Heteroskedasticity-robust Wald / Q (same restriction)    = 50.112
   robust p-value                                        = 0.0000e+00

Staiger–Stock rule of thumb (F > 10)?                    = Yes
   — but see the caveats in the text: the rule is calibrated for the conventional
     homoskedastic F with a single endogenous regressor, not for the robust statistic.

rng_mc = np.random.default_rng(2026)
R = 2_000
I_sim = 500
β_true = 0.5
π_s = [0.02, 0.05, 0.1, 0.3, 0.5]     # grid of instrument strengths, s in [S]

results = {}
F_avg = {}
for π in π_s:
    β_r = np.empty(R)
    F_r = np.empty(R)
    for r in range(R):
        x_i, ysim_i, z_i = simulate_iv(I=I_sim, β=β_true, π=π, rng=rng_mc)
        # 2SLS slope (one regressor, one instrument: just-identified IV)
        β_r[r] = (z_i @ ysim_i) / (z_i @ x_i)
        # First-stage F = t² for the slope of x on z
        z2 = z_i @ z_i
        π_hat = (z_i @ x_i) / z2
        ehat_i = x_i - π_hat * z_i
        σ2_e = (ehat_i @ ehat_i) / (I_sim - 1)
        se_π = np.sqrt(σ2_e / z2)
        F_r[r] = (π_hat / se_π) ** 2
    results[π] = β_r
    F_avg[π] = F_r.mean()

# The weak-instrument distributions have very heavy tails; the plot shows only the window
# [-2, 3] (out-of-window draws are EXCLUDED, not piled at the border) and each density is
# over the displayed window only; the excluded shares are printed first.
for π in π_s:
    out = np.mean((results[π] < -2) | (results[π] > 3))
    print(f'π = {π:>4}:  mean F = {F_avg[π]:7.1f},  share of draws outside the [-2, 3] '
          f'plotting window = {out:6.2%}')

π = 0.02:  mean F =     1.3,  share of draws outside the [-2, 3] plotting window = 27.45%
π = 0.05:  mean F =     2.2,  share of draws outside the [-2, 3] plotting window = 18.20%
π =  0.1:  mean F =     6.3,  share of draws outside the [-2, 3] plotting window =  3.90%
π =  0.3:  mean F =    46.0,  share of draws outside the [-2, 3] plotting window =  0.00%
π =  0.5:  mean F =   126.9,  share of draws outside the [-2, 3] plotting window =  0.00%

fig, ax = plt.subplots(figsize=(9, 4.5))
bins = np.linspace(-2, 3, 100)
for π, β_r in results.items():
    ax.hist(β_r[(β_r >= -2) & (β_r <= 3)], bins=bins, density=True, alpha=0.45,
            label=fr'$\pi = {π}$,  $\overline{{F}} \approx {F_avg[π]:.1f}$')
ax.axvline(β_true, color='k', ls='--', lw=1.0, label=fr'truth $\beta = {β_true}$')
ax.set_xlim(-2, 3)
ax.set_xlabel(r'$\hat\beta_{\rm 2SLS}$')
ax.set_ylabel('density')
ax.set_title(f'Sampling distribution of 2SLS across instrument strengths, I = {I_sim}, R = {R}')
ax.legend(loc='upper right', fontsize=9)
plt.tight_layout()
plt.show()

Instrumental variables and 2SLS¶

Alfred Galichon (NYU)¶

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

With python code examples¶

Learning objectives¶

References¶

Motivation: structural parameters and selection¶

Loading our libraries and data¶

1. The endogeneity problem¶

1.1 Numerical demonstration: omitted-variable bias¶

2. The instrumental variables estimator¶

2.1 Setup and definitions¶

2.2 Consistency¶

3. Two-stage least squares for over-identified models¶

3.1 The estimator¶

3.2 Asymptotic distribution¶

3.3 DIY implementation¶

4. Application: returns to schooling on Mroz (1987)¶

4.1 Set-up¶

4.2 Three estimators side by side¶

4.3 Comparison with `linearmodels.IV2SLS`¶

5. Diagnostic tests¶

5.1 The Sargan test for over-identification¶

5.2 Hausman test for endogeneity¶

6. Weak instruments¶

6.1 First-stage $F$ on Mroz¶

6.2 Monte Carlo demonstration¶

Looking ahead¶

Instrumental variables and 2SLS¶

Alfred Galichon (NYU)¶

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

With python code examples¶

Learning objectives¶

References¶

Motivation: structural parameters and selection¶

Loading our libraries and data¶

1. The endogeneity problem¶

1.1 Numerical demonstration: omitted-variable bias¶

2. The instrumental variables estimator¶

2.1 Setup and definitions¶

2.2 Consistency¶

3. Two-stage least squares for over-identified models¶

3.1 The estimator¶

3.2 Asymptotic distribution¶

3.3 DIY implementation¶

4. Application: returns to schooling on Mroz (1987)¶

4.1 Set-up¶

4.2 Three estimators side by side¶

4.3 Comparison with linearmodels.IV2SLS¶

5. Diagnostic tests¶

5.1 The Sargan test for over-identification¶

5.2 Hausman test for endogeneity¶

6. Weak instruments¶

6.1 First-stage $F$ on Mroz¶

6.2 Monte Carlo demonstration¶

Looking ahead¶

4.3 Comparison with `linearmodels.IV2SLS`¶