RScorer class for causal model selection (#361)

* added rscorer for model selection

* added readme on model selection

Author: vsyrgkanis

README.md

@@ -319,7 +319,7 @@ treatment_effects = est.effect(X_test)
 ```
 </details>
 
- See the <a href="#references">References</a> section for more details.
+See the <a href="#references">References</a> section for more details.
 
 ### Interpretability
 <details>
@@ -370,6 +370,54 @@ treatment_effects = est.effect(X_test)
 
 </details>
 
+
+### Causal Model Selection and Cross-Validation
+
+
+<details>
+  <summary>Causal model selection with the `RScorer` (click to expand)</summary>
+
+  ```Python
+  from econml.score import Rscorer
+
+  # split data in train-validation
+  X_train, X_val, T_train, T_val, Y_train, Y_val = train_test_split(X, T, y, test_size=.4)
+
+  # define list of CATE estimators to select among
+  reg = lambda: RandomForestRegressor(min_samples_leaf=20)
+  clf = lambda: RandomForestClassifier(min_samples_leaf=20)
+  models = [('ldml', LinearDML(model_y=reg(), model_t=clf(), discrete_treatment=True,
+                               linear_first_stages=False, n_splits=3)),
+            ('xlearner', XLearner(models=reg(), cate_models=reg(), propensity_model=clf())),
+            ('dalearner', DomainAdaptationLearner(models=reg(), final_models=reg(), propensity_model=clf())),
+            ('slearner', SLearner(overall_model=reg())),
+            ('drlearner', DRLearner(model_propensity=clf(), model_regression=reg(),
+                                    model_final=reg(), n_splits=3)),
+            ('rlearner', NonParamDML(model_y=reg(), model_t=clf(), model_final=reg(),
+                                     discrete_treatment=True, n_splits=3)),
+            ('dml3dlasso', DML(model_y=reg(), model_t=clf(),
+                               model_final=LassoCV(cv=3, fit_intercept=False),
+                               discrete_treatment=True,
+                               featurizer=PolynomialFeatures(degree=3),
+                               linear_first_stages=False, n_splits=3))
+  ]
+
+  # fit cate models on train data
+  models = [(name, mdl.fit(Y_train, T_train, X=X_train)) for name, mdl in models]
+
+  # score cate models on validation data
+  scorer = RScorer(model_y=reg(), model_t=clf(),
+                   discrete_treatment=True, n_splits=3, mc_iters=2, mc_agg='median')
+  scorer.fit(Y_val, T_val, X=X_val)
+  rscore = [scorer.score(mdl) for _, mdl in models]
+  # select the best model
+  mdl, _ = scorer.best_model([mdl for _, mdl in models])
+  # create weighted ensemble model based on score performance
+  mdl, _ = scorer.ensemble([mdl for _, mdl in models])
+  ```
+
+</details>
+
 ### Inference
 
 Whenever inference is enabled, then one can get a more structure `InferenceResults` object with more elaborate inference information, such

doc/reference.rst

@@ -14,6 +14,7 @@ Public Module Reference
     econml.metalearners
     econml.ortho_forest
     econml.ortho_iv
+    econml.score
     econml.two_stage_least_squares
     econml.utilities
 

econml/__init__.py

@@ -3,5 +3,5 @@
            'cate_interpreter', 'causal_forest',
            'data', 'deepiv', 'dml', 'drlearner', 'inference',
            'metalearners', 'ortho_forest', 'ortho_iv',
-           'sklearn_extensions', 'tree',
+           'score', 'sklearn_extensions', 'tree',
            'two_stage_least_squares', 'utilities']

econml/dml/__init__.py

@@ -7,30 +7,31 @@
 (aka residual outcome and residual treatment).
 Then estimates a CATE model by regressing the residual outcome on the residual treatment
 in a manner that accounts for heterogeneity in the regression coefficient, with respect
-to X.
+to X. For the theoretical foundations of these methods see [dml]_, [rlearner]_, [paneldml]_,
+[lassodml]_, [ortholearner]_.
 
 References
 ----------
 
-\\ V. Chernozhukov, D. Chetverikov, M. Demirer, E. Duflo, C. Hansen, and a. W. Newey.
+.. [dml] V. Chernozhukov, D. Chetverikov, M. Demirer, E. Duflo, C. Hansen, and a. W. Newey.
     Double Machine Learning for Treatment and Causal Parameters.
-    https://arxiv.org/abs/1608.00060, 2016.
+    `<https://arxiv.org/abs/1608.00060>`_, 2016.
 
-\\ X. Nie and S. Wager.
+.. [rlearner] X. Nie and S. Wager.
     Quasi-Oracle Estimation of Heterogeneous Treatment Effects.
-    arXiv preprint arXiv:1712.04912, 2017. URL http://arxiv.org/abs/1712.04912.
+    arXiv preprint arXiv:1712.04912, 2017. URL `<http://arxiv.org/abs/1712.04912>`_.
 
-\\ V. Chernozhukov, M. Goldman, V. Semenova, and M. Taddy.
+.. [paneldml] V. Chernozhukov, M. Goldman, V. Semenova, and M. Taddy.
     Orthogonal Machine Learning for Demand Estimation: High Dimensional Causal Inference in Dynamic Panels.
-    https://arxiv.org/abs/1712.09988, December 2017.
+    `<https://arxiv.org/abs/1712.09988>`_, December 2017.
 
-\\ V. Chernozhukov, D. Nekipelov, V. Semenova, and V. Syrgkanis.
+.. [lassodml] V. Chernozhukov, D. Nekipelov, V. Semenova, and V. Syrgkanis.
     Two-Stage Estimation with a High-Dimensional Second Stage.
-    https://arxiv.org/abs/1806.04823, 2018.
+    `<https://arxiv.org/abs/1806.04823>`_, 2018.
 
-\\ Dylan Foster, Vasilis Syrgkanis (2019).
+.. [ortholearner] Dylan Foster, Vasilis Syrgkanis (2019).
     Orthogonal Statistical Learning.
-    ACM Conference on Learning Theory. https://arxiv.org/abs/1901.09036
+    ACM Conference on Learning Theory. `<https://arxiv.org/abs/1901.09036>`_
 
 """
 

econml/grf/__init__.py

@@ -1,6 +1,16 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
+""" An efficient Cython implementation of Generalized Random Forests [grf]_ and special
+case python classes.
+
+References
+----------
+.. [grf] Athey, Susan, Julie Tibshirani, and Stefan Wager. "Generalized random forests."
+    The Annals of Statistics 47.2 (2019): 1148-1178
+    https://arxiv.org/pdf/1610.01271.pdf
+"""
+
 from ._criterion import LinearMomentGRFCriterion, LinearMomentGRFCriterionMSE
 from .classes import CausalForest, CausalIVForest, RegressionForest, MultiOutputGRF
 

econml/metalearners.py

@@ -3,7 +3,7 @@
 
 """Metalearners for heterogeneous treatment effects in the context of discrete treatments.
 
-For more details on these CATE methods, see <https://arxiv.org/abs/1706.03461>
+For more details on these CATE methods, see `<https://arxiv.org/abs/1706.03461>`_
 (Künzel S., Sekhon J., Bickel P., Yu B.) on Arxiv.
 """
 

econml/score/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""
+A suite of scoring methods for scoring CATE models out-of-sample for the
+purpose of model selection.
+"""
+
+from .rscorer import RScorer
+from .ensemble_cate import EnsembleCateEstimator
+
+__all__ = ['RScorer',
+           'EnsembleCateEstimator']

econml/score/ensemble_cate.py

@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import numpy as np
+from sklearn.utils.validation import check_array
+from .._cate_estimator import BaseCateEstimator, LinearCateEstimator
+
+
+class EnsembleCateEstimator:
+    """ A CATE estimator that represents a weighted ensemble of many
+    CATE estimators. Returns their weighted effect prediction.
+
+    Parameters
+    ----------
+    cate_models : list of BaseCateEstimator objects
+        A list of fitted cate estimator objects that will be used in the ensemble.
+        The models are passed by reference, and not copied internally, because we
+        need the fitted objects, so any change to the passed models will affect
+        the internal predictions (e.g. if the input models are refitted).
+    weights : np.ndarray of shape (len(cate_models),)
+        The weight placed on each model. Weights must be non-positive. The
+        ensemble will predict effects based on the weighted average predictions
+        of the cate_models estiamtors, weighted by the corresponding weight in `weights`.
+    """
+
+    def __init__(self, *, cate_models, weights):
+        self.cate_models = cate_models
+        self.weights = weights
+
+    def effect(self, X=None, *, T0=0, T1=1):
+        return np.average([mdl.effect(X=X, T0=T0, T1=T1) for mdl in self.cate_models],
+                          weights=self.weights, axis=0)
+    effect.__doc__ = BaseCateEstimator.effect.__doc__
+
+    def marginal_effect(self, T, X=None):
+        return np.average([mdl.marginal_effect(T, X=X) for mdl in self.cate_models],
+                          weights=self.weights, axis=0)
+    marginal_effect.__doc__ = BaseCateEstimator.marginal_effect.__doc__
+
+    def const_marginal_effect(self, X=None):
+        if np.any([not hasattr(mdl, 'const_marginal_effect') for mdl in self.cate_models]):
+            raise ValueError("One of the base CATE models in parameter `cate_models` does not support "
+                             "the `const_marginal_effect` method.")
+        return np.average([mdl.const_marginal_effect(X=X) for mdl in self.cate_models],
+                          weights=self.weights, axis=0)
+    const_marginal_effect.__doc__ = LinearCateEstimator.const_marginal_effect.__doc__
+
+    @property
+    def cate_models(self):
+        return self._cate_models
+
+    @cate_models.setter
+    def cate_models(self, value):
+        if (not isinstance(value, list)) or (not np.all([isinstance(model, BaseCateEstimator) for model in value])):
+            raise ValueError('Parameter `cate_models` should be a list of `BaseCateEstimator` objects.')
+        self._cate_models = value
+
+    @property
+    def weights(self):
+        return self._weights
+
+    @weights.setter
+    def weights(self, value):
+        weights = check_array(value, accept_sparse=False, ensure_2d=False, allow_nd=False, dtype='numeric',
+                              force_all_finite=True)
+        if np.any(weights < 0):
+            raise ValueError("All weights in parameter `weights` must be non-negative.")
+        self._weights = weights