causalml/inference/meta/rlearner.py [85:113]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        X, treatment, y = convert_pd_to_np(X, treatment, y)
        check_treatment_vector(treatment, self.control_name)
        self.t_groups = np.unique(treatment[treatment != self.control_name])
        self.t_groups.sort()

        if p is None:
            logger.info('Generating propensity score')
            p = dict()
            p_model = dict()
            for group in self.t_groups:
                mask = (treatment == group) | (treatment == self.control_name)
                treatment_filt = treatment[mask]
                X_filt = X[mask]
                w_filt = (treatment_filt == group).astype(int)
                w = (treatment == group).astype(int)
                p[group], p_model[group] = compute_propensity_score(X=X_filt, treatment=w_filt,
                                                                    X_pred=X, treatment_pred=w)
            self.propensity_model = p_model
            self.propensity = p
        else:
            check_p_conditions(p, self.t_groups)

        if isinstance(p, (np.ndarray, pd.Series)):
            treatment_name = self.t_groups[0]
            p = {treatment_name: convert_pd_to_np(p)}
        elif isinstance(p, dict):
            p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}

        self._classes = {group: i for i, group in enumerate(self.t_groups)}
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



causalml/inference/meta/xlearner.py [96:124]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        X, treatment, y = convert_pd_to_np(X, treatment, y)
        check_treatment_vector(treatment, self.control_name)
        self.t_groups = np.unique(treatment[treatment != self.control_name])
        self.t_groups.sort()

        if p is None:
            logger.info('Generating propensity score')
            p = dict()
            p_model = dict()
            for group in self.t_groups:
                mask = (treatment == group) | (treatment == self.control_name)
                treatment_filt = treatment[mask]
                X_filt = X[mask]
                w_filt = (treatment_filt == group).astype(int)
                w = (treatment == group).astype(int)
                p[group], p_model[group] = compute_propensity_score(X=X_filt, treatment=w_filt,
                                                                    X_pred=X, treatment_pred=w)
            self.propensity_model = p_model
            self.propensity = p
        else:
            check_p_conditions(p, self.t_groups)

        if isinstance(p, (np.ndarray, pd.Series)):
            treatment_name = self.t_groups[0]
            p = {treatment_name: convert_pd_to_np(p)}
        elif isinstance(p, dict):
            p = {treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()}

        self._classes = {group: i for i, group in enumerate(self.t_groups)}
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -