fix incorrect bootstrap mean calculation in meta-learners (#828)

Author: jeongyoonlee

causalml/inference/meta/base.py

@@ -1,4 +1,4 @@
-from abc import ABCMeta, abstractclassmethod
+from abc import ABCMeta, abstractmethod
 import logging
 import numpy as np
 import pandas as pd
@@ -12,11 +12,13 @@
 
 
 class BaseLearner(metaclass=ABCMeta):
-    @abstractclassmethod
+    @classmethod
+    @abstractmethod
     def fit(self, X, treatment, y, p=None):
         pass
 
-    @abstractclassmethod
+    @classmethod
+    @abstractmethod
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
@@ -37,7 +39,8 @@ def fit_predict(
         self.fit(X, treatment, y, p)
         return self.predict(X, treatment, y, p, return_components, verbose)
 
-    @abstractclassmethod
+    @classmethod
+    @abstractmethod
     def estimate_ate(
         self,
         X,

causalml/inference/meta/drlearner.py

@@ -434,7 +434,7 @@ def estimate_ate(
                 cate_b = self.bootstrap(
                     X, treatment, y, p, size=bootstrap_size, seed=seed
                 )
-                ate_bootstraps[:, n] = cate_b.mean()
+                ate_bootstraps[:, n] = cate_b.mean(axis=0)
 
             ate_lower = np.percentile(
                 ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1

causalml/inference/meta/rlearner.py

@@ -324,7 +324,7 @@ def estimate_ate(
                 else:
                     p = self._format_p(p, self.t_groups)
                 cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
-                ate_bootstraps[:, n] = cate_b.mean()
+                ate_bootstraps[:, n] = cate_b.mean(axis=0)
 
             ate_lower = np.percentile(
                 ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1

causalml/inference/meta/slearner.py

@@ -278,7 +278,7 @@ def estimate_ate(
 
             for n in tqdm(range(n_bootstraps)):
                 ate_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
-                ate_bootstraps[:, n] = ate_b.mean()
+                ate_bootstraps[:, n] = ate_b.mean(axis=0)
 
             ate_lower = np.percentile(
                 ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1

causalml/inference/meta/tlearner.py

@@ -280,7 +280,7 @@ def estimate_ate(
 
             for n in tqdm(range(n_bootstraps)):
                 ate_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
-                ate_bootstraps[:, n] = ate_b.mean()
+                ate_bootstraps[:, n] = ate_b.mean(axis=0)
 
             ate_lower = np.percentile(
                 ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1

causalml/inference/meta/xlearner.py

@@ -390,7 +390,7 @@ def estimate_ate(
 
             for n in tqdm(range(n_bootstraps)):
                 cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
-                ate_bootstraps[:, n] = cate_b.mean()
+                ate_bootstraps[:, n] = cate_b.mean(axis=0)
 
             ate_lower = np.percentile(
                 ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "causalml"
-version = "0.15.3"
+version = "0.15.4dev"
 description = "Python Package for Uplift Modeling and Causal Inference with Machine Learning Algorithms"
 readme = { file = "README.md", content-type = "text/markdown" }