Implement (non-)robust residual prediction test (#127)

* Implement (non-)robust residual prediction test

* Use robust in test.

* Fix test

Author: mlondschien

ivmodels/tests/residual_prediction.py

@@ -11,6 +11,7 @@ def residual_prediction_test(
     X,
     y,
     C=None,
+    robust=False,
     nonlinear_model=None,
     fit_intercept=True,
     train_fraction=None,
@@ -36,6 +37,8 @@ def residual_prediction_test(
         Outcomes.
     C: np.ndarray of dimension (n, mc) or None, optional, default = None
         Included exogenous regressors.
+    robust: bool or string, optional, default = False
+        Whether to use heteroskedasticity-robust standard errors.
     nonlinear_model: object, optional, default = None
         Object with a ``fit`` and ``predict`` method. If ``None``, uses an
         ``sklearn.ensemble.RandomForestRegressor()``.
@@ -136,15 +139,22 @@ def residual_prediction_test(
 
     XCb_proj = np.hstack([proj(np.hstack([Zb, Cb]), Xb), Cb])
     XCb = np.hstack([Xb, Cb])
-    # pinv(X) = (X^T @ X)^(-1) @ X^T
-    sigma_sq_hat = (
-        np.mean((wb - np.linalg.pinv(XCb_proj).T @ XCb.T @ wb) ** 2 * residuals_b**2)
-        - np.mean(wb * residuals_b) ** 2
-    )
+
+    if robust:
+        # pinv(X) = (X^T @ X)^(-1) @ X^T
+        sigma_sq_hat = (
+            np.mean(
+                (wb - np.linalg.pinv(XCb_proj).T @ XCb.T @ wb) ** 2 * residuals_b**2
+            )
+            - np.mean(wb * residuals_b) ** 2
+        )
+    else:
+        sigma_sq_hat = np.mean((wb - np.linalg.pinv(XCb_proj).T @ XCb.T @ wb) ** 2)
+        sigma_sq_hat *= np.mean(residuals_b**2)
 
     if sigma_sq_hat < gamma:  # Pre-test for variance
         return -np.inf, 1
 
-    stat = wb.T @ residuals_b / np.sqrt(sigma_sq_hat) / np.sqrt(n)
+    stat = wb.T @ residuals_b / np.sqrt(sigma_sq_hat) / np.sqrt(Xb.shape[0])
     p_value = 1 - scipy.stats.norm.cdf(stat)
     return stat, p_value

tests/tests/test_residual_prediction.py

@@ -1,15 +1,17 @@
 import numpy as np
 import pytest
+import scipy.stats
 from sklearn.ensemble import RandomForestRegressor
 
 from ivmodels.tests import residual_prediction_test
 
 
+@pytest.mark.parametrize("robust", [False, True])
 @pytest.mark.parametrize(
     "n, k, mx, mc, fit_intercept",
-    [(200, 3, 3, 1, True), (200, 3, 1, 1, False), (200, 15, 5, 5, False)],
+    [(500, 3, 3, 1, True), (500, 3, 1, 1, False), (500, 15, 5, 5, True)],
 )
-def test_residual_prediction_test(n, k, mx, mc, fit_intercept):
+def test_residual_prediction_test(n, k, mx, mc, fit_intercept, robust):
     rng = np.random.default_rng(0)
 
     Pi = rng.normal(size=(k, mx))
@@ -19,7 +21,7 @@ def test_residual_prediction_test(n, k, mx, mc, fit_intercept):
     Pi_CX = rng.normal(size=(mc, mx))
     Pi_Cy = rng.normal(size=(mc, 1))
 
-    n_seeds = 50
+    n_seeds = 20
     statistics = np.zeros(n_seeds)
     p_values = np.zeros(n_seeds)
 
@@ -29,20 +31,27 @@ def test_residual_prediction_test(n, k, mx, mc, fit_intercept):
         U = rng.normal(size=(n, mx + 1))
         X = Z @ Pi + U @ gamma + C @ Pi_CX + rng.normal(size=(n, mx))
         X[:, 0] += Z[:, 0] ** 2  # allow for nonlinearity Z -> X
-        y = X @ beta + U[:, 0:1] + U[:, 0:1] ** 3 + C @ Pi_Cy + rng.normal(size=(n, 1))
+        noise = rng.normal(size=(n, 1))
+        if robust:
+            noise *= Z[:, 0:1] ** 2
+        y = X @ beta + U[:, 0:1] + np.sin(U[:, 0:1]) + C @ Pi_Cy + noise
 
         statistics[idx], p_values[idx] = residual_prediction_test(
             Z=Z,
             X=X,
             y=y,
             C=C,
+            robust=robust,
             nonlinear_model=RandomForestRegressor(n_estimators=20, random_state=0),
             fit_intercept=fit_intercept,
-            train_fraction=0.6,
+            train_fraction=0.4,
             seed=0,
         )
 
-    assert np.mean(p_values < 0.1) < 0.05
+    assert (
+        scipy.stats.kstest(p_values, scipy.stats.uniform(loc=0.0, scale=1.0).cdf).pvalue
+        > 0.05
+    )
 
 
 def test_residual_prediction_test_rejects():