Estoy sorprendido por la capacidad y flexibilidad para expresar modelos estadísticos complejos con STAN. Como sabrán, STAN es un lenguage declarativo de programación probabilístico. Los modelos estadísticos son descritos con una sintaxis sencilla para especificar de una forma muy similar a lo que hacemos en el pizarrón o en papel. STAN está basado en principios bayesianos y utiliza varios métodos para muestrear distribuciones posteriores, entre los más conocidos están Hamiltonian Monte Carlo, Non U Turn y Metrópolis-Hastings (no voy a meterme a detalle en esto ahora, en internet hay montón de referencias al respecto).

En este post veremos como definir un modelo logistico jerárquico. Supongamos que para toda observación binaria \(y_i \in \{0,1\} \) existe un nivel associado \(j_i \in \{1,2,…,J\}\) y un vector de covariables (predictores) \(x_i \in \mathbb{R}^K\). Supondremos pues que las observaciones son realizaciones de un proceso bernoulli independendiente cuando se condiciona con un proceso normal jerarquizado. Esto es:

\([Y | S] \sim Bernoulli(S)\)

Ahora, \(S\) puede ser descrita de la siguiente forma:

\(S_{i,j} = \beta^t_j x_i + \epsilon\)

Y se lee de la siguiente forma. La variable S en la observación i del nivel j es igual a \(\beta^t_j x_{i} + \epsilon \). Es decir, el vector \(\beta\) son los coeficientes associados a la respuesta \(S_{i}\) en el nivel j.

En STAN esto es equivalente al siguiente código:

hierarquical_log_model = u"""
data {
  int<lower=0> N; // num obs. 
  int<lower=0> J; // number of levels
  int<lower=1> K;                // num covariates
  matrix[N, K] x;                // design matrix
  int<lower=1,upper=J> level[N];  // type of level (spec)
  int<lower=0,upper=1> y[N];              // observations, in this case is binary.
  
} 
parameters {
  real mu[K];
  vector[K] beta[J];
  real<lower=0,upper=100> sigma; // same sigma

}

model {
    for (k in 1:K) {
        mu[k] ~ normal(0,100);
        for (j in 1:J)
            beta[j,k] ~ normal(mu[k], sigma);
    }
    for (i in 1:N)
        y[i] ~ bernoulli_logit(x[i] * beta[level[i]]);
  
}"""

Para correr el modelo en Python usando PyStan hacer esto.

varying_intercept_model = pystan.StanModel(model_code=hierarquical_log_model)
multilevel_linear_data = {'N': X.shape[0],
                          'J': 2,
                          'K': X.shape[1],
                          'level': data.codes.values, # niveles
                          'x' : X.values,
                          'y': data.obs.values}
%time fit = varying_intercept_model.sampling(data=multilevel_linear_data, iter=1000,chains=1)
sample = fit.extract(permuted=True)

En este caso mis datos están definidos en un Pandas Dataframe llamado X, separado de las observaciones. Para eso use patsy. La muestra de la distribución posterior estará guardada en el objeto ‘sample’