# Buggy code
def estimate_ate(
    self,
    X,
    treatment,
    y,
    p=None,
    bootstrap_ci=False,
    n_bootstraps=1000,
    bootstrap_size=10000,
):
    """Estimate the Average Treatment Effect (ATE).

    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.
        bootstrap_ci (bool): whether run bootstrap for confidence intervals
        n_bootstraps (int): number of bootstrap iterations
        bootstrap_size (int): number of samples per bootstrap
    Returns:
        The mean and confidence interval (LB, UB) of the ATE estimate.
    """
    X, treatment, y = convert_pd_to_np(X, treatment, y)
    te = self.fit_predict(X, treatment, y, p)

    if p is None:
        p = self.propensity
    else:
        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()}

    ate = np.zeros(self.t_groups.shape[0])
    ate_lb = np.zeros(self.t_groups.shape[0])
    ate_ub = np.zeros(self.t_groups.shape[0])

    for i, group in enumerate(self.t_groups):
        w = (treatment == group).astype(int)
        prob_treatment = float(sum(w)) / X.shape[0]
        _ate = te[:, i].mean()

        se = (
            np.sqrt(
                (self.vars_t[group] / prob_treatment)
                + (self.vars_c[group] / (1 - prob_treatment))
                + te[:, i].var()
            )
            / X.shape[0]
        )

        _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
        _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)

        ate[i] = _ate
        ate_lb[i] = _ate_lb
        ate_ub[i] = _ate_ub

    if not bootstrap_ci:
        return ate, ate_lb, ate_ub
    else:
        t_groups_global = self.t_groups
        _classes_global = self._classes
        model_mu_global = deepcopy(self.model_mu)
        models_tau_global = deepcopy(self.models_tau)

        logger.info("Bootstrap Confidence Intervals for ATE")
        ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))

        for n in tqdm(range(n_bootstraps)):
            cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
            ate_bootstraps[:, n] = cate_b.mean()

        ate_lower = np.percentile(ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1)
        ate_upper = np.percentile(
            ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
        )

        # set member variables back to global (currently last bootstrapped outcome)
        self.t_groups = t_groups_global
        self._classes = _classes_global
        self.model_mu = deepcopy(model_mu_global)
        self.models_tau = deepcopy(models_tau_global)
        return ate, ate_lower, ate_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 estimate_ate(
    self,
    X,
    treatment,
    y,
    p=None,
    bootstrap_ci=False,
    n_bootstraps=1000,
    bootstrap_size=10000,
):
    """Estimate the Average Treatment Effect (ATE).

    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.
        bootstrap_ci (bool): whether run bootstrap for confidence intervals
        n_bootstraps (int): number of bootstrap iterations
        bootstrap_size (int): number of samples per bootstrap
    Returns:
        The mean and confidence interval (LB, UB) of the ATE estimate.
    """
    X, treatment, y = convert_pd_to_np(X, treatment, y)
    te = self.fit_predict(X, treatment, y, p)

    if p is None:
        p = self.propensity
    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()}

    ate = np.zeros(self.t_groups.shape[0])
    ate_lb = np.zeros(self.t_groups.shape[0])
    ate_ub = np.zeros(self.t_groups.shape[0])

    for i, group in enumerate(self.t_groups):
        w = (treatment == group).astype(int)
        prob_treatment = float(sum(w)) / X.shape[0]
        _ate = te[:, i].mean()

        se = (
            np.sqrt(
                (self.vars_t[group] / prob_treatment)
                + (self.vars_c[group] / (1 - prob_treatment))
                + te[:, i].var()
            )
            / X.shape[0]
        )

        _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
        _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)

        ate[i] = _ate
        ate_lb[i] = _ate_lb
        ate_ub[i] = _ate_ub

    if not bootstrap_ci:
        return ate, ate_lb, ate_ub
    else:
        t_groups_global = self.t_groups
        _classes_global = self._classes
        model_mu_global = deepcopy(self.model_mu)
        models_tau_global = deepcopy(self.models_tau)

        logger.info("Bootstrap Confidence Intervals for ATE")
        ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))

        for n in tqdm(range(n_bootstraps)):
            cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
            ate_bootstraps[:, n] = cate_b.mean()

        ate_lower = np.percentile(ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1)
        ate_upper = np.percentile(
            ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
        )

        # set member variables back to global (currently last bootstrapped outcome)
        self.t_groups = t_groups_global
        self._classes = _classes_global
        self.model_mu = deepcopy(model_mu_global)
        self.models_tau = deepcopy(models_tau_global)
        return ate, ate_lower, ate_upper