Simulation for zero-one inflated beta regression in Stan

zoib-recovery.R

## DATA GENERATION
# reproducibility
library(bayesplot)
set.seed(1839)

# matches the stan function
inv_logit <- function(x){
  exp(x) / (1 + exp(x))
}

n <- 10000 # sample size
x <- rnorm(n) # predictor value

# coefficients for alpha
a_b0 <- 0
a_b1 <- 0.3

# calculating alpha, adding error
alpha <- inv_logit(a_b0 + a_b1 * x)

# coefficients for gamma
g_b0 <- 0
g_b1 <- 0.3

# calculating gamma, adding error
gamma <- inv_logit(g_b0 + g_b1 * x)

# coefficients for mu
m_b0 <- -0.6
m_b1 <- 0.5

# calculating mu, adding error 
mu <- inv_logit(m_b0 + m_b1 * x)

# coefficients for phi
p_b0 <- -0.5
p_b1 <- 0.7

# calculating phi, adding error
phi <- exp(p_b0 + p_b1 * x)

# calculating shape parameters for beta distribution
p <- mu * phi
q <- phi - mu * phi

# calculate, using alpha, if their score comes from bernoulli trial
y_is_discrete <- rbinom(n, 1, alpha)

# calculate y scores
y <- rep(NA, n) # initialize empty vector of outcomes
for (i in 1:n) {
  if (y_is_discrete[i] == 1) { # if y came from bernoulli...
    y[i] <- rbinom(1, 1, gamma[i]) # being 1 is determined by its gamma
  } else {
    y[i] <- rbeta(1, p[i], q[i]) # if it did not, draw from beta
  }
}

hist(y, breaks = 100)

library(rstan)

stan_d <- list(y = y, n = n, x = x)

m_init <- stan_model('zoib.stan')
m_fit <- sampling(m_init, data = stan_d, 
                  pars = c('coef_a', 'coef_g', 'coef_m', 'coef_p'), 
                  cores = 2)
m_fit

draws <- as.matrix(m_fit)
draws <- draws[, !(colnames(draws) == 'lp__')]

true <- c(a_b0, a_b1, 
          g_b0, g_b1, 
          m_b0, m_b1, 
          p_b0, p_b1)

mcmc_recover_intervals(draws, true)

zoib.stan

data {
  int n;
  vector[n] x; 
  vector<lower=0, upper=1>[n] y;
}
transformed data {
  int<lower=0, upper=1> is_discrete[n];
  int<lower=-1, upper=1> y_discrete[n];

  // create indicator for whether y is discrete 
  // and an integer value to pass to bernoulli_lpmf for discrete y
  for (i in 1:n) {
    if (y[i] == 0) {
      is_discrete[i] = 1;
      y_discrete[i] = 0;
    } else if (y[i] == 1) {
      is_discrete[i] = 1;
      y_discrete[i] = 1;
    } else {
      is_discrete[i] = 0;
      // hack to ensure that throws error if passed to bernoulli_lpmf
      y_discrete[i] = -1;
    }
  }
}
parameters {
  vector[2] coef_a;
  vector[2] coef_g;
  vector[2] coef_m;
  vector[2] coef_p;
}
transformed parameters {
  vector<lower=0, upper=1>[n] alpha;
  vector<lower=0, upper=1>[n] gamma;
  vector[n] mu;
  vector<lower=0>[n] phi;
  vector<lower=0>[n] p;
  vector<lower=0>[n] q;

  for (i in 1:n) {
    alpha[i] = inv_logit(coef_a[1] + coef_a[2] * x[i]);
    gamma[i] = inv_logit(coef_g[1] + coef_g[2] * x[i]);
    mu[i] = inv_logit(coef_m[1] + coef_m[2] * x[i]);
    phi[i] = exp(coef_p[1] + coef_p[2] * x[i]);
    p[i] = mu[i] * phi[i];
    q[i] = phi[i] - mu[i] * phi[i];
  }
}
model {
  coef_a ~ normal(0, 1);
  coef_g ~ normal(0, 1);
  coef_m ~ normal(0, 1);
  coef_p ~ normal(0, 1);
  is_discrete ~ bernoulli(alpha); 
  for (i in 1:n) {
    if (is_discrete[i] == 1) {
      y_discrete[i] ~ bernoulli(gamma[i]);
    } else {
      y[i] ~ beta(p[i], q[i]);
    }
  }
}