robust and reduces the burden of model specification by the researcher. Subsequently, the average difference in the treatment effect within the desired principal strata is calculated. Variance is estimated using Rubin’s combine rule to account for the imputation uncertainty [@heitjan]. Detailed steps for our method are given below.

1.  Generate a bootstrap sample $b$ from the data by sampling the units with replacement.

2.  Estimate the model $X_{z_1^{(b)}}^{(b)}|Z_1^{(b)}=z_1^{(b)},W_1^{(b)},V^{(b)}$. Use this model to compute the counterfactual of $X_{z_1^{(b)}}^{(b)}$ for bootstrap sample $b$.

3.  Estimate the distribution of $Z_1^{(b)}|W_1^{(b)},V^{(b)}$. Use this model to estimate the propensity to be assigned treatment $Z_1^{(b)}=z_1^{(b)}$ as $P^*_{z_1^{(b)}}=Pr(Z_1^{(b)}=z_1^{(b)}|W_1^{(b)},V^{(b)})$. Note that we did not perform a logit transformation to obtain $P^*_{z_1^{(b)}}$ (See PENCOMP Steps 2 and 5). This is because by using PENCOMP modified with BART to predict the outcomes, the non-linear effect of the propensity of assigned treatment will be handled automatically. Hence, any non-linear transformation on the propensity of assigned treatment would not be needed.

4.  Estimate the model $Y_{z_1^{(b)}}^{(b)}|P^*_{z_1^{(b)}},Z_1^{(b)}=z_1^{(b)},X_{z_1^{(b)}}^{(b)},W_1^{(b)},V^{(b)}$. As mentioned, we used PENCOMP modified with BART to estimate this model. The advantage of using BART is the researcher no longer needs to specify the model. BART automatically takes care of any linear or non-linear 