WAIC with hierarchical models

model1-add1-speedup.stan

functions {
  real f(int Y, real[] theta, real x) {
    return(gamma_lpdf(x | theta[2], theta[2]/theta[1]) + poisson_lpmf(Y | x));
  }

  real log_lik_Simpson(int Y, real[] theta, real a, real b, int M) {
    vector[M+1] lp;
    real h;
    h = (b-a)/M;
    lp[1] = f(Y, theta, a);
    for (m in 1:(M/2))
      lp[2*m] = log(4) + f(Y, theta, a + h*(2*m-1));
    for (m in 1:(M/2-1))
      lp[2*m+1] = log(2) + f(Y, theta, a + h*2*m);
    lp[M+1] = f(Y, theta, b);
    return(log(h/3) + log_sum_exp(lp));
  }
}

data {
  int N;
  int Y[N];
}

parameters {
  real<lower=0> theta[2];
  real<lower=0> r[N];
}

model {
  Y ~ poisson(r);
  r ~ gamma(theta[2], theta[2]/theta[1]);
}

generated quantities {
  real log_lik[N];
  real log_lik_table[max(Y)+1];
  for (k in 0:max(Y))
    log_lik_table[k+1] = log_lik_Simpson(k, theta, 0.0001, 50, 200);
  for (n in 1:N)
    log_lik[n] = log_lik_table[Y[n]+1];
}

model1-add1.stan

functions {
  real f(int Y, real[] theta, real x) {
    return(gamma_lpdf(x | theta[2], theta[2]/theta[1]) + poisson_lpmf(Y | x));
  }

  real log_lik_Simpson(int Y, real[] theta, real a, real b, int M) {
    vector[M+1] lp;
    real h;
    h = (b-a)/M;
    lp[1] = f(Y, theta, a);
    for (m in 1:(M/2))
      lp[2*m] = log(4) + f(Y, theta, a + h*(2*m-1));
    for (m in 1:(M/2-1))
      lp[2*m+1] = log(2) + f(Y, theta, a + h*2*m);
    lp[M+1] = f(Y, theta, b);
    return(log(h/3) + log_sum_exp(lp));
  }
}

data {
  int N;
  int Y[N];
}

parameters {
  real<lower=0> theta[2];
  real<lower=0> r[N];
}

model {
  Y ~ poisson(r);
  r ~ gamma(theta[2], theta[2]/theta[1]);
}

generated quantities {
  real log_lik[N];
  for (n in 1:N)
    log_lik[n] = log_lik_Simpson(Y[n], theta, 0.0001, 50, 200);
}

model1-givenG-add1.stan

data {
  int N;
  int Y[N];
}

parameters {
  real<lower=0> mu;
  real<lower=0> phi;
  real<lower=0> r[N];
}

model {
  Y ~ poisson(r);
  r ~ gamma(phi, phi/mu);
}

generated quantities {
  real log_lik[N];
  for (n in 1:N)
    log_lik[n] = poisson_lpmf(Y[n] | r[n]);
}

model1-nbinom.stan

data {
  int N;
  int Y[N];
}

parameters {
  real<lower=0> mu;
  real<lower=0> phi;
}

model {
  Y ~ neg_binomial_2(mu, phi);
}

generated quantities {
  real log_lik[N];
  for (n in 1:N)
    log_lik[n] = neg_binomial_2_lpmf(Y[n] | mu, phi);
}

model2-add1-speedup.stan

functions {
  real f_fix(real[] theta, real x) {
    return(normal_lpdf(x | theta[1], theta[2]));
  }

  real f(real Y, real[] theta, real x) {
    return(normal_lpdf(Y | x, theta[3]));
  }

  vector log_f_fix(real[] theta, real a, real b, int M) {
    vector[M+1] lp;
    real h ;
    h = (b-a)/M;
    lp[1] = f_fix(theta, a);
    for (m in 1:(M/2))
      lp[2*m] = log(4) + f_fix(theta, a + h*(2*m-1));
    for (m in 1:(M/2-1))
      lp[2*m+1] = log(2) + f_fix(theta, a + h*2*m);
    lp[M+1] = f_fix(theta, b);
    return(lp);
  }

  vector log_f(real Y, real[] theta, real a, real b, int M) {
    vector[M+1] lp;
    real h;
    h = (b-a)/M;
    lp[1] = f(Y, theta, a);
    for (m in 1:(M/2))
      lp[2*m] = f(Y, theta, a + h*(2*m-1));
    for (m in 1:(M/2-1))
      lp[2*m+1] = f(Y, theta, a + h*2*m);
    lp[M+1] = f(Y, theta, b);
    return(lp);
  }
}

data {
  int G;
  int N;
  vector[N] Y;
  int N2G[N];
}

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

transformed parameters {
  real theta[3];
  theta[1] = mu;
  theta[2] = s0;
  theta[3] = sigma;
}

model {
  r ~ normal(mu, s0);
  Y ~ normal(r[N2G], sigma);
}

generated quantities {
  vector[N] log_lik;
  int M;
  M = 100;
  {
    vector[M+1] fix;
    vector[M+1] rest;
    fix = log_f_fix(theta, mu-5*s0, mu+5*s0, M);
    for (n in 1:N) {
      rest = log_f(Y[n], theta, mu-5*s0, mu+5*s0, M);
      log_lik[n] = log(10*s0/M/3) + log_sum_exp(fix + rest);
    }
  }
}

model2-add1.stan

functions {
  real f(real Y, real[] theta, real x) {
    return(normal_lpdf(x | theta[1], theta[2]) + normal_lpdf(Y | x, theta[3]));
  }

  real log_lik_Simpson(real Y, real[] theta, real a, real b, int M) {
    vector[M+1] lp;
    real h;
    h = (b-a)/M;
    lp[1] = f(Y, theta, a);
    for (m in 1:(M/2))
      lp[2*m] = log(4) + f(Y, theta, a + h*(2*m-1));
    for (m in 1:(M/2-1))
      lp[2*m+1] = log(2) + f(Y, theta, a + h*2*m);
    lp[M+1] = f(Y, theta, b);
    return(log(h/3) + log_sum_exp(lp));
  }
}

data {
  int N;
  int G;
  vector[N] Y;
  int N2G[N];
}

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

transformed parameters {
  real theta[3];
  theta[1] = mu;
  theta[2] = s0;
  theta[3] = sigma;
}

model {
  r ~ normal(mu, s0);
  Y ~ normal(r[N2G], sigma);
}

generated quantities {
  real log_lik[N];
  for (n in 1:N)
    log_lik[n] = log_lik_Simpson(Y[n], theta, mu-5*s0, mu+5*s0, 100);
}

model2-addG.stan

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

  real log_lik_Simpson(vector Y, real[] theta, real a, real b, int M) {
    vector[M+1] lp;
    real h;
    h = (b-a)/M;
    lp[1] = f(Y, theta, a);
    for (m in 1:(M/2))
      lp[2*m] = log(4) + f(Y, theta, a + h*(2*m-1));
    for (m in 1:(M/2-1))
      lp[2*m+1] = log(2) + f(Y, theta, a + h*2*m);
    lp[M+1] = f(Y, theta, b);
    return(log(h/3) + log_sum_exp(lp));
  }
}

data {
  int G;
  int NbyG[G];
  int N;
  vector[N] Y;
  int N2G[N];
}

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

transformed parameters {
  real theta[3];
  theta[1] = mu;
  theta[2] = s0;
  theta[3] = sigma;
}

model {
  r ~ normal(mu, s0);
  Y ~ normal(r[N2G], sigma);
}

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

model2-givenG-add1.stan

data {
  int G;
  int N;
  vector[N] Y;
  int N2G[N];
}

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

model {
  r ~ normal(mu, s0);
  Y ~ normal(r[N2G], sigma);
}

generated quantities {
  real log_lik[N];
  for (n in 1:N)
    log_lik[n] = normal_lpdf(Y[n] | r[N2G[n]], sigma);
}

run-model1.R

library(rstan)

set.seed(123)
N <- 500
Y <- rnbinom(N, size=1.2, mu=8)
data <- list(N=N, Y=Y)

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)
}

fit1 <- stan(file='model/model1-givenG-add1.stan', data=data, seed=1234)
log_lik <- extract(fit1)$log_lik
waic1_byG <- sapply(1:N, function(n) waic(log_lik[,n,drop=F]))

fit2 <- stan(file='model/model1-nbinom.stan', data=data, seed=1234)
waic2 <- waic(extract(fit2)$log_lik)

# fit3 <- stan(file='model/model1-add1.stan', data=data, seed=1234)
# waic3 <- waic(extract(fit3)$log_lik)

fit3 <- stan(file='model/model1-add1-speedup.stan', data=data, seed=1234)
waic3 <- waic(extract(fit3)$log_lik)

run-model2.R

library(rstan)

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 <- list(N=N, G=G, NbyG=NbyG, N2G=N2G$G, Y=Y)

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)
}

stanmodel1 <- stan_model(file='model/model2-givenG-add1.stan')
fit1 <- sampling(stanmodel1, data=data, seed=1234)
log_lik <- extract(fit1)$log_lik
G2N_list <- split(N2G$N, N2G$G)
waic1_byG <- sapply(G2N_list, function(n) waic(log_lik[,n]))

stanmodel2 <- stan_model(file='model/model2-addG.stan')
fit2 <- sampling(stanmodel2, data=data, seed=1234)
waic2 <- waic(extract(fit2)$log_lik)

# stanmodel3 <- stan_model(file='model/model2-add1.stan')
# fit3 <- sampling(stanmodel3, data=data, seed=1234)
# waic3 <- waic(extract(fit3)$log_lik)

stanmodel3 <- stan_model(file='model/model2-add1-speedup.stan')
fit3 <- sampling(stanmodel3, data=data, seed=1234)
waic3 <- waic(extract(fit3)$log_lik)