Greta test

greta_test.R

#devtools::install_github("wkmor1/msmod")
library(msmod) # for eucs dataset
library(lme4)
library(greta)
sc = function(x) scale(x)[, 1] / 2 # center and scale covariates by 2 sds

# Fit glmer model
m_glmer = glmer(
  present ~ sc(logit_rock) + (1 + sc(logit_rock) | species),
  data = eucs, family = stats::binomial
)

# No rowsum method for greta arrays?
rsum = function(x) {
  ans = rep(0, nrow(x)) 
  for (i in seq_len(ncol(x))) ans = ans + x[, i]
  ans
}

# Prep data
y   = eucs$present
x   = cbind(1, sc(eucs$logit_rock))
g   = as.numeric(factor(eucs$species))
n_g = length(unique(g))
n_b = ncol(as.matrix(x))

# Fit Greta model
m               = normal(0, 1, n_b)
s               = greta_array(dim = c(n_b, n_b))
diag(s)         = normal(0, 1, n_b, c(0, Inf))
s               = s %*% s %*% lkj_correlation(1, n_b)
#b               = multivariate_normal(m, s, n_g)
# Whitening?
b               = t(m)[rep(1, n_g), ] + normal(0, 1, c(n_g, n_b)) %*% chol(s)
distribution(y) = bernoulli(ilogit(rsum(b[g, ] * x)))
m_greta         = model(m, s, b)

samples_greta = 
  mcmc(
    m_greta,
    warmup = 1000#,
    # Inital values from glmer model (doesn't work for whitening trans?)
    #  initial_values =
    #    setNames(
    #     unlist(c(
    #       fixef(m_glmer), 
    #       coef(m_glmer)$species,
    #       diag(summary(m_glmer)$varcor$species),
    #       attr(summary(m_glmer)$varcor$species, "cor")[lower.tri(diag(2))]
    #     )),
    #     names(m_greta$dag$example_parameters())
    #   )
    )

bayesplot::mcmc_trace(samples_greta, regex_pars = "m")