PyStan の実行例 – Eight Schools

!pip install pystan==2.19.1.1

import pystan
import arviz as az

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

model_code = """
data {
  int<lower=0> J;         // number of schools
  real y[J];              // estimated treatment effects
  real<lower=0> sigma[J]; // standard error of effect estimates
}
parameters {
  real mu;                // population treatment effect
  real<lower=0> tau;      // standard deviation in treatment effects
  vector[J] eta;          // unscaled deviation from mu by school
}
transformed parameters {
  vector[J] theta = mu + tau * eta;        // school treatment effects
}
model {
  target += normal_lpdf(eta | 0, 1);       // prior log-density
  target += normal_lpdf(y | theta, sigma); // log-likelihood
}
"""

%%time

sm = pystan.StanModel(model_code=model_code)

INFO:pystan:COMPILING THE C++ CODE FOR MODEL anon_model_67cb7d0f2cb7720776cbeb52007d2dbb NOW.

CPU times: user 1.86 s, sys: 132 ms, total: 1.99 s
Wall time: 1min 10s

%%time

mcmc = sm.sampling(data=data, iter=100000, warmup=500, chains=1, control=dict(adapt_delta=0.99, max_treedepth=10))

WARNING:pystan:2 of 99500 iterations ended with a divergence (0.00201 %).
WARNING:pystan:Try running with adapt_delta larger than 0.99 to remove the divergences.

CPU times: user 7.44 s, sys: 130 ms, total: 7.57 s
Wall time: 7.54 s

%%time

mcmc = sm.sampling(data=data, iter=100000, warmup=500, chains=4, control=dict(adapt_delta=0.99, max_treedepth=10))

WARNING:pystan:10 of 398000 iterations ended with a divergence (0.00251 %).
WARNING:pystan:Try running with adapt_delta larger than 0.99 to remove the divergences.

CPU times: user 7.45 s, sys: 434 ms, total: 7.88 s
Wall time: 38.3 s

idata = az.from_pystan(mcmc)

az.plot_trace(idata);

az.summary(idata)