Add subgroup CATE estimation (econml DML) to lending audit

LendingFairnessAudit.estimate_cate() wraps econml.dml.CausalForestDML to estimate
the conditional average treatment effect of the protected attribute on approval
(as in P1). The treatment defaults to group != reference and the outcome to the
approval labels; nuisance models default to random forests. Causal validity is
conditional on the DML assumptions, stated in the docstring. econml is an optional
dependency (pip install fairscope[lending]); a guarded import raises an informative
ImportError naming the extra. A planted-effect test (importorskip) confirms the
estimator recovers the right sign on the strong stratum.

Author: Rajveer-code

fairscope/lending/audit.py

@@ -39,6 +39,64 @@ def from_outcomes(cls, approved, group, year, *, reference, alpha=0.05):
         """Build an audit from precomputed approval outcomes."""
         return cls(approved=approved, group=group, year=year, reference=reference, alpha=alpha)
 
+    def _binary_treatment(self):
+        return (self.group != self.reference).astype(int)
+
+    def estimate_cate(
+        self,
+        X,
+        *,
+        treatment=None,
+        outcome=None,
+        model_y=None,
+        model_t=None,
+        n_estimators=500,
+        random_state=0,
+    ):
+        """Per-subgroup conditional average treatment effect (CATE) of the protected
+        attribute on approval, via Causal Forest DML (``econml.dml.CausalForestDML``,
+        as in P1).
+
+        The CAUSAL CLAIM IS CONDITIONAL on the DML assumptions (unconfoundedness given
+        the supplied features ``X``, and overlap). This estimates an effect under those
+        assumptions; it does not, on its own, prove discrimination.
+
+        Parameters
+        ----------
+        X : array (n, k) of heterogeneity features.
+        treatment : binary array; defaults to ``group != reference``.
+        outcome : binary array; defaults to the ``approved`` outcomes.
+        model_y, model_t : nuisance estimators; default to random forests.
+
+        Returns ``{"ate", "effect", "effect_interval"}``. Requires the optional
+        dependency: ``pip install fairscope[lending]``.
+        """
+        try:
+            from econml.dml import CausalForestDML
+        except ImportError as exc:  # optional dependency
+            raise ImportError(
+                "Subgroup CATE requires the optional dependency: " "pip install fairscope[lending]"
+            ) from exc
+        from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+        T = self._binary_treatment() if treatment is None else np.asarray(treatment)
+        Y = self.approved if outcome is None else np.asarray(outcome)
+        X = np.asarray(X)
+        est = CausalForestDML(
+            model_y=model_y or RandomForestRegressor(random_state=random_state),
+            model_t=model_t or RandomForestClassifier(random_state=random_state),
+            discrete_treatment=True,
+            n_estimators=n_estimators,
+            random_state=random_state,
+        )
+        est.fit(Y, T, X=X)
+        lo, hi = est.effect_interval(X, alpha=self.alpha)
+        return {
+            "ate": float(est.ate(X)),
+            "effect": est.effect(X),
+            "effect_interval": (lo, hi),
+        }
+
     def run(self) -> LendingReport:
         rows = []
         for yr in sorted(np.unique(self.year).tolist()):

tests/test_lending.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pandas as pd
+import pytest
 
 from fairscope.lending import LendingFairnessAudit, LendingReport
 
@@ -52,3 +53,46 @@ def test_summary_returns_text():
         .summary()
     )
     assert isinstance(text, str) and "minority" in text
+
+
+def test_estimate_cate_missing_econml_raises(monkeypatch):
+    import builtins
+
+    real = builtins.__import__
+
+    def _no_econml(name, *a, **k):
+        if name.startswith("econml"):
+            raise ImportError("no econml")
+        return real(name, *a, **k)
+
+    monkeypatch.setattr(builtins, "__import__", _no_econml)
+    rng = np.random.default_rng(2)
+    n = 200
+    audit = LendingFairnessAudit.from_outcomes(
+        rng.integers(0, 2, n),
+        np.where(rng.random(n) < 0.5, "minority", "reference"),
+        np.full(n, 2021),
+        reference="reference",
+    )
+    with pytest.raises(ImportError, match=r"fairscope\[lending\]"):
+        audit.estimate_cate(np.zeros((n, 2)))
+
+
+def test_estimate_cate_recovers_planted_effect():
+    pytest.importorskip("econml")
+    rng = np.random.default_rng(3)
+    n = 2000
+    X = rng.normal(size=(n, 3))  # heterogeneity covariates
+    T = (rng.random(n) < 0.5).astype(int)  # protected indicator (treatment)
+    tau = 0.8 * (X[:, 0] > 0)  # heterogeneous planted effect on the X0>0 stratum
+    Y = (0.3 + 0.5 * X[:, 0] + tau * T + rng.normal(scale=0.3, size=n) > 0.5).astype(int)
+    audit = LendingFairnessAudit.from_outcomes(
+        Y, np.where(T == 1, "minority", "reference"), np.full(n, 2021), reference="reference"
+    )
+    # defaults: treatment derives from group != reference, outcome from approved (P1 style)
+    res = audit.estimate_cate(X, n_estimators=200)
+    assert res["ate"] != 0
+    lo, hi = res["effect_interval"]
+    assert len(lo) == len(hi) == n
+    strong = X[:, 0] > 0
+    assert np.median(res["effect"][strong]) > 0  # right sign on the strong stratum