class_docstr: A parent class for S-learner classes.
  An S-learner estimates treatment effects with one machine learning model.
  Details of S-learner are available at Kunzel et al. (2018) (https://arxiv.org/abs/1706.03461).
class_docstr_tok: ["A","parent","class","for","S-learner","classes",".","An","S-learner","estimates","treatment","effects","with","one","machine","learning","model",".","Details","of","S-learner","are","available","at","Kunzel","et","al",".","(","2018",")","(","https",":","\/\/arxiv.org\/abs\/1706.03461",")","."]
docstr: Estimate the Average Treatment Effect (ATE).
  
  Args:
      X (np.matrix, 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
      return_ci (bool, optional): whether to return confidence intervals
      bootstrap_ci (bool): whether to return 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",",","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","return_ci","(","bool",",","optional",")",":","whether","to","return","confidence","intervals","bootstrap_ci","(","bool",")",":","whether","to","return","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,
      y,
      p=None,
      return_ci=False,
      bootstrap_ci=False,
      n_bootstraps=1000,
      bootstrap_size=10000,
      pretrain=False,
  ):
      """Estimate the Average Treatment Effect (ATE).
  
      Args:
          X (np.matrix, 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
          return_ci (bool, optional): whether to return confidence intervals
          bootstrap_ci (bool): whether to return 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, yhat_cs, yhat_ts = self.predict(
              X, treatment, y, return_components=True
          )
      else:
          te, yhat_cs, yhat_ts = self.fit_predict(
              X, treatment, y, return_components=True
          )
  
      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]
          y_filt = y[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]
  
          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 return_ci:
          return ate
      elif return_ci and not bootstrap_ci:
          return ate, ate_lb, ate_ub
      else:
          t_groups_global = self.t_groups
          _classes_global = self._classes
          models_global = deepcopy(self.models)
  
          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)):
              ate_b = self.bootstrap(
                  X, treatment, y, size=bootstrap_size
              )
              ate_bootstraps[:, n] = ate_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 = deepcopy(models_global)
  
          return ate, ate_lower, ate_upper
  
imports: import logging
  import numpy
  from tqdm import tqdm
  from scipy.stats import norm
  from sklearn.dummy import DummyRegressor
  import statsmodels.api
  from copy import deepcopy
  from causalml.inference.meta.base import BaseLearner
  from causalml.inference.meta.utils import check_treatment_vector, convert_pd_to_np
  from causalml.metrics import regression_metrics, classification_metrics
label_desc: Use image node_id 6 for calling the BaseSLearner obj's underlying member method code with example usage: obj.estimate_ate(X, treatment, y, p, return_ci, bootstrap_ci, n_bootstraps, bootstrap_size, pretrain) and returns: ate, ate,  ate_lb,  ate_ub, ate,  ate_lower,  ate_upper