feat: add S-Learner, T-Learner, and X-Learner implementations with ut…

foundry/__init__.py

@@ -1 +1 @@
-__version__ = '0.2.6'
+__version__ = '0.2.7'

foundry/uplift/__init__.py

@@ -0,0 +1 @@
+from .meta_learners import SLearner, TLearner, XLearner

foundry/uplift/meta_learners/__init__.py

@@ -0,0 +1,3 @@
+from .t_learner import TLearner
+from .s_learner import SLearner
+from .x_learner import XLearner

foundry/uplift/meta_learners/s_learner.py

@@ -0,0 +1,114 @@
+import warnings
+from typing import Tuple, Union, overload, Literal, Optional
+
+import numpy as np
+import pandas as pd
+from sklearn.base import BaseEstimator, clone
+
+from foundry.util import SliceDict, to_1d, safe_predict
+from ..util import get_qini_curve, get_cumulative_gain_score
+
+
+class SLearner(BaseEstimator):
+    """
+    An S-learner. Please note that current implementation assumes randomized treatment/control!
+
+    :param estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it).
+    :param include_interaction: Whether to include X * treatment interaction terms.
+    """
+    estimator_: Optional[BaseEstimator] = None
+
+    def __init__(self, estimator: BaseEstimator, include_interaction: bool = True) -> None:
+        self.estimator = estimator
+        self.include_interaction = include_interaction
+
+    def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict, **fit_kwargs) -> "SLearner":
+        y_arr, treatment_ind = self._normalize_y(y)
+
+        X_aug = self._augment_with_treatment(X, treatment_ind)
+
+        self.estimator_ = clone(self.estimator).fit(X=X_aug, y=y_arr, **fit_kwargs)
+        return self
+
+
+    @overload
+    def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[False] = ..., **predict_kwargs) -> np.ndarray: ...
+    @overload
+    def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[True], **predict_kwargs) -> Tuple[np.ndarray, np.ndarray]: ...
+
+    def predict(self, X, return_components=False, **predict_kwargs):
+        X_t = self._augment_with_treatment(X, np.ones(len(X), dtype=bool))
+        X_c = self._augment_with_treatment(X, np.zeros(len(X), dtype=bool))
+
+        yhat_t = safe_predict(self.estimator_, X=X_t, **predict_kwargs)
+        yhat_c = safe_predict(self.estimator_, X=X_c, **predict_kwargs)
+
+        if return_components:
+            return yhat_t, yhat_c
+
+        return yhat_t - yhat_c
+
+    def score(
+        self,
+        X: Union[pd.DataFrame, np.ndarray],
+        y: SliceDict,
+        sample_weight: Union[np.ndarray, None] = None,
+        method: str = 'qini',
+        normalize: bool = True,
+        **kwargs,
+    ) -> float:
+        y_arr, treatment_ind = self._normalize_y(y)
+        if sample_weight is not None:
+            raise NotImplementedError
+        pred = self.predict(X=X)
+        if method == 'cumulative_gain':
+            return get_cumulative_gain_score(
+                y_true=y_arr,
+                treatment=treatment_ind,
+                score=pred,
+                **kwargs,
+            )
+        qini = get_qini_curve(
+            y_true=y_arr,
+            treatment=treatment_ind,
+            score=pred,
+            normalize=normalize,
+            **kwargs,
+        )
+        random_area = np.linspace(0, qini[-1], qini.shape[0]).sum()
+        return (np.nansum(qini) - random_area) / qini.shape[0]
+
+    def _augment_with_treatment(
+        self,
+        X: Union[pd.DataFrame, np.ndarray],
+        treatment_ind: np.ndarray,
+    ) -> Union[pd.DataFrame, np.ndarray]:
+        treatment_col = treatment_ind.astype(int)
+
+        if isinstance(X, pd.DataFrame):
+            X_aug = X.copy()
+            X_aug["treatment"] = treatment_col
+
+            if self.include_interaction:
+                for col in X.columns:
+                    X_aug[f"{col}_x_treatment"] = X[col] * treatment_col
+
+            return X_aug
+
+        else:
+            treatment_col = treatment_col.reshape(-1, 1)
+
+            if self.include_interaction:
+                interaction_terms = X * treatment_col
+                return np.hstack([X, treatment_col, interaction_terms])
+
+            return np.hstack([X, treatment_col])
+
+    @staticmethod
+    def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]:
+        y = y.copy()
+        y_arr = to_1d(np.asanyarray(y.pop("value")))
+        treatment_ind = to_1d(np.asanyarray(y.pop("is_treatment")).astype(bool))
+        if len(y.keys()):
+            warnings.warn(f"Unused keys in ``y``: {set(y)}")
+        return y_arr, treatment_ind

foundry/uplift/meta_learners/t_learner.py

@@ -0,0 +1,80 @@
+import warnings
+from typing import Optional, Tuple, Union, overload, Literal
+
+import numpy as np
+import pandas as pd
+from sklearn.base import BaseEstimator, clone
+
+from foundry.util import SliceDict, to_1d, safe_predict
+from ..util import get_qini_curve, get_cumulative_gain_score
+
+
+class TLearner(BaseEstimator):
+    """
+    A T-learner. Please note that current implementation assumes randomized treatment/control!
+
+    :param estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it).
+    """
+    treatment_est_: Optional[BaseEstimator] = None
+    control_est_: Optional[BaseEstimator] = None
+
+    def __init__(self, estimator: BaseEstimator) -> None:
+        self.estimator = estimator
+
+    def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict, **fit_kwargs) -> "TLearner":
+        y_arr, treatment_ind = self._normalize_y(y)
+        self.treatment_est_ = clone(self.estimator).fit(X=X[treatment_ind], y=y_arr[treatment_ind], **fit_kwargs)
+        self.control_est_ = clone(self.estimator).fit(X[~treatment_ind], y_arr[~treatment_ind], **fit_kwargs)
+
+        return self
+
+    @overload
+    def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[False] = ..., **predict_kwargs) -> np.ndarray: ...
+    @overload
+    def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[True], **predict_kwargs) -> Tuple[np.ndarray, np.ndarray]: ...
+
+    def predict(self, X, return_components=False, **predict_kwargs):
+        yhat_t = safe_predict(self.treatment_est_, X=X, **predict_kwargs)
+        yhat_c = safe_predict(self.control_est_, X=X, **predict_kwargs)
+        if return_components:
+            return yhat_t, yhat_c
+        return yhat_t - yhat_c
+
+    def score(
+        self,
+        X: Union[pd.DataFrame, np.ndarray],
+        y: SliceDict,
+        sample_weight: Optional[np.ndarray] = None,
+        method: str = 'qini',
+        normalize: bool = True,
+        **kwargs,
+    ) -> float:
+        y_arr, treatment_ind = self._normalize_y(y)
+        if sample_weight is not None:
+            raise NotImplementedError
+        pred = self.predict(X=X)
+        if method == 'cumulative_gain':
+            return get_cumulative_gain_score(
+                y_true=y_arr,
+                treatment=treatment_ind,
+                score=pred,
+                **kwargs,
+            )
+        qini = get_qini_curve(
+            y_true=y_arr,
+            treatment=treatment_ind,
+            score=pred,
+            normalize=normalize,
+            **kwargs
+        )
+        random_area = np.linspace(0, qini[-1], qini.shape[0]).sum()
+        return (np.nansum(qini) - random_area) / qini.shape[0]
+
+    @staticmethod
+    def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]:
+        y = y.copy()
+        y_arr = to_1d(np.asanyarray(y.pop('value')))
+        treatment_ind = to_1d(np.asanyarray(y.pop('is_treatment')).astype(bool))
+        if len(y.keys()):
+            warnings.warn(f"Unused keys in ``y``: {set(y)}")
+        return y_arr, treatment_ind

foundry/uplift/meta_learners/x_learner.py

@@ -0,0 +1,133 @@
+import warnings
+from typing import Any, Dict, Optional, Tuple, Union, overload, Literal
+
+import numpy as np
+import pandas as pd
+from sklearn.base import BaseEstimator, clone
+
+from foundry.util import SliceDict, to_1d, safe_predict
+from ..util import get_qini_curve, get_cumulative_gain_score
+
+
+class XLearner(BaseEstimator):
+    """
+    An X-learner. Please note that current implementation assumes randomized treatment/control!
+    Adapted from https://matheusfacure.github.io/python-causality-handbook/21-Meta-Learners.html 
+    and the original paper by Kunzel et al. (2019): https://arxiv.org/abs/1706.03461
+
+    :param first_stage_estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it).
+    :param second_stage_estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it).
+    :param propensity_estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it).
+    :param first_stage_fit_params: Optional dict of kwargs passed to ``fit()`` for the first-stage models.
+    :param second_stage_fit_params: Optional dict of kwargs passed to ``fit()`` for the second-stage models.
+    :param propensity_fit_params: Optional dict of kwargs passed to ``fit()`` for the propensity model.
+    """
+    first_treatment_est_: Optional[BaseEstimator] = None
+    first_control_est_: Optional[BaseEstimator] = None
+    second_treatment_est_: Optional[BaseEstimator] = None
+    second_control_est_: Optional[BaseEstimator] = None
+    propensity_est_: Optional[BaseEstimator] = None
+
+    def __init__(
+        self,
+        first_stage_estimator: BaseEstimator,
+        second_stage_estimator: BaseEstimator,
+        propensity_estimator: BaseEstimator,
+        first_stage_fit_params: Optional[Dict[str, Any]] = None,
+        second_stage_fit_params: Optional[Dict[str, Any]] = None,
+        propensity_fit_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self.first_stage_estimator = first_stage_estimator
+        self.second_stage_estimator = second_stage_estimator
+        self.propensity_estimator = propensity_estimator
+        self.first_stage_fit_params = first_stage_fit_params
+        self.second_stage_fit_params = second_stage_fit_params
+        self.propensity_fit_params = propensity_fit_params
+
+    def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict) -> "XLearner":
+        y_arr, treatment_ind = self._normalize_y(y)
+
+        _first_stage_fit_params = self.first_stage_fit_params or {}
+        _second_stage_fit_params = self.second_stage_fit_params or {}
+        _propensity_fit_params = self.propensity_fit_params or {}
+
+        self.first_control_est_ = clone(self.first_stage_estimator).fit(
+            X[~treatment_ind], y_arr[~treatment_ind], **_first_stage_fit_params
+        )
+        self.first_treatment_est_ = clone(self.first_stage_estimator).fit(
+            X[treatment_ind], y_arr[treatment_ind], **_first_stage_fit_params
+        )
+
+        self.propensity_est_ = clone(self.propensity_estimator).fit(
+            X, treatment_ind, **_propensity_fit_params
+        )
+
+        imputed_te = np.where(
+            treatment_ind,
+            y_arr - safe_predict(self.first_control_est_, X),
+            safe_predict(self.first_treatment_est_, X) - y_arr,
+        )
+
+        self.second_control_est_ = clone(self.second_stage_estimator).fit(
+            X[~treatment_ind], imputed_te[~treatment_ind], **_second_stage_fit_params
+        )
+        self.second_treatment_est_ = clone(self.second_stage_estimator).fit(
+            X[treatment_ind], imputed_te[treatment_ind], **_second_stage_fit_params
+        )
+
+        return self
+
+    @overload
+    def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[False] = ..., **predict_kwargs) -> np.ndarray: ...
+    @overload
+    def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[True], **predict_kwargs) -> Tuple[np.ndarray, np.ndarray]: ...
+
+    def predict(self, X, return_components=False, **predict_kwargs):
+        p_treatment = safe_predict(self.propensity_est_, X)
+        p_control = 1 - p_treatment
+
+        tau0 = safe_predict(self.second_control_est_,   X, **predict_kwargs)
+        tau1 = safe_predict(self.second_treatment_est_, X, **predict_kwargs)
+
+        if return_components:
+            return tau0, tau1
+        return p_treatment * tau0 + p_control * tau1
+
+    def score(
+        self,
+        X: Union[pd.DataFrame, np.ndarray],
+        y: SliceDict,
+        sample_weight: Optional[np.ndarray] = None,
+        method: str = 'qini',
+        normalize: bool = True,
+        **kwargs,
+    ) -> float:
+        y_arr, treatment_ind = self._normalize_y(y)
+        if sample_weight is not None:
+            raise NotImplementedError
+        pred = self.predict(X=X)
+        if method == 'cumulative_gain':
+            return get_cumulative_gain_score(
+                y_true=y_arr,
+                treatment=treatment_ind,
+                score=pred,
+                **kwargs,
+            )
+        qini = get_qini_curve(
+            y_true=y_arr,
+            treatment=treatment_ind,
+            score=pred,
+            normalize=normalize,
+            **kwargs,
+        )
+        random_area = np.linspace(0, qini[-1], qini.shape[0]).sum()
+        return (np.nansum(qini) - random_area) / qini.shape[0]
+
+    @staticmethod
+    def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]:
+        y = y.copy()
+        y_arr = to_1d(np.asanyarray(y.pop('value')))
+        treatment_ind = to_1d(np.asanyarray(y.pop('is_treatment')).astype(bool))
+        if len(y.keys()):
+            warnings.warn(f"Unused keys in ``y``: {set(y)}")
+        return y_arr, treatment_ind