Small fixes

README.md

@@ -401,7 +401,7 @@ The reStructuredText files that make up the documentation are stored in the [doc
 
 * June 2019: [Treatment Effects with Instruments paper](https://arxiv.org/pdf/1905.10176.pdf)
 
-* May 2019: [Open Data Science Conference Workshop](https://staging5.odsc.com/training/portfolio/machine-learning-estimation-of-heterogeneous-treatment-effect-the-microsoft-econml-library) 
+* May 2019: [Open Data Science Conference Workshop](https://odsc.com/speakers/machine-learning-estimation-of-heterogeneous-treatment-effect-the-microsoft-econml-library/) 
 
 * 2018: [Orthogonal Random Forests paper](http://proceedings.mlr.press/v97/oprescu19a.html)
 

econml/_ortho_learner.py

@@ -27,9 +27,9 @@ class in this module implements the general logic in a very versatile way
 import numpy as np
 import copy
 from warnings import warn
-from .utilities import (shape, reshape, ndim, hstack, cross_product, transpose,
+from .utilities import (shape, reshape, ndim, hstack, cross_product, transpose, inverse_onehot,
                         broadcast_unit_treatments, reshape_treatmentwise_effects,
-                        StatsModelsLinearRegression, LassoCVWrapper)
+                        StatsModelsLinearRegression, _EncoderWrapper)
 from sklearn.model_selection import KFold, StratifiedKFold, check_cv
 from sklearn.linear_model import LinearRegression, LassoCV
 from sklearn.preprocessing import (PolynomialFeatures, LabelEncoder, OneHotEncoder,
@@ -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 nuisance 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,25 @@ 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,)
+
+        # scores entries should be lists of scores, so make each entry a singleton list
+        scores = tuple([s] for s in 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 +171,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 +256,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::
 
@@ -256,6 +280,10 @@ class _OrthoLearner(TreatmentExpansionMixin, LinearCateEstimator):
     discrete_instrument: bool
         Whether the instrument values should be treated as categorical, rather than continuous, quantities
 
+    categories: 'auto' or list
+        The categories to use when encoding discrete treatments (or 'auto' to use the unique sorted values).
+        The first category will be treated as the control treatment.
+
     n_splits: int, cross-validation generator or an iterable
         Determines the cross-validation splitting strategy.
         Possible inputs for cv are:
@@ -321,7 +349,7 @@ def score(self, Y, T, W=None, nuisances=None):
         est = _OrthoLearner(ModelNuisance(LinearRegression(), LinearRegression()),
                             ModelFinal(),
                             n_splits=2, discrete_treatment=False, discrete_instrument=False,
-                            random_state=None)
+                            categories='auto', random_state=None)
         est.fit(y, X[:, 0], W=X[:, 1:])
 
     >>> est.score_
@@ -383,7 +411,7 @@ def score(self, Y, T, W=None, nuisances=None):
         y = T + W[:, 0] + np.random.normal(0, 0.01, size=(100,))
         est = _OrthoLearner(ModelNuisance(LogisticRegression(solver='lbfgs'), LinearRegression()),
                             ModelFinal(), n_splits=2, discrete_treatment=True, discrete_instrument=False,
-                            random_state=None)
+                            categories='auto', random_state=None)
         est.fit(y, T, W=W)
 
     >>> est.score_
@@ -408,10 +436,12 @@ 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, *,
-                 discrete_treatment, discrete_instrument, n_splits, random_state):
+                 discrete_treatment, discrete_instrument, categories, n_splits, random_state):
         self._model_nuisance = clone(model_nuisance, safe=False)
         self._models_nuisance = None
         self._model_final = clone(model_final, safe=False)
@@ -420,8 +450,9 @@ def __init__(self, model_nuisance, model_final, *,
         self._discrete_instrument = discrete_instrument
         self._random_state = check_random_state(random_state)
         if discrete_treatment:
-            self._label_encoder = LabelEncoder()
-            self._one_hot_encoder = OneHotEncoder(categories='auto', sparse=False)
+            if categories != 'auto':
+                categories = [categories]  # OneHotEncoder expects a 2D array with features per column
+            self._one_hot_encoder = OneHotEncoder(categories=categories, sparse=False, drop='first')
         super().__init__()
 
     @staticmethod
@@ -511,26 +542,25 @@ def _fit_nuisances(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
         stratify = self._discrete_treatment or self._discrete_instrument
 
         if self._discrete_treatment:
-            T = self._label_encoder.fit_transform(T.ravel())
+            T = self._one_hot_encoder.fit_transform(reshape(T, (-1, 1)))
 
         if self._discrete_instrument:
             z_enc = LabelEncoder()
             Z = z_enc.fit_transform(Z.ravel())
 
             if self._discrete_treatment:  # need to stratify on combination of Z and T
-                to_split = T + Z * len(self._label_encoder.classes_)
+                to_split = inverse_onehot(T) + Z * len(self._one_hot_encoder.categories_[0])
             else:
                 to_split = Z  # just stratify on Z
 
-            z_ohe = OneHotEncoder(categories='auto', sparse=False)
-            Z = z_ohe.fit_transform(reshape(Z, (-1, 1)))[:, 1:]
+            z_ohe = OneHotEncoder(categories='auto', sparse=False, drop='first')
+            Z = z_ohe.fit_transform(reshape(Z, (-1, 1)))
             self.z_transformer = FunctionTransformer(
-                func=(lambda Z:
-                      z_ohe.transform(
-                          reshape(z_enc.transform(Z.ravel()), (-1, 1)))[:, 1:]),
+                func=_EncoderWrapper(z_ohe, z_enc).encode,
                 validate=False)
         else:
-            to_split = T  # stratify on T if discrete, and fine to pass T as second arg to KFold.split even when not
+            # stratify on T if discrete, and fine to pass T as second arg to KFold.split even when not
+            to_split = inverse_onehot(T) if self._discrete_treatment else T
             self.z_transformer = None
 
         if self._n_splits == 1:  # special case, no cross validation
@@ -550,19 +580,15 @@ def _fit_nuisances(self, Y, T, X=None, W=None, Z=None, sample_weight=None):
                 folds = splitter.split(np.ones((T.shape[0], 1)), to_split)
 
         if self._discrete_treatment:
-            # drop first column since all columns sum to one
-            T = self._one_hot_encoder.fit_transform(reshape(T, (-1, 1)))[:, 1:]
-
             self._d_t = shape(T)[1:]
             self.transformer = FunctionTransformer(
-                func=(lambda T:
-                      self._one_hot_encoder.transform(
-                          reshape(self._label_encoder.transform(T.ravel()), (-1, 1)))[:, 1:]),
+                func=_EncoderWrapper(self._one_hot_encoder).encode,
                 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):