Add attributes for nuisance scores

econml/_ortho_learner.py

@@ -76,7 +76,7 @@ def _crossfit(model, folds, *args, **kwargs):
     -------
     nuisances : tuple of numpy matrices
         Each entry in the tuple is a nuisance parameter matrix. Each row i-th in the
-        matric corresponds to the value of the nuisancee parameter for the i-th input
+        matrix corresponds to the value of the nuisancee parameter for the i-th input
         sample.
     model_list : list of objects of same type as input model
         The cloned and fitted models for each fold. Can be used for inspection of the
@@ -85,6 +85,8 @@ def _crossfit(model, folds, *args, **kwargs):
         The indices of the arrays for which the nuisance value was calculated. This
         corresponds to the union of the indices of the test part of each fold in
         the input fold list.
+    scores : tuple of list of float or None
+        The out-of-sample model scores for each nuisance model
 
     Examples
     --------
@@ -108,7 +110,7 @@ def predict(self, X, y, W=None):
         y = X[:, 0] + np.random.normal(size=(5000,))
         folds = list(KFold(2).split(X, y))
         model = Lasso(alpha=0.01)
-        nuisance, model_list, fitted_inds = _crossfit(Wrapper(model), folds, X, y, W=y, Z=None)
+        nuisance, model_list, fitted_inds, scores = _crossfit(Wrapper(model), folds, X, y, W=y, Z=None)
 
     >>> nuisance
     (array([-1.105728... , -1.537566..., -2.451827... , ...,  1.106287...,
@@ -121,18 +123,22 @@ def predict(self, X, y, W=None):
     """
     model_list = []
     fitted_inds = []
+    calculate_scores = hasattr(model, 'score')
 
     if folds is None:  # skip crossfitting
         model_list.append(clone(model, safe=False))
         kwargs = {k: v for k, v in kwargs.items() if v is not None}
         model_list[0].fit(*args, **kwargs)
         nuisances = model_list[0].predict(*args, **kwargs)
+        scores = model_list[0].score(*args, **kwargs) if calculate_scores else None
 
         if not isinstance(nuisances, tuple):
             nuisances = (nuisances,)
+        if not isinstance(scores, tuple):
+            scores = (scores,)
 
         first_arr = args[0] if args else kwargs.items()[0][1]
-        return nuisances, model_list, np.arange(first_arr.shape[0])
+        return nuisances, model_list, np.arange(first_arr.shape[0]), scores
 
     for idx, (train_idxs, test_idxs) in enumerate(folds):
         model_list.append(clone(model, safe=False))
@@ -162,7 +168,19 @@ def predict(self, X, y, W=None):
         for it, nuis in enumerate(nuisance_temp):
             nuisances[it][test_idxs] = nuis
 
-    return nuisances, model_list, np.sort(fitted_inds.astype(int))
+        if calculate_scores:
+            score_temp = model_list[idx].score(*args_test, **kwargs_test)
+
+            if not isinstance(score_temp, tuple):
+                score_temp = (score_temp,)
+
+            if idx == 0:
+                scores = tuple([] for _ in score_temp)
+
+            for it, score in enumerate(score_temp):
+                scores[it].append(score)
+
+    return nuisances, model_list, np.sort(fitted_inds.astype(int)), (scores if calculate_scores else None)
 
 
 class _OrthoLearner(TreatmentExpansionMixin, LinearCateEstimator):
@@ -235,6 +253,9 @@ class _OrthoLearner(TreatmentExpansionMixin, LinearCateEstimator):
         methods. If ``discrete_treatment=True``, then the input ``T`` to both above calls will be the
         one-hot encoding of the original input ``T``, excluding the first column of the one-hot.
 
+        If the estimator also provides a score method with the same arguments as fit, it will be used to
+        calculate scores during training.
+
     model_final: estimator for fitting the response residuals to the features and treatment residuals
         Must implement `fit` and `predict` methods that must have signatures::
 
@@ -409,6 +430,8 @@ def score(self, Y, T, W=None, nuisances=None):
         If the model_final has a score method, then `score_` contains the outcome of the final model
         score when evaluated on the fitted nuisances from the first stage. Represents goodness of fit,
         of the final CATE model.
+    nuisance_scores_ : tuple of lists of floats or None
+        The out-of-sample scores from training each nuisance model
     """
 
     def __init__(self, model_nuisance, model_final, *,
@@ -561,9 +584,10 @@ def _fit_nuisances(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                           reshape(self._label_encoder.transform(T.ravel()), (-1, 1)))[:, 1:]),
                 validate=False)
 
-        nuisances, fitted_models, fitted_inds = _crossfit(self._model_nuisance, folds,
-                                                          Y, T, X=X, W=W, Z=Z, sample_weight=sample_weight)
+        nuisances, fitted_models, fitted_inds, scores = _crossfit(self._model_nuisance, folds,
+                                                                  Y, T, X=X, W=W, Z=Z, sample_weight=sample_weight)
         self._models_nuisance = fitted_models
+        self.nuisance_scores_ = scores
         return nuisances, fitted_inds
 
     def _fit_final(self, Y, T, X=None, W=None, Z=None, nuisances=None, sample_weight=None, sample_var=None):

econml/_rlearner.py

@@ -185,6 +185,8 @@ def predict(self, X):
         An instance of the model_final object that was fitted after calling fit.
     score_ : float
         The MSE in the final residual on residual regression, i.e.
+    nuisance_scores_ : list of float
+        The out-of-sample scores for each nuisance model
 
         .. math::
             \\frac{1}{n} \\sum_{i=1}^n (Y_i - \\hat{E}[Y|X_i, W_i]\
@@ -213,6 +215,17 @@ def fit(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 self._model_y.fit(X, W, Y, sample_weight=sample_weight)
                 return self
 
+            def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
+                if hasattr(self._model_y, 'score'):
+                    Y_score = self._model_y.score(X, W, Y, sample_weight=sample_weight)
+                else:
+                    Y_score = None
+                if hasattr(self._model_t, 'score'):
+                    T_score = self._model_t.score(X, W, T, sample_weight=sample_weight)
+                else:
+                    T_score = None
+                return Y_score, T_score
+
             def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 Y_pred = self._model_y.predict(X, W)
                 T_pred = self._model_t.predict(X, W)
@@ -334,3 +347,11 @@ def models_y(self):
     @property
     def models_t(self):
         return [mdl._model_t for mdl in super().models_nuisance]
+
+    @property
+    def nuisance_scores_y(self):
+        return self.nuisance_scores_[0]
+
+    @property
+    def nuisance_scores_t(self):
+        return self.nuisance_scores_[1]

econml/dml.py

@@ -100,6 +100,17 @@ def predict(self, X, W):
         else:
             return self._model.predict(self._combine(X, W, n_samples, fitting=False))
 
+    def score(self, X, W, Target, sample_weight=None):
+        if (not self._is_Y) and self._discrete_treatment:
+            # In this case, the Target is the one-hot-encoding of the treatment variable
+            # We need to go back to the label representation of the one-hot so as to call
+            # the classifier.
+            Target = inverse_onehot(Target)
+        if sample_weight is not None:
+            return self._model.score(self._combine(X, W, Target.shape[0]), Target, sample_weight=sample_weight)
+        else:
+            return self._model.score(self._combine(X, W, Target.shape[0]), Target)
+
 
 class _FinalWrapper:
     def __init__(self, model_final, fit_cate_intercept, featurizer, use_weight_trick):

econml/drlearner.py

@@ -275,6 +275,21 @@ def fit(self, Y, T, X=None, W=None, *, sample_weight=None):
                 self._model_regression.fit(np.hstack([XW, T]), Y, **filtered_kwargs)
                 return self
 
+            def score(self, Y, T, X=None, W=None, *, sample_weight=None):
+                XW = self._combine(X, W)
+                filtered_kwargs = _filter_none_kwargs(sample_weight=sample_weight)
+
+                if hasattr(self._model_propensity, 'score'):
+                    propensity_score = self._model_propensity.score(XW, inverse_onehot(T), **filtered_kwargs)
+                else:
+                    propensity_score = None
+                if hasattr(self._model_regression, 'score'):
+                    regression_score = self._model_regression.fit(np.hstack([XW, T]), Y, **filtered_kwargs)
+                else:
+                    regression_score = None
+
+                return propensity_score, regression_score
+
             def predict(self, Y, T, X=None, W=None, *, sample_weight=None):
                 XW = self._combine(X, W)
                 propensities = np.maximum(self._model_propensity.predict_proba(XW), min_propensity)
@@ -472,6 +487,16 @@ def models_regression(self):
         """
         return [mdl._model_regression for mdl in super().models_nuisance]
 
+    @property
+    def nuisance_scores_propensity(self):
+        """Gets the score for the propensity model on out-of-sample training data"""
+        return self.nuisance_scores_[0]
+
+    @property
+    def nuisance_scores_regression(self):
+        """Gets the score for the regression model on out-of-sample training data"""
+        return self.nuisance_scores_[1]
+
     @property
     def featurizer(self):
         """

econml/ortho_iv.py

@@ -61,6 +61,26 @@ def fit(self, X, *args, sample_weight=None):
         else:
             self._model.fit(self._combine(X, Z, Target.shape[0]), Target)
 
+    def score(self, X, *args, sample_weight=None):
+        if len(args) == 1:
+            Target, = args
+            Z = None
+        else:
+            (Z, Target) = args
+        if self._discrete_target:
+            # In this case, the Target is the one-hot-encoding of the treatment variable
+            # We need to go back to the label representation of the one-hot so as to call
+            # the classifier.
+            if np.any(np.all(Target == 0, axis=0)) or (not np.any(np.all(Target == 0, axis=1))):
+                raise AttributeError("Provided crossfit folds contain training splits that " +
+                                     "don't contain all treatments")
+            Target = inverse_onehot(Target)
+
+        if sample_weight is not None:
+            return self._model.score(self._combine(X, Z, Target.shape[0]), Target, sample_weight=sample_weight)
+        else:
+            return self._model.score(self._combine(X, Z, Target.shape[0]), Target)
+
     def predict(self, X, Z=None):
         n_samples = X.shape[0] if X is not None else (Z.shape[0] if Z is not None else 1)
         if self._discrete_target:
@@ -203,6 +223,21 @@ def fit(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 self._model_Z_X.fit(W, Z, sample_weight=sample_weight)
                 return self
 
+            def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
+                if hasattr(self._model_Y_X, 'score'):
+                    Y_X_score = self._model_Y_X.score(W, Y, sample_weight=sample_weight)
+                else:
+                    Y_X_score = None
+                if hasattr(self._model_T_X, 'score'):
+                    T_X_score = self._model_T_X.score(W, T, sample_weight=sample_weight)
+                else:
+                    T_X_score = None
+                if hasattr(self._model_Z_X, 'score'):
+                    Z_X_score = self._model_Z_X.score(W, Z, sample_weight=sample_weight)
+                else:
+                    Z_X_score = None
+                return Y_X_score, T_X_score, Z_X_score
+
             def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 Y_pred = self._model_Y_X.predict(W)
                 T_pred = self._model_T_X.predict(W)
@@ -239,6 +274,21 @@ def fit(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 self._model_T_XZ.fit(W, Z, T, sample_weight=sample_weight)
                 return self
 
+            def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
+                if hasattr(self._model_Y_X, 'score'):
+                    Y_X_score = self._model_Y_X.score(W, Y, sample_weight=sample_weight)
+                else:
+                    Y_X_score = None
+                if hasattr(self._model_T_X, 'score'):
+                    T_X_score = self._model_T_X.score(W, T, sample_weight=sample_weight)
+                else:
+                    T_X_score = None
+                if hasattr(self._model_T_XZ, 'score'):
+                    T_XZ_score = self._model_T_XZ.score(W, Z, T, sample_weight=sample_weight)
+                else:
+                    T_XZ_score = None
+                return Y_X_score, T_X_score, T_XZ_score
+
             def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 Y_pred = self._model_Y_X.predict(W)
                 TX_pred = self._model_T_X.predict(W)
@@ -344,6 +394,21 @@ def fit(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 self._model_T_XZ.fit(X, Z, T, sample_weight=sample_weight)
                 return self
 
+            def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
+                if hasattr(self._model_Y_X, 'score'):
+                    Y_X_score = self._model_Y_X.score(X, Y, sample_weight=sample_weight)
+                else:
+                    Y_X_score = None
+                if hasattr(self._model_T_X, 'score'):
+                    T_X_score = self._model_T_X.score(X, T, sample_weight=sample_weight)
+                else:
+                    T_X_score = None
+                if hasattr(self._model_T_XZ, 'score'):
+                    T_XZ_score = self._model_T_XZ.score(X, Z, T, sample_weight=sample_weight)
+                else:
+                    T_XZ_score = None
+                return Y_X_score, T_X_score, T_XZ_score
+
             def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 Y_pred = self._model_Y_X.predict(X)
                 TXZ_pred = self._model_T_XZ.predict(X, Z)
@@ -521,6 +586,27 @@ def models_T_XZ(self):
         """
         return [mdl._model for mdl in super().models_T_XZ]
 
+    @property
+    def nuisance_scores_Y_X(self):
+        """
+        Get the scores for Y_X model on the out-of-sample training data
+        """
+        return self.nuisance_scores_[0]
+
+    @property
+    def nuisance_scores_T_X(self):
+        """
+        Get the scores for T_X model on the out-of-sample training data
+        """
+        return self.nuisance_scores_[1]
+
+    @property
+    def nuisance_scores_T_XZ(self):
+        """
+        Get the scores for T_XZ model on the out-of-sample training data
+        """
+        return self.nuisance_scores_[2]
+
     def cate_feature_names(self, input_feature_names=None):
         """
         Get the output feature names.
@@ -989,6 +1075,28 @@ def fit(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 self._prel_model_effect.fit(Y, inverse_onehot(T), inverse_onehot(Z), X=X)
                 return self
 
+            def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
+                if hasattr(self._model_Y_X, 'score'):
+                    Y_X_score = self._model_Y_X.score(X, Y, sample_weight=sample_weight)
+                else:
+                    Y_X_score = None
+                if hasattr(self._model_T_XZ, 'score'):
+                    T_XZ_score = self._model_T_XZ.score(X, Z, T, sample_weight=sample_weight)
+                else:
+                    T_XZ_score = None
+                if hasattr(self._prel_model_effect, 'score'):
+                    # we need to undo the one-hot encoding for calling effect,
+                    # since it expects raw values
+                    # also note that score here is defined on _OrthoLearner
+                    # whereas fit is calling the fit override from this class with reordered arguments,
+                    # which is why the signature used here is different from the call to _prel_model_effect.fit
+                    # above in fit
+                    effect_score = self._prel_model_effect.score(Y, inverse_onehot(T), X=X, Z=inverse_onehot(Z))
+                else:
+                    effect_score = None
+
+                return Y_X_score, T_XZ_score, effect_score
+
             def predict(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 Y_pred = self._model_Y_X.predict(X)
                 T_pred_zero = self._model_T_XZ.predict(X, np.zeros(Z.shape))
@@ -1117,6 +1225,18 @@ def models_Y_X(self):
     def models_T_XZ(self):
         return [mdl._model_T_XZ._model for mdl in super().models_nuisance]
 
+    @property
+    def nuisance_scores_Y_X(self):
+        return self.nuisance_scores_[0]
+
+    @property
+    def nuisance_scores_T_XZ(self):
+        return self.nuisance_scores_[1]
+
+    @property
+    def nuisance_scores_effect(self):
+        return self.nuisance_scores_[2]
+
 
 class LinearIntentToTreatDRIV(StatsModelsCateEstimatorMixin, IntentToTreatDRIV):
     """

econml/tests/test_dml.py

@@ -148,6 +148,12 @@ def make_random(is_discrete, d):
                                             np.testing.assert_array_equal(
                                                 marg_eff if d_x else marg_eff[0:1], const_marg_eff)
 
+                                            assert isinstance(est.score_, float)
+                                            for score in est.nuisance_scores_y:
+                                                assert isinstance(score, float)
+                                            for score in est.nuisance_scores_t:
+                                                assert isinstance(score, float)
+
                                             T0 = np.full_like(T, 'a') if is_discrete else np.zeros_like(T)
                                             eff = est.effect(X, T0=T0, T1=T)
                                             self.assertEqual(shape(eff), effect_shape)

econml/tests/test_ortho_iv.py

@@ -20,7 +20,7 @@
 import econml.tests.utilities  # bugfix for assertWarns
 
 
-class TestDML(unittest.TestCase):
+class TestOrthoIV(unittest.TestCase):
 
     def test_cate_api(self):
         """Test that we correctly implement the CATE API."""