15 - Regresión lineal con {RStan} (Clima en Australia)

En este artículo se muestra como implementar en R y {RStan} algunos de los modelos necesarios en el ejercicio Clima en Australia de la Práctica 4.

Comenzamos cargando las liberías que vamos a usar.

library(dplyr)
library(ggplot2)
library(readr)
library(rstan)

c_orange_dark <- "#e76f51"
c_orange_light <- "#f4a261"
c_yellow <- "#e9c46a"
c_green_blue <-"#2a9d8f"
c_blue_dark <- "#264653"
c_lightblue <- "#219ebc"
df <- read_csv(
    "https://raw.githubusercontent.com/estadisticaunr/estadistica-bayesiana/main/datos/weather_WU.csv"
)

# Explorar los datos
ggplot(df) +
  geom_point(aes(x = temp9am, y = temp3pm, color = location), alpha = 0.7, size = 2) +
  scale_color_manual(values = c(c_orange_dark, c_green_blue)) +
  labs(
    x = "Temperatura a las 9 a.m. (C°)", 
    y = "Temperatura a las 3 p.m. (C°)",
    color = "Ciudad"
  ) +
  theme_bw()

Modelos

A continuación se presentan los modelos que se consideran en el ejercicio. En todos los casos se tiene:

  • \(Y_i\): temperatura a las 3 p.m. de la observación i-ésima.
  • \(X_i\): temperatura a las 9 a.m. de la observacion i-ésima.

Modelo 1

\[ \begin{aligned} Y_i \mid \mu_i, \sigma &\underset{iid}{\sim} \text{Normal}(\mu_i, \sigma) \\ \mu_i &= \beta_0 + \beta_1 x_i \\ \beta_0 &\sim \text{Normal}(15, 8^2)\\ \beta_1 &\sim \text{Normal}(0, 3^2)\\ \sigma &\sim \text{Normal}^+(0, 12^2)\\ \end{aligned} \]

Modelo 2

\[ \begin{aligned} Y_i \mid \mu_i, \sigma &\underset{iid}{\sim}\text{Normal}(\mu_i, \sigma) \\ \mu_i &= \beta_{0, j[i]} \\ \beta_{0, j} &\sim \text{Normal}(20, 10^2) \quad \forall j = 1, 2 \\ \sigma &\sim \text{Normal}^+(0, 15^2) \\ \end{aligned} \]

donde \(j\) indexa a las diferentes ciudades.

Para pensar: ¿qué indican \(\beta_{0, 1}\) y \(\beta_{0, 2}\)?

Modelo 3

\[ \begin{aligned} Y_i \mid \mu_i, \sigma &\underset{iid}{\sim} \text{Normal}(\mu_i, \sigma) \\ \mu_i &= \beta_{0, j[i]} + \beta_1 x_i \\ \beta_{0, j} &\sim \text{Normal}(15, 8^2) \quad \forall j = 1, 2 \\ \beta_1 &\sim \text{Normal}(0, 3^2) \\ \sigma &\sim \text{Normal}^+(0, 12^2) \\ \end{aligned} \]

Para pensar: ¿en qué difieren los \(\beta_{0, 1}\) y \(\beta_{0, 2}\) de este modelo con los del modelo 2? ¿Y en qué difiere \(\beta_1\) de este modelo con el del modelo 1?

Implementación

Modelo 1

data {
  int<lower=0> N;     // Cantidad de observaciones
  vector[N] temp9am;  // Temperatura a las 9 am (predictor)
  vector[N] temp3pm;  // Temperatura a las 3 pm (respuesta)
}
parameters {
  real beta0;           // Intercepto
  real beta1;           // Pendiente
  real<lower=0> sigma;  // Desvío estándar del error
}
model {
  beta0 ~ normal(15, 8);
  beta1 ~ normal(0, 3);
  sigma ~ normal(0, 12);
  temp3pm ~ normal(beta0 + beta1 * temp9am, sigma);
}
generated quantities {
  vector[N] mu;
  vector[N] y_rep;
  vector[N] log_likelihood;
  
  // Calcular 'mu'
  mu = beta0 + beta1 * temp9am;

  for (i in 1:N) {
    // Obtención de muestras de la distribución predictiva a posteriori
    y_rep[i] = normal_rng(mu[i], sigma);

    // Cálculo de la log-verosimilitud
    log_likelihood[i] = normal_lpdf(temp3pm[i] | mu[i], sigma);
  }
}
# Datos que pasamos a Stan
stan_data_1 <- list(
    N = nrow(df), 
    temp9am = df$temp9am, 
    temp3pm = df$temp3pm
)

ruta_modelo_1 <- here::here(
  "recursos", "codigo", "stan", "australia", "modelo_1.stan"
)

# Muestreamos del posterior con Stan
fit_1 <- stan( 
    file = ruta_modelo_1,
    data = stan_data_1,
    model_name = "modelo 1",
    chains = 2,
    refresh = 0,
    seed = 121195
)

# Convertimos las muestras a data.frame, conservando solo los parámetros de interés
df_draws_1 <- as.data.frame(extract(fit_1, c("beta0", "beta1", "sigma")))    

Graficamos la distribución a posteriori marginal de los parámetros.

df_draws_long_1 <- df_draws_1 |>
  tidyr::pivot_longer(c("beta0", "beta1", "sigma"), names_to = "parametro")

ggplot(df_draws_long_1) +
  geom_histogram(
    aes(x = value, y = after_stat(density)), 
    bins = 30, 
    fill = c_orange_light, 
    color = c_orange_light
) +
  facet_wrap(~ parametro, scales = "free") +
  labs(x = "Valor", y = "Densidad") +
  theme_bw()

La recta de regresión está dada por un intercepto y una pendiente. Al contar con múltiples valores de interceptos y pendientes, es posible pensar que se tienen múltiples muestras de rectas de regresión. Para graficar algunas de estas rectas simplemente seleccionamos algunos de los pares de (beta0, beta1) que se obtuvieron del posterior.

ggplot(df) +
  geom_point(aes(x = temp9am, y = temp3pm, color = location), alpha = 0.7, size = 2) +
  geom_abline(
    aes(intercept = beta0, slope = beta1),
    alpha = 0.3,
    color = "gray30",
    data = df_draws_1[sample(4000, 40), ]
  ) +
  scale_color_manual(values = c(c_orange_dark, c_green_blue)) +
  labs(
      x = "Temperatura a las 9 a.m. (C°)", 
      y = "Temperatura a las 3 p.m. (C°)",
      color = "Ciudad"
  ) +
  theme_bw()

Modelo 2

data {
  int<lower=0> N;       // Cantidad de observaciones
  int location_idx[N];  // Índice de la ciudad
  vector[N] temp3pm;    // Temperatura a las 3 pm (respuesta)
}
parameters {
  vector[2] beta0;      // Interceptos
  real<lower=0> sigma;  // Desvío estándar del error
}
model {
  beta0 ~ normal(20, 10);
  sigma ~ normal(0, 15);
  temp3pm ~ normal(beta0[location_idx], sigma);
}
generated quantities {
  vector[N] mu;
  vector[N] y_rep;
  vector[N] log_likelihood;
  
  // Calcular 'mu'
  mu = beta0[location_idx];

  for (i in 1:N) {
    // Obtención de muestras de la distribución predictiva a posteriori
    y_rep[i] = normal_rng(mu[i], sigma);

    // Cálculo de la log-verosimilitud
    log_likelihood[i] = normal_lpdf(temp3pm[i] | mu[i], sigma);
  }
}

Este modelo incorpora la ubicación. Una forma de implementarlo es utilizando índices que indiquen a que ubicación pertence cada observación.

as.factor(df$location)
  [1] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
  [7] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [13] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [19] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [25] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [31] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [37] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [43] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [49] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [55] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [61] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [67] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [73] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [79] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [85] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [91] Uluru      Uluru      Uluru      Uluru      Uluru      Uluru     
 [97] Uluru      Uluru      Uluru      Uluru      Wollongong Wollongong
[103] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[109] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[115] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[121] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[127] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[133] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[139] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[145] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[151] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[157] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[163] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[169] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[175] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[181] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[187] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[193] Wollongong Wollongong Wollongong Wollongong Wollongong Wollongong
[199] Wollongong Wollongong
Levels: Uluru Wollongong
as.integer(as.factor(df$location))
  [1] 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
 [38] 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
 [75] 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 2
[112] 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
[149] 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
[186] 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2
stan_data_2 <- list(
    N = nrow(df), 
    location_idx = as.integer(as.factor(df$location)),
    temp3pm = df$temp3pm
)

ruta_modelo_2 <- here::here(
  "recursos", "codigo", "stan", "australia", "modelo_2.stan"
)

fit_2 <- stan( 
    file = ruta_modelo_2,
    data = stan_data_2,
    model_name = "modelo 2",
    chains = 2,
    refresh = 0,
    seed = 121195
)

# Extraemos las muestras del posterior y renombramos la columnas de beta0
df_draws_2 <- as.data.frame(extract(fit_2, c("beta0", "sigma")))
colnames(df_draws_2) <- c("beta0[Uluru]", "beta0[Wollongong]", "sigma")

El código para visualizar los posteriors marginales es análogo.

df_draws_long_2 <- df_draws_2 |>
  tidyr::pivot_longer(
    c("beta0[Uluru]", "beta0[Wollongong]", "sigma"), 
    names_to = "parametro"
)

ggplot(df_draws_long_2) +
  geom_histogram(
    aes(x = value, y = after_stat(density)), 
    bins = 30, 
    fill = c_orange_light, 
    color = c_orange_light
) +
  facet_wrap(~ parametro, scales = "free") +
  labs(x = "Valor", y = "Densidad") +
  theme_bw()

En este segundo modelo se tienen dos interceptos, uno para cada ciudad, y una pendiente nula.

data_lines <- data.frame(
    beta0 = c(df_draws_2[["beta0[Uluru]"]], df_draws_2[["beta0[Wollongong]"]]),
    location = rep(c("Uluru", "Wollongong"), each = nrow(df_draws_2))
)

ggplot(df) +
  geom_point(aes(x = temp9am, y = temp3pm, color = location), alpha = 0.7, size = 2) +
  geom_abline(
    aes(intercept = beta0, slope = 0, color = location),
    alpha = 0.3,
    data = data_lines[sample(nrow(data_lines), 80), ]
  ) +
  scale_color_manual(values = c(c_orange_dark, c_green_blue)) +
  labs(
    x = "Temperatura a las 9 a.m. (C°)", 
    y = "Temperatura a las 3 p.m. (C°)",
    color = "Ciudad"
  ) +
  theme_bw()

El modelo predice una temperatura única para cada ciudad a las 3 p.m., independiente de la temperatura a las 9 a.m.

Modelo 3

data {
  int<lower=0> N;       // Cantidad de observaciones
  int location_idx[N];  // Índice de la ciudad
  vector[N] temp9am;    // Temperatura a las 9 am (predictor)
  vector[N] temp3pm;    // Temperatura a las 3 pm (respuesta)
}
parameters {
  vector[2] beta0;      // Interceptos
  real beta1;           // Pendiente
  real<lower=0> sigma;  // Desvío estándar del error
}
model {
  beta0 ~ normal(15, 8);
  beta1 ~ normal(0, 3);
  sigma ~ normal(0, 15);
  temp3pm ~ normal(beta0[location_idx] + beta1 * temp9am, sigma);
}
generated quantities {
  vector[N] mu;
  vector[N] y_rep;
  vector[N] log_likelihood;
  
  // Calcular 'mu'
  mu = beta0[location_idx] + beta1 * temp9am;

  for (i in 1:N) {
    // Obtención de muestras de la distribución predictiva a posteriori
    y_rep[i] = normal_rng(mu[i], sigma);

    // Cálculo de la log-verosimilitud
    log_likelihood[i] = normal_lpdf(temp3pm[i] | mu[i], sigma);
  }
}

Este último modelo presenta un intercepto distinto para cada ciudad y una pendiente común a ambas. Por lo tanto, es necesario pasarle a Stan el vector de índices para las ciudades y la temperatura a las 9 a.m.

stan_data_3 <- list(
    N = nrow(df), 
    location_idx = as.integer(as.factor(df$location)),
    temp9am = df$temp9am, 
    temp3pm = df$temp3pm
)

ruta_modelo_3 <- here::here(
  "recursos", "codigo", "stan", "australia", "modelo_3.stan"
)

fit_3 <- stan( 
    file = ruta_modelo_3,
    data = stan_data_3,
    model_name = "modelo 3",
    chains = 2,
    refresh = 0,
    seed = 121195
)

df_draws_3 <- as.data.frame(extract(fit_3, c("beta0", "beta1", "sigma")))
colnames(df_draws_3) <- c("beta0[Uluru]", "beta0[Wollongong]", "beta1", "sigma")
df_draws_long_3 <- df_draws_3 |>
  tidyr::pivot_longer(
    c("beta0[Uluru]", "beta0[Wollongong]", "beta1", "sigma"), 
    names_to = "parametro"
)

ggplot(df_draws_long_3) +
  geom_histogram(
    aes(x = value, y = after_stat(density)), 
    bins = 30, 
    fill = c_orange_light, 
    color = c_orange_light
) +
  facet_wrap(~ parametro, scales = "free") +
  labs(x = "Valor", y = "Densidad") +
  theme_bw()

data_lines <- data.frame(
    beta0 = c(df_draws_3[["beta0[Uluru]"]], df_draws_3[["beta0[Wollongong]"]]),
    beta1 = rep(df_draws_3[["beta1"]], 2),
    location = rep(c("Uluru", "Wollongong"), each = nrow(df_draws_3))
)

ggplot(df) +
  geom_abline(
    aes(intercept = beta0, slope = beta1, color = location),
    alpha = 0.3,
    data = data_lines[sample(nrow(data_lines), 80), ]
  ) +
  geom_point(aes(x = temp9am, y = temp3pm, color = location), alpha = 0.7, size = 2) +
  scale_color_manual(values = c(c_orange_dark, c_green_blue)) +
  labs(
    x = "Temperatura a las 9 a.m. (C°)", 
    y = "Temperatura a las 3 p.m. (C°)",
    color = "Ciudad"
  ) +
  theme_bw()

De manera cuantitativa, es posible ver que este modelo es el que mejor ajusta a la nube de puntos.

Comparación

Posterior predictive checks

La librería {bayesplot} nos permite realizar pruebas predictivas a posteriori mediante el uso de ppc_dens_overlay(). El gráfico que se obtiene permite comparar la distribución empírica de los datos (marginal) con distribuciones estimadas a partir de valores de la distribución predictiva a posteriori.

Para crear este gráfico es necesario haber obtenido muestras de la distribución predictiva a posteriori con Stan. En nuestro caso, las guardamos en y_rep.

library(bayesplot)

# Notar que solo se muestran 200 muestras para que el gráfico quede menos cargado.
y_rep_1 <- extract(fit_1, "y_rep")$y_rep
ppc_dens_overlay(df$temp3pm, y_rep_1[sample(2000, 200), ]) +
  labs(x = "Temperatura a las 3 p.m. (C°)", y = "Densidad") + 
  theme_bw()

y_rep_2 <- extract(fit_2, "y_rep")$y_rep
ppc_dens_overlay(df$temp3pm, y_rep_2[sample(2000, 200), ]) +
  labs(x = "Temperatura a las 3 p.m. (C°)", y = "Densidad") + 
  theme_bw()

y_rep_3 <- extract(fit_3, "y_rep")$y_rep
ppc_dens_overlay(df$temp3pm, y_rep_3[sample(2000, 200), ]) +
  labs(x = "Temperatura a las 3 p.m. (C°)", y = "Densidad") + 
  theme_bw()

Las distribuciones generadas con los modelos 1 y 2 no replican fielmente la distribución empírica de la variable respuesta, mientras que las distribuciones generadas con el modelo 3 proveen una mejor aproximación.

LOO

En esta sección utilizamos la función loo() de la librería {loo} para estimar el ELPD (Expected log pointwise predictive density for a new dataset) mediante una aproximación a la validación cruzada conocido como PSIS-LOO CV.

Lo único que necesitamos es el objeto que obtuvimos al muestrear del posterior e indicar el nombre de la variable donde guardamos el cómputo del log-likelihood punto a punto (log_likelihood en todos los casos).

library(loo)

loo_1 <- loo(fit_1, pars = "log_likelihood")
loo_1

Computed from 2000 by 200 log-likelihood matrix.

         Estimate   SE
elpd_loo   -568.5  8.5
p_loo         2.6  0.4
looic      1137.1 16.9
------
MCSE of elpd_loo is 0.1.
MCSE and ESS estimates assume MCMC draws (r_eff in [0.3, 1.2]).

All Pareto k estimates are good (k < 0.7).
See help('pareto-k-diagnostic') for details.
loo_2 <- loo(fit_2, pars = "log_likelihood")
loo_2

Computed from 2000 by 200 log-likelihood matrix.

         Estimate   SE
elpd_loo   -625.8  9.5
p_loo         3.0  0.3
looic      1251.6 19.0
------
MCSE of elpd_loo is 0.0.
MCSE and ESS estimates assume MCMC draws (r_eff in [0.7, 1.2]).

All Pareto k estimates are good (k < 0.7).
See help('pareto-k-diagnostic') for details.
loo_3 <- loo(fit_3, pars = "log_likelihood")
loo_3

Computed from 2000 by 200 log-likelihood matrix.

         Estimate   SE
elpd_loo   -461.2 22.9
p_loo         7.6  3.5
looic       922.4 45.8
------
MCSE of elpd_loo is NA.
MCSE and ESS estimates assume MCMC draws (r_eff in [0.4, 1.1]).

Pareto k diagnostic values:
                         Count Pct.    Min. ESS
(-Inf, 0.7]   (good)     199   99.5%   377     
   (0.7, 1]   (bad)        1    0.5%   <NA>    
   (1, Inf)   (very bad)   0    0.0%   <NA>    
See help('pareto-k-diagnostic') for details.

loo() devuelve un objeto con varias métricas asociadas al ELPD. Para entender mejor cada una de ellas, se puede echar un vistazo al glosario de loo. elpd_loo es la estimación de ELPD mediante el método PSIS-LOO CV, para la cual también se tiene una medida de la variabilidad con un error estándar.

{loo} también provee la función loo_compare() que toma objetos devueltos por loo() y los ordena de mejor a peor según este criterio. En este caso se obtiene una estimación de la diferencia de ELPDs y un error estándar de la misma.

loo_compare(loo_1, loo_2, loo_3)
       elpd_diff se_diff
model3    0.0       0.0 
model1 -107.3      19.2 
model2 -164.6      22.1 

Según este criterio, el mejor modelo es el modelo 3. La diferencia de ELPD entre el modelo 3 y el modelo 2 es -107.3 y el error estándar de la misma es 19,2. Si se construye un intervalo aproximado del 95% restando y sumando 2 desvíos estándar de la estimación puntual, el mismo no cubre el 0 y se puede concluir que la diferencia de ELPDs entre el modelo 3 y el 2 es significativa.