class_docstr: A parent class for R-learner classes.
  
  An R-learner estimates treatment effects with two machine learning models and the propensity score.
  
  Details of R-learner are available at Nie and Wager (2019) (https://arxiv.org/abs/1712.04912).
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): only needed when pretrain=False, a treatment vector
      y (np.array or pd.Series):only needed when pretrain=False, 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.
      sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
          weight of each observation for `effect_learner`. If None, it assumes equal weight.
      bootstrap_ci (bool): whether run bootstrap for confidence intervals
      n_bootstraps (int): number of bootstrap iterations
      bootstrap_size (int): number of samples per bootstrap
      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",")",":","only","needed","when","pretrain=False",",","a","treatment","vector","y","(","np.array","or","pd.Series",")",":","only","needed","when","pretrain=False",",","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",".","sample_weight","(","np.array","or","pd.Series",",","optional",")",":","an","array","of","sample","weights","indicating","the","weight","of","each","observation","for","`","effect_learner","`",".","If","None",",","it","assumes","equal","weight",".","bootstrap_ci","(","bool",")",":","whether","run","bootstrap","for","confidence","intervals","n_bootstraps","(","int",")",":","number","of","bootstrap","iterations","bootstrap_size","(","int",")",":","number","of","samples","per","bootstrap","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=None,
      y=None,
      p=None,
      sample_weight=None,
      bootstrap_ci=False,
      n_bootstraps=1000,
      bootstrap_size=10000,
      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): only needed when pretrain=False, a treatment vector
          y (np.array or pd.Series):only needed when pretrain=False, 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.
          sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
              weight of each observation for `effect_learner`. If None, it assumes equal weight.
          bootstrap_ci (bool): whether run bootstrap for confidence intervals
          n_bootstraps (int): number of bootstrap iterations
          bootstrap_size (int): number of samples per bootstrap
          pretrain (bool): whether a model has been fit, default False.
      Returns:
          The mean and confidence interval (LB, UB) of the ATE estimate.
      """
      X, treatment, y = convert_pd_to_np(X, treatment, y)
      if pretrain:
          te = self.predict(X, p)
      else:
          if not len(treatment) or not len(y):
              raise ValueError(
                  "treatmeng and y must be provided when pretrain=False"
              )
          te = self.fit_predict(
              X, treatment, y, p, sample_weight, return_ci=False
          )
  
      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):
          w = (treatment == group).astype(int)
          prob_treatment = float(sum(w)) / X.shape[0]
          _ate = te[:, i].mean()
  
          se = (
              np.sqrt(
                  (self.vars_t[group] / prob_treatment)
                  + (self.vars_c[group] / (1 - prob_treatment))
                  + te[:, i].var()
              )
              / X.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
          model_mu_global = deepcopy(self.model_mu)
          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)):
              if p is None:
                  p = self.propensity
              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_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.model_mu = deepcopy(model_mu_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 tqdm import tqdm
  from scipy.stats import norm
  from sklearn.model_selection import cross_val_predict, KFold, train_test_split
  from xgboost import XGBRegressor
  from causalml.inference.meta.base import BaseLearner
  from causalml.inference.meta.utils import check_treatment_vector, get_xgboost_objective_metric, convert_pd_to_np, get_weighted_variance
  from causalml.propensity import compute_propensity_score, ElasticNetPropensityModel
label_desc: Use image node_id 6 for calling the BaseRLearner obj's underlying member method code with example usage: obj.estimate_ate(X, treatment, y, p, sample_weight, bootstrap_ci, n_bootstraps, bootstrap_size, pretrain) and returns: ate,  ate_lb,  ate_ub, ate,  ate_lower,  ate_upper