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,20 @@ 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 hasattr(self._model, 'score'):
+            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)
+        else:
+            return None
+
 
 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):
         """