Modelos de conteos con sobredispersión (con Stan)

Esta entrada muestra cómo afrontar (con Stan) un problema que encontré el otro día en un lugar que no puedo mencionar pero en el que sé que me leen (y los destinatarios sabrán que va por ellos).

El contexto es el siguiente: se hace un test A/B donde la variable de interés son unos conteos. Hay varios grupos (aquí los reduciré a dos) y los datos siguen aproximadamente (aquí omitiré la parte de la inflación de ceros) una distribución de Poisson. Pero solo aproximadamente: existe sobredispersión, es decir, la varianza de los datos excede su media.

Así que la formulación más natural, en la que cada n_i \sim \text{Pois}(\lambda) no vale. Recurrir a la binomial negativa es una opción, la opción de libro viejo, pero destruye la interpretabilidad del modelo y obliga a renunciar a las propiedades más interesantes de la Poisson (p.e., la aditividad).

Una mejor aproximación es interpretar la sobredispersión de la Poisson como efecto de la heterogeneidad de los sujetos. Es decir, que cada sujeto tiene su propia \lambda y crear un modelo jerárquico: uno para las lambdas (p.e., \lambda_i \sim \Gamma(a, b) y otro para los conteos, n_i \sim \text{Pois}(\lambda_i).

Para lo cual, genero datos (obviamente, de acuerdo con mi modelo):

# sujetos por experimento
n <- 5000

# mejora en tratamiento
incr <- .05

# parámetros iniciales de la 
# distribución de heterogeneidad
a <- 3
b <- 4


prop.control <- rgamma(n, a, b)
obs.control <- sapply(prop.control, 
                      function(lambda) rpois(1, lambda))

prop.tratamiento <- rgamma(n, a, b) * (1 + incr)
obs.tratamiento <- sapply(prop.tratamiento, 
                          function(lambda) rpois(1, lambda))

Y ahora modelo en Stan:

data {
  int n;
  int control[n];
  int tratamiento[n];
}

parameters {
  real<lower = 0> a;
  real<lower = 0> b;
  real<lower = -0.01, upper = 0.2> incr;
  real<lower = 0> lambdas_control[n];
  real<lower = 0> lambdas_tratamiento[n];
}

model {
  lambdas_control ~ gamma(a, b);
  lambdas_tratamiento ~ gamma(a, b);
  
  for(i in 1:n){
    control[i] ~ poisson(lambdas_control[i]);
    tratamiento[i] ~ poisson(lambdas_tratamiento[i] * (1 + incr));
  }
}

Y ejecuto:

library(rstan)

fit <- stan("disp_poisson.stan", 
            data = list(n = n, control = obs.control, 
                   tratamiento = obs.tratamiento), 
            iter = 10000, warmup = 2000, 
            chains = 1, thin = 10)

res <- as.data.frame(fit)

El resultado es lo de menos. Lo de más, que tenemos disponible este procedimiento para modelar la sobre dispersión de los conteos, que es más natural que otros.

Nota: para una visión alternativa, esto.