17 - Regresión lineal con `{RStan}`

El siguiente programa muestra el código necesario para con R los ejercicios Mi primer regresión bayesiana y Mejorando mi regresión bayesiana de la Práctica 4.

En primer lugar, se cargan librerías necesarias.

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

Luego se leen y visualizan los datos.

# Leer los datos desde el repositorio
df_sales <- read_csv(
    "https://raw.githubusercontent.com/estadisticaunr/estadistica-bayesiana/main/datos/sales.csv"
)

# Explorar los datos
ggplot(df_sales) +
  geom_point(aes(x = x, y = y), alpha = 0.6, size = 2) +
  labs(x = "Publicidad ($)", y = "Ventas ($)") +
  theme_bw()

Se crea el siguiente modelo de regresión lineal bayesiana: \[ \begin{aligned} \text{ventas}_i &\sim \text{Normal}(\mu_i, \sigma^2) \\ \mu_i &= \beta_0 + \beta_1 \text{publicidad}_i \\ \end{aligned} \]

donde los parámetros \(\beta_0\), \(\beta_1\) y \(\sigma\) siguen distribuciones a priori uniformes. El programa de Stan para implementar el modelo es el siguiente:

stan/regresion_lineal/01_modelo.stan

data {
  int<lower=0> N;  // Cantidad de observaciones
  vector[N] x;     // Valores de la variable predictora
  vector[N] y;     // Valores de la variable respuesta
}
parameters {
  real beta0;           // Intercepto
  real beta1;           // Pendiente
  real<lower=0> sigma;  // Desvio estándar del error
}
transformed parameters {
  vector[N] mu;
  mu = beta0 + beta1 * x;
}
model {
  y ~ normal(mu, sigma);
}
generated quantities {
  // Obtención de muestras de la distribución predictiva a posteriori
  vector[N] y_rep;
  for (i in 1:N) {
    y_rep[i] = normal_rng(mu[i], sigma);
  }
}

Ahora se crea una lista con los datos para el modelo y se obtienen muestras del posterior utilizando la función stan().

stan_data <- list(
  N = nrow(df_sales), # Cantidad de observaciones
  x = df_sales$x,     # Publicidad
  y = df_sales$y      # Ventas
)

ruta_modelo_1 <- here::here(
  "recursos", "codigo", "stan", "regresion_lineal", "01_modelo.stan"
)

# Pasamos la ubicación del archivo con el codigo del modelo en Stan
modelo_1 <- stan_model(file = ruta_modelo_1)

modelo_1_fit <- sampling(
  modelo_1,             # El modelo
  data = stan_data,     # Datos
  chains = 4,           # Cantidad de cadenas
  seed = 1211,          # Para que el resultado sea reproducible
  refresh = 0,          # Cada cuantos pasos mostrar mensajes del sampler (0: nunca)
)

Al imprimir el objeto que devuelve stan() se puede encontrar un resumen del posterior, incluyendo medidas de diagnóstico como el tamaño de muestra efectivo y \(\hat{R}\).

print(modelo_1_fit, pars = c("beta0", "beta1", "sigma"))

Inference for Stan model: anon_model.
4 chains, each with iter=2000; warmup=1000; thin=1; 
post-warmup draws per chain=1000, total post-warmup draws=4000.

      mean se_mean   sd 2.5%  25%  50%  75% 97.5% n_eff Rhat
beta0 5.45    0.01 0.49 4.52 5.12 5.46 5.77  6.40  1288    1
beta1 0.29    0.00 0.04 0.21 0.26 0.29 0.31  0.36  1307    1
sigma 1.21    0.00 0.12 0.99 1.12 1.19 1.28  1.48  1755    1

Samples were drawn using NUTS(diag_e) at Thu Jun 12 22:51:51 2025.
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).

Para continuar explorando el posterior y calcular cantidades de interés, conviene trabajar con las muestras en un data.frame. Para ello resulta fundamental la función extract() que extrae las muestras del posterior y las devuelve en una lista.

df_draws_1 <- as.data.frame(extract(modelo_1_fit, pars = c("beta0", "beta1", "sigma")))
head(df_draws_1)

     beta0     beta1    sigma
1 5.778072 0.2742644 1.383950
2 5.845044 0.2217190 1.476740
3 5.745932 0.2736212 1.319527
4 5.483745 0.2623577 1.335966
5 5.646768 0.2485603 1.180145
6 5.256127 0.2939936 1.198690

Lo siguiente es un enfoque posible, aunque bastante manual, para graficar los posteriors marginales de los parámetros del modelo.

df_draws_long_1 <- df_draws_1 |>
  select(beta0, beta1, sigma) |>
  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 = 40) +
  facet_wrap(~ parametro, scales = "free") +
  labs(x = "Valor", y = "Densidad") +
  theme_bw()

Un primer paso en el análisis de un modelo lineal simple es visualizar la recta de regresión estimada.

# Calcular la media a posteriori del intercepto y la pendiente
intercept_mean <- mean(df_draws_1$beta0)
slope_mean <- mean(df_draws_1$beta1)

# Utilizar estos dos valores para graficar la media de la recta de regresión
ggplot(df_sales) +
  geom_point(aes(x = x, y = y), alpha = 0.6, size = 2) +
  geom_abline(
    intercept = intercept_mean,
    slope = slope_mean,
    linewidth = 1,
    color = "red"
  ) +
  labs(x = "Publicidad ($)", y = "Ventas ($)") +
  theme_bw()

Pero no hay que olvidar que estamos trabajando con un modelo bayesiano. Por lo tanto, más que visualizar una recta basada en la media de las distribuciones marginales, es mejor visualizar las rectas asociadas a muestras del posterior. Esto también brinda una idea de la variabilidad en la estimación de la recta.

ggplot(df_sales) +
  geom_point(aes(x = x, y = y), alpha = 0.6, size = 2) +
  geom_abline(
    aes(intercept = beta0, slope = beta1),
    alpha = 0.3,
    color = "gray30",
1    data = df_draws_1[sample(nrow(df_draws_1), 40), ]
  ) +
  geom_abline(
    intercept = intercept_mean,
    slope = slope_mean,
    linewidth = 1,
    color = "red"
  ) +
  labs(x = "Publicidad ($)", y = "Ventas ($)") +
  theme_bw()

1: Se utiliza sample(nrow(df_draws_1), 40) para seleccionar 40 muestras del posterior al azar. La figura resultaría muy difícil de leer si se intenta visualizar las rectas asociadas a muchas más muestras.

También es posible utilizar todas las muestras del posterior para obtener bandas de credibilidad para la recta de regresión. A continuación, se obtiene la distribución condicional de \(\mu_i\) para valores de \(\text{ventas}_i\) en una grilla que cubre el rango de valores observados. A partir de esas muestras, se calculan intervalos de credibilidad.

1x_grid <- seq(4, 20, length.out = 100)
2mu_matrix <- matrix(nrow = 4000, ncol = 100)

for (i in seq_along(x_grid)) {
3    mu_matrix[, i] <- df_draws_1$beta0 + df_draws_1$beta1 * x_grid[i]
}

4mu_mean <- apply(mu_matrix, 2, mean)
5mu_qts <- t(apply(mu_matrix, 2, function(x) quantile(x, c(0.025, 0.975))))
6mu_qts2 <- t(apply(mu_matrix, 2, function(x) quantile(x, c(0.25, 0.75))))

# Finalmente, se lamacenan los valores calculados en un data frame
data_mu <- data.frame(
  x = x_grid,
  y = mu_mean,
  lower_95 = mu_qts[, 1],
  upper_95 = mu_qts[, 2],
  lower_50 = mu_qts2[, 1],
  upper_50 = mu_qts2[, 2]
)

head(data_mu)

1: Se crea la grilla de valores para \(\text{ventas}\).
2: Se crea una matriz con tantas filas como muestras del posterior y tantas columnas como cantidad de valores en la grilla. Acá almacenamos los valores de \(\mu_i\).
3: Para cada valor de la grilla, se obtiene la distribución condicional de \(\mu_i\).
4: Se calcula la media a posteriori de \(\mu_i\) para cada valor en la grilla.
5: Se calculan los percentiles de \(\mu_i\) que se corresponden con un intervalo de colas iguales del 95%.
6: Se calculan los percentiles de \(\mu_i\) que se corresponden con un intervalo de colas iguales del 50%.

         x        y lower_95 upper_95 lower_50 upper_50
1 4.000000 6.600767 5.950424 7.292322 6.371004 6.824642
2 4.161616 6.647243 6.008847 7.324778 6.421652 6.866754
3 4.323232 6.693718 6.065794 7.361577 6.472019 6.908684
4 4.484848 6.740194 6.121453 7.395583 6.521862 6.951351
5 4.646465 6.786669 6.174731 7.432017 6.572691 6.993651
6 4.808081 6.833145 6.230303 7.468890 6.621086 7.036685

Las bandas de credibilidad se construyen usando geom_ribbon de {ggplot2}.

ggplot(df_sales) +
  geom_ribbon(
    aes(x, ymin = lower_95, ymax = upper_95),
    fill = "grey50",
    alpha = 0.6,
    data = data_mu
  ) +
  geom_ribbon(
    aes(x, ymin = lower_50, ymax = upper_50),
    fill = "grey35",
    alpha = 0.6,
    data = data_mu
  ) +
  geom_point(aes(x = x, y = y), alpha = 0.6, size = 2) +
  geom_line(
    aes(x, y),
    color = "firebrick",
    data = data_mu
  ) +
  labs(x = "Publicidad ($)", y = "Ventas ($)") +
  theme_bw()

También se puede visualizar la distribución condicional de \(\mu_i\) para un valor particular de la variable predictora \(\text{ventas}\).

publicidad <- 15
mu <- df_draws_1$beta0 + df_draws_1$beta1 * publicidad
data.frame(mu = mu) |>
  ggplot() +
  geom_histogram(aes(mu, y = after_stat(density)), bins = 40) +
  labs(
      x = expression(mu[i] ~ " | " ~ publicidad[i] ~ " = 15"),
      y = "Densidad"
  ) +
  theme_bw()

Finalmente, se utiliza la función ppc_intervals() de la librería {bayesplot} para visualizar intervalos de credibilidad de la distribución predictiva a posteriori para cada una de las observaciones. Para ello, primero se debe extraer y_rep del objeto modelo_1_fit.

y_pp <- extract(modelo_1_fit, pars = "y_rep")$y_rep
bayesplot::ppc_intervals(
  y = df_sales$y,
  yrep = y_pp,
  x = df_sales$x,
  prob = 0.5,
  prob_outer = 0.95
) +
  labs(x = "Publicidad ($)", y = "Ventas ($)") +
  theme_bw()

Ahora se muestra como utilizar distribuciones a priori no uniformes para los parámetros del modelo. El mismo se describe: \[ \begin{aligned} \text{ventas}_i &\sim \text{Normal}(\mu_i, \sigma^2) \\ \mu_i &= \beta_0 + \beta_1 \text{publicidad}_i \\ \beta_0 &\sim \text{Normal}(\overline{\text{ventas}}, 10^2) \\ \beta_1 &\sim \text{Normal}(0, 0.5^2) \\ \sigma &\sim \text{Normal}^+(5) \end{aligned} \]

y el programa de Stan que lo implementa es el siguiente

stan/regresion_lineal/02_modelo_con_priors.stan

data {
  int<lower=0> N;  // Cantidad de observaciones
  vector[N] x;     // Valores del predictor
  vector[N] y;     // Valores de la respuesta
  real beta0_mu;   // Media del prior del intercepto
}
parameters {
  real beta0;
  real beta1;
  real<lower=0> sigma;
}
transformed parameters {
  vector[N] mu;
  mu = beta0 + beta1 * x;
}
model {
  beta0 ~ normal(beta0_mu, 10);
  beta1 ~ normal(0, 0.5);
  sigma ~ normal(0, 5);
  y ~ normal(mu, sigma);
}
generated quantities {
  // Obtención de muestras de la distribución predictiva a posteriori
  vector[N] y_rep;
  for (i in 1:N) {
    y_rep[i] = normal_rng(mu[i], sigma);
  }
}

La estructura del código para ajustar el modelo es idéntica a la utilizada anteriormente. La única diferencia es que debemos pasar el valor de la media del prior del intercepto \(\beta_0\).

stan_data <- list(
  N = nrow(df_sales),          # Cantidad de observaciones
  x = df_sales$x,              # Publicidad
  y = df_sales$y,              # Ventas
  beta0_mu = mean(df_sales$y)  # Media del prior del intercepto
)

ruta_modelo_2 <- here::here(
    "recursos", "codigo", "stan", "regresion_lineal", "02_modelo_con_priors.stan"
)

modelo_2 <- stan_model(file = ruta_modelo_2)

modelo_2_fit <- sampling(
  modelo_2,
  chains = 4,         # Cantidad de cadenas
  data = stan_data,   # Datos
  refresh = 0,        # No mostrar mensajes del sampler
  seed = 121195       # Para que el resultado sea reproducible
)

Se puede ver el resumen del posterior:

print(modelo_2_fit, pars = c("beta0", "beta1", "sigma"))

Inference for Stan model: anon_model.
4 chains, each with iter=2000; warmup=1000; thin=1; 
post-warmup draws per chain=1000, total post-warmup draws=4000.

      mean se_mean   sd 2.5%  25%  50%  75% 97.5% n_eff Rhat
beta0 5.43    0.01 0.47 4.50 5.11 5.42 5.73  6.34  1389    1
beta1 0.29    0.00 0.04 0.21 0.26 0.29 0.32  0.36  1426    1
sigma 1.20    0.00 0.13 0.98 1.11 1.19 1.28  1.48  1807    1

Samples were drawn using NUTS(diag_e) at Thu Jun 12 22:52:48 2025.
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).

y también se podrían crear visualizaciones como las anteriores utilizando el mismo código.

Se recomienda echar un vistazo a las funciones disponibles en la librería {bayesplot} aquí. Por ejemplo, mcmc_trace(), que permite obtener un traceplot que sirve para evaluar la convergencia y mezcla de las cadenas.

bayesplot::mcmc_trace(
  modelo_2_fit,
  pars = c("beta0", "beta1", "sigma"),
  facet_args = list(nrow = 3)
)

También hay funciones que permiten calcular el \(\hat{R}\) y estimar el tamaño efectivo de muestra \(N_{\text{eff}}\) (como proporción de la cantidad de muestras).

bayesplot::rhat(modelo_2_fit, pars = c("beta0", "beta1", "sigma"))

   beta0    beta1    sigma 
1.002616 1.001983 1.001746

bayesplot::neff_ratio(modelo_2_fit, pars = c("beta0", "beta1", "sigma"))

    beta0     beta1     sigma 
0.3472752 0.3564939 0.4517207

Por último, también podemos visualizar intervalos de la predictiva a posteriori para cada observación:

y_pp <- extract(modelo_2_fit, pars = "y_rep")$y_rep
bayesplot::ppc_intervals(
  y = df_sales$y,
  yrep = y_pp,
  x = df_sales$x,
  prob = 0.5,
  prob_outer = 0.95
) +
  labs(x = "Publicidad ($)", y = "Ventas ($)") +
  theme_bw()