Merge pull request #242 from paullo0106/propensity_score_check

Fix #241 by putting back pandas Series as an acceptable format for pr…
uber · Oct 22, 2020 · 4323818 · 4323818
2 parents c99ebc5 + ab8d271
commit 4323818
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
@@ -237,7 +237,7 @@ def estimate_ate(self, X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps
         else:
             check_p_conditions(p, self.t_groups)
 
-        if isinstance(p, np.ndarray):
+        if isinstance(p, (np.ndarray, pd.Series)):
             treatment_name = self.t_groups[0]
             p = {treatment_name: convert_pd_to_np(p)}
         elif isinstance(p, dict):

diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
@@ -245,7 +245,7 @@ def fit_predict(self, X, treatment, y, p=None, return_ci=False, n_bootstraps=100
             p = self.propensity
 
         check_p_conditions(p, self.t_groups)
-        if isinstance(p, np.ndarray):
+        if isinstance(p, (np.ndarray, pd.Series)):
             treatment_name = self.t_groups[0]
             p = {treatment_name: convert_pd_to_np(p)}
         elif isinstance(p, dict):
@@ -305,7 +305,7 @@ def estimate_ate(self, X, treatment, y, p=None, bootstrap_ci=False, n_bootstraps
             p = self.propensity
         else:
             check_p_conditions(p, self.t_groups)
-        if isinstance(p, np.ndarray):
+        if isinstance(p, (np.ndarray, pd.Series)):
             treatment_name = self.t_groups[0]
             p = {treatment_name: convert_pd_to_np(p)}
         elif isinstance(p, dict):