Implement WIP versions of ate learners; refactoring

Some refactoring of common code blocks for propensity learners. Setting
and predicting with propensity learners is bundled in the utils to
ensure later on that all propensity learners are properly calibrated
for the task of predicting probabilities.

src/justcause/learners/__init__.py

@@ -4,3 +4,6 @@
 from .meta.xlearner import XLearner  # noqa: F401
 
 from .tree.causal_forest import CausalForest  # noqa: F401
+
+from .ate.double_robust import DoubleRobustEstimator  # noqa: F401
+from .ate.propensity_weighting import PSWEstimator  # noqa: F401

src/justcause/learners/ate/double_robust.py

@@ -1,66 +1,89 @@
 import copy
 
 import numpy as np
-from learners.causal_method import CausalMethod
-from sklearn.calibration import CalibratedClassifierCV
+from sklearn.linear_model import LassoLars
 
-from justcause.utils import get_regressor_name
+from ..utils import fit_predict_propensity, set_propensity_learner
 
 
-class DoubleRobust(CausalMethod):
-    def __init__(self, propensity_regressor, outcome_regressor):
-        super().__init__()
-        self.given_regressor = propensity_regressor
-        self.propensity_regressor = propensity_regressor
-        self.outcome_regressor = outcome_regressor
-        self.outcome_regressor_ctrl = copy.deepcopy(outcome_regressor)
+class DoubleRobustEstimator:
+    """ Implements Double Robust Estmation with generic learners based on the
+        equations of M. Davidian
+
+
+    References:
+        [1] M. Davidian, “Double Robustness in Estimation of Causal Treatment Effects”
+                2007. Presentation
+                http://www.stat.ncsu.edu/∼davidian North
+
+    """
+
+    def __init__(
+        self, propensity_learner=None, learner=None, learner_c=None, learner_t=None
+    ):
+        """
+
+        Args:
+            propensity_learner: propensity regression model
+            learner: generic outcome learner for both outcomes
+            learner_c: specific control outcome learner
+            learner_t: specific treatment outcome learner
+        """
+        self.propensity_learner = set_propensity_learner(propensity_learner)
+
+        if learner is None:
+            if learner_c is None and learner_t is None:
+                self.learner_c = LassoLars()
+                self.learner_t = LassoLars()
+            else:
+                self.learner_c = learner_c
+                self.learner_t = learner_t
+
+        else:
+            self.learner_c = copy.deepcopy(learner)
+            self.learner_t = copy.deepcopy(learner)
+
+        # TODO: This is not very clean here
         self.delta = 0.0001
 
-    def requires_provider(self):
-        return False
+    def __str__(self):
+        """ Simple string representation for logs and outputs"""
+        return "{}(control={}, treated={}, propensity={})".format(
+            self.__class__.__name__,
+            self.learner_c.__class__.__name__,
+            self.learner_t.__class__.__name__,
+            self.propensity_learner.__class__.__name__,
+        )
+
+    def predict_ate(self, x, t, y, propensity=None):
+        """ **Fits** and Predicts average treatment effect of the given population"""
+        # TODO: Out-of-sample prediction makes little sense here
+
+        self.fit(x, t, y)
 
-    def predict_ate(self, x, t=None, y=None):
-        # Predict ATE always for training set, thus test set evaluation is pretty bad
-        if t is not None and y is not None:
-            # Fit for requested set
-            self.fit(x, t, y)
-            self.x = x
-            self.t = t
-            self.y = y
+        # predict propensity
+        if propensity is None:
+            propensity = fit_predict_propensity(self.propensity_learner, x, t)
 
-        prop = self.propensity_regressor.predict_proba(x)[:, 1]
         dr1 = (
             np.sum(
-                ((self.t * self.y) / (prop + self.delta))
-                - ((self.t - prop + self.delta) / (prop + self.delta))
-                * self.outcome_regressor.predict(x)
+                ((t * y) / (propensity + self.delta))
+                - ((t - propensity + self.delta) / (propensity + self.delta))
+                * self.learner_t.predict(x)
             )
             / x.shape[0]
         )
         dr0 = (
             np.sum(
-                ((1 - self.t) * self.y / (1 - prop + self.delta))
-                - ((self.t - prop + self.delta) / (1 - prop + self.delta))
-                * self.outcome_regressor_ctrl.predict(x)
+                ((1 - t) * y / (1 - propensity + self.delta))
+                - ((t - propensity + self.delta) / (1 - propensity + self.delta))
+                * self.learner_c.predict(x)
             )
             / x.shape[0]
         )
         return dr1 - dr0
 
-    def fit(self, x, t, y, refit=False) -> None:
-        # Fit propensity score model
-        self.t = t
-        self.y = y
-        self.propensity_regressor = CalibratedClassifierCV(self.given_regressor)
-        self.propensity_regressor.fit(x, t)
-        # Fit the two outcome models
-        self.outcome_regressor.fit(x[t == 1], y[t == 1])
-        self.outcome_regressor_ctrl.fit(x[t == 0], y[t == 0])
-
-    def __str__(self):
-        return (
-            "DoubleRobustEstimator - P: "
-            + get_regressor_name(self.given_regressor)
-            + " O: "
-            + get_regressor_name(self.outcome_regressor)
-        )
+    def fit(self, x, t, y) -> None:
+        """ Fits the outcome models on treated and control separately"""
+        self.learner_c.fit(x[t == 0], y[t == 0])
+        self.learner_t.fit(x[t == 1], y[t == 1])

src/justcause/learners/ate/propensity_weighting.py

@@ -1,50 +1,36 @@
 import numpy as np
-from learners.causal_method import CausalMethod
-from sklearn.calibration import CalibratedClassifierCV
-
-from justcause.utils import get_regressor_name
-
-
-class PropensityScoreWeighting(CausalMethod):
-    def __init__(self, propensity_regressor):
-        super().__init__()
-        self.given_regressor = propensity_regressor
-        self.propensity_regressor = propensity_regressor
-        self.delta = 0.0001
-
-    def requires_provider(self):
-        return False
-
-    def predict_ate(self, x, t=None, y=None):
-        # Predict ATE always for training set, thus test set evaluation is pretty bad
-        if t is not None and y is not None:
-            # Fit for requested set
-            # self.fit(x, t, y)
-            self.x = x
-            self.t = t
-            self.y = y
-
-        prop = self.propensity_regressor.predict_proba(self.x)[:, 1]
-        m1 = np.sum(
-            ((self.t * self.y + self.delta) / (prop + self.delta)) / self.x.shape[0]
-        )
-        m0 = np.sum(
-            (((1 - self.t) * self.y + self.delta) / (1 - prop + self.delta))
-            / self.x.shape[0]
-        )
-        return m1 - m0
 
-    def fit(self, x, t, y, refit=False) -> None:
-        # Fit propensity score model
-        self.x = x
-        self.t = t
-        self.y = y
-        self.propensity_regressor = CalibratedClassifierCV(self.given_regressor)
-        self.propensity_regressor.fit(x, t)
+from ..utils import fit_predict_propensity, set_propensity_learner
+
+
+class PSWEstimator:
+    """ Implements the simple propensity score weighting estimator """
+
+    def __init__(self, propensity_learner=None):
+        self.propensity_learner = set_propensity_learner(propensity_learner)
+
+        # TODO: Not clean here
+        self.delta = 0.001
 
     def __str__(self):
-        return "PropensityScoreWeighting - " + get_regressor_name(self.given_regressor)
+        return "{}(p_learner={})".format(
+            self.__class__.__name__, self.propensity_learner.__class__.__name__
+        )
+
+    def predict_ate(self, x, t, y, propensity=None):
+        """ Fits the models on given population and calculates the ATE"""
+        # TODO: Discuss: Out-of-sample prediction makes little sense here
+
+        num_samples = x.shape[0]
+
+        if propensity is None:
+            propensity = fit_predict_propensity(self.propensity_learner, x, t)
+
+        m1 = np.sum(((t * y + self.delta) / (propensity + self.delta)))
+        m0 = np.sum((((1 - t) * y + self.delta) / (1 - propensity + self.delta)))
+        return (m1 - m0) / num_samples
 
-    def predict_ite(self, x, t=None, y=None):
-        # Broadcast ATE to all instances
-        return np.full(x.shape[0], self.predict_ate(x, t, y))
+    def fit(self, x, t, y) -> None:
+        """Shell method to avoid errors"""
+        # TODO: Discuss use and convention of ATE learners
+        pass

src/justcause/learners/meta/slearner.py

@@ -1,9 +1,12 @@
 from typing import Optional
 
 import numpy as np
-from causalml.propensity import ElasticNetPropensityModel
 
-from ..utils import replace_factual_outcomes
+from ..utils import (
+    fit_predict_propensity,
+    replace_factual_outcomes,
+    set_propensity_learner,
+)
 
 
 class BaseSLearner(object):
@@ -96,7 +99,6 @@ class WeightedSLearner(BaseSLearner):
     [S-Learner](https://arxiv.org/pdf/1706.03461.pdf)
     """
 
-    # TODO: Add CalibratedCV in order to ensure correct probabilites
     def __init__(self, learner, propensity_learner=None):
         """
         Checks if the given propenstiy_regressor has the predict_proba function
@@ -107,15 +109,7 @@ def __init__(self, learner, propensity_learner=None):
             propensity_learner: the propensity learner fitting x -> p(T | X)
         """
         super().__init__(learner)
-
-        if propensity_learner is None:
-            self.propensity_learner = ElasticNetPropensityModel()
-        else:
-            assert hasattr(
-                propensity_learner, "predict_proba"
-            ), "propensity learner must have predict_proba method"
-
-            self.propensity_learner = propensity_learner
+        self.propensity_learner = set_propensity_learner(propensity_learner)
 
     def __str__(self):
         """ Simple String Representation for logs and outputs"""
@@ -143,17 +137,10 @@ def fit(
             y: factual outcomes
             propensity: propensity scores to be used
         """
-        if propensity is not None:
-            assert len(propensity) == len(t)
-            ipt = 1 / propensity
-        else:
-            if isinstance(self.propensity_learner, ElasticNetPropensityModel):
-                # Use special API of Elastic Net
-                ipt = self.propensity_learner.fit_predict(x, t)
-            else:
-                # Use predict_proba of sklearn classifier
-                self.propensity_learner.fit(x, t)
-                ipt = 1 / self.propensity_learner.predict_proba(x)[:, 1]
+        if propensity is None:
+            propensity = fit_predict_propensity(self.propensity_learner, x, t)
+
+        ipt = 1 / propensity
 
         train = np.c_[x, t]
         self.learner.fit(train, y, sample_weight=ipt)

src/justcause/learners/meta/tlearner.py

@@ -2,10 +2,13 @@
 from typing import Optional, Tuple, Union
 
 import numpy as np
-from causalml.propensity import ElasticNetPropensityModel
 from sklearn.linear_model import LassoLars
 
-from ..utils import replace_factual_outcomes
+from ..utils import (
+    fit_predict_propensity,
+    replace_factual_outcomes,
+    set_propensity_learner,
+)
 
 #: Type alias for predict_ite return type
 SingleComp = Union[Tuple[np.array, np.array, np.array], np.array]
@@ -136,14 +139,7 @@ def __init__(
             random_state:
         """
         super().__init__(learner, learner_c, learner_t, random_state)
-        if propensity_learner is None:
-            self.propensity_learner = ElasticNetPropensityModel()
-        else:
-            assert hasattr(
-                propensity_learner, "predict_proba"
-            ), "propensity learner must have predict_proba method"
-
-            self.propensity_learner = propensity_learner
+        self.propensity_learner = set_propensity_learner(propensity_learner)
 
     def __str__(self):
         """ Simple string representation for logs and outputs"""
@@ -173,11 +169,7 @@ def fit(
             propensity: propensity scores to be used
         """
         if propensity is None:
-            if isinstance(self.propensity_learner, ElasticNetPropensityModel):
-                propensity = self.propensity_learner.fit_predict(x, t)
-            else:
-                self.propensity_learner.fit(x, t)
-                propensity = self.propensity_learner.predict_proba(x)[:, 1]
+            propensity = fit_predict_propensity(self.propensity_learner, x, t)
 
         ipt = 1 / propensity
         self.learner_c.fit(x[t == 0], y[t == 0], sample_weight=ipt[t == 0])

src/justcause/learners/nn/dragonnet_wrapper.py

@@ -1,35 +1,23 @@
-from learners.causal_method import CausalMethod
-
 from justcause.contrib.dragonnet import dragonnet
 
 
-class DragonNetWrapper(CausalMethod):
+class DragonNetWrapper:
     def __init__(
-        self,
-        seed=0,
-        learning_rate=0.001,
-        num_epochs=50,
-        num_covariates=25,
-        batch_size=512,
+        self, learning_rate=0.001, num_epochs=50, num_covariates=25, batch_size=512,
     ):
-        super().__init__(seed)
         self.learning_rate = learning_rate
         self.num_epochs = num_epochs
         self.batch_size = batch_size
         self.num_covariates = num_covariates
 
-    def fit(self, x, t, y, refit=False) -> None:
+    def fit(self, x, t, y) -> None:
         self.model = dragonnet.train_dragon(
             t, y, x, num_epochs=self.num_epochs, batch_size=self.batch_size
         )
 
     def predict_ite(self, x, t=None, y=None):
-        # returns a 4-tupel for each instance : [y_0, y_1, t_pred, epislon]
         res = self.model.predict(x)
         return res[:, 1] - res[:, 0]
 
-    def requires_provider(self):
-        return False
-
     def __str__(self):
         return "DragonNet" + str(self.num_epochs)