eight schools Turing / Stan

es.jl

using Turing, FlexiChains
import Mooncake, Enzyme

#=
(test) pkg> st
Status `~/ppl/test/Project.toml`
  [7da242da] Enzyme v0.13.129
  [4a37a8b9] FlexiChains v0.3.4
  [da2b9cff] Mooncake v0.5.17
  [fce5fe82] Turing v0.43.0
=#

ad_mc = AutoMooncake()
ad_en = AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse))
spl(ad) = NUTS(10000, 0.65; adtype=ad)
kwargs = (; num_warmup=10000, thinning=10, chain_type=VNChain)

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

@model function esc(J, y, sigma)
    mu ~ Normal(0, 5)
    tau ~ truncated(Cauchy(0, 5); lower=0)
    theta ~ MvNormal(fill(mu, J), tau^2 * I)
    for i in 1:J
        y[i] ~ Normal(theta[i], sigma[i])
    end
end

model_esc = esc(J, y, sigma)
@info "Centred Enzyme"   # 0.36 s
@time chain = sample(model_esc, spl(ad_en), 1000; kwargs...)
@info "Centred Mooncake" # 0.49 s
@time chain = sample(model_esc, spl(ad_mc), 1000; kwargs...)

@model function esnc(J, y, sigma)
    mu ~ Normal(0, 5)
    tau ~ truncated(Cauchy(0, 5); lower=0)
    theta_trans ~ MvNormal(zeros(J), I)
    theta := theta_trans .* tau .+ mu
    for i in 1:J
        y[i] ~ Normal(theta[i], sigma[i])
    end
end
model_esnc = esnc(J, y, sigma)
@info "Non-centred Enzyme"   # 0.29 s
@time chain = sample(model_esnc, spl(ad_en), 1000; kwargs...)
@info "Non-centred Mooncake" # 0.41 s
@time chain = sample(model_esnc, spl(ad_mc), 1000; kwargs...)

esc.stan

data {
  int<lower=0> J; // number of schools
  array[J] real y; // estimated treatment
  array[J] real<lower=0> sigma; // std of estimated effect
}
parameters {
  array[J] real theta; // treatment effect in school j
  real mu; // hyper-parameter of mean
  real<lower=0> tau; // hyper-parameter of sdv
}
model {
  tau ~ cauchy(0, 5); // a non-informative prior
  theta ~ normal(mu, tau);
  y ~ normal(theta, sigma);
  mu ~ normal(0, 5);
}

esnc.stan

data {
  int<lower=0> J; // number of schools
  array[J] real y; // estimated treatment
  array[J] real<lower=0> sigma; // std of estimated effect
}
parameters {
  vector[J] theta_trans; // transformation of theta
  real mu; // hyper-parameter of mean
  real<lower=0> tau; // hyper-parameter of sd
}
transformed parameters {
  vector[J] theta;
  // original theta
  theta = theta_trans * tau + mu;
}
model {
  theta_trans ~ normal(0, 1);
  y ~ normal(theta, sigma);
  mu ~ normal(0, 5); // a non-informative prior
  tau ~ cauchy(0, 5);
}

stan_driver.py

from cmdstanpy import CmdStanModel, install_cmdstan
from pathlib import Path
import time
from sys import argv

"""
Stan 2.38.0

$ uv pip list
Using Python 3.14.2 environment at: venv
Package         Version
--------------- -----------
cmdstanpy       1.3.0
numpy           2.4.3
pandas          3.0.1
pip             25.3
python-dateutil 2.9.0.post0
six             1.17.0
stanio          0.5.1
tqdm            4.67.3
"""

install_cmdstan()

DATA = {
    "y": [28, 8, -3, 7, -1, 1, 18, 12],
    "sigma": [15, 10, 16, 11, 9, 11, 10, 18],
    "J": 8,
}

def main():
    if len(argv) != 2:
        print("Usage: python a.py <stan_file>")
        return
    stan_file = Path(__file__).parent / argv[1]
    model = CmdStanModel(stan_file=stan_file)
    x = time.time()
    fit = model.sample(data=DATA, chains=1,
                       iter_warmup=10000, save_warmup=False,
                       iter_sampling=10000, thin=10)
    y = time.time()
    print(fit.summary())
    print(f"Time taken: {y - x} seconds")


if __name__ == "__main__":
    main()