16.3_Smart_Drug_two_groups.Rmd

---
title: "16.3: Smart Drug (two groups)"
output: html_notebook
---


## Introduction

This is the example from Section 16.3 of Kruschke on the (fictitious) effect of the "smart drug", this time comparing the "smart drug" group against the placebo group. Make sure the "TwoGroupIQ.csv" file is in your project directory.


## Preliminaries

Load necessary packages:

```{r}
library(tidyverse)
library(rstan)
library(tidybayes)
library(bayesplot)
```

Set Stan to save compiled code.

```{r}
rstan_options(auto_write = TRUE)
```

Set Stan to use parallel processing where possible.

```{r}
options(mc.cores = parallel::detectCores())
```


## Data

```{r}
IQ_data <- read_csv("TwoGroupIQ.csv")
IQ_data
```

Since we are now comparing two groups, we can use the whole data set (120 total observations: 63 experimental and 57 control):

```{r}
table(IQ_data$Group)
```

We will, however, need to give the groups numerical values to pass to Stan (rather than the strings "Placebo" and "Smart Drug"). The easiest way to do this is to `factor` the `Group` variable and then convert the factor variable to a numeric variable. This will change the underlying structure to 1s and 2s. Observe:

```{r}
IQ_data <- IQ_data %>%
    mutate(Group_num = as.numeric(factor(Group)))
IQ_data
```

Now "Placebo" is 1 and "Smart Drug" is 2.

We will also calculate the total number of groups and pass it to Stan as data. While we only have two groups, we can write the code to allow for any number of groups.

```{r}
G <- IQ_data %>%
    summarize(n_distinct(Group)) %>%
    pull()
G
```


Bundle everything together into a list:

```{r}
N <- NROW(IQ_data)
y <- IQ_data$Score
G <- G
group <- IQ_data$Group_num
stan_data <- list(N = N, y = y, G = G, group = group)
```


## Prior

### Simulation code

Almost the same as in the other Chapter 16 examples, except that now we have a group index, `g`. There will be a $\mu_{g}$ and $\sigma_{g}$ for each group, but only one choice of `nu` (for convenience). 

```{stan, output.var = "IQ2_prior", cache = TRUE}
data {
    int<lower = 0> N;
    array[N] real y;
    int<lower = 1> G;
    array[N] int<lower = 1, upper = G> group;
}
transformed data {
    real M;
    real<lower = 0> S;
    real<lower = 0> R_nu_minus_one;
    real<lower = 0> R_sigma;
    
    M = 100;   // mean centered at 100
    S = 50;    // mean between 0 and 200
    R_nu_minus_one = 1.0/29;  // df of 30
    R_sigma = 1.0/50;         // sd of 50
}
generated quantities {
    real<lower = 0> nu_minus_one;
    real<lower = 1> nu;
    array[G] real mu;
    array[G] real<lower = 0> sigma;
    array[N] real y_sim;
    
    nu_minus_one = exponential_rng(R_nu_minus_one);
    nu = nu_minus_one + 1;
    
    for(g in 1:G) {
        mu[g] = normal_rng(M, S);
        sigma[g] = exponential_rng(R_sigma);
    }
    
    for(n in 1:N) {
        y_sim[n] = student_t_rng(nu, mu[group[n]], sigma[group[n]]);
    }
}
```

```{r, cache = TRUE}
fit_IQ2_prior <- sampling(IQ2_prior,
                                data = stan_data,
                                chains = 1,
                                algorithm = "Fixed_param",
                                seed = 11111,
                                refresh = 0)
```

```{r}
samples_IQ2_prior <- tidy_draws(fit_IQ2_prior)
samples_IQ2_prior
```

### Examine prior

```{r}
mcmc_hist(samples_IQ2_prior,
          pars = "nu")
```

```{r}
mcmc_hist(samples_IQ2_prior,
          pars = vars(starts_with("mu")))
```

```{r}
mcmc_hist(samples_IQ2_prior,
          pars = vars(starts_with("sigma")))
```

```{r}
mcmc_pairs(fit_IQ2_prior,
           pars = vars("nu", starts_with(c("mu", "sigma"))))
```

### Prior predictive distribution

With multiple groups, this is a little tricky. For example, the column `y_sim[1]` is sampled using the group parameters for the 1st patient, who happens to be in the Smart Drug group, group 2. In other words, that data depends on $\mu_{2}$ and $\sigma_{2}$. However, `y_sim[120]` is sampled using $\mu_{1}$ and $\sigma_{1}$ since patient 120 was in the Placebo group.

While some `ppd_` commands have grouped versions, others do not. If it were easy to do so, we would `select` out only the `y_sim[j]` columns for all values of j belonging to one group. This would require quite a bit of code, however, so we'll just take what we can get.

At least for this data set, all the patients in each group appear consecutively in the data. Therefore, we get lucky in that we can just grab entries from each group separately: 1--63 for group 2 (Smart Drug) and 64--120 for group 1 (Placebo).

```{r}
y_sim_IQ2_prior <- samples_IQ2_prior %>%
    dplyr::select(starts_with("y_sim")) %>%
    as.matrix()
```

Group 1 (Placebo):

```{r}
ppd_hist(ypred = y_sim_IQ2_prior[1:20, 64:120])
```

```{r}
ppd_boxplot(ypred = y_sim_IQ2_prior[1:10, 64:120],
            notch = FALSE)
```

Group 2 (Smart Drug):

```{r}
ppd_hist(ypred = y_sim_IQ2_prior[1:20, 1:63])
```

```{r}
ppd_boxplot(ypred = y_sim_IQ2_prior[1:10, 1:63],
            notch = FALSE)
```

Some `ppd_` functions have a `grouped` version for comparing groups.

```{r}
ppd_intervals_grouped(ypred = y_sim_IQ2_prior,
                      group = group)
```


## Model

The only change here is the addition of a `diff` variable to record the difference between the two group means. (If G is larger than 2, this code could be changed to calculate the difference between any two parameters of interest.)

```{stan, output.var = "IQ2", cache = TRUE}
data {
    int<lower = 0> N;
    array[N] real y;
    int<lower = 1> G;
    array[N] int<lower = 1, upper = G> group;
}
transformed data {
    real M;
    real<lower = 0> S;
    real<lower = 0> R_nu_minus_one;
    real<lower = 0> R_sigma;
    
    M = 100;   // mean centered at 100
    S = 50;    // mean between 0 and 200
    R_nu_minus_one = 1.0/29;  // df of 30
    R_sigma = 1.0/50;         // sd of 50
}
parameters {
    real<lower = 0> nu_minus_one;
    array[G] real mu;
    array[G] real<lower = 0> sigma;
}
transformed parameters {
    real<lower = 1> nu;
    
    nu = nu_minus_one + 1;
}
model {
    nu_minus_one ~ exponential(R_nu_minus_one);
    mu ~ normal(M, S);
    sigma ~ exponential(R_sigma);
    y ~ student_t(nu, mu[group], sigma[group]);
}
generated quantities {
    array[N] real y_sim;
    real diff;
    
    for (n in 1:N) {
        y_sim[n] = student_t_rng(nu, mu[group[n]], sigma[group[n]]);
    }
    
    diff = mu[1] - mu[2];
}
```

```{r, cache = TRUE}
fit_IQ2 <- sampling(IQ2,
                   data = stan_data,
                   seed = 11111,
                   refresh = 0)
```

```{r}
samples_IQ2 <- tidy_draws(fit_IQ2)
samples_IQ2
```

## Model diagnostics

```{r}
stan_trace(fit_IQ2,
           pars = c("nu", "mu", "sigma"))
```

```{r}
mcmc_acf(fit_IQ2,
         pars = vars("nu", starts_with(c("mu","sigma"))))
```

```{r}
mcmc_rhat(rhat(fit_IQ2))
```

```{r}
mcmc_neff(neff_ratio(fit_IQ2))
```


## Model summary

```{r}
print(fit_IQ2,
      pars = c("nu", "mu", "sigma"))
```


## Model visualization

```{r}
mcmc_areas(fit_IQ2, pars = "nu")
```

```{r}
mcmc_areas(fit_IQ2, pars = vars(starts_with("mu")))
```

```{r}
mcmc_areas(fit_IQ2, pars = vars(starts_with("sigma")))
```

```{r}
pairs(fit_IQ2,
      pars = c("nu", "mu", "sigma"))
```


## Posterior predictive check

```{r}
y_sim_IQ2 <- samples_IQ2 %>%
    dplyr::select(starts_with("y_sim")) %>%
    as.matrix()
```

Same issue as before: it's hard to make sure we look at representative samples from both groups in general, but it's easy for this data set.


Group 1 (Placebo):

```{r}
ppc_hist(y = y[64:120],
         yrep = y_sim_IQ2[1:19, 64:120])
```

```{r}
ppc_boxplot(y = y[64:120],
            yrep = y_sim_IQ2[1:10, 64:120],
            notch = FALSE)
```

```{r}
ppc_scatter(y = y[64:120],
            yrep = y_sim_IQ2[1:9, 64:120])
```

```{r}
ppc_stat_2d(y = y[64:120],
            yrep = y_sim_IQ2[, 64:120])
```

Group 2 (Smart Drug):

```{r}
ppc_hist(y = y[1:63],
         yrep = y_sim_IQ2[1:19, 1:63])
```

```{r}
ppc_boxplot(y = y[1:63],
            yrep = y_sim_IQ2[1:10, 1:63],
            notch = FALSE)
```

```{r}
ppc_scatter(y = y[1:63],
            yrep = y_sim_IQ2[1:9, 1:63])
```

```{r}
ppc_stat_2d(y = y[1:63],
            yrep = y_sim_IQ2[, 1:63])
```

The following PPC graphs have grouped versions.

```{r}
ppc_stat_grouped(y = y,
                 yrep = y_sim_IQ2,
                 group = group)
```

```{r}
ppc_dens_overlay_grouped(y = y,
                         yrep = y_sim_IQ2[1:50, ],
                         group = group)
```

```{r}
ppc_scatter_avg_grouped(y = y,
                        yrep = y_sim_IQ2,
                        group = group)
```

```{r}
ppc_intervals_grouped(y = y,
                      yrep = y_sim_IQ2[1:19, ],
                      group = group)
```