Introducción modelos jerárquicos (II)

Modified

Friday, March 29, 2024

knitr::opts_chunk$set(echo = TRUE, warning = F, message = F)

paquetes <- c("rethinking", "tidyverse", "magrittr", 'patchwork',
              'cmdstanr', 'loo', "MASS", 'ellipse')

sapply(paquetes, library, character.only = T)

options(mc.cores = parallel::detectCores())

source('functions_mod_diagnostics.R')

En el capítulo anterior expuse cómo los modelos jerárquicos usan la información disponible en cada agrupamiento (cluster) de datos para generar estimaciones más precisas de los parámetros (partial pooling). Presenté las ventajas de esta aproximación, y su utilidad cuando trabajamos con observaciones desbalanceadas entre grupos. Usar esta aproximación suponía considerar distribuciones de probabilidad de la población de parámetros (previa adaptativa), hiper previas (distribución previa de la previa adaptativa) y la reparametrización de los modelos para mejorar la eficiencia computacional y la convergencia de las cadenas. Bien, en esta sección aumentaremos el nivel de abstracción. Emplearemos partial pooling para cruzar la información ya no solo entre clusters, sino también entre parámetros. Este método aprovecha la covariación entre parámetros para extraer una mayor cantidad de información de los datos, realizar estimaciones más precisas, y es la base para extender los modelos jerárquicos a áreas más especializadas (e.g. modelos filogenéticos, modelos espacio-estado, análisis de redes).

El material de esta sección es avanzado, requiere la revisión del capítulo anterior, algunas nociones adicionales de programación en Stan y muuuucha paciencia. Si usted es una persona “normal”, como yo, el entendimiento/intuición para definir y ajustar este tipo de modelos no ocurrirá de inmediato, solo llegará con la práctica y repetición.

1 Parámetros correlacionados y distribución normal multivariada

Acabo de mencionar que este tipo de modelos aprovecha la covariación entre parámetros para intercambiar información y producir mejores estimaciones. En este caso, el modelo emplea un vector de promedios para cada parámetros –cada promedio es en realidad una distribución de probabilidad– y una matriz de varianza-covarianza para parametrizar una distribución normal multivariada que asocia los parámetros y realizar su estimación. Veamos en qué consiste:

Supongamos los parámetros \(\theta \sim normal(\mu=2, \sigma = 1.2)\) y \(\psi \sim normal(\mu = 1.1, \sigma=0.8)\), cuyo coeficiente de correlación \(r = 0.7\). Para generar la distribución normal multivariada que relaciona ambos parámetros requerimos:

  1. Vector de promedios \(\mu = [\mu_\theta, \mu_\psi]\):
theta <- 2
psi <- 1.1
mu <- c(theta, psi)
  1. Matriz de varianza-covarianza. Esta es una matriz cuadrada cuya diagonal corresponde a la varianza de los parámetros (i.e. \(\sigma^2_\theta\), \(\sigma^2_\psi\)), mientras que los demás elementos de la matriz corresponden a la covariación \(\rho_{\theta \psi}\) entre ambos parámetros. Dicha covarianza se calcula a partir de la fórmula: \(COV(\theta, \psi) = \sigma^2_\theta \times \sigma^2_\psi \times \rho_{\theta \psi}\). Calculemos \(COV\) y definamos nuestra matriz de varianza-covarianza:
sigma_theta <- 1.2
sigma_psi <- 0.8
r <- 0.7
cov_theta_psi <- sigma_theta * sigma_psi * r

cov_m <- matrix(c(sigma_theta^2, cov_theta_psi, cov_theta_psi, sigma_psi^2), 
                ncol = 2) # tantas columnas como parámetros
cov_m
      [,1]  [,2]
[1,] 1.440 0.672
[2,] 0.672 0.640

Tenemos los elementos necesarios para parametrizar la distribución normal multivariada que describe \(\theta\) y \(\psi\). Usaremos MASS::mvrnorm() para generar valores aleatorios de la distribución:

multi_norm <- mvrnorm(1e3, mu = mu, Sigma = cov_m)

par(mfrow = c(1, 2), mar = c(4, 4, 2, 1))
plot(density(multi_norm[, 1]), main = '', xlab = expression(theta),
     lwd = 2, col = 2)
lines(density(rnorm(1e3, theta, sigma_theta)),  
     lwd = 2, col = 3)
plot(density(multi_norm[, 2]), main = '', xlab = expression(psi), 
     lwd = 3, col = 2)
lines(density(rnorm(1e3, psi, sigma_psi)),  
     lwd = 3, col = 3)

par(mfrow = c(1, 1))

MASS::mvrnorm() genera, en este caso, una distribución normal con dos dimensiones, una por parámetro (Figura anterior). Grafiquemos la distribución multivariada:

plot(multi_norm[, 1], multi_norm[, 2], main = '', 
     ylab = expression(theta), xlab = expression(psi))
for (i in seq(0.1, 0.9, by = 0.2)) {
  lines(ellipse(cov_m, centre = mu, level = i), col = 4, lwd = 1.5)
}

Será justamente este tipo de distribución normal multivariada la que usaremos como distribución previa en nuestros modelos. Ajustarlos requerirá algunas consideraciones adicionales respecto a la parametrización y reparametrización de las distribuciones previas, nuevas distribuciones de probabilidad, operaciones con matrices y algunos elementos nuevos de programación en Stan.

veamos un ejemplo:

2 Ejemplo 1: datos sintéticos de depredación de huevos

Para contextualizar el ajuste de este tipo de modelos generemos un conjunto de datos sintéticos sobre depredación de nidadas del ave X. Supongamos que una investigadora y su equipo monitorearon 1000 nidadas con N huevos cada una; los huevos estuvieron distribuidos en 20 parches de bosque y bajo dos condiciones de estrato vegetal, arbustivo y arbóreo. El objetivo es calcular la probabilidad de que los huevos sean depredados en cada condición y, por lo tanto, se debe “controlar” la variación asociada a los nidos y los parches –por “controlar” me refiero a incluir dichas variables en el modelo para estimar parámetros para cada nido y parche. Además, lo que la investigadora descubrirá solo una vez ajuste los modelos, es que generaremos los datos de manera que los parches y la condición de estrato vegetal covaríen.

Primero, siguiendo el procedimiento descrito arriba, generemos la estructura de covariación entre parches (\(\tau\)) y estrato (\(\beta\)).

alpha <- 0.05
tau <- 0.1
beta <- 1.2
sigma_alpha <- 0.8
sigma_tau <- 1
sigma_beta <- 0.9
rho_alpha_beta <- 0.005
rho_tau_beta <- 0.8

mu1 <- c(alpha, beta)
mu2 <- c(tau, beta)

cov_alpha_beta <- sigma_alpha*sigma_beta*rho_alpha_beta
cov_tau_beta <- sigma_tau*sigma_beta*rho_tau_beta

sig_alpha <- c(sigma_alpha, sigma_beta)
mat_VC_alpha <- matrix(c(1, 
                         cov_alpha_beta, 
                         cov_alpha_beta, 
                         1), ncol = 2) # matriz de correlación
mat_VC_alpha <- diag(sig_alpha) %*% mat_VC_alpha %*% diag(sig_alpha)

sig_tau <- c(sigma_tau, sigma_beta)
mat_VC_tau <- matrix(c(1, 
                       cov_tau_beta, 
                       cov_tau_beta, 
                       1), ncol = 2) # matriz de correlación
mat_VC_tau <- diag(sig_tau) %*% mat_VC_tau %*% diag(sig_tau)

La única diferencia con el ejemplo anterior es que en este caso la matriz de varianza-covarianza está construida a partir del producto de tres matrices: dos con la varianza de cada parámetro en la diagonal y la matriz de correlación (líneas 24-28 y 17-21). En sintaxis de R el símbolo %*% se usa para calcular el producto de dos matrices.

Veamos gráficamente la correlación entre ambos parámetros:

par(mfrow = c(1, 2), mar = c(4, 4, 1, 1))
set.seed(123)
nidos <- rmvnorm(1000, mu1, sigma = mat_VC_alpha)
plot(nidos[, 1], nidos[, 2], xlab = expression(alpha), 
     ylab = expression(beta))
for (i in seq(0.1, 0.9, by = 0.1)) {
  l <- ellipse(mat_VC_alpha, centre = mu1, level = i)
  lines(l[, 1], l[, 2], col = 'tomato', lwd = 2)
}

set.seed(123)
parche <- rmvnorm(1000, mu2, sigma = mat_VC_tau)
plot(parche[, 1], parche[, 2], xlab = expression(tau), 
     ylab = expression(beta))
for (i in seq(0.1, 0.9, by = 0.1)) {
  l <- ellipse(mat_VC_tau, centre = mu2, level = i)
  lines(l[, 1], l[, 2], col = 'tomato', lwd = 2)
}

par(mfrow = c(1, 1))

Ahora generemos los datos:

N <- 1000
set.seed(123)
nidada <- rpois(N, 10)
nido_id <- 1:N
parches <- rep(1:20, each = 10 , length.out = N)
bosque <- rep(0:1, each = 5, length.out = N)
alpha_nido <- nidos[, 1]
tau_parche <- parche[, 1]
beta_bosque <- parche[, 2]
mu <- alpha_nido[nido_id] + tau_parche[parches] + beta_bosque[parches]*bosque
mu <- inv_logit(mu)
set.seed(123)
huevos <- rbinom(N, size = nidada, prob = mu)

dat <- 
  list(
    N = N, 
    N_parche = 20,
    N_bosque = 2,
    N_nidos = 1000,
    nido = nido_id,
    nidada = nidada,
    bosque = bosque, 
    huevos = huevos,
    parche = parches
  )

Con el fin de comparar ambas aproximaciones, primero ajustemos un modelo jerárquico con partial pooling entre clusters. Definamos el modelo en notación matemática y luego en lenguaje Stan.

$$ \[\begin{align} & Huevos~sobrevivientes_i \sim binomial(huevos~observados_i,~ p_i) \\ & logit(p_i) = \alpha_{nido_i} + \tau_{parche_i} + \beta_{parche_i} \times bosque_i \\ & z_\alpha ~ normal(0, 1) \\ & \mu_\alpha ~ normal(0, 1) \\ & \sigma_\alpha ~ exponential(1) \\ & z_\tau ~ normal(0, 1) \\ & \mu_\tau ~ normal(5, 2) \\ & \sigma_\tau ~ exponential(1) \\ & z_\beta ~ normal(0, 1) \\ & \mu_\beta ~ normal(0, 1) \\ & \sigma_\beta ~ exponential(1) \\ & \alpha = \mu_\alpha + z_\alpha \times \sigma_\alpha \\ & \tau = \mu_\tau + z_\tau \times \sigma_\tau \\ & \beta = \mu_\beta + z_\beta \times \sigma_\beta \\ \end{align}\] $$

cat(file = 'multilevel_mods_II/eg_cov_0.stan', 
    '
    data{
      int N;
      int N_parche;
      int N_bosque;
      int N_nidos;
      array[N] int nido;
      array[N] int nidada;
      array[N] int bosque;
      array[N] int huevos;
      array[N] int parche;
    }
    
    parameters {
      vector[N_nidos] z_alpha;
      real mu_alpha;
      real<lower = 0> sigma_alpha;
      vector[N_parche] z_tau;
      real mu_tau;
      real<lower = 0> sigma_tau;
      vector[N_parche] z_beta;
      real mu_beta;
      real<lower = 0> sigma_beta;
    }
    
    transformed parameters {
      vector[N_nidos] alpha;
      vector[N_parche] tau;
      vector[N_parche] beta;
      alpha = mu_alpha + z_alpha*sigma_alpha;
      tau = mu_tau + z_tau*sigma_tau;
      beta = mu_beta + z_beta*sigma_beta;
    }
    
    model{
      vector[N] p;
      z_alpha ~ normal(0, 1);
      mu_alpha ~ normal(0, 1);
      sigma_alpha ~ exponential(1);
      z_tau ~ normal(0, 1);
      mu_tau ~ normal(5, 2);
      sigma_tau ~ exponential(1);
      z_beta ~ normal(0, 1);
      mu_beta ~ normal(0, 1);
      sigma_beta ~ exponential(1);
    
      for (i in 1:N) {
        p[i] = alpha[nido[i]] + tau[parche[i]] + beta[parche[i]] * bosque[i];
        p[i] = inv_logit(p[i]);
      }
    
      huevos ~ binomial(nidada, p);
    }
    
    generated quantities{
      vector[N] log_lik;
      vector[N] p;
      array[N] int ppcheck;
    
      for (i in 1:N) {
        p[i] = alpha[nido[i]] + tau[parche[i]] + beta[parche[i]] * bosque[i];
        p[i] = inv_logit(p[i]);
      }
    
      for (i in 1:N) log_lik[i] = binomial_lpmf(huevos[i] | nidada[i], p[i]);
    
      ppcheck = binomial_rng(nidada, p);
    }
    ')

Se deben notar varias cosas. Usé una parametrización no centrada para ajustar el modelo, pero además incluimos un nuevo bloque de código en Stan y algunas funciones adicionales. Veamos:

  1. bloque generated quantities: lo usamos para realizar operaciones con parámetros cuyo resultado estará disponible para usarlo en el bloque model y generated quantities.

  2. binomial_lpmf: las siglas _lpmf denotan log probability mass function lo calculo para estimar WAIC, PSIS o las distancias Pareto K.

  3. binomial_rng: las siglas _rng denotan random number generator y lo uso para generar números pseudo aleatorios (ahorra algunas líneas de código en R para calcular las distribuciones predictivas posteriores).

file <- paste(getwd(), '/multilevel_mods_II/eg_cov_0.stan', sep = '')
fit_m_noCOR <- cmdstan_model(file, compile = T)

m_noCORR <- 
  fit_m_noCOR$sample(
    data = dat,
    iter_sampling = 4e3, 
    iter_warmup = 500,
    chains = 3,
    parallel_chains = 3, 
    thin = 3,
    refresh = 500
  )

m_noCORR$save_object('multilevel_mods_II/eg_cov_0.rds')

Verifiquemos el ajuste del modelo y detectemos potenciales problemas:

m_noCORR <- readRDS('eg_cov_0.rds')

out_noCOR <- m_noCORR$summary()

par(mfrow = c(3, 3), mar = c(4, 4, 1, 1))
for (i in 1:9) trace_plot(m_noCORR, out_noCOR$variable[i], 3)

par(mfrow = c(1, 1))

mod_diagnostics(m_noCORR, out_noCOR)

Tenemos algunos valores de Pareto k > 0.7, lo que puede sugerir que el modelo es demasiado optimista en término de la estimación de los parámetros. Por lo demás, todos los parámetros tienen un ess > 1000 y Rhat < 1.01. Ya volveremos a esto cuando ajustemos los modelos incluyendo la estructura de covariación entre los parámetros.

Ahora veamos las distribuciones predictivas posteriores:

ppcheck <- as.matrix(m_noCORR$draws(variables = 'ppcheck', format = 'df'))

plot(density(ppcheck[1, ]), main = '', lwd = 0.1, ylim = c(0, 0.14), 
     xlab = 'Huevos sobrevivientes')
for (i in 1:500) lines(density(ppcheck[i, ]), lwd = 0.1)
lines(density(dat$huevos), col = 'red', lwd = 2)

Todo se ve en orden. Grafiquemos los parámetros correlacionados:

post <- m_noCORR$draws(variable = c('alpha', 'tau', 'beta'), format = 'df')
post_noCOR <- 
  list(
    alpha = post[, grep('alpha', colnames(post))],
    tau = post[, grep('tau', colnames(post))],
    beta = post[, grep('beta', colnames(post))]
  )


par(mar = c(4, 4, 1, 1))
plot(apply(post_noCOR$tau, 2, mean), 
     apply(post_noCOR$beta, 2, mean), 
     main = '', xlab = expression(tau), ylab = expression(beta))

Claramente existe una asociación entre los parámetros, pero el algoritmo no puede tomar ventaja de ella porque no la incluimos en la estructura del modelo. Reformulemos el modelo:

La función de likelihood y la función lineal son los mismos:

$$ \[\begin{align} & Huevos~sobrevivientes_i \sim binomial(huevos~observados_i,~ p_i) \\ & logit(p_i) = \alpha_{nido_i} + \tau_{parche_i} + \beta_{parche_i} \times bosque_i \\ \end{align}\] $$ La diferencia entre los modelos ajustados anteriormente y este, radica en que la consideración de la covariación entre parámetros requiere el uso de una distribución normal multivariada como previa adaptativa y una previa para la matriz de correlaciones de los parámetros. Veamos la continuación del modelo:

$$ \[\begin{align} & [\tau_{parche_i}, ~ \beta_{parche_i}] \sim MVNormal([\overline{\tau}, ~ \overline{\beta}],~R_{\tau \beta}, [\overline{\sigma_{\tau}}, \overline{\sigma_{\beta}}]) \\ & [\alpha_{parche_i}, ~ \beta_{parche_i}] \sim MVNormal([\overline{\alpha}, ~ \overline{\beta}],~R_{\alpha \beta}, [\overline{\sigma_{\alpha}}, \overline{\sigma_{\beta}}]) \\ & R_{\tau \beta} \sim LKJCorr(2) \\ & R_{\alpha \beta} \sim LKJCorr(2) \\ & \overline{\tau} \sim Normal(0, 1) \\ & \overline{\alpha} \sim Normal(0, 1) \\ & \overline{\beta} \sim Normal(0, 1) \\ & \overline{\sigma_{\tau}} \sim exponential(1) \\ & \overline{\sigma_{\alpha}} \sim exponential(1) \\ & \overline{\sigma_{\beta}} \sim exponential(1) \\ \end{align}\] $$ \(MVNormal([\overline{\tau}, ~ \overline{\beta}],~R, [\overline{\sigma_{\tau}}, \overline{\sigma_{\beta}}])\) corresponde la previa adaptativa de la poblacion de los parámetros cuya correlación nos interesa incorporar (i.e. \([\alpha_{nido_i}, ~ \beta_{nido_i}]\)); \(R\) corresponde a la matriz de correlación de los parámetros, cuya previa proviene de una distribución de probabilidad de matrices de correlación \(LKJCorr()\) (ver abajo las consecuencias de reducir o aumentar el valor de su parámetro); \([\overline{\sigma_{\tau}}, \overline{\sigma_{\beta}}]\) corresponde a un vector de párametros de dispersión para \(\tau_{parche}\) y \(\beta_{parche}\). La misma estructura de previas se repite para \([\alpha_{parche_i}, ~ \beta_{parche_i}]\).

par(mfrow = c(3, 3), mar = c(4, 4, 1, 1))

for (i in 1:9) {
  
  R <- rlkjcorr(1e3, 
                K = 2, # número de parámetros (dimension de la matriz)
                eta = i) # forma de la distribución
  
  plot(density(R[, 1, 2]),
       main = '', 
       xlab = 'Correlacion'); title(paste('eta=', i))
}

par(mfrow = c(1, 1))

Ajustemos el modelo en Stan empleando una parametrización centrada y discutamos algunos elementos de la sintaxis.

cat(file = 'multilevel_mods_II/eg_cov_par1.stan', 
    '
    data{
      int N;
      int N_parche;
      int N_bosque;
      int N_nidos;
      array[N] int nido;
      array[N] int nidada;
      array[N] int bosque;
      array[N] int huevos;
      array[N] int parche;
    }
    
    parameters{
      matrix[N_parche, N_bosque] M_tau;
      corr_matrix[N_bosque] rho_parche;
      vector<lower = 0>[N_bosque] sigma_parche;
      vector[N_bosque] tau_mu;
    
      matrix[N_nidos, N_bosque] M_alpha;
      corr_matrix[N_bosque] rho_nido;
      vector<lower = 0>[N_bosque] sigma_nido;
    }
    
    transformed parameters{
      vector[N_parche] tau;
      vector[N_parche] beta;
      vector[N_nidos] alpha;
      
      alpha = M_alpha[, 1];
      tau = M_tau[, 1];
      beta = M_tau[, 2];
    }
    
    model{
      vector[N] p;
      sigma_parche ~ exponential(1);
      tau_mu ~ normal(0, 1);
      rho_parche ~ lkj_corr(1);
      sigma_nido ~ exponential(1);
      rho_nido ~ lkj_corr(1);
      for (i in 1:N_parche) M_tau[i, :] ~ multi_normal(tau_mu, 
                                                       quad_form_diag(rho_parche, 
                                                       sigma_parche));
      for (i in 1:N_nidos) M_alpha[i, :] ~ multi_normal(rep_vector(0, N_bosque),
                                                       quad_form_diag(rho_nido,
                                                       sigma_nido));
    
      for (i in 1:N) {
        p[i] = alpha[nido[i]] + tau[parche[i]] + beta[parche[i]]*bosque[i];
        p[i] = inv_logit(p[i]);
      }
      
      huevos ~ binomial(nidada, p);
    }
    
    generated quantities{
      vector[N] log_lik;
      vector[N] p;
      array[N] int ppcheck;
    
      for (i in 1:N) {
        p[i] = alpha[nido[i]] + tau[parche[i]] + beta[parche[i]]*bosque[i];
        p[i] = inv_logit(p[i]);
      }
    
      for (i in 1:N) log_lik[i] = binomial_lpmf(huevos[i] | nidada[i], p[i]);
    
      ppcheck = binomial_rng(nidada, p);
    }
    ')

La novedad es la asignación de previas para cada nivel de la variable de grupo (i.e. nido o parche). Esto es, iteramos sobre cada fila de la matriz (i.e. \([grupo, ~condicion~bosque]\)) para asignar a cada nivel de factor una distribución previa \(MVNormal([\overline{\tau}, ~ \overline{\beta}],~R, [\overline{\sigma_{\tau}}, \overline{\sigma_{\beta}}])\) o \(MVNormal([\overline{\alpha}, ~ \overline{\beta}],~R_{\alpha \beta}, [\overline{\sigma_{\alpha}}, \overline{\sigma_{\beta}}])\). Stan usa la función quad_form_diag para convertir el vector de los parámetros de dispersión en la diagonal de una matriz y factorizarla con la matriz de correlación de los parámetros (un poco sobre esto acá y acá).

Ahora ajustemos el modelo:

file <- paste(getwd(), '/multilevel_mods_II/eg_cov_par1.stan', sep = '')
fit <- cmdstan_model(file, compile = T)

m <- 
  fit$sample(
    data = dat,
    iter_sampling = 4e3,
    iter_warmup = 500, 
    thin = 3, 
    chains = 3, 
    parallel_chains = 3, 
    refresh = 200, 
    seed = 123
  )

m$save_object(paste(getwd(), '/multilevel_mods_II/eg_cov_par1.rds', sep = ''))

Verifiquemos la calidad de su ajuste

m <- readRDS('eg_cov_par1.rds')

m_out <- m$summary()

par(mfrow = c(3, 3), mar = c(4, 4, 1, 1))
for (i in 1:9) trace_plot(m, m_out[m_out$rhat >1.01,][-c(1:5), ]$variable[i+2], 3)

par(mforw = c(1, 1))

mod_diagnostics(m, m_out)

m$diagnostic_summary()
$num_divergent
[1]  0 32  0

$num_max_treedepth
[1]  2  5 53

$ebfmi
[1] 0.01269898 0.01247557 0.02439522

Mmmmmm… Las cadenas de varios parámetros exploraron distintas regiones del espacio (del parámetro :) ), tenemos valores de Rhat > 1.01 y un ess < 1000 para muchos parámetros.

Para recordar: Las transiciones divergentes son un lugar común en modelos jerárquicos dado que tenemos parámetros cuya distribución es función de otros parámetros. Esto resulta en ineficiencia del algoritmo para explorar el espacio del parámetro, tiempos de cómputo más prolongados y bajo número de estimados estadísticamente independientes para cada distribución posterior. En el mejor de los casos ajustar el modelo con mayor número de iteraciones puede resultar en la convergencia de las cadenas; en el peor, las cadenas simplemente no van a converger y la reparametrización es indispensable.

Cómo ya es rutina con los modelos jerárquicos debemos reparametrizar el modelo:

$$ \[\begin{align} & Huevos~sobrevivientes_i \sim binomial(huevos~observados_i,~ p_i) \\ & logit(p_i) = \alpha_{nido_i} + \tau_{parche_i} + \beta_{parche_i} \times bosque_i \\ & \alpha_{nido_i} = \overline{\alpha} + \omega[, 1] \\ & \tau_{parche_i} = \overline{\tau} + \theta[, 1] \\ & \beta_{parche_i} = \overline{\beta} + \theta[, 2] \\ & \omega = (diag(\sigma_\alpha), R1_{cholesky}Z_\alpha)^T \\ & \theta = (diag(\sigma_\tau), R2_{cholesky}Z_\tau)^T \\ & \overline{\alpha} \sim Normal(0, 1) \\ & \overline{\tau} \sim Normal(0, 1) \\ & \sigma_\alpha \sim exponential(1) \\ & \sigma_\tau \sim exponential(1) \\ & Z_{\alpha,\beta} \sim Normal(0, 1) \\ & Z_{\tau,\beta} \sim Normal(0, 1) \\ & R1_{cholesky} \sim LKJCorr(2) \\ & R2_{cholesky} \sim LKJCorr(2) \\ \end{align}\] $$ Discutamos un poco el modelo reparametrizado. Reemplazamos la distribución normal multivariada por las matrices \(Z_{\alpha,\beta}\) y \(Z_{\tau,\beta}\) a cuyos elementos fue asignada una previa \(Normal(0,~1)\). Luego tenemos que las matrices de correlación, \(R1_{cholesky}\) y \(R2_{cholesky}\), son factorizadas a través de una descomposición de Cholesky (necesaria para que funcione la reparametrización). La expresión \((diag(\sigma_\alpha), R1_{cholesky}Z_\alpha)^T\) denota operaciones de álgebra lineal para factorizar las matrices de correlación, estandarizada (\(Z\)) y la diagonal de los parámetros de dispersión. Se usa la transpuesta (\(^T\)) de esta matriz para generar los parámetros de la función lineal.

¡Si! La reparametrización de estos modelos es ligeramente más compleja. Sin embargo, convengamos que Stan se encarga del trabajo duro (i.e. factorización de matrices), mientras que nosotros sólo debemos prestar atención a la asignación de las dimensiones de las matrices y a la indexación. Veamos el modelo en Stan:

cat(file = 'multilevel_mods_II/eg_cov_par2.stan', 
    "
    data{
      int N;
      int N_parche;
      int N_bosque;
      int N_nidos;
      array[N] int nido;
      array[N] int nidada;
      array[N] int bosque;
      array[N] int huevos;
      array[N] int parche;
    }
    
    parameters{
      matrix[N_bosque, N_parche] Z_tau;
      cholesky_factor_corr[N_bosque] rho_tau;
      vector<lower = 0>[N_bosque] sigma_tau;
      vector[N_bosque] tau_bar;
    
      matrix[N_bosque, N_nidos] Z_alpha;
      cholesky_factor_corr[N_bosque] rho_alpha;
      vector<lower = 0>[N_bosque] sigma_alpha;
      vector[N_bosque] alpha_bar;
    }
    
    transformed parameters{
      vector[N_nidos] alpha;
      vector[N_nidos] beta2;
      vector[N_parche] tau;
      vector[N_parche] beta;
      matrix[N_nidos, N_bosque] M_alpha;
      matrix[N_parche, N_bosque] M_tau;
      M_alpha = (diag_pre_multiply(sigma_alpha, rho_alpha) * Z_alpha)';
      M_tau = (diag_pre_multiply(sigma_tau, rho_tau) * Z_tau)';
      alpha = alpha_bar[1] + M_alpha[, 1];
      beta2 = alpha_bar[2] + M_alpha[, 2];
      tau = tau_bar[1] + M_tau[, 1];
      beta = tau_bar[2] + M_tau[, 2];
    }
    
    model{
      vector[N] p;
      vector[N] p2;
      sigma_alpha ~ exponential(1);
      sigma_tau ~ exponential(1);
      tau_bar ~ normal(0, 1);
      alpha_bar ~ normal(0, 1);
      to_vector(Z_alpha) ~ normal(0, 1);
      to_vector(Z_tau) ~ normal(0, 1);
      rho_alpha ~ lkj_corr_cholesky(2);
      rho_tau ~ lkj_corr_cholesky(2);
    
      for (i in 1:N) {
        p[i] = alpha[nido[i]] + tau[parche[i]] + beta[parche[i]]*bosque[i];
        p[i] = inv_logit(p[i]);
      }
      
      huevos ~ binomial(nidada, p);
    
    for (i in 1:N) {
        p2[i] = alpha[nido[i]] + tau[parche[i]] + beta2[parche[i]]*bosque[i];
        p2[i] = inv_logit(p2[i]);
      }
      
      huevos ~ binomial(nidada, p2);
    
    }
    
    generated quantities{
      vector[N] log_lik;
      vector[N] p;
      array[N] int ppcheck;
      matrix[N_bosque, N_bosque] rho_parche;
      matrix[N_bosque, N_bosque] rho_nido;
      rho_parche = multiply_lower_tri_self_transpose(rho_tau);
      rho_nido = multiply_lower_tri_self_transpose(rho_alpha);
    
      for (i in 1:N) {
        p[i] = alpha[nido[i]] + tau[parche[i]] + beta[parche[i]]*bosque[i];
        p[i] = inv_logit(p[i]);
      }
    
      for (i in 1:N) log_lik[i] = binomial_lpmf(huevos[i] | nidada[i], p[i]);
    
      ppcheck = binomial_rng(nidada, p);
    }
    ")

Elementos nuevos de la sintaxis: cholesky_factor_corr indica una matriz de correlación factorizada con una descomposición de Cholesky; diag_pre_multiply realiza la factorización de las matrices de Cholesky, estandarizada y diagonal de los parámetros de dispersión; ' indica la transpuesta de la matriz; lkj_corr_cholesky es el equivalente de lkj_corr() para matrices de correlaciones factorizadas con descomposición de Cholesky; la función multiply_lower_tri_self_transpose convierte la matriz de Cholesky a una matriz de correlación que podemos interpretar en rangos entre -1 y 1.

Ajustemos el modelo:

file <- paste(getwd(), '/multilevel_mods_II/eg_cov_par2.stan', sep = '')
fit <- cmdstan_model(file, compile = T)

m2 <- 
  fit$sample(
    data = dat,
    iter_sampling = 4e3,
    iter_warmup = 500, 
    thin = 3, 
    chains = 3, 
    parallel_chains = 3, 
    refresh = 200, 
    seed = 123
  )

Verifiquemos la “sanidad” del ajuste

m2 <- readRDS('eg_cov_par2.rds')

m2_out <- m2$summary()

mod_diagnostics(m2, m2_out)

Ahora tenemos suficientes estimados estadísticamente independientes, Rhat < 1.01 y redujimos los valores de Pareto-k. Grafiquemos las distribuciones predictivas posteriores:

ppcheck <- as.matrix(m2$draws(variables = 'ppcheck', format = 'df'))

plot(density(ppcheck[1, ]), xlab = 'Huevos sobrevivientes', 
     ylab = 'Density', main = '', lwd = 0.01, ylim = c(0, 0.12))
for (i in 1:500) lines(density(ppcheck[i, ]), lwd = 0.1)
lines(density(dat$huevos), col = 'red', lwd = 2)

Todo En orden.

Extraigamos muestras de las distribuciones posteriores de los parámetros y comprobemos que el modelo pudo estimar la correlación entre \(\beta\) y \(\tau\). También, comparemos los parámetros generadores de los datos (i.e. los “reales”), los estimados por el partial pooling entre clusters, y el modelo con partial pooling entre parámetros/clusters.

post <- m2$draws(variables = c('alpha', 'tau', 'beta', 'beta2'), format = 'df')
post <- 
  list(alpha = apply(post[, grep('alpha', colnames(post))], 2, mean),
       tau = apply(post[, grep('tau', colnames(post))], 2, mean), 
       beta = apply(post[, grep('beta\\[', colnames(post))], 2, mean), 
       beta2 = apply(post[, grep('beta2', colnames(post))], 2, mean))

rho_parche <- m2$draws(variables = c('rho_parche'), 
                       format = 'df')

sigma_parche <- m2$draws(variables = 'sigma_tau', format = 'df')

sigmas_parche <- 
  matrix(c(mean(sigma_parche$`sigma_tau[1]`)^2,
             mean(rho_parche$`rho_parche[2,1]`),
             mean(rho_parche$`rho_parche[2,1]`),
             mean(sigma_parche$`sigma_tau[2]`)^2), ncol = 2)

mu_parche <- c(mean(post$tau), mean(post$beta))

mu_nido <- c(mean(post$alpha), mean(post$beta2))

sigmas_nido <- m2$draws(variables = 'sigma_alpha', format = 'df')
sigmas_nido <- apply(sigmas_nido, 2, mean)

rho_nido <- m2$draws(variables = 'rho_nido', format = 'df')

sigmas_nido <- 
  matrix(
    c(sigmas_nido[1]^2, 
      mean(rho_nido$`rho_nido[2,1]`), 
      mean(rho_nido$`rho_nido[2,1]`), 
      sigmas_nido[2]^2), ncol = 2
  )


tau_noCORR <- apply(post_noCOR$tau, 2, mean)
beta_noCORR <- apply(post_noCOR$beta, 2, mean)
obs_tau_beta <- parche[1:20, ]

par(mfrow = c(1, 2), mar = c(4, 4, 2.5, 1))
post %$% plot(tau, beta, xlab = expression(tau), 
              ylab = expression(beta), col = 2, 
              xlim = c(-3, 3), ylim = c(-1, 4.5), 
              main = 'Modelo\n par. pool. pars/clusters')
for (i in seq(0.1, 0.9, by = 0.2))
  lines(ellipse(sigmas_parche, centre = mu_parche, level = i), lwd = 0.5)

post %$% plot(tau, beta, xlab = expression(tau), 
              ylab = expression(beta), col = 2, 
              xlim = c(-3, 3), ylim = c(-1, 4.5), 
              main = 'Pars. estimados y reales')
points(tau_noCORR, beta_noCORR, col = 4)
points(obs_tau_beta[, 1], obs_tau_beta[, 2], pch = 16)

par(mfrow = c(1, 1))

Los círculos rojos corresponden a los parámetros estimados con el modelo que incorpora la covarianza entre parámetros y las elipses denotan intervalos de credibilidad, desde 10 hasta 90%. En la derecha vemos los parámetros estimados de cada modelo (círculos azules = partial pooling entre clusters) y sus valores reales (puntos negros). A simple vista notamos que los puntos negros y los círculos rojos tienden a estar más solapados, lo que sugiere una mayor precisión del modelo que usa la información entre parámetros y clusters. Calculemos estos errores de estimación y realicemos algunos gráficos para enfatizar las ventajas de un modelo sobre otro.

parametros <- 
  tibble(beta_cor = post$beta, 
       tau_cor = post$tau,
       beta_noCORR = beta_noCORR, 
       tau_noCORR = tau_noCORR, 
       real_beta = obs_tau_beta[, 2], 
       real_tau = obs_tau_beta[, 1]) 

parametros <- 
  parametros |> 
  mutate(error_corr_beta = real_beta - beta_cor, 
         error_corr_tau = real_tau - tau_cor, 
         error_noCorr_beta = real_beta - beta_noCORR, 
         error_noCorr_tau = real_tau - tau_noCORR, )

par(mfrow = c(2, 2), mar = c(4, 4, 2.5, 1))
parametros %$% 
  plot(tau_noCORR, beta_noCORR, xlim = c(-1.5, 2.5), ylim = c(-1, 4), 
       main = 'Error de estimación\n(par. pool. clusters)', 
       xlab = expression(tau), ylab = expression(beta), col = 4)
parametros %$%
  points(real_tau, real_beta, pch = 16)
parametros %$%
  segments(x0 = tau_noCORR, x1 = real_tau, 
           y0 = beta_noCORR, y1 = real_beta, lty = 3, lwd = 0.8)

parametros %$% 
  plot(tau_cor, beta_cor, xlim = c(-1.5, 1.8), ylim = c(-1, 4),
       main = 'Error de estimación\n(par. pool. pars/clusters)', 
       col = 2, xlab = expression(tau), ylab = expression(beta))
parametros %$%
  points(real_tau, real_beta, pch = 16)
parametros %$%
  segments(x0 = tau_cor, x1 = real_tau, 
           y0 = beta_cor, y1 = real_beta, lty = 3, lwd = 0.8)

parametros %$% plot(density(error_corr_beta), col = 2, lwd = 3, 
                    main = 'Error (observado - estimado)', 
                    xlab = expression(beta))
abline(v = median(parametros$error_corr_beta), lty = 3, col = 2, 
       lwd = 3)
parametros %$%
  lines(density(error_noCorr_beta), col = 4, lwd = 3)
abline(v = median(parametros$error_noCorr_beta), lty = 3, col = 4, 
       lwd = 3)

parametros %$% plot(density(error_corr_tau), col = 2, lwd = 3, 
                    main = 'Error (observado - estimado)', 
                    xlab = expression(tau), 
                    xlim = c(-1.5, 0.45))
abline(v = median(parametros$error_corr_tau), lty = 3, col = 2, 
       lwd = 3)
parametros %$%
  lines(density(error_noCorr_tau), col = 4, lwd = 3)
abline(v = median(parametros$error_noCorr_tau), lty = 3, col = 4, 
       lwd = 3)

par(mfrow = c(1, 1))

El color azul corresponde al modelo con partial pooling entre clusters, el rojo al modelo con partial polling entre parámetros/clusters. Ahora es claro que que ambos modelos estiman de manera similar \(\beta\) –aunque sólo estamos comparando promedios sin la dispersión de las distribuciones posteriores, \(\tau\) es otra historia. El modelo que usa la covariación entre parámetros estima \(\tau\) de manera bastante precisa, mientras que el modelo que solo emplea partial pooling entre clusters subestima el valor real de los \(\tau_i\) en casi una unidad. En efecto, incluir la estructura de covarianza de los parámetros en el modelo resulta en una mayor capacidad predictiva del modelo (out of sample prediction); comprobémoslo usando leave-one-out cross-validation:

loo_compare(m_noCORR$loo(), m2$loo())
       elpd_diff se_diff
model2    0.0       0.0 
model1 -467.8      15.9 

El 2do modelo es el mejor ranqueado.

Finalmente, grafiquemos la distribución posterior de la correlación entre \(\tau\) y \(\beta\):

plot(density(rho_parche$`rho_parche[2,1]`), col =4, lwd = 3, 
     main = expression(rho[beta~tau]), xlab = expression(italic('r')))
abline(v = median(rho_parche$`rho_parche[2,1]`), lty = 3, col = 4, 
       lwd = 3)
abline(v = rho_tau_beta, lty = 3, col = 1, 
       lwd = 3)

El valor real de la correlación \(\rho_{\beta - \tau}\) se encuentra contenido en la distribución posterior de nuestro parámetro, WELL DONE!.

3 Ejemplo 2: pingüinos y rasgos funcionales

Usaremos los datos palmerpenguins para analizar la relación del tamaño del ala/aleta con el tamaño corporal del pingüino. Los datos son interesantes porque nos permitirán extender el ejemplo anterior ya que involucran variables continuas y están estratificados en cuatro niveles anidados: especie (niveles = 3), sexo (niveles = 2), isla (niveles = 3) y año (niveles = 2). La idea es modelar esta relación considerando la potencial covariación entre los efectos de grupo (i.e. año e isla) y efectos poblacionales (i.e. especies y sexo). Además, permitiremos que el algoritmo explore valores estratificados de la pendiente por especie/sexo, empleando partial pooling entre clusters.

Carguemos los datos y preparémoslos en el formato que requerimos:

d <- na.omit(palmerpenguins::penguins)

d <- d |> dplyr::select(species, year, island, sex, flipper_length_mm, body_mass_g)

d$year <- as.factor(d$year)

d$species_sex <- as.factor(paste(d$species, d$sex, sep = '_'))

dat <- lapply(d, function(x) if(!is.double(x)) as.integer(x) else(x))

dat$N <- length(dat$species)
dat$N_spp <- max(dat$species)
dat$N_sex <- 2
dat$N_island <- max(dat$island)
dat$N_year <- max(dat$year)
dat$N_spp_sex <- max(dat$species_sex)

Definamos el modelo en notación matemática y luego en lenguaje Stan.

$$ \[\begin{align} & largo~aleta_{i}~(mm) \sim normal(\mu_i, ~ \sigma_i) \\ & \mu_i = \theta_{species/sex_i} + \alpha_{[island_i, ~species/sex_i]} + \\ &~~~~~~~~~ \tau_{[year_i, ~species/sex_i]} + \beta_{species/sex_i}\times~masa~corporal~(g)_i \\ & \alpha_{species/sex_i} = (diag(\sigma_\alpha), R1_{cholesky}Z_\alpha)^T \\ & \tau_{[year_i, ~species/sex_i]} = (diag(\sigma_\tau), R1_{cholesky}Z_\tau)^T \\ & \beta_{species/sex_i} = \mu_\beta + Z_\beta \times \sigma_\beta \\ & \theta_{species/sex_i} \sim Normal(200, ~50) \\ & \sigma_\alpha \sim exponential(1) \\ & \sigma_\tau \sim exponential(1) \\ & R1_{cholesky} \sim LKJCorr(2) \\ & R2_{cholesky} \sim LKJCorr(2) \\ & Z_\alpha \sim normal(0, 1) \\ & Z_\tau \sim normal(0, 1) \\ & \mu_\beta \sim nrormal(0.5, 0.5) \\ & Z_\beta \sim normal(0, 1) \\ & \sigma_\beta \sim expoenential(1) \\ & \sigma_i \sim exponential(1) \\ \end{align}\] $$ Este modelo comprende varios de los procedimientos para cruzar información en la estimación de los parámetros del modelo: no pooling, partial polling entre clusters y partial pooling entre parámetros y clusters. Por ejemplo, no empleé partial pooling para estimar los valores del intercepto \(\theta_{species/sex_i}\), es discutible. No lo hice por motivos más conceptuales que metodológicos, pero perfectamente podría incluirse. También, notemos que usé una parametrización no centrada y que los parámetros \(\alpha\) y \(\tau\) son incluidos en el modelo como matrices. En el caso de \(\beta\) definí una parametrización no centrada y partial pooling sin efectos correlacionados.

En lenguaje Stan:

cat(file = 'multilevel_mods_II/penguins_par_pool.stan', 
    "
    data{
      int N;
      int N_spp_sex;
      int N_island;
      int N_year;
      vector[N] flipper_length_mm;
      vector[N] body_mass_g;
      array[N] int species_sex;
      array[N] int year;
      array[N] int island;
    }
    
    parameters{
      
      vector[N_spp_sex] z_beta;
      real mu_beta;
      real<lower = 0> sigma_beta;
      
      matrix[N_spp_sex, N_island] z_island;  
      cholesky_factor_corr[N_spp_sex] Rho_island;
      vector<lower = 0>[N_spp_sex] sigma_island;
      
      matrix[N_spp_sex, N_year] z_year;
      cholesky_factor_corr[N_spp_sex] Rho_year;
      vector<lower = 0>[N_spp_sex] sigma_year;
      
      vector[N_spp_sex] theta;
      real<lower = 0> sigma;
    }
    
    transformed parameters{
    
      vector[N_spp_sex] beta;
      matrix[N_island, N_spp_sex] alpha;
      matrix[N_year, N_spp_sex] tau;
      alpha = (diag_pre_multiply(sigma_island, Rho_island) * z_island)';
      tau = (diag_pre_multiply(sigma_year, Rho_year) * z_year)';
      beta = mu_beta + z_beta*sigma_beta;
    }
    
    model{
      vector[N] mu;
      sigma_island ~ exponential(1);
      sigma_year ~ exponential(1);
      Rho_island ~ lkj_corr_cholesky(2);
      Rho_year ~ lkj_corr_cholesky(2);
      to_vector(z_island) ~ normal(0, 1);
      to_vector(z_year) ~ normal(0, 1);
      theta ~ normal(200, 50);
      mu_beta ~ normal(0.5, 0.5);
      z_beta ~ normal(0, 1);
      sigma_beta ~ exponential(1);
      sigma ~ exponential(1);
     
      for (i in 1:N) {
        mu[i] = theta[species_sex[i]] +  
                alpha[island[i], species_sex[i]] + tau[year[i], species_sex[i]] + 
                beta[species_sex[i]]*body_mass_g[i];
      }
      
      flipper_length_mm ~ normal(mu, sigma);
    }
    
    generated quantities{
      vector[N] log_lik;
      vector[N] mu_;
      array[N] real ppcheck;
      matrix[N_spp_sex, N_spp_sex] Rho_isla_corr;
      matrix[N_spp_sex, N_spp_sex] Rho_year_corr;
      
      Rho_isla_corr = multiply_lower_tri_self_transpose(Rho_island);
      Rho_year_corr = multiply_lower_tri_self_transpose(Rho_year);
    
      for (i in 1:N) {
        mu_[i] = theta[species_sex[i]] +  
                alpha[island[i], species_sex[i]] + tau[year[i], species_sex[i]] + 
                beta[species_sex[i]]*body_mass_g[i];
      }
      
      for (i in 1:N) log_lik[i] = normal_lpdf(flipper_length_mm[i] | mu_[i], sigma);
    
      ppcheck = normal_rng(mu_, sigma);
    }")

No hay muchos elementos nuevos en el modelo. Sugeriría prestar particular atención a la manera en que indexé los parámetros \(\alpha\) y \(\tau\). También, nótese que sus previas son el resultado de la factorización de tres matrices.

Ejecutemos el algoritmo y verifiquemos la calidad del ajuste:

file <- paste(getwd(), '/penguins_par_pool.stan', sep = '')

fit_m1 <- cmdstan_model(file, compile = T)

m1 <- 
  fit_m1$sample(
    data = dat,
    iter_sampling = 4e3,
    iter_warmup = 500,
    chains = 3,
    parallel_chains = 3,
    thin = 3,
    refresh = 500,
    seed = 123
  )
Running MCMC with 3 parallel chains...

Chain 1 Iteration:    1 / 4500 [  0%]  (Warmup) 
Chain 2 Iteration:    1 / 4500 [  0%]  (Warmup) 
Chain 3 Iteration:    1 / 4500 [  0%]  (Warmup) 
Chain 3 Iteration:  500 / 4500 [ 11%]  (Warmup) 
Chain 3 Iteration:  501 / 4500 [ 11%]  (Sampling) 
Chain 2 Iteration:  500 / 4500 [ 11%]  (Warmup) 
Chain 2 Iteration:  501 / 4500 [ 11%]  (Sampling) 
Chain 1 Iteration:  500 / 4500 [ 11%]  (Warmup) 
Chain 1 Iteration:  501 / 4500 [ 11%]  (Sampling) 
Chain 3 Iteration: 1000 / 4500 [ 22%]  (Sampling) 
Chain 2 Iteration: 1000 / 4500 [ 22%]  (Sampling) 
Chain 1 Iteration: 1000 / 4500 [ 22%]  (Sampling) 
Chain 3 Iteration: 1500 / 4500 [ 33%]  (Sampling) 
Chain 2 Iteration: 1500 / 4500 [ 33%]  (Sampling) 
Chain 1 Iteration: 1500 / 4500 [ 33%]  (Sampling) 
Chain 3 Iteration: 2000 / 4500 [ 44%]  (Sampling) 
Chain 2 Iteration: 2000 / 4500 [ 44%]  (Sampling) 
Chain 3 Iteration: 2500 / 4500 [ 55%]  (Sampling) 
Chain 1 Iteration: 2000 / 4500 [ 44%]  (Sampling) 
Chain 2 Iteration: 2500 / 4500 [ 55%]  (Sampling) 
Chain 3 Iteration: 3000 / 4500 [ 66%]  (Sampling) 
Chain 1 Iteration: 2500 / 4500 [ 55%]  (Sampling) 
Chain 2 Iteration: 3000 / 4500 [ 66%]  (Sampling) 
Chain 3 Iteration: 3500 / 4500 [ 77%]  (Sampling) 
Chain 1 Iteration: 3000 / 4500 [ 66%]  (Sampling) 
Chain 3 Iteration: 4000 / 4500 [ 88%]  (Sampling) 
Chain 2 Iteration: 3500 / 4500 [ 77%]  (Sampling) 
Chain 1 Iteration: 3500 / 4500 [ 77%]  (Sampling) 
Chain 3 Iteration: 4500 / 4500 [100%]  (Sampling) 
Chain 3 finished in 70.5 seconds.
Chain 2 Iteration: 4000 / 4500 [ 88%]  (Sampling) 
Chain 1 Iteration: 4000 / 4500 [ 88%]  (Sampling) 
Chain 2 Iteration: 4500 / 4500 [100%]  (Sampling) 
Chain 2 finished in 80.5 seconds.
Chain 1 Iteration: 4500 / 4500 [100%]  (Sampling) 
Chain 1 finished in 84.9 seconds.

All 3 chains finished successfully.
Mean chain execution time: 78.6 seconds.
Total execution time: 85.0 seconds.
out_m1 <- m1$summary()

par(mfrow = c(3, 3), mar = c(4, 4, 1, 1))
for (i in 1:9) trace_plot(m1, out_m1$variable[i], 3)

par(mfrow = c(1, 1))

mod_diagnostics(m1, out_m1)

Las cadenas convergieron a la misma distribución estacionaria, tenemos valores de Rhat < 1.01, ess > 1000 para todos los parámetros y valores de Pareto-k < 0.7. Todo en orden.

Grafiquemos algunas distribuciones predictivas posteriores y veamos cómo el modelo predice los datos que usamos para ajustarlo:

ppcheck_m1 <- m1$draws(variables = 'ppcheck', format = 'matrix')

plot(NULL, main = '', 
     xlim = c(160, 245), ylim = c(0, 0.04), xlab = 'longitud de la aleta (mm)', 
     ylab = 'Density')
for (i in 1:100) lines(density(ppcheck_m1[i, ]), lwd = 0.1)
lines(density(dat$flipper_length_mm), col = 'red')

Perfect!

Ahora grafiquemos los efectos principales, marginalizando el efecto del año y la isla.

post_m1 <- m1$draws(variables = c('theta', 'alpha', 'tau', 'beta', 'sigma'), 
                    format = 'df')

post_m1 <- 
  list(spp = post_m1[, grepl('theta', colnames(post_m1))],
       I = post_m1[, grepl('^alpha', colnames(post_m1))],
       Y = post_m1[, grepl('^tau', colnames(post_m1))], 
       beta = post_m1[, grepl('beta', colnames(post_m1))], 
       sigma = post_m1[, grepl('sigma', colnames(post_m1))])

d_temp <- d

d_temp$spp_sex <- dat$species_sex

d_temp <- split(d_temp, list(d_temp$species, d_temp$sex))

seq_x <- lapply(d_temp, FUN = 
                  function(x) {
                    
                    z <- x$body_mass_g
                    
                    z <- seq(min(x$body_mass_g), 
                           max(x$body_mass_g), 
                           length.out = 100)
                    
                    tibble(x_seq = z, 
                           spp_sex = x$spp_sex[[1]])
                    
                  })
seq_x <- do.call('rbind', seq_x)

sex <- c('H', 'M')


colnames(post_m1$spp) <- levels(d$species_sex)

nombres <- colnames(post_m1$spp)

est_cor <- lapply(1:ncol(post_m1$spp), FUN = 
                    function(i) {
                      
                      x_seq <- seq_x[seq_x$spp_sex == i, ]$x_seq
                      
                      df <- 
                        sapply(x_seq, FUN = 
                                 function(z) {
                                   
                                   post_m1$spp[, i, drop = T] + 
                                     apply(post_m1$I, 1, mean) +
                                     apply(post_m1$Y, 1, mean) +
                                     post_m1$beta[, i, drop = T]*z
                                   
                                 })
                      
                      df <- do.call('rbind',
                                    apply(df, 2, FUN = 
                                            function(j) {
                                              tibble(mu_est = mean(j), 
                                                     li_est = quantile(j, 0.025), 
                                                     ls_est = quantile(j, 0.975))
                                            }, simplify = 'list'))
                      
                      df_pred <- 
                        sapply(x_seq, FUN = 
                                 function(z) {
                                   
                                   mu <- post_m1$spp[, i, drop = T] + 
                                     apply(post_m1$I, 1, mean) +
                                     apply(post_m1$Y, 1, mean) +
                                     post_m1$beta[, i, drop = T]*z
                                   
                                   rnorm(1e3, mu, post_m1$sigma$sigma)
                                   
                                 })
                      
                      df_pred <- do.call('rbind',
                                    apply(df_pred, 2, FUN = 
                                            function(j) {
                                              tibble(li_pred = quantile(j, 0.025), 
                                                     ls_pred = quantile(j, 0.975))
                                            }, simplify = 'list'))
                      
                      df <- cbind(df, df_pred)
                      df$x <- x_seq
                      df$spp <- gsub('^(.*)(_)(.*)$', '\\1', nombres[i])
                      df$sex <- gsub('^(.*)(_)(.*)$', '\\3', nombres[i])
                      df

                    })

est_cor <- do.call('rbind', est_cor)

est_cor <- as_tibble(est_cor)
colnames(est_cor)[7] <- 'species'

ggplot() +
  geom_point(data = d, aes(body_mass_g, flipper_length_mm, color = sex)) +
  geom_line(data = est_cor, 
            aes(x, mu_est, color = sex)) +
  geom_ribbon(data = est_cor, 
              aes(x, mu_est, fill = sex, 
                  ymin = li_pred, ymax = ls_pred), alpha = 0.3) +
  facet_wrap(~ species) +
  theme_bw() +
  labs(x = 'Masa corporal (g)', y = 'Longitud de la aleta (mm)') +
  theme(
    panel.grid = element_blank(),
    legend.position = 'top',
    legend.title = element_blank()
  )

Marginalicemos el efecto del tamaño corporal y veamos la variación de la longitud de la aleta entre especie, sexo, año e isla.

d1_ <- unique(d[, c("species", "year", "island", "sex")])

d1 <- apply(d1_, 2, FUN = 
              function(x) {
                
                x <- factor(x, labels = 1:length(unique(x)))
                as.integer(x)
                
              })

d1 <- as_tibble(d1)

sigmas <- m1$draws(format = 'df')
sigmas <- sigmas[, grepl('sigma', colnames(sigmas))]


d$spp_sex <- dat$species_sex

est_islas_year <- lapply(1:nrow(d1), FUN = 
                           function(x) {
                             
                             se <- d1$sex[[x]]
                             s_b <- ifelse(se == 1, d1$species[[x]] + 0, 
                                           d1$species[[x]] + 3)
                             is <- d1$island[[x]]
                             y <- d1$year[[x]]
                             
                             p <- with(post_m1, 
                                       {
                                         spp[, s_b, drop = T] + 
                                           I[, is, drop = T] +
                                           Y[, y, drop = T] +
                                           beta[, s_b, drop = T] *
                                           mean(d[d$spp_sex == s_b, ]$body_mass_g)
                                         
                                       })
                             
                             p <- rnorm(1e3, p, sigmas$sigma)
                             
                             cbind(tibble(val = p), d1_[x, ])
                             
                           })


est_islas_year <- as_tibble(do.call('rbind', est_islas_year))

est_islas_year |> 
  ggplot() +
  geom_boxplot(aes(year, val, fill = species)) +
  facet_wrap( ~ island + sex) +
  labs(x = NULL, y = 'Longitud aleta (mm)') +
  theme_bw() +
  theme(legend.position = 'top')

Grafiquemos las distribuciones posteriores de los \(\sigma\) para cada especie/sexo entre años e islas para mostrar otra característica útil e interesante de los modelos jerárquicos:

sigmas_islas <- as.matrix(sigmas[, grep('sigma_is', colnames(sigmas))])
sigmas_year <- as.matrix(sigmas[, grep('sigma_y', colnames(sigmas))])

nombres <- unique(d[, c("species", "spp_sex", "sex")])

nombres$labels <- paste(nombres$species, nombres$sex, sep = "_")

nombres <- split(nombres, nombres$species)

par(mfrow = c(3, 1), mar = c(4, 3.5, 1, 0.5))

for (i in 1:3) {
  
  n1 <- nombres[[i]]$spp_sex[[1]]
  n2 <- nombres[[i]]$spp_sex[[2]]
  nom1 <- nombres[[i]]$sex[[1]]
  nom2 <- nombres[[i]]$sex[[2]]
  sp <- nombres[[i]]$species[[1]]
  
  plot(density(sigmas_islas[, n1]), col = 2, lwd = 3, 
       xlim = c(-0.5, 8), ylim = c(0, 1.3), main = sp, lty = 1,
       xlab = expression(sigma))
  lines(density(sigmas_year[, n1]), col = 2, lwd = 3, lty = 3)
  lines(density(sigmas_islas[, n2]), col = 4, lwd = 3, lty = 1)
  lines(density(sigmas_year[, n2]), col = 4, lwd = 3, lty = 3)
  legend(x = 4, y = 1.1, 
         legend = c(paste(nom1, '- isla'), 
                    paste(nom1, '- año'),
                    paste(nom2, '- isla'), 
                    paste(nom2, '- año')),
         col = c(2, 2, 4, 4), lty = rep(c(1, 3), 2), cex = 0.8, lwd = 3)
}

par(mfrow = c(1, 1))

Los modelos jerárquicos son ideales para modelar variación entre clusters. En nuestro caso, podemos notar que la variación de la longitud de la aleta se comporta distinto entre especies/sexo, años e islas. Éste, en mi opinión, y a pesar de no ser el objetivo principal, es el resultado más interesante del modelo. Lo es porque pone de manifiesto que si bien cada estudio puede buscar la estimación de un parámetro particular –en este caso \(\beta\), una completa exploración del modelo puede resultar en nuevas y más profundas preguntas acerca del sistema.

Para finalizar un inevitablemente extenso capítulo, grafiquemos las distribuciones posteriores de los coeficientes \(\rho_{año-species/sex}\) y \(\rho_{isla-species/sex}\) y verifiquemos si, en efecto, los parámetros co-varían.

rho_isla <- m1$draws(variables = 'Rho_isla_corr', format = 'df')
rho_year <- m1$draws(variables = 'Rho_year_corr', format = 'df')

rho_isla <- lapply(2:6, FUN = 
                     function(i) {
                       
                       df <- sapply(i:6, FUN = 
                                      function(j) {
                                        
                                        rho_isla[, grep(paste('^(.*)', i-1, ',', j,'.$', sep = ''), 
                                                        colnames(rho_isla)), drop = T]
                                        
                                      })
                       name <- i:6
                       df <- as_tibble(df)
                       
                       for (z in seq_along(name)) colnames(df)[z] <- 
                         paste('rho', i-1, name[z], sep = '_')
                       df
                     })
rho_isla <- as_tibble(do.call('cbind', rho_isla))

rho_year <- lapply(2:6, FUN = 
                     function(i) {
                       
                       df <- sapply(i:6, FUN = 
                                      function(j) {
                                        
                                        rho_year[, grep(paste('^(.*)', i-1, ',', j,'.$', sep = ''), 
                                                        colnames(rho_year)), drop = T]
                                        
                                      })
                       name <- i:6
                       df <- as_tibble(df)
                       
                       for (z in seq_along(name)) colnames(df)[z] <- 
                         paste('rho', i-1, name[z], sep = '_')
                       df
                     })
rho_year <- as_tibble(do.call('cbind', rho_year))
  
rho_isla
# A tibble: 4,002 × 15
    rho_1_2  rho_1_3 rho_1_4 rho_1_5  rho_1_6  rho_2_3 rho_2_4 rho_2_5 rho_2_6
      <dbl>    <dbl>   <dbl>   <dbl>    <dbl>    <dbl>   <dbl>   <dbl>   <dbl>
 1  0.00925 -0.457    0.0569  0.0496  0.782    0.458   -0.324  -0.270   0.366 
 2  0.275    0.289   -0.113  -0.184  -0.761   -0.392    0.0339  0.459  -0.559 
 3 -0.217    0.0122  -0.305  -0.246   0.161    0.159    0.0835  0.0822 -0.611 
 4 -0.0366   0.219    0.0866  0.247  -0.180    0.00772 -0.225   0.116   0.351 
 5 -0.174   -0.458   -0.279   0.186   0.273    0.0427   0.0416 -0.353  -0.532 
 6  0.0602  -0.0239  -0.647   0.0822  0.212   -0.415   -0.287   0.176  -0.217 
 7  0.240   -0.238   -0.329  -0.535   0.00311 -0.181    0.635   0.209  -0.703 
 8  0.468   -0.00103 -0.278   0.714   0.430   -0.456   -0.745   0.499  -0.0609
 9 -0.168   -0.384    0.549   0.115  -0.272    0.201    0.427   0.126  -0.234 
10  0.289    0.113    0.0286  0.170   0.307   -0.410    0.589  -0.155   0.229 
# … with 3,992 more rows, and 6 more variables: rho_3_4 <dbl>, rho_3_5 <dbl>,
#   rho_3_6 <dbl>, rho_4_5 <dbl>, rho_4_6 <dbl>, rho_5_6 <dbl>
par(mfrow = c(1, 2), mar = c(4, 4, 1, 1))
plot(NULL, xlim = c(-1, 1), ylim = c(0, 1.5), xlab = expression(rho['isla']), ylab = 'Density')
for (i in 1:ncol(rho_isla)) {
  lines(density(rho_isla[, i, drop = T]), col = i)
  abline(v = mean(rho_isla[, i, drop = T]), lty = 3, col = i)
}

plot(NULL, xlim = c(-1, 1), ylim = c(0, 1.5), xlab = expression(rho['year']), ylab = 'Density')
for (i in 1:ncol(rho_year)) {
  lines(density(rho_year[, i, drop = T]), col = i)
  abline(v = mean(rho_year[, i, drop = T]), lty = 3, col = i)
}

par(mfrow = c(1, 1))

Las distribuciones posteriores están más o menos centradas en \(\rho = 0\), así que no tenemos evidencia que indique una co-variación de los parámetros –tal vez un poco para algunos niveles de \(\rho_{isla-species/sex}\), pero en cualquier caso nada contundente.

Iris (derecha) y Ziggy (izquierda) viéndo, complacidos, cómo llegaste hasta el final.
sessionInfo()
R version 4.2.1 (2022-06-23)
Platform: x86_64-apple-darwin17.0 (64-bit)
Running under: macOS Big Sur ... 10.16

Matrix products: default
BLAS:   /Library/Frameworks/R.framework/Versions/4.2/Resources/lib/libRblas.0.dylib
LAPACK: /Library/Frameworks/R.framework/Versions/4.2/Resources/lib/libRlapack.dylib

locale:
[1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8

attached base packages:
[1] parallel  stats     graphics  grDevices utils     datasets  methods  
[8] base     

other attached packages:
 [1] ellipse_0.4.3     MASS_7.3-57       loo_2.5.1         cmdstanr_0.5.3   
 [5] patchwork_1.1.2   magrittr_2.0.3    lubridate_1.9.2   forcats_1.0.0    
 [9] stringr_1.5.0     dplyr_1.1.1       purrr_1.0.2       readr_2.1.4      
[13] tidyr_1.3.0       tibble_3.2.1      ggplot2_3.4.1     tidyverse_2.0.0  
[17] rethinking_2.13.2

loaded via a namespace (and not attached):
 [1] mvtnorm_1.1-3        lattice_0.20-45      palmerpenguins_0.1.1
 [4] ps_1.7.2             digest_0.6.33        utf8_1.2.3          
 [7] R6_2.5.1             backports_1.4.1      evaluate_0.17       
[10] coda_0.19-4          pillar_1.8.1         rlang_1.1.2         
[13] data.table_1.14.2    rstudioapi_0.15.0    checkmate_2.1.0     
[16] rmarkdown_2.17       labeling_0.4.2       htmlwidgets_1.5.4   
[19] munsell_0.5.0        compiler_4.2.1       xfun_0.33           
[22] pkgconfig_2.0.3      shape_1.4.6          htmltools_0.5.3     
[25] tidyselect_1.2.0     tensorA_0.36.2       matrixStats_0.63.0  
[28] fansi_1.0.4          tzdb_0.3.0           withr_2.5.0         
[31] grid_4.2.1           distributional_0.3.1 jsonlite_1.8.7      
[34] gtable_0.3.1         lifecycle_1.0.3      posterior_1.3.1     
[37] scales_1.2.1         cli_3.6.2            stringi_1.7.8       
[40] farver_2.1.1         ellipsis_0.3.2       generics_0.1.3      
[43] vctrs_0.6.5          tools_4.2.1          glue_1.6.2          
[46] hms_1.1.2            processx_3.8.0       abind_1.4-5         
[49] fastmap_1.1.0        yaml_2.3.7           timechange_0.2.0    
[52] colorspace_2.1-0     knitr_1.40