greta-dev / greta

simple and scalable statistical modelling in R
https://greta-stats.org
Other
528 stars 63 forks source link

better column names for mcmc outputs #143

Closed goldingn closed 6 years ago

goldingn commented 6 years ago

When MCMC sampling greta arrays with 2+ dimensions, the columns names in the coda::mcmc objects returned are a bit cryptic; it's not clear to which element of the array they belong. For example:

library (greta)
z <- normal(0, 1, dim = c(2, 2))
m <- model(z)
draws <- mcmc(m, warmup = 5, n_samples = 5)[[1]]
draws
Markov Chain Monte Carlo (MCMC) output:
Start = 1 
End = 5 
Thinning interval = 1 
            z1         z2         z3        z4
[1,] 0.5028811 0.04083608 -0.1917491 0.6928466
[2,] 0.5888769 0.11905159 -0.2655282 0.6975463
[3,] 0.5836352 0.11562616 -0.2206950 0.6682547
[4,] 0.5548779 0.19657053 -0.2389385 0.6436528
[5,] 0.4782978 0.12375486 -0.2113774 0.7031416

greta returns the elements columnwise (like you usually get in R by coercing to a vector), though that isn't clear to users.

In JAGS, the coda mcmc output has the same format, but the column names are more informative:

library (rjags)
mod <-"
model {
  for (i in 1:2) {
    for (j in 1:2) {
      z[i, j] ~ dnorm(0, 1)
    }
  }
}
"
writeLines(mod, tmp <- tempfile())
jags <- jags.model(tmp, n.chains = 1)
coda.samples(jags,
             "z",
             n.iter = 5)[[1]]
Markov Chain Monte Carlo (MCMC) output:
Start = 1 
End = 5 
Thinning interval = 1 
          z[1,1]     z[2,1]     z[1,2]     z[2,2]
[1,]  1.44686351  1.6917898  0.1123499 -1.7850698
[2,] -0.82897144 -1.8866577  0.9759450  1.1456189
[3,] -1.04369933  0.2572350 -1.8803257 -0.1053721
[4,] -0.08653975  0.1263241 -0.7157092  1.3224209
[5,] -0.24756447  1.6448110 -0.3366253  0.5714593

greta should add useful column names like these

goldingn commented 6 years ago

This function will do the bulk of the work:

# given a vector of dimensions, return a vector of strings giving the row,
# column, etc. indices for each element, columnwise, in the format "[i,j]"
get_indices_text <- function (dims, name = "") {
  vec <- seq_len(prod(dims))
  indices <- arrayInd(vec, dims)
  mid_text <- apply(indices, 1, paste, collapse = ",")
  paste0(name, "[", mid_text, "]")
}
a <- greta:::ones(2, 3)
get_indices_text(dim(a), "a")
[1] "a[1,1]" "a[2,1]" "a[1,2]" "a[2,2]" "a[1,3]" "a[2,3]"
goldingn commented 6 years ago

Done on the dev branch