import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats, optimize
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,                       # endogenous regressor
    mroz_w['exper'].values,
    mroz_w['exper'].values ** 2,
])
Z_i_l = np.column_stack([
    np.ones(I),
    mroz_w['motheduc'].values,                   # instrument 1 (excluded)
    mroz_w['fatheduc'].values,                   # instrument 2 (excluded)
    mroz_w['exper'].values,
    mroz_w['exper'].values ** 2,
])
xnames = ['const', 'educ', 'exper', 'exper^2']
K, L = X_i_k.shape[1], Z_i_l.shape[1]
print(f'I = {I}, K = {K} (parameters), L = {L} (moments); over-identification = L - K = {L - K}')

I = 428, K = 4 (parameters), L = 5 (moments); over-identification = L - K = 1

def linear_gmm(X_i_k, y_i, Z_i_l, W_l_l):
    XZ_k_l = X_i_k.T @ Z_i_l
    Zy_l = Z_i_l.T @ y_i
    β_k  = np.linalg.solve(XZ_k_l @ W_l_l @ XZ_k_l.T, XZ_k_l @ W_l_l @ Zy_l)
    εhat_i = y_i - X_i_k @ β_k
    return SimpleNamespace(β_k=β_k, εhat_i=εhat_i)

def linear_gmm_vcov(X_i_k, Z_i_l, W_l_l, εhat_i):
    # Heteroskedasticity-consistent GMM variance.
    # Aggregates are unscaled sums (S_I = Σ_i ε_i² z_i z_i', etc.) to avoid
    # confusion between per-observation averages and totals.
    XZ_k_l = X_i_k.T @ Z_i_l
    g_i_l = Z_i_l * εhat_i[:, None]                   # I × L matrix of score vectors
    S_l_l = g_i_l.T @ g_i_l                     # = Σ_i ε_i² z_i z_i' (unscaled)
    bread_k_k   = np.linalg.inv(XZ_k_l @ W_l_l @ XZ_k_l.T)
    middle_k_k  = XZ_k_l @ W_l_l @ S_l_l @ W_l_l @ XZ_k_l.T
    V_k_k       = bread_k_k @ middle_k_k @ bread_k_k
    return V_k_k, S_l_l

# Step 1: 2SLS, i.e., W_1 = (Z'Z)^{-1}
W1_l_l = np.linalg.inv(Z_i_l.T @ Z_i_l)
fit1 = linear_gmm(X_i_k, y_i, Z_i_l, W1_l_l)
V1_k_k, S1_l_l = linear_gmm_vcov(X_i_k, Z_i_l, W1_l_l, fit1.εhat_i)
se1_k = np.sqrt(np.diag(V1_k_k))

# Step 2: efficient GMM, W_2 = S_1^{-1}
W2_l_l = np.linalg.inv(S1_l_l)
fit2 = linear_gmm(X_i_k, y_i, Z_i_l, W2_l_l)
V2_k_k, _ = linear_gmm_vcov(X_i_k, Z_i_l, W2_l_l, fit2.εhat_i)
se2_k = np.sqrt(np.diag(V2_k_k))

print(f"{'name':<10s}{'2SLS (W₁)':>22s}{'efficient GMM (W₂)':>26s}{'SE ratio':>12s}")
for name, b1, b2, s1, s2 in zip(xnames, fit1.β_k, fit2.β_k, se1_k, se2_k):
    out1 = f'{b1:>10.4f} ({s1:.4f})'
    out2 = f'{b2:>10.4f} ({s2:.4f})'
    print(f'{name:<10s}{out1:>22s}{out2:>26s}{s2/s1:>12.3f}')

name                   2SLS (W₁)        efficient GMM (W₂)    SE ratio
const            0.0481 (0.4278)           0.0477 (0.4277)       1.000
educ             0.0614 (0.0332)           0.0611 (0.0332)       1.000
exper            0.0442 (0.0155)           0.0451 (0.0154)       0.997
exper^2         -0.0009 (0.0004)          -0.0009 (0.0004)       0.996

data = mroz_w.copy()
data['expersq'] = data['exper'] ** 2

mod_2step = iv.IVGMM.from_formula(
    'lwage ~ 1 + exper + expersq + [educ ~ motheduc + fatheduc]',
    data=data,
    weight_type='robust',
)
res_2step = mod_2step.fit(iter_limit=2)        # two-step

lm_index = ['Intercept', 'educ', 'exper', 'expersq']
smβ_k  = res_2step.params.reindex(lm_index).values
smse_k = res_2step.std_errors.reindex(lm_index).values

print(f"{'name':<10s}{'DIY GMM β̂':>14s}{'lm β̂':>12s}{'DIY se':>12s}{'lm se':>12s}")
for name, b1, b2, s1, s2 in zip(xnames, fit2.β_k, smβ_k, se2_k, smse_k):
    print(f'{name:<10s}{b1:>14.6f}{b2:>12.6f}{s1:>12.6f}{s2:>12.6f}')

name          DIY GMM β̂       lm β̂      DIY se       lm se
const           0.047654    0.047654    0.427730    0.427730
educ            0.061053    0.061053    0.033170    0.033170
exper           0.045135    0.045135    0.015421    0.015421
exper^2        -0.000931   -0.000931    0.000426    0.000426

# Compute J at the efficient estimator.
# J is *the minimized objective itself*, scaled by I, so it must use the same weight matrix
# the second step actually minimized — Ω̂₁ built from the STEP-1 residuals — not a fresh one
# recomputed at θ̂₂. The recomputed version is asymptotically equivalent and is reported
# below as a check, but only the first is the value of the objective at its own minimum.
gbar_l = (Z_i_l * fit2.εhat_i[:, None]).sum(axis=0) / I        # = (1/I) Σ_i ε̂_i z_i
Ω̂_1 = S1_l_l / I                                              # the weight used in step 2
J = I * gbar_l @ np.linalg.solve(Ω̂_1, gbar_l)
df_J = L - K
p_J = 1 - stats.chi2.cdf(J, df=df_J)
print(f'Hansen 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}')

# J equals I times the minimized GMM objective, by construction:
obj_min = gbar_l @ W2_l_l @ gbar_l * I**2      # W2_l_l = S1_l_l^{-1}, an unscaled sum
print(f'\nI x minimized objective    = {obj_min:.4f}   (identical to J above)')

# The asymptotically equivalent variant, reweighting at θ̂₂ rather than θ̂₁:
S2_l_l = (Z_i_l * fit2.εhat_i[:, None]).T @ (Z_i_l * fit2.εhat_i[:, None])
J_recomputed = I * gbar_l @ np.linalg.solve(S2_l_l / I, gbar_l)
print(f'J with Ω̂ recomputed at θ̂₂  = {J_recomputed:.4f}  (same limit, different statistic)')

# Compare with Sargan from Lecture 10 (homoskedastic version)
σ2_hat = (fit2.εhat_i @ fit2.εhat_i) / I
PZε_i = Z_i_l @ np.linalg.solve(Z_i_l.T @ Z_i_l, Z_i_l.T @ fit2.εhat_i)
Sargan = (fit2.εhat_i @ PZε_i) / σ2_hat
print(f'\nSargan statistic (Lecture 10) = {Sargan:.4f}')
print(f'Difference                   = {J - Sargan:.4f}')

Hansen J-statistic         = 0.4435
χ²_1 95% critical      = 3.8415
p-value                    = 0.5055

I x minimized objective    = 0.4435   (identical to J above)
J with Ω̂ recomputed at θ̂₂  = 0.4433  (same limit, different statistic)

Sargan statistic (Lecture 10) = 0.3847
Difference                   = 0.0588

def simulate_heteroskedastic_iv(I, β_k=0.5, π1=0.4, π2=0.3, α=1.5, rng=None):
    if rng is None:
        rng = np.random.default_rng()
    z1_i = rng.standard_normal(I)
    z2_i = rng.standard_normal(I)
    w_i  = rng.standard_normal(I)                       # source of endogeneity
    v_i  = rng.standard_normal(I)
    x_i  = π1 * z1_i + π2 * z2_i + 0.5 * w_i + np.sqrt(1 - 0.5**2) * v_i
    σ_i = np.sqrt(np.exp(α * z1_i))                     # strong heteroskedasticity in z1_i
    u_i = σ_i * rng.standard_normal(I)
    ε_i  = 0.7 * w_i + u_i
    y_i  = β_k * x_i + ε_i
    return x_i, y_i, z1_i, z2_i

rng = np.random.default_rng(2026)
M = 2_000
I_s = 1_000
β_true = 0.5

β2sls_m = np.empty(M)
βgmm_m  = np.empty(M)

for m in range(M):
    x_i, ys_i, z1_i, z2_i = simulate_heteroskedastic_iv(I_s, β_k=β_true, rng=rng)
    Xs_i_k = x_i[:, None]
    Zs_i_l = np.column_stack([z1_i, z2_i])
    # 2SLS = GMM with W_l_l = (Z_i_l'Z_i_l)^{-1}
    W1_l_l = np.linalg.inv(Zs_i_l.T @ Zs_i_l)
    fit1 = linear_gmm(Xs_i_k, ys_i, Zs_i_l, W1_l_l)
    # Efficient two-step with W_l_l = S_l_l^{-1}
    S_l_l = (Zs_i_l * fit1.εhat_i[:, None]).T @ (Zs_i_l * fit1.εhat_i[:, None])
    W2_l_l = np.linalg.inv(S_l_l)
    fit2 = linear_gmm(Xs_i_k, ys_i, Zs_i_l, W2_l_l)
    β2sls_m[m] = fit1.β_k[0]
    βgmm_m[m]  = fit2.β_k[0]

print(f'Heteroskedastic DGP (α = 1.5):')
print(f"{'':<14s}{'mean β̂':>10s}{'std β̂':>10s}{'rel. eff.':>12s}")
print(f"{'2SLS':<14s}{β2sls_m.mean():>10.4f}{β2sls_m.std():>10.4f}{1.0:>12.3f}")
print(f"{'two-step GMM':<14s}{βgmm_m.mean():>10.4f}{βgmm_m.std():>10.4f}{β2sls_m.std()/βgmm_m.std():>12.3f}")

Heteroskedastic DGP (α = 1.5):
                 mean β̂    std β̂   rel. eff.
2SLS              0.5030    0.1784       1.000
two-step GMM      0.5020    0.1545       1.155

fig, ax = plt.subplots(figsize=(9, 4.2))
bins_g = np.linspace(min(β2sls_m.min(), βgmm_m.min()),
                     max(β2sls_m.max(), βgmm_m.max()), 60)
ax.hist(β2sls_m, bins=bins_g, density=True, alpha=0.45, edgecolor='white',
        label=f'2SLS  (sd {β2sls_m.std():.4f})')
ax.hist(βgmm_m, bins=bins_g, density=True, alpha=0.45, edgecolor='white',
        label=f'two-step GMM  (sd {βgmm_m.std():.4f})')
ax.axvline(β_true, color='k', ls='--', lw=1.2, label=r'true $\beta_0$')
ax.set_xlabel(r'$\hat\beta$')
ax.set_ylabel('density')
ax.set_title(f'Efficient weighting under heteroskedasticity (M = {M} replications)')
ax.legend()
plt.tight_layout()
plt.show()

Estimator	Moment $g_i(\theta)$	Weight matrix $W$	$V$ at the efficient choice
OLS	$x_i(y_i - x_i^\top\beta)$	any (just-identified)	$\sigma^2\bigl(\mathbb{E}[xx^\top]\bigr)^{-1}$ under homosk.; sandwich under heterosk.
IV (just-identified)	$z_i(y_i - x_i^\top\beta)$	any (just-identified)	$\bigl(\mathbb{E}[zx^\top]\bigr)^{-1}\Omega\bigl(\mathbb{E}[xz^\top]\bigr)^{-1}$
2SLS	$z_i(y_i - x_i^\top\beta)$	$(Z^\top Z/I)^{-1}$	efficient under conditional homoskedasticity only
Efficient GMM-IV	$z_i(y_i - x_i^\top\beta)$	$\hat\Omega^{-1}$	$(G^\top\Omega^{-1}G)^{-1}$
MLE	$s_i(\theta) = \nabla\log f(z_i;\theta)$	any (just-identified)	$\mathcal{I}(\theta_0)^{-1}$ by the information identity

Generalized Method of Moments¶

Alfred Galichon (NYU)¶

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

With python code examples¶

Learning objectives¶

References¶

Motivation: moment conditions as the primitive¶

Loading our libraries and data¶

1. Method of moments and GMM¶

1.1 Moment conditions¶

1.2 The just-identified case ($L = K$): method of moments¶

1.3 The over-identified case ($L > K$): GMM¶

2. Asymptotic theory of GMM¶

2.1 Consistency¶

2.2 Asymptotic normality¶

3. The efficient weight matrix¶

3.1 Special cases¶

3.2 Two-step GMM¶

4. DIY GMM for linear IV moment conditions¶

4.1 Implementation¶

4.2 Comparison with `linearmodels.IVGMM`¶

5. Hansen's $J$-test for over-identification¶

6. Efficiency gain in a heteroskedastic DGP¶

7. GMM as a unifying framework¶

Looking ahead¶

Generalized Method of Moments¶

Alfred Galichon (NYU)¶

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

With python code examples¶

Learning objectives¶

References¶

Motivation: moment conditions as the primitive¶

Loading our libraries and data¶

1. Method of moments and GMM¶

1.1 Moment conditions¶

1.2 The just-identified case ($L = K$): method of moments¶

1.3 The over-identified case ($L > K$): GMM¶

2. Asymptotic theory of GMM¶

2.1 Consistency¶

2.2 Asymptotic normality¶

3. The efficient weight matrix¶

3.1 Special cases¶

3.2 Two-step GMM¶

4. DIY GMM for linear IV moment conditions¶

4.1 Implementation¶

4.2 Comparison with linearmodels.IVGMM¶

5. Hansen's $J$-test for over-identification¶

6. Efficiency gain in a heteroskedastic DGP¶

7. GMM as a unifying framework¶

Looking ahead¶

4.2 Comparison with `linearmodels.IVGMM`¶