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 inference model.
  
  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.
      seed (int): random seed for cross-fitting
docstr_tok: ["Fit","the","inference","model",".","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",".","seed","(","int",")",":","random","seed","for","cross-fitting"]
code: def fit(
      self,
      X,
      assignment,
      treatment,
      y,
      p=None,
      pZ=None,
      seed=None,
      calibrate=True,
  ):
      """Fit the inference model.
  
      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.
          seed (int): random seed for cross-fitting
      """
      X, treatment, assignment, y = convert_pd_to_np(
          X, treatment, assignment, y
      )
      check_treatment_vector(treatment, self.control_name)
      self.t_groups = np.unique(
          treatment[treatment != self.control_name]
      )
      self.t_groups.sort()
      self._classes = {
          group: i for i, group in enumerate(self.t_groups)
      }
  
      # The estimator splits the data into 3 partitions for cross-fit on the propensity score estimation,
      # the outcome regression, and the treatment regression on the doubly robust estimates. The use of
      # the partitions is rotated so we do not lose on the sample size.  We do not cross-fit the assignment
      # score estimation as the assignment process is usually simple.
      cv = KFold(n_splits=3, shuffle=True, random_state=seed)
      split_indices = [index for _, index in cv.split(y)]
  
      self.models_mu_c = {
          group: [
              deepcopy(self.model_mu_c),
              deepcopy(self.model_mu_c),
              deepcopy(self.model_mu_c),
          ]
          for group in self.t_groups
      }
      self.models_mu_t = {
          group: [
              deepcopy(self.model_mu_t),
              deepcopy(self.model_mu_t),
              deepcopy(self.model_mu_t),
          ]
          for group in self.t_groups
      }
      self.models_tau = {
          group: [
              deepcopy(self.model_tau),
              deepcopy(self.model_tau),
              deepcopy(self.model_tau),
          ]
          for group in self.t_groups
      }
  
      if p is None:
          self.propensity_1 = {
              group: np.zeros(y.shape[0]) for group in self.t_groups
          }  # propensity scores for those assigned
          self.propensity_0 = {
              group: np.zeros(y.shape[0]) for group in self.t_groups
          }  # propensity scores for those not assigned
      if pZ is None:
          self.propensity_assign, _ = compute_propensity_score(
              X=X,
              treatment=assignment,
              X_pred=X,
              treatment_pred=assignment,
              calibrate_p=calibrate,
          )
      else:
          self.propensity_assign = pZ
  
      for ifold in range(3):
          treatment_idx = split_indices[ifold]
          outcome_idx = split_indices[(ifold + 1) % 3]
          tau_idx = split_indices[(ifold + 2) % 3]
  
          treatment_treat, treatment_out, treatment_tau = (
              treatment[treatment_idx],
              treatment[outcome_idx],
              treatment[tau_idx],
          )
          assignment_treat, assignment_out, assignment_tau = (
              assignment[treatment_idx],
              assignment[outcome_idx],
              assignment[tau_idx],
          )
          y_out, y_tau = y[outcome_idx], y[tau_idx]
          X_treat, X_out, X_tau = (
              X[treatment_idx],
              X[outcome_idx],
              X[tau_idx],
          )
          pZ_tau = self.propensity_assign[tau_idx]
  
          if p is None:
              logger.info("Generating propensity score")
              cur_p_1 = dict()
              cur_p_0 = dict()
  
              for group in self.t_groups:
                  mask = (treatment_treat == group) | (
                      treatment_treat == self.control_name
                  )
                  mask_1, mask_0 = (
                      mask & (assignment_treat == 1),
                      mask & (assignment_treat == 0),
                  )
                  cur_p_1[group], _ = compute_propensity_score(
                      X=X_treat[mask_1],
                      treatment=(
                          treatment_treat[mask_1] == group
                      ).astype(int),
                      X_pred=X_tau,
                      treatment_pred=(treatment_tau == group).astype(
                          int
                      ),
                  )
                  if (treatment_treat[mask_0] == group).sum() == 0:
                      cur_p_0[group] = np.zeros(X_tau.shape[0])
                  else:
                      cur_p_0[group], _ = compute_propensity_score(
                          X=X_treat[mask_0],
                          treatment=(
                              treatment_treat[mask_0] == group
                          ).astype(int),
                          X_pred=X_tau,
                          treatment_pred=(
                              treatment_tau == group
                          ).astype(int),
                      )
                  self.propensity_1[group][tau_idx] = cur_p_1[group]
                  self.propensity_0[group][tau_idx] = cur_p_0[group]
          else:
              cur_p_1 = dict()
              cur_p_0 = dict()
              if isinstance(p[0], (np.ndarray, pd.Series)):
                  cur_p_0 = {
                      self.t_groups[0]: convert_pd_to_np(p[0][tau_idx])
                  }
              else:
                  cur_p_0 = {
                      g: prop[tau_idx] for g, prop in p[0].items()
                  }
              check_p_conditions(cur_p_0, self.t_groups)
  
              if isinstance(p[1], (np.ndarray, pd.Series)):
                  cur_p_1 = {
                      self.t_groups[0]: convert_pd_to_np(p[1][tau_idx])
                  }
              else:
                  cur_p_1 = {
                      g: prop[tau_idx] for g, prop in p[1].items()
                  }
              check_p_conditions(cur_p_1, self.t_groups)
  
          logger.info("Generate outcome regressions")
          for group in self.t_groups:
              mask = (treatment_out == group) | (
                  treatment_out == self.control_name
              )
              mask_1, mask_0 = (
                  mask & (assignment_out == 1),
                  mask & (assignment_out == 0),
              )
              self.models_mu_c[group][ifold].fit(
                  X_out[mask_0], y_out[mask_0]
              )
              self.models_mu_t[group][ifold].fit(
                  X_out[mask_1], y_out[mask_1]
              )
  
          logger.info("Fit pseudo outcomes from the DR formula")
  
          for group in self.t_groups:
              mask = (treatment_tau == group) | (
                  treatment_tau == self.control_name
              )
              treatment_filt = treatment_tau[mask]
              X_filt = X_tau[mask]
              y_filt = y_tau[mask]
              w_filt = (treatment_filt == group).astype(int)
              p_1_filt = cur_p_1[group][mask]
              p_0_filt = cur_p_0[group][mask]
              z_filt = assignment_tau[mask]
              pZ_filt = pZ_tau[mask]
              mu_t = self.models_mu_t[group][ifold].predict(X_filt)
              mu_c = self.models_mu_c[group][ifold].predict(X_filt)
              dr = (
                  z_filt * (y_filt - mu_t) / pZ_filt
                  - (1 - z_filt) * (y_filt - mu_c) / (1 - pZ_filt)
                  + mu_t
                  - mu_c
              )
              weight = (
                  z_filt * (w_filt - p_1_filt) / pZ_filt
                  - (1 - z_filt) * (w_filt - p_0_filt) / (1 - pZ_filt)
                  + p_1_filt
                  - p_0_filt
              )
              dr /= weight
              self.models_tau[group][ifold].fit(
                  X_filt, dr, sample_weight=weight**2
              )
  
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 3 for calling the BaseDRIVLearner obj's underlying member method code with example usage: obj.fit(X, assignment, treatment, y, p, pZ, seed, calibrate) without return types