class_docstr: A parent class for R-learner classifier classes.
docstr: Fit the treatment effect and outcome models of the R learner.
  
  Args:
      X (np.matrix or 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
      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.
      verbose (bool, optional): whether to output progress logs
docstr_tok: ["Fit","the","treatment","effect","and","outcome","models","of","the","R","learner",".","Args",":","X","(","np.matrix","or","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","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",".","verbose","(","bool",",","optional",")",":","whether","to","output","progress","logs"]
code: def fit(
      self, X, treatment, y, p=None, sample_weight=None, verbose=True
  ):
      """Fit the treatment effect and outcome models of the R learner.
  
      Args:
          X (np.matrix or 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
          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.
          verbose (bool, optional): whether to output progress logs
      """
      X, treatment, y = convert_pd_to_np(X, treatment, y)
      check_treatment_vector(treatment, self.control_name)
      if sample_weight is not None:
          assert len(sample_weight) == len(
              y
          ), "Data length must be equal for sample_weight and the input data"
          sample_weight = convert_pd_to_np(sample_weight)
      self.t_groups = np.unique(
          treatment[treatment != self.control_name]
      )
      self.t_groups.sort()
  
      if p is None:
          self._set_propensity_models(X=X, treatment=treatment, y=y)
          p = self.propensity
      else:
          p = self._format_p(p, self.t_groups)
  
      self._classes = {
          group: i for i, group in enumerate(self.t_groups)
      }
      self.models_tau = {
          group: deepcopy(self.model_tau) for group in self.t_groups
      }
      self.vars_c = {}
      self.vars_t = {}
  
      if verbose:
          logger.info("generating out-of-fold CV outcome estimates")
      yhat = cross_val_predict(
          self.model_mu,
          X,
          y,
          cv=self.cv,
          method="predict_proba",
          n_jobs=-1,
      )[:, 1]
  
      for group in self.t_groups:
          mask = (treatment == group) | (treatment == self.control_name)
          treatment_filt = treatment[mask]
          X_filt = X[mask]
          y_filt = y[mask]
          yhat_filt = yhat[mask]
          p_filt = p[group][mask]
          w = (treatment_filt == group).astype(int)
  
          weight = (w - p_filt) ** 2
          diff_c = y_filt[w == 0] - yhat_filt[w == 0]
          diff_t = y_filt[w == 1] - yhat_filt[w == 1]
          if sample_weight is not None:
              sample_weight_filt = sample_weight[mask]
              sample_weight_filt_c = sample_weight_filt[w == 0]
              sample_weight_filt_t = sample_weight_filt[w == 1]
              self.vars_c[group] = get_weighted_variance(
                  diff_c, sample_weight_filt_c
              )
              self.vars_t[group] = get_weighted_variance(
                  diff_t, sample_weight_filt_t
              )
              weight *= sample_weight_filt  # update weight
          else:
              self.vars_c[group] = diff_c.var()
              self.vars_t[group] = diff_t.var()
  
          if verbose:
              logger.info(
                  "training the treatment effect model for {} with R-loss".format(
                      group
                  )
              )
          self.models_tau[group].fit(
              X_filt,
              (y_filt - yhat_filt) / (w - p_filt),
              sample_weight=weight,
          )
  
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 2 for calling the BaseRClassifier obj's underlying member method code with example usage: obj.fit(X, treatment, y, p, sample_weight, verbose) without return types