Add STAN example to documentation

[florianhartig]

Example from Amael Le Squin for residual checks with STAN - to be integrated with the help

#### Example from https://cran.r-project.org/web/packages/DHARMa/vignettes/DHARMaForBayesians.html ran with Stan

library(data.table)
library(cmdstanr)
library(DHARMa)

#### Create data
set.seed(123)

dat = DHARMa::createData(200, overdispersion = 0.2)

Data = as.list(dat)
Data$nobs = nrow(dat)
Data$nGroups = length(unique(dat$group))

#### Fit data
## Common variables
n_chains = 3
maxIter = 3000

## Compile model
model = cmdstan_model("./florian_eg.stan")

## Run model
results = model$sample(data = Data, parallel_chains = n_chains, refresh = 50, chains = n_chains,
    iter_warmup = round(maxIter/2), iter_sampling = round(maxIter/2), save_warmup = FALSE)

n_sampling = results$metadata()$iter_sampling

#### Compute residuals
## Get parameters
intercept_array = as.vector(results$draws("intercept")) # dimensions: iter_sampling * n_chains * 1
env_array = as.vector(results$draws("env")) # dimensions: iter_sampling * n_chains * 1

## Posterior distributions
# Posterior lambda 
posteriorPredDistr = matrix(data = 0, nrow = n_sampling*n_chains, ncol = nrow(dat))
for (i in 1:nrow(dat))
    posteriorPredDistr[, i] = exp(intercept_array + env_array*Data$Environment1[i])

# Posterior response variable
posteriorPredSim = results$draws("observedResponseSim") # dimension: iter_sampling * n_chains * nobs

# Reshape posteriorPredSim as a matrix nobs x n_repetition
posteriorPredSim_matrix = matrix(data = 0, nrow = nrow(dat), ncol = n_sampling*n_chains)
for (obs in 1:nrow(posteriorPredSim_matrix))
{
    col_start = 1
    for (chain in 1:n_chains)
    {
        col_end = chain*n_sampling
        posteriorPredSim_matrix[obs, col_start:col_end] = posteriorPredSim[, chain, obs]
        col_start = col_end + 1
    }
}

## Create dharma data for the residuals
sim = createDHARMa(simulatedResponse = posteriorPredSim_matrix,
    observedResponse = dat$observedResponse,
    fittedPredictedResponse = apply(posteriorPredDistr, 2, median),
    integerResponse = TRUE)

## Plot residuals
plot(sim)

########################################################
###############    RANDOM EFFECT PART    ###############
########################################################
#### Add a random effect.
## Explanations:
# By default (see ?createData), the intercept is 0, the fixed effect is 1, the number of groups is 10, the random effect is on
# the intercept and its variance is 1. We create the expected result below and then run a stan model accounting for the random effect.

expected_results = c(intercept_mean = 0, env = 1, randEff_var = 1)

#### Create stan data
## Create indexing for groups (used in Stan to compute likelihood)
# Convert data to data table and order it by group
setDT(dat)
setorder(dat, group)

# Index
ls_groups = as.integer(dat[, unique(group)])
indices = integer(Data$nGroups + 1)

for (i in 1:Data$nGroups)
    indices[i] = min(which(dat[, group] == ls_groups[i]))
indices[Data$nGroups + 1] = Data$nobs + 1

## Data list for stan
Data = as.list(dat)
Data$nobs = nrow(dat)
Data$nGroups = length(unique(dat$group))
Data$indices = indices

#### Fit data
## Common variables
n_chains = 3
maxIter = 3000

## Compile model
model = cmdstan_model("./florian_eg_randEff.stan")

## Run model
results = model$sample(data = Data, parallel_chains = n_chains, refresh = 50, chains = n_chains,
    iter_warmup = round(maxIter/2), iter_sampling = round(maxIter/2), save_warmup = FALSE)

results$cmdstan_diagnose()

results$print(names(expected_results))
print(expected_results)

n_sampling = results$metadata()$iter_sampling

#### Compute residuals
## Posterior distributions
# Posterior lambda 
posteriorPredDistr = results$draws("lambda") # dimension: iter_sampling * n_chains * nobs
posteriorPredDistr_matrix = matrix(data = 0, nrow = n_sampling*n_chains, ncol = nrow(dat)) # One obs per column

for (i in 1:nrow(dat))
    posteriorPredDistr_matrix[, i] = as.vector(posteriorPredDistr[,,i])

# Posterior response variable
posteriorPredSim = results$draws("observedResponseSim") # dimension: iter_sampling * n_chains * nobs

# Reshape posteriorPredSim as a matrix nobs x n_repetition
posteriorPredSim_matrix = matrix(data = 0, nrow = nrow(dat), ncol = n_sampling*n_chains)
for (obs in 1:nrow(posteriorPredSim_matrix))
{
    col_start = 1
    for (chain in 1:n_chains)
    {
        col_end = chain*n_sampling
        posteriorPredSim_matrix[obs, col_start:col_end] = posteriorPredSim[, chain, obs]
        col_start = col_end + 1
    }
}

## Create dharma data for the residuals
sim = createDHARMa(simulatedResponse = posteriorPredSim_matrix,
    observedResponse = dat$observedResponse,
    fittedPredictedResponse = apply(posteriorPredDistr_matrix, 2, median),
    integerResponse = TRUE)

## Plot residuals
plot(sim)

florian_eg.stan

data {
    // Number of data
    int<lower = 0> nobs;

    // Observations
    int observedResponse [nobs];

    // Explanatory variables
    vector [nobs] Environment1;
}

parameters {
    real intercept;
    real env;
}

model {
    // Declare variables
    vector [nobs] lambda = exp(intercept + env*Environment1);

    // Priors
    target += normal_lpdf(intercept | 0, 1e4);
    target += normal_lpdf(env | 0, 1e4);

    // Likelihood
    target += poisson_lpmf(observedResponse | lambda);
}

generated quantities {
    int<lower = 0> observedResponseSim [nobs] = poisson_rng(exp(intercept + env*Environment1));
}

florian_eg_randEff.stan

data {
    // Number of data
    int<lower = 1> nobs;
    int<lower = 1> nGroups;

    // Indexing (index of the beginning of each group)
    array [nGroups + 1] int indices; // This syntax is compatible only from Stan 2.29!

    // Observations
    array [nobs] int observedResponse;

    // Explanatory variables
    vector [nobs] Environment1;
}

parameters {
    real intercept_mean;
    array [nGroups] real intercepts; // This syntax is compatible only from Stan 2.29!

    real env;

    // Random effect
    real<lower = 0> randEff_var;
}

model {
    // Declare variables
    vector [nobs] lambda;
    for (i in 1:nGroups)
        lambda[indices[i]:(indices[i + 1] - 1)] = exp(intercepts[i] + env*Environment1[indices[i]:(indices[i + 1] - 1)]);

    // Priors
    target += normal_lpdf(intercept_mean | 0, 1e4);
    target += normal_lpdf(intercepts | intercept_mean, randEff_var);
    target += normal_lpdf(env | 0, 1e4);

    target += gamma_lpdf(randEff_var | 1.0/1000, 1.0/1000); // Give an average of 1, and a variance of 1000

    // Likelihood
    target += poisson_lpmf(observedResponse | lambda);
}

generated quantities {
    array [nobs] int<lower = 0> observedResponseSim;
    vector<lower = 0> [nobs] lambda;

    for (i in 1:nGroups)
    {
        observedResponseSim[indices[i]:(indices[i + 1] - 1)] = poisson_rng(exp(intercepts[i] + env*Environment1[indices[i]:(indices[i + 1] - 1)]));
        lambda[indices[i]:(indices[i + 1] - 1)] = exp(intercepts[i] + env*Environment1[indices[i]:(indices[i + 1] - 1)]);
    }
}