Separating fixed and random components in a correlated effects model

correlated-effects-greta.R

library(tidyverse)
library(greta)
library(mvtnorm)
library(bayesplot)


# ---- set up ----
J <- 40  # number of groups
n_j <- sample.int(15, J, replace = TRUE)  # number of observations per group

# ---- generating values ----
beta_true <- c(b0 = 2, b1 = -1.5, b2 = 0.2)  # fixed effects
Omega_true <- rbind(  # correlations
  c(1, -0.4, 0.7),
  c(-0.4, 1, 0.25),
  c(0.7, 0.25, 1)
)
sigma_true <- c(1, 3, 2)  # standard deviations
Sigma_true <- diag(sigma_true) %*% Omega_true %*% diag(sigma_true)  # covmat
sigma_y_true <- 1  # observation variance

eps_j_true <- rmvnorm(J, rep(0, length(beta_true)), Sigma_true, 
                      method = 'svd')
coefs_true <- sweep(eps_j_true, 2, beta_true, "+")

# ---- (simulated) data ----
dimnames(coefs_true)[[2]] <- names(beta_true) 
mu_true <- as_tibble(coefs_true) %>% 
  mutate(x1 = rnorm(n()), x2 = runif(n(), -1, 1),
         lp = b0 + b1 * x1 + b2 * x2,  # the linear predictor
         j = 1:n()) %>% 
  left_join(tibble(j = 1:J, n_j), by = 'j')

d <- bind_rows(lapply(mu_true, rep, mu_true$n_j)) %>% 
  transmute(y = rnorm(n(), lp, sigma_y_true), j, x1, x2)

# model matrix
X <- model.matrix(~ x1 + x2, d) 


# --- the model ----

M <- ncol(X) # number of varying coefficients

# priors on fixed effect coefficients
beta_mu <- normal(0, 10, dim = M)

# prior on the standard deviation of the varying coefficient
sd <- cauchy(0, 3, truncation = c(0, Inf), dim = M)

# prior on the correlation between the varying coefficients
Omega <- lkj_correlation(eta = 1, M) 

# optimization of the varying coefficient sampling through
# cholesky factorization
Omega_U <- chol(Omega)
Sigma_U <- sweep(Omega_U, 2, sd, "*")

Sigma <- greta:::chol2symm(Sigma_U)
eps <- multivariate_normal(mean = zeros(1, M),
                           Sigma = Sigma,
                           dimension = M,
                           n_realisations = J)
ab <- sweep(eps, 2, beta_mu, "+")

# the linear predictor
mu <- rowSums(ab[d$j, ] * X)

# the residual variance
sigma_e <- cauchy(0, 3, truncation = c(0, Inf))

# model
y <- d$y
distribution(y) <- normal(mu, sigma_e)

m <- model(Omega, sigma_e, beta_mu, ab)
draws <- greta::mcmc(m, warmup = 2000, n_samples = 2000, one_by_one = TRUE, 
                     chains = 2)

# ---- validate ----
mcmc_hist(draws, regex_pars = 'sigma_e') + 
  geom_vline(xintercept = sigma_y_true)

d_beta <- tibble(beta_true) %>% 
  mutate(Parameter = sprintf('beta_mu[%s,1]', 1:n()))
mcmc_hist(draws, regex_pars = 'beta_mu', 
          facet_args = list(ncol = 1, scales = 'fixed')) +
  geom_vline(data = d_beta, aes(xintercept = beta_true)) +
  geom_vline(xintercept = 0, color = 'gray', linetype = 'dashed')

A <- cor(eps_j_true)  #Omega_true
d_Omega <- tibble(Omega_empirical = A[upper.tri(A)]) %>% 
  mutate(Parameter = sprintf('Omega[%s,%s]', 
                             row(A)[upper.tri(A)], col(A)[upper.tri(A)]))
mcmc_hist(draws, regex_pars = 'Omega\\[1,2\\]|Omega\\[1,3\\]|Omega\\[2,3\\]',
          facet_args = list(ncol = 1, scales = 'fixed')) +
  geom_vline(data = d_Omega, aes(xintercept = Omega_empirical)) +
  geom_vline(xintercept = 0, color = 'gray', linetype = 'dashed')