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