import numpy as np, pymc as pm, pytensor.tensor as pt, arviz as az
from seeds_gibbs import seeds_data, seeds_gibbs
r, n, X = seeds_data(); I = len(r); k = X.shape[1]
with pm.Model() as m:
    beta = pm.Normal('beta', 0, 10, shape=k)
    sigma = pm.HalfNormal('sigma', 1.0); z = pm.Normal('z', 0, 1, shape=I); b = pm.Deterministic('b', sigma*z)
    p = pm.math.invlogit(pt.dot(X, beta) + b)
    pm.Binomial('r', n=n.astype(int), p=p, observed=r.astype(int))
    idata = pm.sample(2000, tune=2000, chains=4, target_accept=0.95, random_seed=1, progressbar=False)
po = idata.posterior; nm = ['intercept','seed type','root extract','interaction']; bb = po['beta'].values.reshape(-1,k)
for j in range(k): print('%-13s %7.3f [%.2f, %.2f]' % (nm[j], bb[:,j].mean(), np.percentile(bb[:,j],2.5), np.percentile(bb[:,j],97.5)))
print('sigma = %.3f [%.2f, %.2f]   max r_hat %.3f' % (float(po['sigma'].mean()), float(po['sigma'].quantile(.025)), float(po['sigma'].quantile(.975)), float(az.summary(idata,var_names=['beta','sigma'])['r_hat'].max())))
g = seeds_gibbs(r, n, X, R=30000, burn=8000, seed=1)
print('\nfrom-scratch vs PyMC: extract %.2f/%.2f  interaction %.2f/%.2f  sigma %.2f/%.2f'
      % (g['beta'].mean(0)[2], bb[:,2].mean(), g['beta'].mean(0)[3], bb[:,3].mean(), g['sigma'].mean(), float(po['sigma'].mean())))

Initializing NUTS using jitter+adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [beta, sigma, z]

Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 6 seconds.

intercept      -0.551 [-0.98, -0.14]
seed type       0.062 [-0.62, 0.72]
root extract    1.363 [0.79, 1.97]
interaction    -0.835 [-1.79, 0.09]
sigma = 0.346 [0.08, 0.67]   max r_hat 1.000

from-scratch vs PyMC: extract 1.34/1.36  interaction -0.80/-0.84  sigma 0.13/0.35

Seeds — PyMC cross-check (binomial GLMM)¶

Results¶