gelman_schools.py

import numpy as np

from arviz import loo
from pymc3 import HalfCauchy, Model, Normal, sample

"""Original Stan model

data {
  int<lower=0> J; // number of schools
  real y[J]; // estimated treatment effects
  real<lower=0> sigma[J]; // s.e. of effect estimates
}
parameters {
  real mu;
  real<lower=0> tau;
  real eta[J];
}
transformed parameters {
  real theta[J];
  for (j in 1:J)
    theta[j] <- mu + tau * eta[j];
}
model {
  eta ~ normal(0, 1);
  y ~ normal(theta, sigma);
}
"""

J = 8
y = np.array([28, 8, -3, 7, -1, 1, 18, 12])
sigma = np.array([15, 10, 16, 11, 9, 11, 10, 18])

with Model() as schools:

    eta = Normal("eta", 0, 1, shape=J)
    mu = Normal("mu", 0, sigma=1e6)
    tau = HalfCauchy("tau", 25)

    theta = mu + tau * eta

    obs = Normal("obs", theta, sigma=sigma, observed=y)


def run(n=1000):
    if n == "short":
        n = 50
    with schools:
        tr = sample(n)
        loo(tr)


if __name__ == "__main__":
    run()