Temporary fix for causal trees missing values support #733 (#734)

* Add temp missing values placeholder for causal trees * Apply black codestyle
uber · Jan 29, 2024 · 9e1f892 · 9e1f892
1 parent 0040ac6
commit 9e1f892
Show file tree

Hide file tree

Showing 8 changed files with 82,960 additions and 76 deletions.
diff --git a/causalml/dataset/synthetic.py b/causalml/dataset/synthetic.py
@@ -383,30 +383,30 @@ def get_synthetic_preds_holdout(
                 # fit the model on training data only
                 learner.fit(X=X_train, treatment=w_train, y=y_train)
                 try:
-                    preds_dict_train[
-                        "{} Learner ({})".format(label_l, label_m)
-                    ] = learner.predict(X=X_train, p=p_hat_train).flatten()
-                    preds_dict_valid[
-                        "{} Learner ({})".format(label_l, label_m)
-                    ] = learner.predict(X=X_val, p=p_hat_val).flatten()
+                    preds_dict_train["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(X=X_train, p=p_hat_train).flatten()
+                    )
+                    preds_dict_valid["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(X=X_val, p=p_hat_val).flatten()
+                    )
                 except TypeError:
-                    preds_dict_train[
-                        "{} Learner ({})".format(label_l, label_m)
-                    ] = learner.predict(
-                        X=X_train, treatment=w_train, y=y_train
-                    ).flatten()
-                    preds_dict_valid[
-                        "{} Learner ({})".format(label_l, label_m)
-                    ] = learner.predict(X=X_val, treatment=w_val, y=y_val).flatten()
+                    preds_dict_train["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(
+                            X=X_train, treatment=w_train, y=y_train
+                        ).flatten()
+                    )
+                    preds_dict_valid["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(X=X_val, treatment=w_val, y=y_val).flatten()
+                    )
             else:
                 learner = base_learner(model())
                 learner.fit(X=X_train, p=p_hat_train, treatment=w_train, y=y_train)
-                preds_dict_train[
-                    "{} Learner ({})".format(label_l, label_m)
-                ] = learner.predict(X=X_train).flatten()
-                preds_dict_valid[
-                    "{} Learner ({})".format(label_l, label_m)
-                ] = learner.predict(X=X_val).flatten()
+                preds_dict_train["{} Learner ({})".format(label_l, label_m)] = (
+                    learner.predict(X=X_train).flatten()
+                )
+                preds_dict_valid["{} Learner ({})".format(label_l, label_m)] = (
+                    learner.predict(X=X_val).flatten()
+                )
 
     return preds_dict_train, preds_dict_valid
 

diff --git a/causalml/feature_selection/filters.py b/causalml/feature_selection/filters.py
@@ -303,9 +303,11 @@ def _GetNodeSummary(
                 if smooth:
                     results[ti].update(
                         {
-                            ci: results_series[ti, ci]
-                            if results_series.index.isin([(ti, ci)]).any()
-                            else 1
+                            ci: (
+                                results_series[ti, ci]
+                                if results_series.index.isin([(ti, ci)]).any()
+                                else 1
+                            )
                         }
                     )
                 else:

diff --git a/causalml/inference/meta/explainer.py b/causalml/inference/meta/explainer.py
@@ -192,9 +192,9 @@ def get_shap_values(self):
         for group, mod in self.models_tau.items():
             explainer = shap.TreeExplainer(mod)
             if self.r_learners is not None:
-                explainer.model.original_model.params[
-                    "objective"
-                ] = None  # hacky way of running shap without error
+                explainer.model.original_model.params["objective"] = (
+                    None  # hacky way of running shap without error
+                )
             shap_values = explainer.shap_values(self.X)
             shap_dict[group] = shap_values