class_docstr: A parent class for DR-learner regressor classes.
  
  A DR-learner estimates treatment effects with machine learning models.
  
  Details of DR-learner are available at Kennedy (2020) (https://arxiv.org/abs/2004.14497).
class_docstr_tok: ["A","parent","class","for","DR-learner","regressor","classes",".","A","DR-learner","estimates","treatment","effects","with","machine","learning","models",".","Details","of","DR-learner","are","available","at","Kennedy","(","2020",")","(","https",":","\/\/arxiv.org\/abs\/2004.14497",")","."]
docstr: Estimate the Average Treatment Effect (ATE).
  
  Args:
      X (np.matrix or np.array or pd.Dataframe): a feature matrix
      treatment (np.array or pd.Series): a treatment vector
      y (np.array or pd.Series): an outcome vector
      p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
          single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
          float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
      bootstrap_ci (bool): whether run bootstrap for confidence intervals
      n_bootstraps (int): number of bootstrap iterations
      bootstrap_size (int): number of samples per bootstrap
      seed (int): random seed for cross-fitting
      pretrain (bool): whether a model has been fit, default False.
  Returns:
      The mean and confidence interval (LB, UB) of the ATE estimate.
docstr_tok: ["Estimate","the","Average","Treatment","Effect","(","ATE",")",".","Args",":","X","(","np.matrix","or","np.array","or","pd.Dataframe",")",":","a","feature","matrix","treatment","(","np.array","or","pd.Series",")",":","a","treatment","vector","y","(","np.array","or","pd.Series",")",":","an","outcome","vector","p","(","np.ndarray","or","pd.Series","or","dict",",","optional",")",":","an","array","of","propensity","scores","of","float","(","0,1",")","in","the","single-treatment","case",";","or",",","a","dictionary","of","treatment","groups","that","map","to","propensity","vectors","of","float","(","0,1",")",";","if","None","will","run","ElasticNetPropensityModel","(",")","to","generate","the","propensity","scores",".","bootstrap_ci","(","bool",")",":","whether","run","bootstrap","for","confidence","intervals","n_bootstraps","(","int",")",":","number","of","bootstrap","iterations","bootstrap_size","(","int",")",":","number","of","samples","per","bootstrap","seed","(","int",")",":","random","seed","for","cross-fitting","pretrain","(","bool",")",":","whether","a","model","has","been","fit",",","default","False",".","Returns",":","The","mean","and","confidence","interval","(","LB",",","UB",")","of","the","ATE","estimate","."]
code: def estimate_ate(
      self,
      X,
      treatment,
      y,
      p=None,
      bootstrap_ci=False,
      n_bootstraps=1000,
      bootstrap_size=10000,
      seed=None,
      pretrain=False,
  ):
      """Estimate the Average Treatment Effect (ATE).
  
      Args:
          X (np.matrix or np.array or pd.Dataframe): a feature matrix
          treatment (np.array or pd.Series): a treatment vector
          y (np.array or pd.Series): an outcome vector
          p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
              single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
              float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
          bootstrap_ci (bool): whether run bootstrap for confidence intervals
          n_bootstraps (int): number of bootstrap iterations
          bootstrap_size (int): number of samples per bootstrap
          seed (int): random seed for cross-fitting
          pretrain (bool): whether a model has been fit, default False.
      Returns:
          The mean and confidence interval (LB, UB) of the ATE estimate.
      """
      if pretrain:
          te, yhat_cs, yhat_ts = self.predict(
              X, treatment, y, p, return_components=True
          )
      else:
          te, yhat_cs, yhat_ts = self.fit_predict(
              X, treatment, y, p, return_components=True, seed=seed
          )
      X, treatment, y = convert_pd_to_np(X, treatment, y)
  
      if p is None:
          p = self.propensity
      else:
          check_p_conditions(p, self.t_groups)
      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):
          p = {
              treatment_name: convert_pd_to_np(_p)
              for treatment_name, _p in p.items()
          }
  
      ate = np.zeros(self.t_groups.shape[0])
      ate_lb = np.zeros(self.t_groups.shape[0])
      ate_ub = np.zeros(self.t_groups.shape[0])
  
      for i, group in enumerate(self.t_groups):
          _ate = te[:, i].mean()
  
          mask = (treatment == group) | (treatment == self.control_name)
          treatment_filt = treatment[mask]
          w = (treatment_filt == group).astype(int)
          prob_treatment = float(sum(w)) / w.shape[0]
  
          yhat_c = yhat_cs[group][mask]
          yhat_t = yhat_ts[group][mask]
          y_filt = y[mask]
  
          # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
          # "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
          se = np.sqrt(
              (
                  (y_filt[w == 0] - yhat_c[w == 0]).var()
                  / (1 - prob_treatment)
                  + (y_filt[w == 1] - yhat_t[w == 1]).var()
                  / prob_treatment
                  + (yhat_t - yhat_c).var()
              )
              / y_filt.shape[0]
          )
  
          _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
          _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
  
          ate[i] = _ate
          ate_lb[i] = _ate_lb
          ate_ub[i] = _ate_ub
  
      if not bootstrap_ci:
          return ate, ate_lb, ate_ub
      else:
          t_groups_global = self.t_groups
          _classes_global = self._classes
          models_mu_c_global = deepcopy(self.models_mu_c)
          models_mu_t_global = deepcopy(self.models_mu_t)
          models_tau_global = deepcopy(self.models_tau)
  
          logger.info("Bootstrap Confidence Intervals for ATE")
          ate_bootstraps = np.zeros(
              shape=(self.t_groups.shape[0], n_bootstraps)
          )
  
          for n in tqdm(range(n_bootstraps)):
              cate_b = self.bootstrap(
                  X, treatment, y, p, size=bootstrap_size, seed=seed
              )
              ate_bootstraps[:, n] = cate_b.mean()
  
          ate_lower = np.percentile(
              ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
          )
          ate_upper = np.percentile(
              ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
          )
  
          # set member variables back to global (currently last bootstrapped outcome)
          self.t_groups = t_groups_global
          self._classes = _classes_global
          self.models_mu_c = deepcopy(models_mu_c_global)
          self.models_mu_t = deepcopy(models_mu_t_global)
          self.models_tau = deepcopy(models_tau_global)
          return ate, ate_lower, ate_upper
  
imports: from copy import deepcopy
  import logging
  import numpy
  import pandas
  from scipy.stats import norm
  from sklearn.model_selection import cross_val_predict, KFold
  from tqdm import tqdm
  from xgboost import XGBRegressor
  from causalml.inference.meta.base import BaseLearner
  from causalml.inference.meta.utils import check_treatment_vector, check_p_conditions, convert_pd_to_np
  from causalml.metrics import regression_metrics, classification_metrics
  from causalml.propensity import compute_propensity_score
label_desc: Use image node_id 6 for calling the BaseDRLearner obj's underlying member method code with example usage: obj.estimate_ate(X, treatment, y, p, bootstrap_ci, n_bootstraps, bootstrap_size, seed, pretrain) and returns: ate,  ate_lb,  ate_ub, ate,  ate_lower,  ate_upper