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: Fit the treatment effect and outcome models of the R learner and predict treatment effects.
  
  Args:
      X (np.matrix or np.array or pd.Dataframe): a feature matrix
      assignment (np.array or pd.Series): a (0,1)-valued 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.
      return_ci (bool): whether to return confidence intervals
      n_bootstraps (int): number of bootstrap iterations
      bootstrap_size (int): number of samples per bootstrap
      return_components (bool, optional): whether to return outcome for treatment and control seperately
      verbose (str): whether to output progress logs
      seed (int): random seed for cross-fitting
  Returns:
      (numpy.ndarray): Predictions of treatment effects for compliers, , i.e. those individuals
          who take the treatment only if they are assigned. Output dim: [n_samples, n_treatment]
          If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
          UB [n_samples, n_treatment]
docstr_tok: ["Fit","the","treatment","effect","and","outcome","models","of","the","R","learner","and","predict","treatment","effects",".","Args",":","X","(","np.matrix","or","np.array","or","pd.Dataframe",")",":","a","feature","matrix","assignment","(","np.array","or","pd.Series",")",":","a","(","0,1",")","-valued","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",".","return_ci","(","bool",")",":","whether","to","return","confidence","intervals","n_bootstraps","(","int",")",":","number","of","bootstrap","iterations","bootstrap_size","(","int",")",":","number","of","samples","per","bootstrap","return_components","(","bool",",","optional",")",":","whether","to","return","outcome","for","treatment","and","control","seperately","verbose","(","str",")",":","whether","to","output","progress","logs","seed","(","int",")",":","random","seed","for","cross-fitting","Returns",":","(","numpy.ndarray",")",":","Predictions","of","treatment","effects","for","compliers",",",",","i.e",".","those","individuals","who","take","the","treatment","only","if","they","are","assigned",".","Output","dim",":","[","n_samples",",","n_treatment","]","If","return_ci",",","returns","CATE","[","n_samples",",","n_treatment","]",",","LB","[","n_samples",",","n_treatment","]",",","UB","[","n_samples",",","n_treatment","]"]
code: def fit_predict(
      self,
      X,
      assignment,
      treatment,
      y,
      p=None,
      pZ=None,
      return_ci=False,
      n_bootstraps=1000,
      bootstrap_size=10000,
      return_components=False,
      verbose=True,
      seed=None,
      calibrate=True,
  ):
      """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
  
      Args:
          X (np.matrix or np.array or pd.Dataframe): a feature matrix
          assignment (np.array or pd.Series): a (0,1)-valued 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.
          return_ci (bool): whether to return confidence intervals
          n_bootstraps (int): number of bootstrap iterations
          bootstrap_size (int): number of samples per bootstrap
          return_components (bool, optional): whether to return outcome for treatment and control seperately
          verbose (str): whether to output progress logs
          seed (int): random seed for cross-fitting
      Returns:
          (numpy.ndarray): Predictions of treatment effects for compliers, , i.e. those individuals
              who take the treatment only if they are assigned. Output dim: [n_samples, n_treatment]
              If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
              UB [n_samples, n_treatment]
      """
      X, assignment, treatment, y = convert_pd_to_np(
          X, assignment, treatment, y
      )
      self.fit(X, assignment, treatment, y, p, seed, calibrate)
  
      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()
              },
          )
  
      if pZ is None:
          pZ = self.propensity_assign
  
      te = self.predict(
          X,
          treatment=treatment,
          y=y,
          return_components=return_components,
      )
  
      if not return_ci:
          return te
      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)
          te_bootstraps = np.zeros(
              shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
          )
  
          logger.info("Bootstrap Confidence Intervals")
          for i in tqdm(range(n_bootstraps)):
              te_b = self.bootstrap(
                  X,
                  assignment,
                  treatment,
                  y,
                  p,
                  pZ,
                  size=bootstrap_size,
                  seed=seed,
              )
              te_bootstraps[:, :, i] = te_b
  
          te_lower = np.percentile(
              te_bootstraps, (self.ate_alpha / 2) * 100, axis=2
          )
          te_upper = np.percentile(
              te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
          )
  
          # 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 (te, te_lower, te_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 5 for calling the BaseDRIVLearner obj's underlying member method code with example usage: obj.fit_predict(X, assignment, treatment, y, p, pZ, return_ci, n_bootstraps, bootstrap_size, return_components, verbose, seed, calibrate) and returns: te, te,  te_lower,  te_upper