Implement adaptive clipping switch

Quantco · Jul 10, 2024 · 8fa0c10 · 8fa0c10
1 parent 1580c96
commit 8fa0c10
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/metalearners/drlearner.py b/metalearners/drlearner.py
@@ -268,6 +268,7 @@ def _pseudo_outcome(
         is_oos: bool,
         oos_method: OosMethod = OVERALL,
         epsilon: float = _EPSILON,
+        adaptive_clipping: bool = False,
     ) -> np.ndarray:
         """Compute the DR-Learner pseudo outcome."""
         validate_valid_treatment_variant_not_control(treatment_variant, self.n_variants)
@@ -317,4 +318,12 @@ def _pseudo_outcome(
             - y0_estimate
         )
 
+        if adaptive_clipping:
+            t_pseudo_outcome = y1_estimate - y0_estimate
+            pseudo_outcome = np.where(
+                propensity_estimates.min(axis=1) < epsilon,
+                t_pseudo_outcome,
+                pseudo_outcome,
+            )
+
         return pseudo_outcome