Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Separating fixed and random components in a correlated effects model
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')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.