Use sandwich covariance for penalized fits

src/irae_risk/logistic.py

@@ -711,35 +711,85 @@ class LogisticPolyRegression:
 
         return H
 
-    """
-        Calculation of asymptotic variance-covariance matrix of regression
-        parameters theta
-
-           cov_{asymp}[theta] = H^{-1}
+    def get_cov(self, x, y, theta, method="model_sandwich"):
+        """Estimate the covariance matrix of fitted parameters.
+
+        Parameters
+        ----------
+        x, y
+            Predictor and binary response arrays.
+        theta
+            Fitted regression parameters.
+        method
+            ``"model_sandwich"`` (default) uses
+
+                A^+ [sum_i p_i(1-p_i) g_i g_i.T] A^+,
+
+            where ``A`` is the observed Hessian of the penalized objective and
+            ``g_i = d eta_i / d theta``.  This accounts for the fact that the
+            deterministic penalty changes the Hessian but does not contribute
+            sampling variability.
+
+            ``"robust_sandwich"`` replaces the middle matrix by the empirical
+            outer product of unpenalized scores,
+
+                sum_i [(y_i-p_i) g_i][(y_i-p_i) g_i].T.
+
+            ``"inverse_hessian"`` returns the previous approximation ``A^+``
+            for comparison.  It should not be treated as the frequentist
+            covariance of a penalized estimator.
+
+        Notes
+        -----
+        ``+`` denotes a symmetric Moore-Penrose pseudoinverse.  These are local
+        normal approximations.  For the monotonic cubic model, they can be
+        unreliable near ``epsilon=0``, where the estimate is on the boundary
+        and the parameter map loses rank.  Bootstrap inference is preferred in
+        that case.  A nonzero L1 penalty is not supported because its Hessian is
+        undefined at zero.
+        """
+        x, y = self._validate_data(x, y)
+        theta = np.asarray(theta, dtype=float)
 
-        where H is hessian of nllf
+        valid_methods = {
+            "model_sandwich",
+            "robust_sandwich",
+            "inverse_hessian",
+        }
+        if method not in valid_methods:
+            raise ValueError(
+                f"Unknown covariance method {method!r}; expected one of "
+                f"{sorted(valid_methods)}."
+            )
 
-           H = [d^2(nllf)/(d(theta)_a d(theta)_b ]_{a,b}
+        if self.lam is not None and self.lam[0] != 0:
+            raise ValueError(
+                "Covariance estimation does not support a nonzero L1 penalty."
+            )
 
-        for the logistic regression of the polynomial model:
+        # Bread: exact observed Hessian of the smooth penalized objective.
+        bread = self.get_cost_hessian(x, y, theta)
+        bread = 0.5 * (bread + bread.T)
+        bread_inv = np.linalg.pinv(bread, hermitian=True)
 
-            log(f(x)/(1 - f(x))) ~ sum_{i=0}^degree b_i(theta) x^i
+        if method == "inverse_hessian":
+            return bread_inv
 
-        Input:
-            x: array of n floats
-            theta: array of r = degree+1 floats, model parameters
-
-        Return:
-            array of rxr floats; r = degree + 1
+        beta = self.get_beta(theta)
+        X = np.column_stack([x**i for i in range(len(beta))])
+        J = self.get_jac_beta(theta)
+        gradient_eta = X @ J
+        probabilities = self.model(x, theta)
 
-        Ref:
-          * https://stats.stackexchange.com/questions/89484/how-to-compute-the-standard-errors-of-a-logistic-regressions-coefficients
-          * https://goodboychan.github.io/machine_learning/2020/09/14/02-Regularized-likelihood-methods.html
-    """
-    def get_cov(self, x, y, theta):
+        if method == "model_sandwich":
+            weights = probabilities * (1.0 - probabilities)
+            meat = (gradient_eta.T * weights) @ gradient_eta
+        else:
+            scores = (y - probabilities)[:, None] * gradient_eta
+            meat = scores.T @ scores
 
-        # covariance matrix C_params = H^-1
-        return np.linalg.pinv(self.get_cost_hessian(x, y, theta))
+        cov = bread_inv @ meat @ bread_inv
+        return 0.5 * (cov + cov.T)
 
     """
         Calculating standard errors fo model parameters for normal distribution of parameters:
@@ -1067,4 +1117,3 @@ class LogisticPolyRegression:
             m,
             max_attempts,
         )
-

tests/test_logistic.py

@@ -97,3 +97,47 @@ def test_goodness_of_fit_reports_reproducible_bootstrap_deviance(regression_data
     ]
     assert 0 < first["p-value(deviance_bootstrap)"] <= 1
     assert first["deviance_bootstrap_samples"] == 25
+
+
+def test_covariance_methods_match_sandwich_formulas(regression_data):
+    x, y, theta = regression_data
+    ridge = 0.5
+    model = LogisticPolyRegression(degree=1, lam=(0.0, ridge))
+
+    design = np.column_stack([np.ones_like(x), x])
+    probabilities = model.model(x, theta)
+    weights = probabilities * (1 - probabilities)
+    information = (design.T * weights) @ design
+    bread = information + 2 * ridge * np.eye(2)
+    bread_inv = np.linalg.pinv(bread, hermitian=True)
+
+    scores = (y - probabilities)[:, None] * design
+    robust_meat = scores.T @ scores
+
+    expected_model = bread_inv @ information @ bread_inv
+    expected_robust = bread_inv @ robust_meat @ bread_inv
+
+    np.testing.assert_allclose(
+        model.get_cov(x, y, theta),
+        expected_model,
+    )
+    np.testing.assert_allclose(
+        model.get_cov(x, y, theta, method="robust_sandwich"),
+        expected_robust,
+    )
+    np.testing.assert_allclose(
+        model.get_cov(x, y, theta, method="inverse_hessian"),
+        bread_inv,
+    )
+
+
+def test_covariance_rejects_unknown_method_and_l1_penalty(regression_data):
+    x, y, theta = regression_data
+    model = LogisticPolyRegression(degree=1)
+
+    with pytest.raises(ValueError, match="Unknown covariance method"):
+        model.get_cov(x, y, theta, method="not-a-method")
+
+    l1_model = LogisticPolyRegression(degree=1, lam=(0.1, 0.0))
+    with pytest.raises(ValueError, match="nonzero L1"):
+        l1_model.get_cov(x, y, theta)