a couple of remaining small fix (#419)

* copy the attribute for inference results
* pass through the raw feature name when parse const_marginal_effect
* fix the broadcasting issue in drlearner score function
* fix final model cate prediction shape when multitask is True for DRLearner and add test

Author: heimengqi

econml/_shap.py

@@ -29,13 +29,13 @@ def _shap_explain_cme(cme_model, X, d_t, d_y,
     cme_models: function
         const_marginal_effect function.
     X: (m, d_x) matrix
-        Features for each sample. Should be in the same shape of fitted X in final stage.
+        Features for each sample. Should be in the same shape of X during fit.
     d_t: tuple of int
         Tuple of number of treatment (exclude control in discrete treatment scenario).
     d_y: tuple of int
         Tuple of number of outcome.
     feature_names: optional None or list of strings of length X.shape[1] (Default=None)
-        The names of input features.
+        The names of raw input features.
     treatment_names: optional None or list (Default=None)
         The name of treatment. In discrete treatment scenario, the name should not include the name of
         the baseline treatment (i.e. the control treatment, which by default is the alphabetically smaller)
@@ -54,8 +54,9 @@ def _shap_explain_cme(cme_model, X, d_t, d_y,
         and the shap_values explanation object as value.
 
     """
-    (dt, dy, treatment_names, output_names, feature_names) = _define_names(d_t, d_y, treatment_names, output_names,
-                                                                           feature_names, input_names)
+    (dt, dy, treatment_names, output_names, feature_names, _) = _define_names(d_t, d_y, treatment_names,
+                                                                              output_names, feature_names,
+                                                                              input_names, None)
     # define masker by using entire dataset, otherwise Explainer will only sample 100 obs by default.
     bg_samples = X.shape[0] if background_samples is None else min(background_samples, X.shape[0])
     background = shap.maskers.Independent(X, max_samples=bg_samples)
@@ -108,7 +109,7 @@ def _shap_explain_model_cate(cme_model, models, X, d_t, d_y, featurizer=None, fe
     featurizer: optional None or instance of featurizer
         Fitted Featurizer of feature X.
     feature_names: optional None or list of strings of length X.shape[1] (Default=None)
-        The names of input features.
+        The names of raw input features.
     treatment_names: optional None or list (Default=None)
         The name of treatment. In discrete treatment scenario, the name should not include the name of
         the baseline treatment (i.e. the control treatment, which by default is the alphabetically smaller)
@@ -129,8 +130,12 @@ def _shap_explain_model_cate(cme_model, models, X, d_t, d_y, featurizer=None, fe
     d_t_, d_y_ = d_t, d_y
     feature_names_, treatment_names_ = feature_names, treatment_names,
     output_names_, input_names_ = output_names, input_names
-    (dt, dy, treatment_names, output_names, feature_names) = _define_names(d_t, d_y, treatment_names, output_names,
-                                                                           feature_names, input_names)
+    (dt, dy, treatment_names, output_names, feature_names, transformed_feature_names) = _define_names(d_t, d_y,
+                                                                                                      treatment_names,
+                                                                                                      output_names,
+                                                                                                      feature_names,
+                                                                                                      input_names,
+                                                                                                      featurizer)
     if featurizer is not None:
         F = featurizer.transform(X)
     else:
@@ -146,11 +151,11 @@ def _shap_explain_model_cate(cme_model, models, X, d_t, d_y, featurizer=None, fe
     for i in range(dt):
         try:
             explainer = shap.Explainer(models[i], background,
-                                       feature_names=feature_names)
+                                       feature_names=transformed_feature_names)
         except Exception as e:
             print("Final model can't be parsed, explain const_marginal_effect() instead!", repr(e))
             return _shap_explain_cme(cme_model, X, d_t_, d_y_,
-                                     feature_names=None,
+                                     feature_names=feature_names_,
                                      treatment_names=treatment_names_,
                                      output_names=output_names_,
                                      input_names=input_names_,
@@ -183,16 +188,17 @@ def _shap_explain_joint_linear_model_cate(model_final, X, d_t, d_y, fit_cate_int
     model_final: a single estimator
         the model's final stage model.
     X: matrix
-        Featurized X
+        Featurized X.
     d_t: tuple of int
         Tuple of number of treatment (exclude control in discrete treatment scenario).
     d_y: tuple of int
         Tuple of number of outcome.
     fit_cate_intercept: bool
         Whether the first entry of the coefficient of the joint linear model associated with
         each treatment, is an intercept.
-    feature_names: optional None or list of strings of length X.shape[1] (Default=None)
-        The names of input features.
+    feature_names: optional None or list of strings of length X.shape[1] or X.shape[1]-1 (Default=None)
+        The name of featurized X (exclude intercept). Length is X.shape[1] if fit_cate_intercpet=False, otherwise
+        length is X.shape[1]-1.
     treatment_names: optional None or list (Default=None)
         The name of treatment. In discrete treatment scenario, the name should not include the name of
         the baseline treatment (i.e. the control treatment, which by default is the alphabetically smaller)
@@ -210,8 +216,11 @@ def _shap_explain_joint_linear_model_cate(model_final, X, d_t, d_y, fit_cate_int
         each treatment name (e.g. "T0" when `treatment_names=None`) as key
         and the shap_values explanation object as value.
     """
-    (d_t, d_y, treatment_names, output_names, feature_names) = _define_names(d_t, d_y, treatment_names, output_names,
-                                                                             feature_names, input_names)
+    # input feature name is already updated by cate_feature_names.
+    (d_t, d_y, treatment_names, output_names, _, _) = _define_names(d_t, d_y, treatment_names,
+                                                                    output_names,
+                                                                    feature_names,
+                                                                    input_names, None)
     X, T = broadcast_unit_treatments(X, d_t)
     X = cross_product(X, T)
     d_x = X.shape[1]
@@ -226,7 +235,7 @@ def _shap_explain_joint_linear_model_cate(model_final, X, d_t, d_y, fit_cate_int
         # define masker by using entire dataset, otherwise Explainer will only sample 100 obs by default.
         bg_samples = X_sub.shape[0] if background_samples is None else min(background_samples, X_sub.shape[0])
         background = shap.maskers.Independent(X_sub, max_samples=bg_samples)
-        explainer = shap.Explainer(model_final, background)
+        explainer = shap.Explainer(model_final, background, feature_names=feature_names)
         shap_out = explainer(X_sub)
 
         data = shap_out.data[:, ind_x[i]]
@@ -236,14 +245,14 @@ def _shap_explain_joint_linear_model_cate(model_final, X, d_t, d_y, fit_cate_int
                 main_effects = None if shap_out.main_effects is None else shap_out.main_effects[..., ind_x[i], j]
                 values = shap_out.values[..., ind_x[i], j]
                 shap_out_new = shap.Explanation(values, base_values=base_values, data=data, main_effects=main_effects,
-                                                feature_names=feature_names)
+                                                feature_names=shap_out.feature_names)
                 shap_outs[output_names[j]][treatment_names[i]] = shap_out_new
         else:
             values = shap_out.values[..., ind_x[i]]
             main_effects = shap_out.main_effects[..., ind_x[i], 0]
             shap_out_new = shap.Explanation(values, base_values=shap_out.base_values, data=data,
                                             main_effects=main_effects,
-                                            feature_names=feature_names)
+                                            feature_names=shap_out.feature_names)
             shap_outs[output_names[0]][treatment_names[i]] = shap_out_new
 
     return shap_outs
@@ -274,7 +283,7 @@ def _shap_explain_multitask_model_cate(cme_model, multitask_model_cate, X, d_t,
     featurizer: optional None or instance of featurizer
         Fitted Featurizer of feature X.
     feature_names: optional None or list of strings of length X.shape[1] (Default=None)
-        The names of input features.
+        The names of raw input features.
     treatment_names: optional None or list (Default=None)
         The name of treatment. In discrete treatment scenario, the name should not include the name of
         the baseline treatment (i.e. the control treatment, which by default is the alphabetically smaller)
@@ -295,8 +304,12 @@ def _shap_explain_multitask_model_cate(cme_model, multitask_model_cate, X, d_t,
     d_t_, d_y_ = d_t, d_y
     feature_names_, treatment_names_ = feature_names, treatment_names,
     output_names_, input_names_ = output_names, input_names
-    (dt, dy, treatment_names, output_names, feature_names) = _define_names(d_t, d_y, treatment_names, output_names,
-                                                                           feature_names, input_names)
+    (dt, dy, treatment_names, output_names, feature_names, transformed_feature_names) = _define_names(d_t, d_y,
+                                                                                                      treatment_names,
+                                                                                                      output_names,
+                                                                                                      feature_names,
+                                                                                                      input_names,
+                                                                                                      featurizer)
     if featurizer is not None:
         F = featurizer.transform(X)
     else:
@@ -311,11 +324,11 @@ def _shap_explain_multitask_model_cate(cme_model, multitask_model_cate, X, d_t,
     for j in range(dy):
         try:
             explainer = shap.Explainer(multitask_model_cate[j], background,
-                                       feature_names=feature_names)
+                                       feature_names=transformed_feature_names)
         except Exception as e:
             print("Final model can't be parsed, explain const_marginal_effect() instead!", repr(e))
             return _shap_explain_cme(cme_model, X, d_t_, d_y_,
-                                     feature_names=None,
+                                     feature_names=feature_names_,
                                      treatment_names=treatment_names_,
                                      output_names=output_names_,
                                      input_names=input_names_,
@@ -336,7 +349,7 @@ def _shap_explain_multitask_model_cate(cme_model, multitask_model_cate, X, d_t,
     return shap_outs
 
 
-def _define_names(d_t, d_y, treatment_names, output_names, feature_names, input_names):
+def _define_names(d_t, d_y, treatment_names, output_names, feature_names, input_names, featurizer):
     """
     Helper function to get treatment and output names
 
@@ -355,28 +368,33 @@ def _define_names(d_t, d_y, treatment_names, output_names, feature_names, input_
         The user provided names of the features
     input_names: dicitionary
         The names of the features, outputs and treatments parsed from the fit input at fit time.
+    featurizer: optional None or instance of featurizer
+        Fitted Featurizer of feature X.
 
     Returns
     -------
     d_t: int
     d_y: int
     treament_names: List
     output_names: List
-    feature_names: List or None
+    feature_names: List
+    transformed_feature_names: List or None
     """
 
     d_t = d_t[0] if d_t else 1
     d_y = d_y[0] if d_y else 1
+
     if treatment_names is None:
-        if (input_names is None) or (input_names['treatment_names'] is None):
-            treatment_names = [f"T{i}" for i in range(d_t)]
-        else:
-            treatment_names = input_names['treatment_names']
+        treatment_names = input_names['treatment_names']
     if output_names is None:
-        if (input_names is None) or (input_names['output_names'] is None):
-            output_names = [f"Y{i}" for i in range(d_y)]
-        else:
-            output_names = input_names['output_names']
-    if (feature_names is None) and (input_names is not None):
+        output_names = input_names['output_names']
+    if feature_names is None:
         feature_names = input_names['feature_names']
-    return (d_t, d_y, treatment_names, output_names, feature_names)
+    if featurizer is None:
+        transformed_feature_names = feature_names
+    elif featurizer is not None and hasattr(featurizer, 'get_feature_names'):
+        transformed_feature_names = featurizer.get_feature_names(feature_names)
+    else:
+        transformed_feature_names = None
+
+    return (d_t, d_y, treatment_names, output_names, feature_names, transformed_feature_names)

econml/dml/causal_forest.py

@@ -555,8 +555,6 @@ def feature_importances(self, max_depth=4, depth_decay_exponent=2.0):
         return imps.reshape(self._d_y + (-1,))
 
     def shap_values(self, X, *, feature_names=None, treatment_names=None, output_names=None, background_samples=100):
-        feature_names = self.cate_feature_names(feature_names)
-
         return _shap_explain_multitask_model_cate(self.const_marginal_effect, self.model_cate.estimators_, X,
                                                   self._d_t, self._d_y, featurizer=self.featurizer_,
                                                   feature_names=feature_names,

econml/dml/dml.py

@@ -1177,8 +1177,6 @@ def refit_final(self, *, inference='auto'):
     refit_final.__doc__ = _OrthoLearner.refit_final.__doc__
 
     def shap_values(self, X, *, feature_names=None, treatment_names=None, output_names=None, background_samples=100):
-        feature_names = self.cate_feature_names(feature_names)
-
         return _shap_explain_model_cate(self.const_marginal_effect, self.model_cate, X, self._d_t, self._d_y,
                                         featurizer=self.featurizer_,
                                         feature_names=feature_names,

econml/dr/_drlearner.py

@@ -127,6 +127,7 @@ def __init__(self, model_final, featurizer, multitask_model_final):
     def fit(self, Y, T, X=None, W=None, *, nuisances, sample_weight=None, sample_var=None):
         Y_pred, = nuisances
         self.d_y = Y_pred.shape[1:-1]  # track whether there's a Y dimension (must be a singleton)
+        self.d_t = Y_pred.shape[-1] - 1  # track # of treatment (exclude baseline treatment)
         if (X is not None) and (self._featurizer is not None):
             X = self._featurizer.fit_transform(X)
         filtered_kwargs = filter_none_kwargs(sample_weight=sample_weight, sample_var=sample_var)
@@ -145,7 +146,7 @@ def predict(self, X=None):
         if (X is not None) and (self._featurizer is not None):
             X = self._featurizer.transform(X)
         if self._multitask_model_final:
-            pred = self.model_cate.predict(X)
+            pred = self.model_cate.predict(X).reshape((-1, self.d_t))
             if self.d_y:  # need to reintroduce singleton Y dimension
                 return pred[:, np.newaxis, :]
             return pred
@@ -158,13 +159,21 @@ def score(self, Y, T, X=None, W=None, *, nuisances, sample_weight=None, sample_v
             X = self._featurizer.transform(X)
         Y_pred, = nuisances
         if self._multitask_model_final:
-            return np.mean(np.average((Y_pred[..., 1:] - Y_pred[..., [0]] - self.model_cate.predict(X))**2,
-                                      weights=sample_weight, axis=0))
+            Y_pred_diff = Y_pred[..., 1:] - Y_pred[..., [0]]
+            cate_pred = self.model_cate.predict(X).reshape((-1, self.d_t))
+            if self.d_y:
+                cate_pred = cate_pred[:, np.newaxis, :]
+            return np.mean(np.average((Y_pred_diff - cate_pred)**2, weights=sample_weight, axis=0))
+
         else:
-            return np.mean([np.average((Y_pred[..., t] - Y_pred[..., 0] -
-                                        self.models_cate[t - 1].predict(X))**2,
-                                       weights=sample_weight, axis=0)
-                            for t in np.arange(1, Y_pred.shape[-1])])
+            scores = []
+            for t in np.arange(1, Y_pred.shape[-1]):
+                # since we only allow single dimensional y, we could flatten the prediction
+                Y_pred_diff = (Y_pred[..., t] - Y_pred[..., 0]).flatten()
+                cate_pred = self.models_cate[t - 1].predict(X).flatten()
+                score = np.average((Y_pred_diff - cate_pred)**2, weights=sample_weight, axis=0)
+                scores.append(score)
+            return np.mean(scores)
 
 
 class DRLearner(_OrthoLearner):
@@ -637,8 +646,6 @@ def fitted_models_final(self):
         return self.ortho_learner_model_final_.models_cate
 
     def shap_values(self, X, *, feature_names=None, treatment_names=None, output_names=None, background_samples=100):
-        feature_names = self.cate_feature_names(feature_names)
-
         if self.ortho_learner_model_final_._multitask_model_final:
             return _shap_explain_multitask_model_cate(self.const_marginal_effect, self.multitask_model_cate, X,
                                                       self._d_t, self._d_y,

econml/inference/_inference.py

@@ -559,7 +559,7 @@ def __init__(self, d_t, d_y, pred, inf_type, fname_transformer=None,
         # For effect summaries, d_t is None, but the result arrays behave as if d_t=1
         self._d_t = d_t or 1
         self.d_y = d_y
-        self.pred = pred
+        self.pred = np.copy(pred) if pred is not None and not np.isscalar(pred) else pred
         self.inf_type = inf_type
         self.fname_transformer = fname_transformer
         self.feature_names = feature_names
@@ -848,7 +848,8 @@ class NormalInferenceResults(InferenceResults):
 
     def __init__(self, d_t, d_y, pred, pred_stderr, inf_type, fname_transformer=None,
                  feature_names=None, output_names=None, treatment_names=None):
-        self.pred_stderr = pred_stderr
+        self.pred_stderr = np.copy(pred_stderr) if pred_stderr is not None and not np.isscalar(
+            pred_stderr) else pred_stderr
         super().__init__(d_t, d_y, pred, inf_type, fname_transformer, feature_names, output_names, treatment_names)
 
     @property
@@ -948,7 +949,7 @@ class EmpiricalInferenceResults(InferenceResults):
 
     def __init__(self, d_t, d_y, pred, pred_dist, inf_type, fname_transformer=None,
                  feature_names=None, output_names=None, treatment_names=None):
-        self.pred_dist = pred_dist
+        self.pred_dist = np.copy(pred_dist) if pred_dist is not None and not np.isscalar(pred_dist) else pred_dist
         super().__init__(d_t, d_y, pred, inf_type, fname_transformer, feature_names, output_names, treatment_names)
 
     @property