fix(iv): 3 critical + 6 high-priority bugs from code review

CRITICAL fixes:
- KP rk: Kronecker order inverted in robust meat (kron(v⊗v,z⊗z) → kron(z⊗z,v⊗v))
- KP rk LM: use SVD min singular value instead of tr(A²) for correct chi² df
- ivmte_lp: remove dead-code tautology in ATT weight path

HIGH fixes:
- IJIVE: remove incorrect D-shrinkage (AD09 differs only in SE, not point est)
- NPIV: multi-Z now uses linear projection index instead of ‖Z‖₂
- Bayesian IV: avoid n×n PZ matrix, use ZtZ_inv for O(nk) computation
- weak_iv_ci: avoid n×n M_Z in CLR/K-test loops via YD'YD - (Zs'YD)'(Zs'YD)
- KP rk: add try/pinv fallback for singular Z'Z
- SW: correct df_denom to n-kW-k-(p-1) per SW (2016) eq. 7

MEDIUM fixes:
- plausibly_exogenous: validate gamma_grid length divisibility
- __init__.py: move numpy import to top of file

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "StatsPAI"
-version = "0.7.1"
+version = "0.8.0"
 description = "The Agent-Native Causal Inference & Econometrics Toolkit for Python"
 readme = "README.md"
 license = {text = "MIT"}

src/statspai/__init__.py

@@ -22,7 +22,7 @@
 >>> sp.outreg2(result, filename="results.xlsx")
 """
 
-__version__ = "0.7.1"
+__version__ = "0.8.0"
 __author__ = "Bryce Wang"
 __email__ = "bryce@copaper.ai"
 

src/statspai/iv/__init__.py

@@ -50,6 +50,8 @@
 
 from typing import Any, Dict, Optional
 
+import numpy as np
+
 # ─── Core estimators (re-exports) ───────────────────────────────────────
 from ..regression.iv import iv, ivreg, IVRegression
 from ..regression.advanced_iv import liml, jive as jive_legacy, lasso_iv
@@ -293,8 +295,6 @@ def _attach_augmented_diagnostics(model, result, opts: Dict[str, Any]):
         result.diagnostics["augmented_diagnostics_error"] = str(e)
 
 
-# np import only needed for the _attach helper; keep local
-import numpy as np  # noqa: E402
 
 
 __all__ = [

src/statspai/iv/bayesian_iv.py

@@ -190,25 +190,26 @@ def bayesian_iv(
     Zt = _residualize(Z, W)
     k = Zt.shape[1]
 
-    # Projection matrix for instruments
-    PZ = Zt @ np.linalg.solve(Zt.T @ Zt + 1e-10 * np.eye(k), Zt.T)
-    MZ = np.eye(n) - PZ
+    # Precompute ZtZ_inv for fast u'PZu = u'Z(Z'Z)^{-1}Z'u = ||Z'u||²_{inv}
+    ZtZ = Zt.T @ Zt + 1e-10 * np.eye(k)
+    ZtZ_inv = np.linalg.inv(ZtZ)
 
     def log_posterior(beta: float) -> float:
         u = Yt - beta * Dt
-        u_Pz = float(u @ PZ @ u)
-        u_Mz = float(u @ MZ @ u)
+        Ztu = Zt.T @ u  # (k,)
+        u_Pz = float(Ztu @ ZtZ_inv @ Ztu)
+        u_u = float(u @ u)
+        u_Mz = u_u - u_Pz
         if u_Mz <= 0:
             return -np.inf
         # AR quasi-log-likelihood: -AR/2 where AR = n * u'PZu / u'MZu
-        # Maximised when u orthogonal to Z (correct β)
         ar = n * u_Pz / u_Mz
         log_lik = -0.5 * ar
         log_prior = -0.5 * (beta / prior_sd) ** 2
         return log_lik + log_prior
 
     # --- Initialize at 2SLS estimate ---
-    D_hat = PZ @ Dt
+    D_hat = Zt @ (ZtZ_inv @ (Zt.T @ Dt))
     denom = float(D_hat @ Dt)
     if abs(denom) > 1e-12:
         beta_init = float(D_hat @ Yt) / denom

src/statspai/iv/ivmte_lp.py

@@ -215,18 +215,11 @@ def _target_weights(
         return c
 
     if target == "att":
-        # ATT = ∫_0^1 MTE(u) · Pr(U≤u | D=1) du  ≈  ∫_0^1 (m_1-m_0) * (p̄≥u) ...
-        # We approximate using observed propensity distribution among treated.
-        # Concretely: weight_k = ∫_0^1 u^k * F̄_P|D=1(u) du
-        # F̄_P|D=1(u) = Pr(P ≥ u | D=1).
-        if target == "att":
-            p_t = p_hat  # treated filter passed externally
-        p_sample = late_bounds if late_bounds is not None else p_hat
-        # simplification: use uniform plug-in via empirical CDF samples
+        # ATT = ∫_0^1 MTE(u) · Pr(P ≥ u | D=1) du  / E[P | D=1]
+        # p_hat is already the treated subsample (caller passes p_hat[D==1])
         u_grid = np.linspace(0.01, 0.99, 101)
-        w = np.array([(p_sample >= u).mean() for u in u_grid])
+        w = np.array([(p_hat >= u).mean() for u in u_grid])
         w /= max(np.trapezoid(w, u_grid), 1e-12)
-        # For each k, ∫ u^k * w(u) du
         wk = np.array([np.trapezoid(u_grid ** k * w, u_grid) for k in range(K + 1)])
         return np.concatenate([wk, -wk])
 

src/statspai/iv/jive_variants.py

@@ -185,14 +185,11 @@ def _jive_estimate(
             D_hat = ((Pz @ D) - hz[:, None] * D) / (1 - hz)[:, None]
 
     elif method == "ijive":
-        # Ackerberg-Devereux 2009 IJIVE: drop own-observation and re-scale
+        # Ackerberg-Devereux 2009 IJIVE: leave-one-out projection
+        # identical to JIVE1 in point estimate; differs in SE.
         D_hat = np.empty_like(D)
         for j in range(p):
             D_hat[:, j] = (D_hat_full[:, j] - h * D[:, j]) / (1 - h)
-        # IJIVE additional degrees-of-freedom correction:
-        # scale by (n - k - W.shape[1]) / n  (AD09 eq 10 variant)
-        scale = (n - k - W.shape[1]) / n if n > k + W.shape[1] else 1.0
-        D_hat = scale * D_hat + (1 - scale) * D
 
     elif method == "rjive":
         if ridge <= 0:

src/statspai/iv/npiv.py

@@ -197,8 +197,14 @@ def npiv(
     Zt = _residualize(Z, W)
 
     # === Stage 1: sieve regression of D on Z ===
-    Phi_Z = _build_basis(Zt[:, 0] if Zt.shape[1] == 1 else np.linalg.norm(Zt, axis=1),
-                         k_z, basis)
+    if Zt.shape[1] == 1:
+        z_scalar = Zt[:, 0]
+    else:
+        # Multi-instrument: build tensor-product basis would be complex;
+        # use first-stage fitted value from linear projection as the index.
+        pi_lin, *_ = np.linalg.lstsq(Zt, Dt, rcond=None)
+        z_scalar = Zt @ pi_lin  # scalar index, preserves instrument info
+    Phi_Z = _build_basis(z_scalar, k_z, basis)
     pi_hat, *_ = np.linalg.lstsq(Phi_Z, Dt, rcond=None)
     Dt_hat = Phi_Z @ pi_hat
     resid_fs = Dt - Dt_hat

src/statspai/iv/plausibly_exogenous.py

@@ -153,7 +153,15 @@ def plausibly_exogenous_uci(
 
     gamma_arr = np.asarray(list(gamma_grid), dtype=float)
     if gamma_arr.ndim == 1:
-        gamma_arr = gamma_arr.reshape(-1, Z.shape[1]) if Z.shape[1] > 1 else gamma_arr.reshape(-1, 1)
+        if Z.shape[1] > 1:
+            if len(gamma_arr) % Z.shape[1] != 0:
+                raise ValueError(
+                    f"gamma_grid length {len(gamma_arr)} is not divisible "
+                    f"by n_instruments={Z.shape[1]}"
+                )
+            gamma_arr = gamma_arr.reshape(-1, Z.shape[1])
+        else:
+            gamma_arr = gamma_arr.reshape(-1, 1)
 
     beta_hat, _, var0 = _tsls(Y, D, Z, W)
     # When p_endog > 1, focus on the "target_endog" coefficient (default: column 0).

src/statspai/iv/weak_identification.py

@@ -235,7 +235,11 @@ def kleibergen_paap_rk(
     Z_tilde = _residualize(Z, W)
 
     # Reduced form coefficients: D_tilde = Z_tilde @ Pi + V
-    ZZ_inv = np.linalg.inv(Z_tilde.T @ Z_tilde)
+    ZtZ = Z_tilde.T @ Z_tilde
+    try:
+        ZZ_inv = np.linalg.inv(ZtZ)
+    except np.linalg.LinAlgError:
+        ZZ_inv = np.linalg.pinv(ZtZ)
     Pi = ZZ_inv @ (Z_tilde.T @ D_tilde)  # k x p
     V = D_tilde - Z_tilde @ Pi
 
@@ -246,13 +250,16 @@ def kleibergen_paap_rk(
         cov_vec = np.kron(Sigma, ZZ_inv)
         cov_label = "nonrobust"
     elif cov_type == "robust":
-        # Meat: sum_i (Z_i Z_i') ⊗ (V_i V_i')
+        # Meat: Σ_i  kron(z_i z_i', v_i v_i')  — KP (2006) eq. 13
+        # Convention: vec(Pi) stacks columns of Pi (k×p), so Var is (kp × kp)
+        # with blocks (Z'Z)^{-1} ⊗ Σ_VV under homoskedasticity.
+        # Robust sandwich: bread = kron(ZZ_inv, I_p) on both sides.
         meat = np.zeros((k * p, k * p))
         for i in range(n):
-            zi = Z_tilde[i:i + 1].T  # k x 1
-            vi = V[i:i + 1].T  # p x 1
-            meat += np.kron(vi @ vi.T, zi @ zi.T)
-        bread = np.kron(np.eye(p), ZZ_inv)
+            zi = Z_tilde[i]   # (k,)
+            vi = V[i]         # (p,)
+            meat += np.kron(np.outer(zi, zi), np.outer(vi, vi))
+        bread = np.kron(ZZ_inv, np.eye(p))
         cov_vec = bread @ meat @ bread
         cov_label = "HC robust"
     elif cov_type == "cluster":
@@ -264,17 +271,14 @@ def kleibergen_paap_rk(
         meat = np.zeros((k * p, k * p))
         for cid in groups:
             idx = np.where((g == cid).values)[0]
-            Zc = Z_tilde[idx]
-            Vc = V[idx]
-            score = np.kron(Vc, Zc).sum(axis=0)
-            # score is p*k (row-major kron of per-obs kron sum); rebuild properly
+            # Cluster score: Σ_{t∈cluster} kron(z_t, v_t) → vectorised form
             score_mat = np.zeros((k, p))
             for t in idx:
                 score_mat += np.outer(Z_tilde[t], V[t])
-            v = score_mat.flatten(order='F')  # matches vec(Pi) stacking cols
+            v = score_mat.flatten(order='F')  # vec(score) with col-stacking
             meat += np.outer(v, v)
         meat *= G / max(G - 1, 1)
-        bread = np.kron(np.eye(p), ZZ_inv)
+        bread = np.kron(ZZ_inv, np.eye(p))
         cov_vec = bread @ meat @ bread
         cov_label = f"cluster ({G} groups)"
     else:
@@ -291,19 +295,26 @@ def kleibergen_paap_rk(
     rk_f = rk_wald / df_num
     rk_wald_pvalue = float(1 - stats.chi2.cdf(rk_wald, df=k - p + 1))
 
-    # KP rk LM statistic (Kleibergen 2002 / Kleibergen-Paap 2006 eq. 17)
-    # Equivalent to Wald when homoskedastic, more stable under weak identification
-    # KP rk LM = n * tr( (D'M_W D)^{-1} D'P_Z D )  after robust-ification
+    # KP rk LM statistic (Kleibergen-Paap 2006, Theorem 1)
+    # Tests H0: rank(Pi) <= p-1 vs H1: rank(Pi) = p.
+    # The rk LM is based on the *smallest* canonical correlation /
+    # singular value of the whitened reduced-form matrix A = Zs' Ds.
+    # Under H0, rk_lm ~ chi²((k - p + 1)).
     Sigma = (V.T @ V) / n
     try:
         Sigma_half_inv = np.linalg.inv(np.linalg.cholesky(Sigma))
     except np.linalg.LinAlgError:
         Sigma_half_inv = np.linalg.pinv(_sqrtm_sym(Sigma))
-    Zs = Z_tilde @ np.linalg.cholesky(ZZ_inv)  # orthonormalised instruments
+    try:
+        ZZ_chol = np.linalg.cholesky(Z_tilde.T @ Z_tilde / n)
+        Zs = Z_tilde @ np.linalg.inv(ZZ_chol.T)  # orthonormalised instruments
+    except np.linalg.LinAlgError:
+        Zs = Z_tilde
     Ds = D_tilde @ Sigma_half_inv.T  # whitened endog
-    A = Zs.T @ Ds  # k x p
-    rk_lm = float(np.sum(A ** 2))
-    rk_lm_pvalue = float(1 - stats.chi2.cdf(rk_lm, df=(k - p + 1) * p))
+    A = Zs.T @ Ds / np.sqrt(n)  # k x p
+    sv = np.linalg.svd(A, compute_uv=False)
+    rk_lm = float(n * sv[-1] ** 2)   # smallest sv², scaled by n
+    rk_lm_pvalue = float(1 - stats.chi2.cdf(rk_lm, df=(k - p + 1)))
 
     return KleibergenPaapResult(
         rk_wald=rk_wald,
@@ -412,7 +423,7 @@ def sanderson_windmeijer(
         tss = float(y_j @ y_j)
 
         df1 = k - (p - 1)  # SW adjusted numerator df
-        df2 = n - n_W - k  # denominator df (net of all controls + instruments)
+        df2 = n - n_W - k - (p - 1)  # SW (2016) eq. 7 denominator df
         if df1 <= 0:
             raise ValueError(
                 f"Not enough instruments: k - (p-1) = {df1} for endogenous '{names[j]}'."

src/statspai/iv/weak_iv_ci.py

@@ -250,13 +250,18 @@ def conditional_lr_ci(
     stat_arr = np.empty(len(beta_grid))
     crit_arr = np.empty(len(beta_grid))
 
+    df_r = max(n - kW - k, 1)
+    # Precompute Zs'Dt (invariant across loop)
+    ZsDt = Zs.T @ Dt  # (k,)
+    DtDt = float(Dt @ Dt)
+    DtZsZsDt = float(ZsDt @ ZsDt)
+
     for i, b0 in enumerate(beta_grid):
         ustar = Yt - b0 * Dt
-        # Reduced-form Sigma under H0
-        M_Z = np.eye(n) - Zs @ Zs.T
+        # Sigma = YD' M_Z YD / df, avoiding n×n M_Z via YD'YD - (Zs'YD)'(Zs'YD)
         YD = np.column_stack([ustar, Dt])
-        df_r = max(n - kW - k, 1)
-        Sigma = YD.T @ M_Z @ YD / df_r
+        ZsYD = Zs.T @ YD  # (k, 2)
+        Sigma = (YD.T @ YD - ZsYD.T @ ZsYD) / df_r
         suu = float(Sigma[0, 0])
         svv = float(Sigma[1, 1])
         suv = float(Sigma[0, 1])
@@ -333,13 +338,13 @@ def k_test_ci(
     crit = stats.chi2.ppf(level, df=1)
 
     stat_arr = np.empty(len(beta_grid))
-    M_Z = np.eye(n) - Zs @ Zs.T
+    df_r = max(n - kW - k, 1)
 
     for i, b0 in enumerate(beta_grid):
         ustar = Yt - b0 * Dt
         YD = np.column_stack([ustar, Dt])
-        df_r = max(n - kW - k, 1)
-        Sigma = YD.T @ M_Z @ YD / df_r
+        ZsYD = Zs.T @ YD
+        Sigma = (YD.T @ YD - ZsYD.T @ ZsYD) / df_r
         suu = float(Sigma[0, 0])
         svv = float(Sigma[1, 1])
         suv = float(Sigma[0, 1])