Implement adaptive clipping switch (#59)

* Implement adaptive clipping switch

* Add adaptive clipping parameter to init and test

* Add docstring

* Update CHANGELOG

CHANGELOG.rst

@@ -10,6 +10,10 @@ Changelog
 0.6.1 (2024-07-xx)
 ------------------
 
+**New features**
+
+* Add optional ``adaptive_clipping`` parameter to :class:`metalearners.DRLearner`.
+
 **Other changes**
 
 * Changed the index columns order in ``MetaLearnerGridSearch.results_``.

metalearners/drlearner.py

@@ -6,7 +6,16 @@
 from joblib import Parallel, delayed
 from typing_extensions import Self
 
-from metalearners._typing import Matrix, OosMethod, Scoring, Vector
+from metalearners._typing import (
+    Features,
+    Matrix,
+    ModelFactory,
+    OosMethod,
+    Params,
+    Scoring,
+    Vector,
+    _ScikitModel,
+)
 from metalearners._utils import (
     clip_element_absolute_value_to_epsilon,
     get_one,
@@ -15,7 +24,7 @@
     index_matrix,
     validate_valid_treatment_variant_not_control,
 )
-from metalearners.cross_fit_estimator import OVERALL
+from metalearners.cross_fit_estimator import OVERALL, CrossFitEstimator
 from metalearners.metalearner import (
     NUISANCE,
     PROPENSITY_MODEL,
@@ -50,6 +59,9 @@ class DRLearner(_ConditionalAverageOutcomeMetaLearner):
 
         * ``"treatment_model"`` which estimates :math:`\mathbb{E}[Y(k) - Y(0) | X]`
 
+    If ``adaptive_clipping`` is set to ``True``, then the pseudo outcomes are computed using
+    adaptive propensity clipping described in section 4.1, equation *DR-Switch* of
+    `Mahajan et al. (2024) <https://arxiv.org/pdf/2211.01939>`_.
     """
 
     @classmethod
@@ -82,6 +94,40 @@ def _supports_multi_treatment(cls) -> bool:
     def _supports_multi_class(cls) -> bool:
         return False
 
+    def __init__(
+        self,
+        is_classification: bool,
+        n_variants: int,
+        nuisance_model_factory: ModelFactory | None = None,
+        treatment_model_factory: ModelFactory | None = None,
+        propensity_model_factory: type[_ScikitModel] | None = None,
+        nuisance_model_params: Params | dict[str, Params] | None = None,
+        treatment_model_params: Params | dict[str, Params] | None = None,
+        propensity_model_params: Params | None = None,
+        fitted_nuisance_models: dict[str, list[CrossFitEstimator]] | None = None,
+        fitted_propensity_model: CrossFitEstimator | None = None,
+        feature_set: Features | dict[str, Features] | None = None,
+        n_folds: int | dict[str, int] = 10,
+        random_state: int | None = None,
+        adaptive_clipping: bool = False,
+    ):
+        super().__init__(
+            nuisance_model_factory=nuisance_model_factory,
+            is_classification=is_classification,
+            n_variants=n_variants,
+            treatment_model_factory=treatment_model_factory,
+            propensity_model_factory=propensity_model_factory,
+            nuisance_model_params=nuisance_model_params,
+            treatment_model_params=treatment_model_params,
+            propensity_model_params=propensity_model_params,
+            fitted_nuisance_models=fitted_nuisance_models,
+            fitted_propensity_model=fitted_propensity_model,
+            feature_set=feature_set,
+            n_folds=n_folds,
+            random_state=random_state,
+        )
+        self.adaptive_clipping = adaptive_clipping
+
     def fit(
         self,
         X: Matrix,
@@ -317,4 +363,12 @@ def _pseudo_outcome(
             - y0_estimate
         )
 
+        if self.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

tests/test_drlearner.py

@@ -0,0 +1,20 @@
+# Copyright (c) QuantCo 2024-2024
+# SPDX-License-Identifier: BSD-3-Clause
+
+from sklearn.linear_model import LinearRegression, LogisticRegression
+
+from metalearners.drlearner import DRLearner
+
+
+def test_adaptive_clipping_smoke(dummy_dataset):
+    X, y, w = dummy_dataset
+    ml = DRLearner(
+        False,
+        2,
+        LinearRegression,
+        LinearRegression,
+        LogisticRegression,
+        n_folds=2,
+        adaptive_clipping=True,
+    )
+    ml.fit(X, y, w)