class_docstr: A parent class for DRIV-learner regressor classes.
  
  A DRIV-learner estimates endogenous treatment effects for compliers with machine learning models.
  
  Details of DR-learner are available at Kennedy (2020) (https://arxiv.org/abs/2004.14497).
  The DR moment condition for LATE comes from Chernozhukov et al (2018) (https://academic.oup.com/ectj/article/21/1/C1/5056401).
docstr: Estimate the Average Treatment Effect (ATE) for compliers.
  
  Args:
      X (np.matrix or np.array or pd.Dataframe): a feature matrix
      assignment (np.array or pd.Series): an assignment vector. The assignment is the
          instrumental variable that does not depend on unknown confounders. The assignment status
          influences treatment in a monotonic way, i.e. one can only be more likely to take the
          treatment if assigned.
      treatment (np.array or pd.Series): a treatment vector
      y (np.array or pd.Series): an outcome vector
      p (2-tuple of np.ndarray or pd.Series or dict, optional): The first (second) element corresponds to
          unassigned (assigned) units. Each is 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.
      pZ (np.array or pd.Series, optional): an array of assignment probability of float (0,1); if None
          will run ElasticNetPropensityModel() to generate the assignment probability score.
      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
  Returns:
      The mean and confidence interval (LB, UB) of the ATE estimate.
docstr_tok: ["Estimate","the","Average","Treatment","Effect","(","ATE",")","for","compliers",".","Args",":","X","(","np.matrix","or","np.array","or","pd.Dataframe",")",":","a","feature","matrix","assignment","(","np.array","or","pd.Series",")",":","an","assignment","vector",".","The","assignment","is","the","instrumental","variable","that","does","not","depend","on","unknown","confounders",".","The","assignment","status","influences","treatment","in","a","monotonic","way",",","i.e",".","one","can","only","be","more","likely","to","take","the","treatment","if","assigned",".","treatment","(","np.array","or","pd.Series",")",":","a","treatment","vector","y","(","np.array","or","pd.Series",")",":","an","outcome","vector","p","(","2-tuple","of","np.ndarray","or","pd.Series","or","dict",",","optional",")",":","The","first","(","second",")","element","corresponds","to","unassigned","(","assigned",")","units",".","Each","is","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",".","pZ","(","np.array","or","pd.Series",",","optional",")",":","an","array","of","assignment","probability","of","float","(","0,1",")",";","if","None","will","run","ElasticNetPropensityModel","(",")","to","generate","the","assignment","probability","score",".","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","Returns",":","The","mean","and","confidence","interval","(","LB",",","UB",")","of","the","ATE","estimate","."]
code: def estimate_ate(
      self,
      X,
      assignment,
      treatment,
      y,
      p=None,
      pZ=None,
      bootstrap_ci=False,
      n_bootstraps=1000,
      bootstrap_size=10000,
      seed=None,
      calibrate=True,
  ):
      """Estimate the Average Treatment Effect (ATE) for compliers.
  
      Args:
          X (np.matrix or np.array or pd.Dataframe): a feature matrix
          assignment (np.array or pd.Series): an assignment vector. The assignment is the
              instrumental variable that does not depend on unknown confounders. The assignment status
              influences treatment in a monotonic way, i.e. one can only be more likely to take the
              treatment if assigned.
          treatment (np.array or pd.Series): a treatment vector
          y (np.array or pd.Series): an outcome vector
          p (2-tuple of np.ndarray or pd.Series or dict, optional): The first (second) element corresponds to
              unassigned (assigned) units. Each is 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.
          pZ (np.array or pd.Series, optional): an array of assignment probability of float (0,1); if None
              will run ElasticNetPropensityModel() to generate the assignment probability score.
          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
      Returns:
          The mean and confidence interval (LB, UB) of the ATE estimate.
      """
      te, yhat_cs, yhat_ts = self.fit_predict(
          X,
          assignment,
          treatment,
          y,
          p,
          return_components=True,
          seed=seed,
          calibrate=calibrate,
      )
      X, assignment, treatment, y = convert_pd_to_np(
          X, assignment, treatment, y
      )
  
      if p is None:
          p = (self.propensity_0, self.propensity_1)
      else:
          check_p_conditions(p[0], self.t_groups)
          check_p_conditions(p[1], self.t_groups)
  
      if isinstance(p[0], (np.ndarray, pd.Series)):
          treatment_name = self.t_groups[0]
          p = (
              {treatment_name: convert_pd_to_np(p[0])},
              {treatment_name: convert_pd_to_np(p[1])},
          )
      elif isinstance(p[0], dict):
          p = (
              {
                  treatment_name: convert_pd_to_np(_p)
                  for treatment_name, _p in p[0].items()
              },
              {
                  treatment_name: convert_pd_to_np(_p)
                  for treatment_name, _p in p[1].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)
          mask_1, mask_0 = mask & (assignment == 1), mask & (
              assignment == 0
          )
          Gamma = (treatment[mask_1] == group).mean() - (
              treatment[mask_0] == group
          ).mean()
  
          y_filt_1, y_filt_0 = y[mask_1], y[mask_0]
          yhat_0 = yhat_cs[group][mask_0]
          yhat_1 = yhat_ts[group][mask_1]
          treatment_filt_1, treatment_filt_0 = (
              treatment[mask_1],
              treatment[mask_0],
          )
          prob_treatment_1, prob_treatment_0 = (
              p[1][group][mask_1],
              p[0][group][mask_0],
          )
          w = (assignment[mask]).mean()
  
          part_1 = (
              (y_filt_1 - yhat_1).var()
              + _ate**2 * (treatment_filt_1 - prob_treatment_1).var()
              - 2
              * _ate
              * (
                  y_filt_1 * treatment_filt_1
                  - yhat_1 * prob_treatment_1
              ).mean()
          )
          part_0 = (
              (y_filt_0 - yhat_0).var()
              + _ate**2 * (treatment_filt_0 - prob_treatment_0).var()
              - 2
              * _ate
              * (
                  y_filt_0 * treatment_filt_0
                  - yhat_0 * prob_treatment_0
              ).mean()
          )
          part_2 = np.mean(
              (
                  yhat_ts[group][mask]
                  - yhat_cs[group][mask]
                  - _ate * (p[1][group][mask] - p[0][group][mask])
              )
              ** 2
          )
  
          # SE formula is based on the lower bound formula (9) from Frölich, Markus. 2006.
          # "Nonparametric IV estimation of local average treatment effects wth covariates."
          # Journal of Econometrics.
          se = np.sqrt((part_1 / w + part_0 / (1 - w)) + part_2) / Gamma
  
          _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,
                  assignment,
                  treatment,
                  y,
                  p,
                  pZ,
                  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: import logging
  from copy import deepcopy
  import numpy
  import pandas
  from causalml.inference.meta.explainer import Explainer
  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
  from scipy.stats import norm
  from sklearn.model_selection import cross_val_predict, KFold
  from tqdm import tqdm
  from xgboost import XGBRegressor
label_desc: Use image node_id 6 for calling the BaseDRIVLearner obj's underlying member method code with example usage: obj.estimate_ate(X, assignment, treatment, y, p, pZ, bootstrap_ci, n_bootstraps, bootstrap_size, seed, calibrate) and returns: ate,  ate_lb,  ate_ub, ate,  ate_lower,  ate_upper