---
title: "detect_idiosyncratic() Tutorial"
author: "Bertling, M., Ding, P., Feller, A., and Miratrix, L."
date: "`r Sys.Date()`"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{detect_idiosyncratic() Tutorial}
  %\VignetteEncoding{UTF-8}
  %\VignetteEngine{knitr::rmarkdown}
editor_options: 
  chunk_output_type: console
---

```{r, echo=FALSE, message=FALSE, warning=FALSE}
par( mgp=c(1.8,0.8,0), mar=c(2.5, 2.5,2,0.1) )
knitr::opts_chunk$set(fig.width = 8)
```

# Introduction
This document demonstrates how to perform permutation test inference for heterogeneous treatment effects. We use a simulated dataset to illustrate. This data is provided with the package as \verb|ToyData|.
Overall, this vignette illustrates the different tests for ideosyncratic variation one might make using our package, and shows how different forms of covariate adjustment can increase power or target inferential questions differently.



---
#, with the data generating process being roughly calibrated to the National Head Start Impact Study (HSIS), a large-scale randomized evaluation of a Federal pre-school programm.  
# Unfortunately, we cannot directly release the HSIS data. 

# SAY MORE ABOUT DGP: what Xs, etc.  What is true model?
---


We first load our package along with some other useful packages.
```{r, message=FALSE, warning=FALSE}
library( mvtnorm )
library( ggplot2 )
library( hettx )
library( dplyr )
library( tidyr )
library( purrr )
data( ToyData )
```

# The illustrative Dataset

We begin with exploring a toy dataset with `r nrow( ToyData )` observations that we included with the package for illustration. `ToyData` has an outcome, four covariates, and a treatment indicator. 
They were generated with the following model:
\[
Y_i(0) =      1 +   x_{1i} + 2 x_{2i} + 4x_{3i} + \epsilon_i
\]
with $\epsilon_i \sim N( 0, 1^2 )$. 
The treatment model is all systematic:
\[
  \tau_i =     2 + 2x_{1i} + x_{2i}
\]
So $x_1$ and $x_2$ are both predictive of treatment impact.  $x_3$ also predicts control side variation, and should therefore be useful for increasing precision.  $x_4$ is useless.
$Y_i(1) = Y_i(0) + \tau_i$, so there is no ideosyncratic variation if we control for both $x_1$ and $x_2$.

As we generated these data, we know the true individual treatment effects.  

```{r, echo=TRUE}
data( ToyData )
head( ToyData )
td <- pivot_longer( ToyData, cols=c(x1, x2, x3, x4), names_to="X", values_to="value" )
td <- pivot_longer( td, cols=c(Y, tau), names_to="outcome", values_to="value2" )
ggplot( td, aes( x=value, y=value2, col=as.factor(Z) ) ) +
        facet_grid( outcome ~ X, scales="free" ) +
        geom_point( alpha=0.5, size=0.5) + 
        geom_smooth( method="loess", se=FALSE ) +
        labs( x="Covariates", y="" )
```

As the data is simulated, we have the true finite-sample ATE:
```{r, echo=TRUE}
mean( ToyData$tau )
```

Before testing, we quickly look at the marginal and residual CDFs of treatment and control. We see heterogeniety at left, but after controlling for observed covariates we have no visibile ideosyncratic heterogeniety left over.

```{r, echo=TRUE}
par( mfrow=c(1,2) )
ll0 = lm( Y ~ Z, data=ToyData )
plot( ecdf( resid(ll0)[ToyData$Z==1] ), pch=".", main="Marginal CDFs of \n treatment and control")
plot( ecdf( resid(ll0)[ToyData$Z==0] ), pch=".", col="red", add=TRUE )

ll1 = lm( Y ~ Z + x1 + x2 + x3 + x4, data=ToyData )
plot( ecdf( resid(ll1)[ToyData$Z==1] ), pch=".", main="Residual CDFs of \n treatment and control" )
plot( ecdf( resid(ll1)[ToyData$Z==0] ), pch=".", col="red", add=TRUE )
```

A simple linear model should give us our parameters from above, since it is a correct specification in this case.

```{r, echo=TRUE}
M0 <- lm( Y ~ Z * (x1+x2+x3), data=ToyData )
round( coef( M0 ), digits=1 )
```



# Testing for ideosyncratic treatment effect variation 


## Basic case: no covariate adjustment

To do our tests, we first must specify some parameters determining the resolution of the grid search and number of permutations at each grid point. Note that these values should be increased for a real analysis. We chose these particular values to reduce computation time for illustration.

```{r, echo=TRUE, results='asis'}
B <- 20
grid.size = 11
```

The basic test for ideosyncratic treatment effect variation is as follows (with no adjustment for covariates):

```{r, echo=TRUE, cache=FALSE}
tst1 = detect_idiosyncratic( Y ~ Z, data=ToyData, B=B, grid.size = grid.size, verbose=FALSE )
summary( tst1 )
```



# Adjusting for covariates

We can increase the power by adjusting for covariates. Please note that we do not include any interaction terms at this point.

We specify a `control.formula` which will be used to generate a matrix to hand to the linear regression function. This will convert factors to dummy variables as needed.

```{r, echo=TRUE, cache=FALSE}
tst2 = detect_idiosyncratic( Y ~ Z, data=ToyData, control.formula = ~ x1 + x2 + x3 + x4, B=B, 
                             test.stat="SKS_stat_cov",  verbose=FALSE )
summary( tst2 )
```

Let's explore how the results might differ when non-treatment covariates are used.

```{r, echo=TRUE, cache=FALSE}
tst2b = detect_idiosyncratic( Y ~ Z, data=ToyData, control.formula = ~ x3 + x4, B=B, 
                              test.stat="SKS_stat_cov", verbose=FALSE )
summary( tst2b )
```

We can compare to when we correct for some useless covariates not related to outcome.

```{r, echo=TRUE, cache=FALSE}
N = nrow(ToyData)
ToyData$x4 = rnorm( N )
tst1b = detect_idiosyncratic( Y ~ Z, data=ToyData, control.formula = ~ x4, B=B, 
                              test.stat="SKS_stat_cov", verbose=FALSE )
summary( tst1b )
```


## Ideosyncratic variation beyond systematic variation

To test for ideosyncratic variation beyond systematic, we pass an `interaction.formula` to the method.

We first test for ideosyncratic variation beyond x1 (and we should get high $p$-value).

```{r ideo_beyond_systematic, echo=TRUE, cache=FALSE}
B = 20 

tst3a1 = detect_idiosyncratic( Y ~ Z, data=ToyData, interaction.formula = ~ x1, B=B, 
                               test.stat="SKS_stat_int_cov", verbose=FALSE )
summary( tst3a1 )
```

Include additional terms to increase power.  We are correcting for x3 and x4.
```{r, echo=TRUE, cache=FALSE}
tst3a2 <- detect_idiosyncratic( Y ~ Z, data=ToyData, 
                        interaction.formula = ~ x1, 
                        control.formula = ~ x3 + x4,
                        B=B, test.stat="SKS_stat_int_cov", 
                        verbose=FALSE )
summary( tst3a2 )
```

For comparison, we next include all terms.
```{r, echo=TRUE, cache=FALSE}
tst3a2b <- detect_idiosyncratic( Y ~ Z, data=ToyData, control.formula = ~ x2 + x3 + x4, 
                         interaction.formula = ~ x1, B=B, test.stat="SKS_stat_int_cov", 
                         verbose=FALSE )
summary( tst3a2b )
```


### Testing for ideosyncratic variation.

Now correct for the other covariates as well, but still have correct heterogeneous treatment model. Note that you should still get high $p$-value.

Start testing for variation beyond x1 and x2. 
```{r beyond_x1_x2, echo=TRUE, cache=FALSE}
tst3b <- detect_idiosyncratic( Y ~ Z, data=ToyData,
                         interaction.formula = ~ x1 + x2, B=B, test.stat="SKS_stat_int_cov", 
                         verbose=FALSE )
summary( tst3b )
```


Continue to test for ideosyncratic variation beyond x1 and x2, adjusting for x3 and x4. Note that you should still get high $p$-value.

```{r, echo=TRUE}
tst3c <- detect_idiosyncratic( Y ~ Z, data=ToyData,
                         interaction.formula = ~ x1 + x2, 
                       control.formula = ~ x3 + x4,
                       B=B, test.stat="SKS_stat_int_cov", 
                       verbose=FALSE )
summary( tst3c )
```


Finally, test for ideosyncratic variation beyond all covariates, even irrelevant ones. Again, you should expect to get high $p$-value.

```{r full_correction, echo=TRUE, cache=FALSE}
tst3d <- detect_idiosyncratic( Y ~ Z, data=ToyData,
                         interaction.formula = ~ x1 + x2 + x3 + x4, 
                       B=B, test.stat="SKS_stat_int_cov", 
              verbose=FALSE )
summary( tst3d )
```

# Comparing the tests

We can easily compare the results by simultaneously displaying the outputs from all tested models by using the `get_p_value()` method.
This method extracts some core summary statistics:
```{r}
get_p_value( tst1b )
```

We can thus collect all our models from above, and make a dataframe of the overall results:
```{r display, echo=TRUE}
tests = list( no_cov=tst1, useless_cov=tst1b, all_covariates=tst2, 
              non_tx_covariates_only=tst2b, het_beyond_x1 = tst3a1,
              het_beyond_x1_with_x3_x4_cov=tst3a2, het_beyond_x1_with_all_cov=tst3a2,
              het_beyond_x1_x2=tst3b, 
              het_beyond_x1_x2_with_cov=tst3c, het_beyond_all=tst3d )

agg.res <- purrr::map( tests, get_p_value ) %>%
  purrr::map( as.list )
agg.res <- bind_rows( agg.res, .id = "test" )
agg.res
```


# Cautionary Tale: A linear model with no treatment interaction.

Let's fit a model that allows for no systematic treatment impact heterogeniety.
This means that all variation would have to be considered ideosyncratic.
The key, however, is we control for covariates to increase precision.
```{r cautionary_tale, echo=TRUE}
ll1 = lm( Y ~ Z + x1 + x2 + x3 + x4, data=ToyData )
print( summary( ll1 ) )
```
The estimated ATE is close to the truth, as expected considering the random assignment.

Next plot residual CDFS of treatment and control groups.
```{r cautionary_tale_2, echo=TRUE}
plot( ecdf( resid(ll1)[ToyData$Z==1] ), pch=".", main="Residual CDFs of treatment and control" )
plot( ecdf( resid(ll1)[ToyData$Z==0] ), pch=".", col="red", add=TRUE )
```


Note the residual ECDFs from above are quite aligned.  The Tx effect variation has been picked up by main effects which means we would not detect ideosyncratic variation even though there is such variation.
In other words, the treatment variation has been distributed across the residuals of the control group as well as treatment, and thus when we compare the distributions they are the same.
This is why the choice of test statistic is delicate, and why for the covariate adjusted SKS statistic, we need to fit the model on the control units only, and then extract the residuals.


# A simple variance ratio test

We also offer an adjusted variance ratio test:
```{r}
variance_ratio_test( ToyData$Y, ToyData$Z )
```

This does not use permutation inference.

# The variety of test statistics

We offer several test statistics one might use.  We also have a method to print out some info on what is available:
```{r}
test_stat_info()
```
