def split_rhat(x):
m, n, d = x.shape
half = n // 2
split = np.concatenate([x[:, :half], x[:, half:2 * half]], axis=0)
within = split.var(axis=1, ddof=1).mean(axis=0)
between = half * split.mean(axis=1).var(axis=0, ddof=1)
variance = (half - 1) * within / half + between / half
return np.sqrt(variance / within)
def autocorrelation(x):
x = x - x.mean()
n = len(x)
spectrum = np.fft.rfft(x, n=2 * n)
acov = np.fft.irfft(spectrum * np.conj(spectrum))[:n]
return acov / acov[0]
def teaching_ess(x):
m, n, d = x.shape
result = []
for j in range(d):
rho = np.mean([autocorrelation(x[c, :, j]) for c in range(m)], axis=0)
positive_pairs = 0.0
for lag in range(1, n - 1, 2):
pair = rho[lag] + rho[lag + 1]
if pair < 0:
break
positive_pairs += pair
result.append(m * n / (1 + 2 * positive_pairs))
return np.asarray(result)
d2l.use_svg_display()
flat_chains = chains.reshape(-1, 2)
rhat = split_rhat(chains)
ess = teaching_ess(chains)
mcmc_mean = flat_chains.mean(axis=0)
mcmc_sd = flat_chains.std(axis=0, ddof=1)
print('split R-hat:', rhat.round(4))
print('ESS :', ess.round(0).astype(int))
print('MCMC mean :', mcmc_mean.round(4),
' grid mean:', posterior_mean.round(4))
print('MCSE :', (mcmc_sd / np.sqrt(ess)).round(4))
fig, axes = d2l.plt.subplots(1, 2, figsize=(10, 3.2))
for c in range(len(chains)):
axes[0].plot(chains[c, :600, 1], linewidth=0.8, alpha=0.8)
axes[0].set(xlabel='post-warmup iteration', ylabel=r'slope $\theta_1$',
title='four chain traces')
for c in range(len(chains)):
axes[1].plot(autocorrelation(chains[c, :, 1])[:80], alpha=0.8)
axes[1].axhline(0, color='black', linewidth=0.7)
axes[1].set(xlabel='lag', ylabel='autocorrelation', title='serial dependence')
d2l.plt.tight_layout()
d2l.plt.show()