Indices of random effects (RE) in stan output for mixture models like beta_gamma model

[xsswang]

For a mixture model, it currently allows separate random effects for each component of the two-part model. 'tmb_params$RE' is set to be two columns. However, when passing this setting to tmbstan, its indices change into one dimension, i.e. RE[1], RE[2], ...., instead of RE[1,1], RE[1,2],... This can be achieved either by stacking two columns (RE[1,1], RE[2,1]...) or by expanding through rows (RE[1,1], RE[1,2]...). Could you please provide a description about how it is re-indexed? It'll help extract sampling from stan outputs.
Thanks in advance!

[seananderson]

The random intercepts appear to get appended one column after the other into a vector. Here's an example:

library(sdmTMB)
library(tmbstan)

pcod$fyear <- as.factor(pcod$year)
fit <- sdmTMB(
  density ~ 1 + (1 | fyear),
  data = pcod,
  family = delta_gamma(),
  spatial = "off" # fast example
)
fit
#> Model fit by ML ['sdmTMB']
#> Formula: density ~ 1 + (1 | fyear)
#> Mesh: NULL (isotropic covariance)
#> Data: pcod
#> Family: delta_gamma(link1 = 'logit', link2 = 'log')
#> 
#> Delta/hurdle model 1: -----------------------------------
#> Family: binomial(link = 'logit') 
#>             coef.est coef.se
#> (Intercept)    -0.15    0.09
#> 
#> Random intercepts:
#>       Std. Dev.
#> fyear      0.25
#> 
#> 
#> Delta/hurdle model 2: -----------------------------------
#> Family: Gamma(link = 'log') 
#>             coef.est coef.se
#> (Intercept)     4.39    0.12
#> 
#> Random intercepts:
#>       Std. Dev.
#> fyear      0.34
#> 
#> Dispersion parameter: 0.60
#> 
#> ML criterion at convergence: 6732.914
#> 
#> See ?tidy.sdmTMB to extract these values as a data frame.

# update sdmTMB first (as of 2022-03-08)
(re1 <- tidy(fit, "ran_vals", model = 1))
#> # A tibble: 9 × 3
#>   term       estimate std.error
#>   <chr>         <dbl>     <dbl>
#> 1 fyear_2003  -0.0859     0.138
#> 2 fyear_2004   0.158      0.139
#> 3 fyear_2005   0.297      0.144
#> 4 fyear_2007  -0.164      0.137
#> 5 fyear_2009  -0.174      0.140
#> 6 fyear_2011  -0.233      0.139
#> 7 fyear_2013   0.303      0.142
#> 8 fyear_2015   0.184      0.139
#> 9 fyear_2017  -0.281      0.143
(re2 <- tidy(fit, "ran_vals", model = 2))
#> # A tibble: 9 × 3
#>   term       estimate std.error
#>   <chr>         <dbl>     <dbl>
#> 1 fyear_2003   -0.164     0.161
#> 2 fyear_2004    0.401     0.156
#> 3 fyear_2005    0.373     0.155
#> 4 fyear_2007   -0.569     0.166
#> 5 fyear_2009   -0.340     0.166
#> 6 fyear_2011    0.294     0.161
#> 7 fyear_2013   -0.104     0.153
#> 8 fyear_2015    0.215     0.154
#> 9 fyear_2017   -0.158     0.165

fit_stan <- tmbstan(fit$tmb_obj, iter = 200, chains = 4, seed = 2938)
#> 
#> SAMPLING FOR MODEL 'tmb_generic' NOW (CHAIN 1).
#> Chain 1: 
#> Chain 1: Gradient evaluation took 6.9e-05 seconds
#> Chain 1: 1000 transitions using 10 leapfrog steps per transition would take 0.69 seconds.
#> Chain 1: Adjust your expectations accordingly!
#> ... [cut]
print(
  fit_stan,
  pars = c("b_j", "ln_phi", "RE")
)
#> Inference for Stan model: sdmTMB.
#> 4 chains, each with iter=200; warmup=100; thin=1; 
#> post-warmup draws per chain=100, total post-warmup draws=400.
#> 
#>         mean se_mean   sd  2.5%   25%   50%   75% 97.5% n_eff Rhat
#> b_j    -0.15    0.01 0.11 -0.37 -0.22 -0.15 -0.09  0.07   135 1.01
#> ln_phi -0.50    0.00 0.03 -0.57 -0.52 -0.50 -0.48 -0.44   516 0.99
#> RE[1]  -0.08    0.01 0.14 -0.37 -0.18 -0.07  0.02  0.19   228 1.00
#> RE[2]   0.17    0.01 0.15 -0.13  0.07  0.17  0.26  0.49   202 1.01
#> RE[3]   0.30    0.01 0.15  0.01  0.20  0.30  0.40  0.61   171 1.00
#> RE[4]  -0.16    0.01 0.15 -0.49 -0.27 -0.15 -0.06  0.14   304 1.00
#> RE[5]  -0.17    0.01 0.14 -0.47 -0.27 -0.18 -0.07  0.10   271 1.00
#> RE[6]  -0.23    0.01 0.15 -0.54 -0.33 -0.23 -0.12  0.05   313 1.00
#> RE[7]   0.30    0.01 0.15  0.03  0.20  0.30  0.41  0.60   166 1.01
#> RE[8]   0.19    0.01 0.16 -0.11  0.08  0.18  0.28  0.47   178 1.02
#> RE[9]  -0.29    0.01 0.16 -0.61 -0.39 -0.28 -0.18 -0.01   268 1.00
#> RE[10] -0.16    0.02 0.18 -0.53 -0.27 -0.16 -0.04  0.18    74 1.04
#> RE[11]  0.41    0.02 0.18  0.05  0.29  0.41  0.53  0.75    74 1.04
#> RE[12]  0.38    0.02 0.18  0.05  0.27  0.39  0.50  0.75    79 1.03
#> RE[13] -0.56    0.02 0.20 -1.01 -0.68 -0.55 -0.41 -0.18    83 1.04
#> RE[14] -0.33    0.02 0.19 -0.74 -0.45 -0.33 -0.23  0.05    85 1.03
#> RE[15]  0.31    0.02 0.18 -0.07  0.19  0.31  0.44  0.68    77 1.03
#> RE[16] -0.10    0.02 0.19 -0.49 -0.22 -0.10  0.03  0.25    72 1.04
#> RE[17]  0.23    0.02 0.18 -0.16  0.12  0.22  0.35  0.59    76 1.04
#> RE[18] -0.16    0.02 0.19 -0.49 -0.27 -0.16 -0.03  0.20   100 1.02
#> 
#> Samples were drawn using NUTS(diag_e) at Wed Mar  8 17:15:27 2023.
#> For each parameter, n_eff is a crude measure of effective sample size,
#> and Rhat is the potential scale reduction factor on split chains (at 
#> convergence, Rhat=1).

post <- rstan::extract(fit_stan)
dim(post$RE)
#> [1] 400  18

n_re <- nrow(re1)
re1_mcmc <- apply(post$RE[,seq(1, n_re)], 2, mean)
re2_mcmc <- apply(post$RE[,seq(n_re + 1, ncol(post$RE))], 2, mean)

plot(re1$estimate, re1_mcmc);abline(0, 1)

plot(re2$estimate, re2_mcmc);abline(0, 1)

^{Created on 2023-03-08 with reprex v2.0.2}