Calculate WAIC using Simpson's rule and Monte Carlo integration (2D version)

calculate-WAIC.R

library(cmdstanr)

waic <- function(log_likelihood) {
  training_error <- - mean(log(colMeans(exp(log_likelihood))))
  functional_variance_div_N <- mean(colMeans(log_likelihood^2) - colMeans(log_likelihood)^2)
  waic <- training_error + functional_variance_div_N
  return(waic)
}

model1_Si <- cmdstan_model('model1-addG-Simpson.stan')
model1_MC <- cmdstan_model('model1-addG-MC.stan')
model2_Si <- cmdstan_model('model2-addG-Simpson.stan')
model2_MC <- cmdstan_model('model2-addG-MC.stan')


## Data 1
## with a random effect on mean but not SD
set.seed(123)
G <- 10
NbyG <- sample.int(100, size = G)
N <- sum(NbyG)
n2g <- data.frame(N = 1:N, G = rep(1:G, times = NbyG))
gr <- rnorm(G, mean = 2, sd = 10)
Y <- rnorm(N, mean = gr[n2g$G], sd = 3)
data_Si <- list(N = N, G = G, NbyG = NbyG, n2g = n2g$G, Y = Y, M = 300)
data_MC <- list(N = N, G = G, NbyG = NbyG, n2g = n2g$G, Y = Y, M = 100000)

fit1_Si <- model1_Si$sample(data = data_Si, seed = 123, parallel_chains = 4, refresh = 200)
fit1_MC <- model1_MC$sample(data = data_MC, seed = 123, parallel_chains = 4, refresh = 200)
fit2_Si <- model2_Si$sample(data = data_Si, seed = 123, parallel_chains = 4, refresh = 200)
fit2_MC <- model2_MC$sample(data = data_MC, seed = 123, parallel_chains = 4, refresh = 200)
log_lik1_Si <- fit1_Si$draws('log_lik', format = 'matrix')
log_lik1_MC <- fit1_MC$draws('log_lik', format = 'matrix')
log_lik2_Si <- fit2_Si$draws('log_lik', format = 'matrix')
log_lik2_MC <- fit2_MC$draws('log_lik', format = 'matrix')

waic(log_lik1_Si) #=> 127.7375
waic(log_lik1_MC) #=> 127.7276
waic(log_lik2_Si) #=> 130.081
waic(log_lik2_MC) #=> 129.3979


## Data 2
## with a random effect on SD in addition to mean
set.seed(123)
G <- 10
NbyG <- sample.int(100, size = G)
N <- sum(NbyG)
n2g <- data.frame(N = 1:N, G = rep(1:G, times = NbyG))
gr <- rnorm(G, mean = 2, sd = 10)
gs <- abs(rnorm(G, mean = 0, sd = 5))
Y <- rnorm(N, mean = gr[n2g$G], sd = gs[n2g$G])
data2_Si <- list(N = N, G = G, NbyG = NbyG, n2g = n2g$G, Y = Y, M = 300)
data2_MC <- list(N = N, G = G, NbyG = NbyG, n2g = n2g$G, Y = Y, M = 100000)

fit1_Si <- model1_Si$sample(data = data2_Si, seed = 123, parallel_chains = 4, refresh = 200)
fit1_MC <- model1_MC$sample(data = data2_MC, seed = 123, parallel_chains = 4, refresh = 200)
fit2_Si <- model2_Si$sample(data = data2_Si, seed = 123, parallel_chains = 4, refresh = 200)
fit2_MC <- model2_MC$sample(data = data2_MC, seed = 123, parallel_chains = 4, refresh = 200)
log_lik1_Si <- fit1_Si$draws('log_lik', format = 'matrix')
log_lik1_MC <- fit1_MC$draws('log_lik', format = 'matrix')
log_lik2_Si <- fit2_Si$draws('log_lik', format = 'matrix')
log_lik2_MC <- fit2_MC$draws('log_lik', format = 'matrix')

waic(log_lik1_Si) #=> 157.7277
waic(log_lik1_MC) #=> 157.6832
waic(log_lik2_Si) #=> 141.7331
waic(log_lik2_MC) #=> 141.0446

model1-addG-MC.stan

functions {
  real f(vector y, vector theta, real r) {
    return normal_lpdf(r | theta[1], theta[2]) 
           + normal_lpdf(y | r, theta[3]);
  }
}

data {
  int G;
  int N;
  array[G] int NbyG;
  array[N] int n2g;
  vector[N] Y;
  int M;  // number of iterations for Monte Carlo integration
}

parameters {
  real mu;
  vector[G] r;
  real<lower=0> s_r;
  real<lower=0> s_y;
}

transformed parameters {
   vector[3] theta = [mu, s_r, s_y]';
}

model {
  r[1:G] ~ normal(mu, s_r);
  Y[1:N] ~ normal(r[n2g], s_y);
  // priors
  mu ~ normal(0, 20);
  s_r ~ normal(0, 20);
  s_y ~ normal(0, 20);
}

generated quantities {
  vector[G] log_lik;
  {
    matrix[G, M] lp;

    // Monte Carlo integration
    int pos = 1;
    for (g in 1:G) {
      for (m in 1:M) {
        real r_rnd = normal_rng(mu, s_r);
        lp[g, m] = normal_lpdf(Y[pos:(pos + NbyG[g] - 1)] | r_rnd, s_y);
      }
      log_lik[g] = - log(M) + log_sum_exp(lp[g]);
      pos = pos + NbyG[g];
    }
  }
}

model1-addG-Simpson.stan

functions {
  real f(vector y, vector theta, real r) {
    return normal_lpdf(r | theta[1], theta[2]) 
           + normal_lpdf(y | r, theta[3]);
  }

  real log_lik_Simpson(vector y, vector theta, 
                       real x_lower, real x_upper, 
                       vector W, int M) {
    vector[M+1] lp;
    real h = (x_upper - x_lower) / M;
    
    for (i in 1:(M+1)) {
      lp[i] = log(W[i]) + f(y, theta, x_lower + h*(i-1));
    }
    return log(h/3) + log_sum_exp(lp);
  }
}

data {
  int G;
  int N;
  array[G] int NbyG;
  array[N] int n2g;
  vector[N] Y;
  int M;  // number of divisions for Simpson's rule
}

transformed data {
  vector[M+1] W_Simpson;
  
  W_Simpson[1] = 1;
  for (m in 1:(M %/% 2)) {
    W_Simpson[2*m] = 4;
  }
  for (m in 1:(M %/% 2 - 1)) {
    W_Simpson[2*m+1] = 2;
  }
  W_Simpson[M+1] = 1;
}

parameters {
  real mu;
  vector[G] r;
  real<lower=0> s_r;
  real<lower=0> s_y;
}

transformed parameters {
  vector[3] theta = [mu, s_r, s_y]';
}

model {
  r[1:G] ~ normal(mu, s_r);
  Y[1:N] ~ normal(r[n2g], s_y);
  // priors
  mu ~ normal(0, 20);
  s_r ~ normal(0, 20);
  s_y ~ normal(0, 20);
}

generated quantities {
  vector[G] log_lik;
  {
    int pos = 1;
    for (g in 1:G) {
      log_lik[g] = log_lik_Simpson(Y[pos:(pos + NbyG[g] - 1)], theta,
                                   mu - 5*s_r, mu + 5*s_r,
                                   W_Simpson, M);
      pos = pos + NbyG[g];
    }
  }
}

model2-addG-MC.stan

functions {
  real f(vector y, vector theta, real r, real s) {
    return normal_lpdf(r | theta[1], theta[2])
           + normal_lpdf(s | 0, theta[3])
           + normal_lpdf(y | r, s);
  }
}

data {
  int G;
  int N;
  array[G] int NbyG;
  array[N] int n2g;
  vector[N] Y;
  int M;  // number of iterations for Monte Carlo integration
}

parameters {
  real mu;
  vector[G] r;
  vector<lower=0>[G] s;
  real<lower=0> s_r;
  real<lower=0> s_s;

}

 transformed parameters {
   vector[3] theta = [mu, s_r, s_s]';
}

model {
  r[1:G] ~ normal(mu, s_r);
  s[1:G] ~ normal(0, s_s);
  Y[1:N] ~ normal(r[n2g], s[n2g]);
  // priors
  mu ~ normal(0, 20);
  s_r ~ normal(0, 20);
  s_s ~ normal(0, 20);
}

generated quantities {
  vector[G] log_lik;
  {
    matrix[G, M] lp;

    // Monte Carlo integration
    int pos = 1;
    for (g in 1:G) {
      for (m in 1:M) {
        real r_rnd = normal_rng(mu, s_r);
        real s_rnd = fabs(normal_rng(0, s_s));
        lp[g, m] = normal_lpdf(Y[pos:(pos + NbyG[g] - 1)] | r_rnd, s_rnd);
      }
      log_lik[g] = - log(M) + log_sum_exp(lp[g]);
      pos = pos + NbyG[g];
    }
  }
}

model2-addG-Simpson.stan

functions {
  real f(vector y, vector theta, real r, real s) {
    return normal_lpdf(r | theta[1], theta[2])
           + normal_lpdf(s | 0, theta[3])
           + normal_lpdf(y | r, s);
  }

  real log_lik_Simpson_2d(vector y, vector theta,
                          real x_lower, real x_upper, real y_lower, real y_upper, 
                          matrix W, int M) {
    matrix[M+1, M+1] lp;
    real h_x = (x_upper - x_lower) / M;
    real h_y = (y_upper - y_lower) / M;

    for (i in 1:(M+1)) {
      for (j in 1:(M+1)) {
        lp[i,j] = log(W[i,j]) + f(y, theta, x_lower + h_x*(i-1), y_lower + h_y*(j-1));
      }
    }
    return log(h_x/3) + log(h_y/3) + log_sum_exp(lp);
  }
}

data {
  int G;
  int N;
  array[G] int NbyG;
  array[N] int n2g;
  vector[N] Y;
  int M;  // number of divisions for Simpson's rule
}

transformed data {
  vector[M+1] V_Simpson;
  matrix[M+1, M+1] W_Simpson;
  
  V_Simpson[1] = 1;
  for (m in 1:(M %/% 2)) {
    V_Simpson[2*m] = 4;
  }
  for (m in 1:(M %/% 2 - 1)) {
    V_Simpson[2*m+1] = 2;
  }
  V_Simpson[M+1] = 1;
  W_Simpson = V_Simpson * V_Simpson';
}

parameters {
  real mu;
  vector[G] r;
  vector<lower=0>[G] s;
  real<lower=0> s_r;
  real<lower=0> s_s;
}

transformed parameters {
  vector[3] theta = [mu, s_r, s_s]';
}

model {
  r[1:G] ~ normal(mu, s_r);
  s[1:G] ~ normal(0, s_s);
  Y[1:N] ~ normal(r[n2g], s[n2g]);
  // priors
  mu ~ normal(0, 20);
  s_r ~ normal(0, 20);
  s_s ~ normal(0, 20);
}

generated quantities {
  vector[G] log_lik;
  {
    int pos = 1;
    for (g in 1:G) {
      log_lik[g] = log_lik_Simpson_2d(Y[pos:(pos + NbyG[g] - 1)], theta,
                                      mu - 5*s_r, mu + 5*s_r, 1e-8, 5*s_s,
                                      W_Simpson, M);
      pos = pos + NbyG[g];
    }
  }
}