2021-intercept-correlations.Rmd

---
title: "Correlations between overall intercept and varying intercept"
date: 2021-12-08
tags: ["R","Stan","Mixed models"]
draft: TRUE
---



```{r setup}
library(cmdstanr)
library(bayesplot)
options(mc.cores = parallel::detectCores())

# Includes sd, intercept, and the first two and last two varying intercepts
my_pars <- function(group_basename, N_groups, intercept_name = "intercept") {
  c("group_sd", intercept_name, paste0(group_basename, "[", c(1,2,N_groups - 1, N_groups), "]"))
}

my_transformations <- list(group_sd = "log")
```

```{r}
set.seed(5422335)
N <- 40
N_groups <- 20

groups <- rep(1:N_groups, length.out = N)

intercept <- 3 
group_sd <- 1
group_r <- rnorm(N_groups, sd = group_sd)

mu <- intercept + group_r[groups]
y <- rpois(N, exp(mu))

data_stan <- list(N = N,
                  N_groups = N_groups,
                  groups = groups,
                  y = y)
```

## Non-centered

```{r}
model_code_nc <- "
data {
  int<lower=0> N;
  int<lower=1> N_groups;
  int<lower=1,upper=N_groups> groups[N];
  int<lower=0> y[N];
}

parameters {
  real intercept;
  vector[N_groups] group_z; 
  real<lower=0> group_sd;
}

transformed parameters {
  vector[N_groups] group_r = group_z * group_sd;
}

model {
  intercept ~ normal(3,1);
  group_z ~ std_normal();
  group_sd ~ normal(0, 1);
  y ~ poisson_log(intercept + group_r[groups]);
}

"

model_nc <- cmdstan_model(write_stan_file(model_code_nc))
```


```{r}
fit_nc <- model_nc$sample(data = data_stan, refresh = 0)
```

```{r}
fit_nc$summary()
```

```{r}
mcmc_pairs(fit_nc$draws(), pars = my_pars("group_z", N_groups), transformations = my_transformations, np = nuts_params(fit_nc))
```

## Centered



```{r}
model_code_c <- "
data {
  int<lower=0> N;
  int<lower=1> N_groups;
  int<lower=1,upper=N_groups> groups[N];
  int<lower=0> y[N];
}

parameters {
  real intercept;
  vector[N_groups] group_r; 
  real<lower=0> group_sd;
}

model {
  intercept ~ normal(3,1);
  group_r ~ normal(0, group_sd);
  group_sd ~ normal(0, 1);
  y ~ poisson_log(intercept + group_r[groups]);
}

"

model_c <- cmdstan_model(write_stan_file(model_code_c))
```
```{r}
fit_c <- model_c$sample(data = data_stan, refresh = 0)
```
```{r}
fit_c$summary()
```

```{r}
mcmc_pairs(fit_c$draws(), pars = my_pars("group_r", N_groups), transformations = my_transformations, np = nuts_params(fit_c))
```

## Centered - sum to zero

```{r}
model_code_c_sz <- "
functions {
  vector Q_sum_to_zero_QR(int N) {
    vector [2*N] Q_r;

    for(i in 1:N) {
      Q_r[i] = -sqrt((N-i)/(N-i+1.0));
      Q_r[i+N] = inv_sqrt((N-i) * (N-i+1));
    }
    Q_r = Q_r * inv_sqrt(1 - inv(N));
    return Q_r;
  }

  vector sum_to_zero_QR(vector x_raw, vector Q_r) {
    int N = num_elements(x_raw) + 1;
    vector [N] x;
    real x_aux = 0;

    for(i in 1:N-1){
      x[i] = x_aux + x_raw[i] * Q_r[i];
      x_aux = x_aux + x_raw[i] * Q_r[i+N];
    }
    x[N] = x_aux;
    return x;
  }
}

data {
  int<lower=0> N;
  int<lower=1> N_groups;
  int<lower=1,upper=N_groups> groups[N];
  int<lower=0> y[N];
}

transformed data {
  vector[2 * N_groups] groups_Q_r = Q_sum_to_zero_QR(N_groups);
}

parameters {
  real intercept;
  vector[N_groups - 1] group_r_raw; 
  real<lower=0> group_sd;
}

transformed parameters {
  vector[N_groups] group_r = sum_to_zero_QR(group_r_raw, groups_Q_r);
}

model {
  intercept ~ normal(3,1);
  group_r_raw ~ normal(0, group_sd);
  group_sd ~ normal(0, 1);
  y ~ poisson_log(intercept + group_r[groups]);
}

"

model_c_sz <- cmdstan_model(write_stan_file(model_code_c_sz))
```
```{r}
fit_c_sz <- model_c_sz$sample(data = data_stan, refresh = 0)
```
```{r}
fit_c_sz$summary()
```

```{r}
mcmc_pairs(fit_c_sz$draws(), pars = my_pars("group_r_raw", N_groups - 1), transformations = my_transformations, np = nuts_params(fit_c_sz))
```

## Non-centered - sum to zero

```{r}
model_code_nc_sz <- "
functions {
  vector Q_sum_to_zero_QR(int N) {
    vector [2*N] Q_r;

    for(i in 1:N) {
      Q_r[i] = -sqrt((N-i)/(N-i+1.0));
      Q_r[i+N] = inv_sqrt((N-i) * (N-i+1));
    }
    Q_r = Q_r * inv_sqrt(1 - inv(N));
    return Q_r;
  }

  vector sum_to_zero_QR(vector x_raw, vector Q_r) {
    int N = num_elements(x_raw) + 1;
    vector [N] x;
    real x_aux = 0;

    for(i in 1:N-1){
      x[i] = x_aux + x_raw[i] * Q_r[i];
      x_aux = x_aux + x_raw[i] * Q_r[i+N];
    }
    x[N] = x_aux;
    return x;
  }
}

data {
  int<lower=0> N;
  int<lower=1> N_groups;
  int<lower=1,upper=N_groups> groups[N];
  int<lower=0> y[N];
}

transformed data {
  vector[2 * N_groups] groups_Q_r = Q_sum_to_zero_QR(N_groups);
}

parameters {
  real intercept;
  vector[N_groups - 1] group_r_raw; 
  real<lower=0> group_sd;
}

transformed parameters {
  vector[N_groups] group_r = sum_to_zero_QR(group_r_raw, groups_Q_r) * group_sd;
}

model {
  intercept ~ normal(3,1);
  group_r_raw ~ normal(0, 1);
  group_sd ~ normal(0, 1);
  y ~ poisson_log(intercept + group_r[groups]);
}

"

model_nc_sz <- cmdstan_model(write_stan_file(model_code_nc_sz))
```
```{r}
fit_nc_sz <- model_nc_sz$sample(data = data_stan)
```
```{r}
fit_nc_sz$summary()
```

```{r}
mcmc_pairs(fit_nc_sz$draws(), pars = my_pars("group_r_raw", N_groups - 1), transformations = my_transformations, np = nuts_params(fit_nc_sz))
```

## Sufficient statistics, centered


```{r}
model_suff_c <- cmdstan_model("intercept_suff_centered.stan")
```


```{r}
fit_suff_c <- model_suff_c$sample(data = data_stan, refresh = 0)
```

```{r}
fit_suff_c$summary()
```

```{r}
mcmc_pairs(fit_suff_c$draws(), pars = c("group_sd", "intercept_sweep", paste0("group_r_raw[", c(1, 2, N_groups - 1), "]")), transformations = my_transformations, np = nuts_params(fit_suff_c), max_treedepth = 10)

# mcmc_pairs(fit_suff_c$draws(), pars = c("group_sd", "intercept_sweep", paste0("group_r_raw[", c(1), "]"), "mean_group_r_raw"), transformations = my_transformations, np = nuts_params(fit_suff_c), max_treedepth = 10)

```

```{r}
mcmc_pairs(fit_suff_c$draws(), regex_pars = c("m_R", "sigma_R"))

```


### SBC test

```{r setupSBC}
library(SBC)
library(future)
plan(multisession)
options(SBC.min_chunk_size = 5)

SBC_cache_dir <- "_intercept-correlations-sbc-cache"
if(!dir.exists(SBC_cache_dir)) {
  dir.create(SBC_cache_dir)
}

```


```{r}

generator <- function(N, N_groups) {

  groups <- rep(1:N_groups, length.out = N)

  intercept <- rnorm(1, mean = 3, sd = 1) 
  group_sd <- abs(rnorm(1))
  group_r <- rnorm(N_groups, sd = group_sd)

  mu <- intercept + group_r[groups]
  y <- rpois(N, exp(mu))

  list(parameters = list(
    intercept = intercept,
    group_sd = group_sd,
    group_r = group_r
  ),
       generated = list(N = N,
                    N_groups = N_groups,
                    groups = groups,
                    y = y)
  )
}

set.seed(5465524)

generator_c <- SBC_generator_function(generator, N = 50, N_groups = 3)
datasets_c <- generate_datasets(generator_c, n_datasets = 500)

suff_c_backend <- SBC_backend_cmdstan_sample(model_suff_c, chains = 2)
SBC_res_suff_c <- compute_results(datasets_c, suff_c_backend, 
                                  gen_quants = generated_quantities(
                                    mu_r = mean(group_r),
                                    group_r_c = group_r - mean(group_r),
                                    intercept_c = intercept + mean(group_r),
                                    sd_c = sd(group_r),
                                    log_lik = sum(dpois(y, exp(intercept + group_r[groups]), log = TRUE))),
                                  keep_fits = FALSE,
                                  cache_mode = "results",
                                  cache_location = file.path(SBC_cache_dir, "suff_c_datasets_c.rds"))
```



```{r}
SBC_res_suff_c$backend_diagnostics
```


```{r}
# fit_id <- 1
# mcmc_pairs(SBC_res_suff_c$fits[[fit_id]]$draws(), pars = c("group_sd", "intercept_sweep", paste0("group_r_raw[", c(1, 2), "]")), transformations = my_transformations, np = nuts_params(SBC_res_suff_c$fits[[fit_id]]))

```


```{r}
plot_ecdf_diff(SBC_res_suff_c)
plot_ecdf_diff(SBC_res_suff_c[SBC_res_suff_c$backend_diagnostics$n_divergent == 0])
plot_rank_hist(SBC_res_suff_c)
plot_rank_hist(SBC_res_suff_c[SBC_res_suff_c$backend_diagnostics$n_divergent == 0])
```

## Sufficient statistics, non-centered


```{r}
model_suff_nc <- cmdstan_model("intercept_suff_noncentered.stan")
```


```{r}
fit_suff_nc <- model_suff_nc$sample(data = data_stan, refresh = 0)
```

```{r}
fit_suff_nc$summary()
```

```{r}
mcmc_pairs(fit_suff_nc$draws(), pars = my_pars("group_r_raw", N_groups - 1, "intercept_sweep"), transformations = my_transformations, np = nuts_params(fit_nc))

# mcmc_pairs(fit_suff_nc$draws(), pars = c("group_sd", "intercept_sweep", paste0("group_r_raw[", c(1), "]")), transformations = my_transformations, np = nuts_params(fit_suff_c), max_treedepth = 10)

```


### SBC Test

```{r}
suff_nc_backend <- SBC_backend_cmdstan_sample(model_suff_nc, chains = 2)
generator_nc <- SBC_generator_function(generator, N = 6, N_groups = 3)
datasets_nc <- generate_datasets(generator_nc, n_datasets = 50)

SBC_res_suff_nc <- compute_results(datasets_c, suff_nc_backend, 
                                  gen_quants = generated_quantities(
                                    mu_r = mean(group_r),
                                    group_r_c = group_r - mean(group_r),
                                    intercept_c = intercept + mean(group_r),
                                    sd_c = sd(group_r),
                                    log_lik = sum(dpois(y, exp(intercept + group_r[groups]), log = TRUE))),
                                  keep_fits = FALSE,
                                  cache_mode = "results",
                                  cache_location = file.path(SBC_cache_dir, "suff_nc_datasets_nc.rds"))

```


```{r}
plot_ecdf_diff(SBC_res_suff_nc)
plot_ecdf_diff(SBC_res_suff_nc[SBC_res_suff_nc$backend_diagnostics$n_divergent == 0])
plot_rank_hist(SBC_res_suff_nc)
plot_rank_hist(SBC_res_suff_nc[SBC_res_suff_nc$backend_diagnostics$n_divergent == 0])
```

```{r}
SBC_res_suff_nc$backend_diagnostics
```
```{r}
fit <- model_suff_nc$sample(datasets_nc$generated[[27]])
mcmc_pairs(fit$draws(), pars = my_pars("group_r_raw", 2, "intercept_sweep"), transformations = my_transformations, np = nuts_params(fit))
```