# Buggy code
def fit_predict(
    self,
    X,
    treatment,
    y,
    p=None,
    return_ci=False,
    n_bootstraps=1000,
    bootstrap_size=10000,
    return_components=False,
    verbose=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
        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.
        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
    Returns:
        (numpy.ndarray): Predictions of treatment effects. 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, treatment, y = convert_pd_to_np(X, treatment, y)
    self.fit(X, treatment, y, p)

    if p is None:
        p = self.propensity

    check_p_conditions(p, self.t_groups)
    if isinstance(p, np.ndarray):
        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()}

    te = self.predict(
        X, treatment=treatment, y=y, p=p, 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_c_global = deepcopy(self.models_tau_c)
        models_tau_t_global = deepcopy(self.models_tau_t)
        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, treatment, y, p, size=bootstrap_size)
            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_c = deepcopy(models_tau_c_global)
        self.models_tau_t = deepcopy(models_tau_t_global)

        return (te, te_lower, te_upper)


# Traceback (IndexError)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-59-76ffa2ace99a> in <module>
1 learner_x = BaseXRegressor(learner=XGBRegressor(random_state=42))
----> 2 te, lb, ub = learner_x.estimate_ate(X=X, p=df['p'], treatment=treatment, y=y)
3 print('Average Treatment Effect (BaseXRegressor using XGBoost): {:.2f} ({:.2f}, {:.2f})'.format(te[0], lb[0], ub[0]))

/dsw/snapshots/720ce1e8-1e96-45bf-a192-8bef9122fb73/python3/lib/python3.6/site-packages/causalml/inference/meta/xlearner.py in estimate_ate(self, X, treatment, y, p, bootstrap_ci, n_bootstraps, bootstrap_size)
326             dhat_c = dhat_cs[group][mask]
327             dhat_t = dhat_ts[group][mask]
--> 328             p_filt = p[group][mask]
329
330             # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.

IndexError: invalid index to scalar variable.

# Fixed code
def fit_predict(
    self,
    X,
    treatment,
    y,
    p=None,
    return_ci=False,
    n_bootstraps=1000,
    bootstrap_size=10000,
    return_components=False,
    verbose=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
        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.
        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
    Returns:
        (numpy.ndarray): Predictions of treatment effects. 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, treatment, y = convert_pd_to_np(X, treatment, y)
    self.fit(X, treatment, y, p)

    if p is None:
        p = self.propensity

    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()}

    te = self.predict(
        X, treatment=treatment, y=y, p=p, 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_c_global = deepcopy(self.models_tau_c)
        models_tau_t_global = deepcopy(self.models_tau_t)
        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, treatment, y, p, size=bootstrap_size)
            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_c = deepcopy(models_tau_c_global)
        self.models_tau_t = deepcopy(models_tau_t_global)

        return (te, te_lower, te_upper)