"""Numerical checks of the derivations in Bishop (2006), Chapter 11 (Sampling Methods).

Each check prints PASS or FAIL with the numbers behind it; the last line counts them.
Run from this folder: python3 checks_ch11.py (under a minute). The samplers checked are the
ones in ch11_sampling.py, which also produce the figures. Monte Carlo checks use fixed seeds
and tolerances of about four standard errors, so they pass deterministically.
"""
import math

import numpy as np
from scipy import integrate, special, stats

import ch11_sampling as m

RESULTS = []


def check(name, ok, detail=""):
    RESULTS.append(bool(ok))
    print("%s  %s%s" % ("PASS" if ok else "FAIL", name, ("  [" + detail + "]") if detail else ""))


def within(est, truth, se, k=4.0):
    return abs(est - truth) <= k * se


# ----------------------------------------------------------------------------- 11.1.1
print("== Transformation method and Box-Muller (Section 11.1.1)")
rng = np.random.default_rng(1)
u = rng.random(200000)
y = m.exponential_from_uniform(u, 2.0)
ks = stats.kstest(y, "expon", args=(0, 0.5))
check("inverse CDF (11.7): exponential, rate 2: mean 1/2, variance 1/4, KS test",
      within(y.mean(), 0.5, 0.5 / math.sqrt(len(y))) and abs(y.var() - 0.25) < 0.01 and ks.pvalue > 0.01,
      "mean %.4f, var %.4f, KS p = %.3f" % (y.mean(), y.var(), ks.pvalue))
c = m.cauchy_from_uniform(u)
ks = stats.kstest(c, "cauchy")
q = np.quantile(c, [0.25, 0.5, 0.75])
check("inverse CDF (11.8), Exercise 11.3: y = tan(pi(u - 1/2)) is Cauchy (quartiles -1, 0, 1; KS test)",
      np.allclose(q, [-1, 0, 1], atol=0.02) and ks.pvalue > 0.01, "quartiles %s, KS p = %.3f" % (np.round(q, 3), ks.pvalue))
c2 = m.cauchy_from_uniform(u, c=9.0, b=math.sqrt(19))
ks = stats.kstest(c2, "cauchy", args=(9.0, math.sqrt(19)))
check("Exercise 11.7: y = b tan(pi(u - 1/2)) + c has the Cauchy density (11.16) with location c, scale b", ks.pvalue > 0.01,
      "KS p = %.3f" % ks.pvalue)
yb = m.inverse_cdf(rng.random(20000), m.bimodal_cdf, -12, 12, tol=1e-10)
ks = stats.kstest(yb, m.bimodal_cdf)
check("inverse CDF by bisection for the bimodal mixture (Fig 11.2): KS test", ks.pvalue > 0.01, "KS p = %.3f" % ks.pvalue)
g, tries = m.box_muller(np.random.default_rng(2), 200000)
y1, y2 = g[0::2], g[1::2]
ks = stats.kstest(g, "norm")
pairs = len(g) / 2
check("Box-Muller (11.10)-(11.12): mean 0, variance 1, pairs uncorrelated, KS test",
      abs(g.mean()) < 4 / math.sqrt(len(g)) and abs(g.var() - 1) < 0.015 and abs(np.corrcoef(y1, y2)[0, 1]) < 0.012 and ks.pvalue > 0.01,
      "mean %.4f, var %.4f, corr %.4f, KS p = %.3f" % (g.mean(), g.var(), np.corrcoef(y1, y2)[0, 1], ks.pvalue))
acc = pairs / tries
check("Box-Muller keeps a fraction pi/4 of the pairs in the square", within(acc, math.pi / 4, math.sqrt(acc * (1 - acc) / tries)),
      "%.4f versus %.4f" % (acc, math.pi / 4))
S = np.array([[2.0, 0.8, 0.1], [0.8, 1.0, -0.3], [0.1, -0.3, 0.5]])
Y = m.gaussian_cholesky(np.random.default_rng(3), np.array([1.0, -1.0, 0.5]), S, 400000)
check("Exercise 11.5: mu + L z has mean mu and covariance L L^T = Sigma", np.allclose(Y.mean(0), [1, -1, 0.5], atol=0.01) and np.allclose(np.cov(Y.T), S, atol=0.015),
      "max cov error %.4f" % np.abs(np.cov(Y.T) - S).max())

# ----------------------------------------------------------------------------- 11.1.2 and 11.1.3
print("== Rejection sampling and adaptive rejection sampling (Sections 11.1.2, 11.1.3)")
env = m.gamma_cauchy_envelope(10.0)
zz = np.linspace(1e-3, 200, 400001)
check("Fig 11.5 envelope: k q(z) >= Gam(z|10, 1) everywhere, touching it (min ratio 1)",
      np.all(env["log_k"] + env["log_q"](zz) >= m.gamma_logpdf(zz, 10) - 1e-9),
      "k = %.4f, max of ln p - ln kq = %.2e" % (env["k"], np.max(m.gamma_logpdf(zz, 10) - env["log_k"] - env["log_q"](zz))))
xs, tries, _ = m.rejection_sample(np.random.default_rng(4), 40000, lambda z: float(m.gamma_logpdf(z, 10)), env["sample"],
                                  lambda z: float(env["log_q"](z)), env["log_k"])
acc = len(xs) / tries
check("(11.14): acceptance rate = (1/k) int p_tilde dz (gamma example, p_tilde normalised so the integral is 1)",
      within(acc, 1 / env["k"], math.sqrt(acc * (1 - acc) / tries)), "%.4f versus 1/k = %.4f (%d tries)" % (acc, 1 / env["k"], tries))
ks = stats.kstest(xs, "gamma", args=(10,))
check("Exercise 11.6: the accepted samples follow Gam(z|10, 1) (KS test)", ks.pvalue > 0.01, "KS p = %.3f" % ks.pvalue)
S4 = m.fig4_setup()
xs, tries, _ = m.rejection_sample(np.random.default_rng(5), 30000, lambda z: float(S4["log_pt"](np.array(z))), S4["sample"],
                                  lambda z: float(S4["log_q"](np.array(z))), S4["log_k"])
acc = len(xs) / tries
truth = m.BIMODAL_Z / S4["k"]
check("(11.14) with an unnormalised p_tilde (Fig 11.4): acceptance = (1/k) int p_tilde = Z_p / k",
      within(acc, truth, math.sqrt(acc * (1 - acc) / tries)) and stats.kstest(xs, m.bimodal_cdf).pvalue > 0.01,
      "%.4f versus %.4f, KS p = %.3f" % (acc, truth, stats.kstest(xs, m.bimodal_cdf).pvalue))
rng = np.random.default_rng(6)
rows = []
ok = True
for D in (1, 2, 5, 10, 20):
    n = 400000
    Z = 1.1 * rng.standard_normal((n, D))
    a = np.mean(rng.random(n) < np.exp(-0.5 * (Z ** 2).sum(1) * (1 - 1 / 1.1 ** 2)))
    th = 1.1 ** -D
    ok &= within(a, th, math.sqrt(th * (1 - th) / n))
    rows.append("D=%d: %.4f/%.4f" % (D, a, th))
check("Section 11.1.3: rejection of N(0, I) by k N(0, 1.1^2 I) in D dimensions (full D-dim draws) accepts (sigma_p/sigma_q)^D", ok, ", ".join(rows))
a = m.highdim_acceptance(np.random.default_rng(7), 1000, 1.01, 4_000_000)
th = 1.01 ** -1000
check("the same at D = 1000 with sigma_q = 1.01 sigma_p: about 1/20 000 (book p. 531), via |z|^2 ~ sigma_q^2 chi^2_D",
      within(a, th, math.sqrt(th / 4e6)), "%.3g versus %.3g = 1/%d" % (a, th, round(1 / th)))
h, dh = m.gamma_ars_functions(10.0)
H = m.ars_hull([4.0, 8.5, 14.0, 20.0], h, dh, 0.0)
zz = np.linspace(1e-3, 80, 200001)
mass_num = integrate.quad(lambda t: math.exp(float(m.ars_upper(H, t))), 0, 200, points=list(H["zb"][1:-1]), limit=400)[0]
cont = max(abs(H["hx"][i] + H["s"][i] * (H["zb"][i + 1] - H["x"][i]) - H["hx"][i + 1] - H["s"][i + 1] * (H["zb"][i + 1] - H["x"][i + 1]))
           for i in range(len(H["x"]) - 1))
check("ARS envelope (11.17): above ln p_tilde everywhere, continuous at the intersections, piece masses (Exercise 11.8) sum to its integral",
      np.all(m.ars_upper(H, zz) >= h(zz) - 1e-12) and cont < 1e-10 and abs(np.exp(H["logm"]).sum() / mass_num - 1) < 1e-8,
      "continuity error %.1e, mass %.6f versus quadrature %.6f" % (cont, np.exp(H["logm"]).sum(), mass_num))
draws = np.array([m.ars_draw(H, a_, b_) for a_, b_ in np.random.default_rng(8).random((100000, 2))])
cdf_env = lambda t: np.array([integrate.quad(lambda x: math.exp(float(m.ars_upper(H, x))), 0, ti, points=[p for p in H["zb"][1:-1] if p < ti] or None, limit=200)[0]
                              for ti in np.atleast_1d(t)]) / mass_num
grid = np.array([3.0, 6.0, 9.0, 12.0, 18.0, 30.0])
emp = np.array([(draws <= g_).mean() for g_ in grid])
check("Exercise 11.9: draws from the piecewise exponential envelope have its CDF", np.allclose(emp, cdf_env(grid), atol=0.006),
      "max CDF difference %.4f" % np.abs(emp - cdf_env(grid)).max())
xs, tries, hist = m.adaptive_rejection_sample(np.random.default_rng(9), 20000, h, dh, [4.0, 8.5, 14.0], lo=0.0)
ks = stats.kstest(xs, "gamma", args=(10,))
check("adaptive rejection sampling of Gam(z|10, 1): KS test, and the envelope tightens (acceptance of the final envelope > 0.98)",
      ks.pvalue > 0.01 and math.exp(special.gammaln(10)) / np.exp(m.ars_hull(hist[-1], h, dh, 0.0)["logm"]).sum() > 0.98,
      "KS p = %.3f, %d rejections in %d tries" % (ks.pvalue, tries - len(xs), tries))

# ----------------------------------------------------------------------------- 11.1.4 to 11.1.6
print("== Importance sampling, SIR (Sections 11.1.4, 11.1.5)")
truth = integrate.quad(lambda t: m.fig8_f(t) * m.bimodal_pdf(t), -15, 15, limit=200)[0]
log_pt = lambda z: np.log(np.maximum(m.bimodal_p_tilde(z), 1e-300))
errs = []
for L in (100, 1000, 10000, 100000):
    e = []
    for rep in range(40):
        z = 2.0 * np.random.default_rng(1000 * L + rep).standard_normal(L)
        e.append(m.importance_sampling(z, log_pt, lambda t: np.log(m.gauss_pdf(t, 0, 2)), m.fig8_f)["estimate"] - truth)
    errs.append(math.sqrt(np.mean(np.square(e))))
slope = np.polyfit(np.log([100, 1000, 10000, 100000]), np.log(errs), 1)[0]
check("self-normalised importance sampling (11.22)-(11.23) is consistent: RMS error falls like L^(-1/2)",
      abs(slope + 0.5) < 0.1 and errs[-1] < 0.005, "RMS errors %s, slope %.3f" % (", ".join("%.4f" % x for x in errs), slope))
z = 2.0 * np.random.default_rng(11).standard_normal(400000)
res = m.importance_sampling(z, log_pt, lambda t: np.log(m.gauss_pdf(t, 0, 2)), m.fig8_f)
r = np.exp(log_pt(z) - np.log(m.gauss_pdf(z, 0, 2)))
check("(11.21): the mean importance ratio estimates Z_p / Z_q (Z_q = 1 here)", within(res["ratio"], m.BIMODAL_Z, r.std() / math.sqrt(len(r))),
      "%.5f versus Z_p = %.5f" % (res["ratio"], m.BIMODAL_Z))
check("effective sample size: 1/sum w^2 = L for equal weights; between 1 and L otherwise",
      abs(m.importance_sampling(z[:1000], lambda t: -0.5 * t * t, lambda t: -0.5 * t * t, m.fig8_f)["ess"] - 1000) < 1e-9 and 1 <= res["ess"] <= len(z),
      "ESS/L = %.3f for q = N(0, 2^2)" % (res["ess"] / len(z)))
zs = 2.0 * np.random.default_rng(12).standard_normal(200000)
w = m.importance_sampling(zs, log_pt, lambda t: np.log(m.gauss_pdf(t, 0, 2)), m.fig8_f)["w"]
rs = m.sir(np.random.default_rng(13), zs, w, 20000)
ks = stats.kstest(rs, m.bimodal_cdf)
check("sampling-importance-resampling (11.25)-(11.26): resampled values follow p (KS test, L = 200 000 -> 20 000)", ks.pvalue > 0.01, "KS p = %.3f" % ks.pvalue)

# ----------------------------------------------------------------------------- 11.2
print("== Markov chains and Metropolis-Hastings (Section 11.2)")
mom = m.random_walk_second_moment(50)
check("Exercise 11.10: the lazy random walk (11.34)-(11.36) has E[z^2] = tau/2 exactly", np.allclose(mom, np.arange(51) / 2, atol=1e-12),
      "E[z^2] at tau = 50: %.12f" % mom[-1])
p = np.array([1.0, 3.0, 0.5, 2.0, 4.0, 0.2])
Qr = np.random.default_rng(14).random((6, 6))
np.fill_diagonal(Qr, 0)
Q = Qr / Qr.sum(1, keepdims=True)          # an asymmetric proposal
T = m.mh_transition_matrix(p, Q)
F = p[:, None] * T
pi = p / p.sum()
check("Metropolis-Hastings kernel (11.44) on 6 states with an asymmetric proposal: detailed balance (11.40), (11.45) holds exactly",
      np.abs(F - F.T).max() < 1e-15 and np.abs(pi @ T - pi).max() < 1e-15 and np.allclose(T.sum(1), 1),
      "max |p(a)T(a,b) - p(b)T(b,a)| = %.1e, max |pT - p| = %.1e" % (np.abs(F - F.T).max(), np.abs(pi @ T - pi).max()))
T0 = m.mh_transition_matrix(p, Q, hastings=False)
F0 = p[:, None] * T0
check("negative control: the plain Metropolis ratio with the same asymmetric proposal breaks detailed balance and invariance",
      np.abs(F0 - F0.T).max() > 1e-3 and np.abs(pi @ T0 - pi).max() > 1e-3, "max imbalance %.3f, max |pT - p| = %.3f" % (np.abs(F0 - F0.T).max(), np.abs(pi @ T0 - pi).max()))
ev = np.sort(np.abs(np.linalg.eigvals(T)))[::-1]
Tn = np.linalg.matrix_power(T, 200)
check("the chain is ergodic: second eigenvalue modulus < 1 and every row of T^200 equals p", ev[1] < 1 and np.abs(Tn - pi).max() < 1e-10,
      "|lambda_2| = %.4f" % ev[1])
Sig = m.FIG9_SIGMA
lp = m._gauss2_logpdf(m.FIG9_MU, Sig)
out, accr = m.metropolis_parallel(np.random.default_rng(15), lp, m.gaussian_cholesky(np.random.default_rng(16), m.FIG9_MU, Sig, 400), 3000, 0.2)
X = out[500:].reshape(-1, 2)
check("Metropolis (11.33) on the Fig 11.9 target: chain mean and covariance match the target (400 chains started in equilibrium)",
      np.allclose(X.mean(0), m.FIG9_MU, atol=0.01) and np.allclose(np.cov(X.T), Sig, atol=0.01),
      "mean %s, cov %s, acceptance %.3f" % (np.round(X.mean(0), 4), np.round(np.cov(X.T), 4).tolist(), accr))
res = m.fig9_run()
check("Figure 11.9 run: 150 candidates, 43 rejected (as in the book)", len(res["acc"]) == 150 and (~res["acc"]).sum() == 43, "%d rejected" % (~res["acc"]).sum())
# Metropolis-Hastings with an asymmetric (multiplicative log-normal) proposal on a gamma target
lpg = lambda z: float(m.gamma_logpdf(z[0], 3.0, 2.0))
prop = lambda rng, z: z * np.exp(0.8 * rng.standard_normal(1))
lq = lambda a, b: float(-math.log(a[0]) - 0.5 * (math.log(a[0] / b[0]) / 0.8) ** 2)
mh = m.metropolis_hastings(np.random.default_rng(17), lpg, np.array([1.0]), 60000, prop, lq)
zc = mh["chain"][1000:, 0]
se = zc.std() * math.sqrt(m.integrated_autocorr_time(zc) / len(zc))
check("Metropolis-Hastings (11.44) with a log-normal proposal (asymmetric q): Gam(3, 2) mean 1.5 recovered",
      within(zc.mean(), 1.5, se), "mean %.4f (se %.4f), acceptance %.3f" % (zc.mean(), se, mh["acc"].mean()))

# ----------------------------------------------------------------------------- 11.3
print("== Gibbs sampling and over-relaxation (Section 11.3)")
r = 0.9
for alpha in (0.0, -0.5, -0.98):
    A = m.overrelaxed_sweep_matrix(r, alpha)
    S2 = np.array([[1.0, r], [r, 1.0]])
    # the noise covariance of a sweep: z1 update adds s^2 (1 - alpha^2) to z1, then z2 = ... + (1 - alpha) r z1' + noise
    s2 = (1 - r * r) * (1 - alpha * alpha)
    A2 = np.array([[1.0, 0.0], [(1 - alpha) * r, alpha]])
    Qn = A2 @ np.diag([s2, 0.0]) @ A2.T + np.diag([0.0, s2])
    check("Gibbs sweep%s on a bivariate Gaussian (r = 0.9): A Sigma A^T + Q = Sigma, so N(0, Sigma) is invariant (exact)" %
          ("" if alpha == 0 else " over-relaxed with alpha = %g (11.50)" % alpha),
          np.abs(A @ S2 @ A.T + Qn - S2).max() < 1e-14, "max error %.1e" % np.abs(A @ S2 @ A.T + Qn - S2).max())
mu3 = np.array([1.0, -2.0, 0.5])
S3 = np.array([[2.0, 0.9, 0.3], [0.9, 1.0, -0.4], [0.3, -0.4, 1.5]])
g3 = m.gibbs_gaussian(np.random.default_rng(18), mu3, S3, np.zeros(3), 60000)["chain"][500:]
check("Gibbs sampling (11.46)-(11.48) of a 3-D Gaussian: sample mean and covariance match", np.allclose(g3.mean(0), mu3, atol=0.04) and np.allclose(np.cov(g3.T), S3, atol=0.05),
      "max cov error %.3f" % np.abs(np.cov(g3.T) - S3).max())
Pz = np.array([[0.10, 0.05, 0.20], [0.15, 0.25, 0.05], [0.02, 0.08, 0.10]])   # p(z1, z2) on 3 x 3 states
n1, n2 = Pz.shape
Tg = np.zeros((9, 9))
for a in range(n1):
    for b in range(n2):
        for a2 in range(n1):
            Tg[a * n2 + b, a2 * n2 + b] = Pz[a2, b] / Pz[:, b].sum()
pv = Pz.ravel()
F = pv[:, None] * Tg
check("Exercise 11.11: a single-site Gibbs update satisfies detailed balance (exact, 3 x 3 discrete distribution)",
      np.abs(F - F.T).max() < 1e-16 and np.abs(pv @ Tg - pv).max() < 1e-16, "max imbalance %.1e" % np.abs(F - F.T).max())
mu_i, s_i = 0.7, 1.3
for alpha in (-0.9, 0.5):
    zi = mu_i + s_i * np.random.default_rng(19).standard_normal(400000)
    nu = np.random.default_rng(20).standard_normal(400000)
    zp = mu_i + alpha * (zi - mu_i) + s_i * math.sqrt(1 - alpha * alpha) * nu
    check("Exercise 11.14: over-relaxation (11.50) with alpha = %g keeps the conditional N(mu_i, sigma_i^2): mean, variance, KS" % alpha,
          within(zp.mean(), mu_i, s_i / math.sqrt(len(zp))) and abs(zp.var() / s_i ** 2 - 1) < 0.01 and stats.kstest(zp, "norm", args=(mu_i, s_i)).pvalue > 0.01,
          "mean %.4f, var %.4f (target %.2f), corr(z, z') = %.3f" % (zp.mean(), zp.var(), s_i ** 2, np.corrcoef(zi, zp)[0, 1]))
check("exact IAT of Gibbs on a bivariate Gaussian equals (1 + r^2)/(1 - r^2) = 2 (L/l)^2 - 1, (L/l)^2 = 1/(1 - r^2)",
      all(abs(m.exact_iat_z1(x) - (1 + x * x) / (1 - x * x)) < 1e-9 for x in (0.3, 0.9, 0.99)), "r = 0.99: %.2f" % m.exact_iat_z1(0.99))
iat_g, iat_o = m.gibbs_iat(0.99), m.gibbs_iat(0.99, alpha=-0.9)
check("measured IAT at r = 0.99 matches the exact values: Gibbs %.1f, over-relaxed alpha = -0.9 %.2f" % (m.exact_iat_z1(0.99), m.exact_iat_z1(0.99, -0.9)),
      abs(iat_g / m.exact_iat_z1(0.99) - 1) < 0.1 and abs(iat_o / m.exact_iat_z1(0.99, -0.9) - 1) < 0.1, "measured %.1f and %.2f" % (iat_g, iat_o))
oo = m.ordered_overrelaxed_gibbs(np.random.default_rng(21), np.zeros(2), np.array([[1, 0.9], [0.9, 1]]), np.zeros(2), 30000, 15)["chain"][200:]
check("ordered over-relaxation (Neal 1999, K = 15) leaves the Gaussian invariant (moments)",
      np.allclose(np.cov(oo.T), [[1, 0.9], [0.9, 1]], atol=0.05) and np.abs(oo.mean(0)).max() < 0.05, "cov %s" % np.round(np.cov(oo.T), 3).tolist())

# ----------------------------------------------------------------------------- 11.4
print("== Slice sampling (Section 11.4)")
lpt = lambda z: float(np.log(max(m.bimodal_p_tilde(z), 1e-300)))
sl = m.slice_sample(np.random.default_rng(22), lpt, 0.0, 60000, 1.0)
zc = sl["chain"][100:]
mean_t = float((m.BIMODAL_W * m.BIMODAL_SD * m.BIMODAL_MU).sum() * math.sqrt(2 * math.pi) / m.BIMODAL_Z)
ez2 = float((m.BIMODAL_W * m.BIMODAL_SD * (m.BIMODAL_MU ** 2 + m.BIMODAL_SD ** 2)).sum() * math.sqrt(2 * math.pi) / m.BIMODAL_Z)
se = zc.std() * math.sqrt(m.integrated_autocorr_time(zc) / len(zc))
thin = zc[::10]
check("slice sampling (11.51)-(11.52) with stepping out and shrinking: mean, E[z^2] and KS test (thinned) for the bimodal target",
      within(zc.mean(), mean_t, se) and abs((zc ** 2).mean() / ez2 - 1) < 0.03 and stats.kstest(thin, m.bimodal_cdf).pvalue > 0.001,
      "mean %.4f (truth %.4f, se %.4f), E[z^2] %.3f (truth %.3f), %.2f evaluations per sample" % (zc.mean(), mean_t, se, (zc ** 2).mean(), ez2, sl["evals"].mean()))
sg = m.slice_sample(np.random.default_rng(23), lambda z: -0.5 * z * z / 9.0, 0.0, 40000, 0.3)["chain"][100:]
check("slice sampling adapts to the scale: N(0, 3^2) with w = 0.3 (ten times too small) still has variance 9", abs(sg.var() / 9 - 1) < 0.06,
      "variance %.3f" % sg.var())

# ----------------------------------------------------------------------------- 11.5
print("== Leapfrog and hybrid Monte Carlo (Section 11.5)")
E4 = lambda z: float(0.25 * np.sum(z ** 4) + 0.5 * z[0] * z[1] + np.sum(np.cos(z)))
g4 = lambda z: z ** 3 + 0.5 * z[::-1] - np.sin(z)
def flow(x, eps=0.17, L=13):
    zs, rs = m.leapfrog(g4, x[:2], x[2:], eps, L)
    return np.r_[zs[-1], rs[-1]]
x0 = np.array([0.3, -0.7, 0.9, 0.2])
J = np.zeros((4, 4))
h_ = 1e-6
for i in range(4):
    e = np.zeros(4)
    e[i] = h_
    J[:, i] = (flow(x0 + e) - flow(x0 - e)) / (2 * h_)
check("leapfrog (11.64)-(11.66) preserves phase-space volume: Jacobian determinant of 13 steps = 1 (non-quadratic 2-D energy)",
      abs(np.linalg.det(J) - 1) < 1e-7, "det = %.10f" % np.linalg.det(J))
y = flow(x0)
back = flow(np.r_[y[:2], -y[2:]])
check("leapfrog is time-reversible: integrate, negate r, integrate again: back to the start", np.abs(np.r_[back[:2], -back[2:]] - x0).max() < 1e-12,
      "error %.1e" % np.abs(np.r_[back[:2], -back[2:]] - x0).max())
yneg = m.leapfrog(g4, x0[:2], x0[2:], -0.17, 13)
fwd = m.leapfrog(g4, y[:2], y[2:], -0.17, 13)
check("integrating with step -eps from the end point retraces the trajectory (the reversibility used by HMC)",
      np.abs(np.r_[fwd[0][-1], fwd[1][-1]] - x0).max() < 1e-12, "error %.1e" % np.abs(np.r_[fwd[0][-1], fwd[1][-1]] - x0).max())
errs = []
epss = np.array([0.2, 0.1, 0.05, 0.025])
for e_ in epss:
    zs, rs = m.leapfrog(g4, x0[:2], x0[2:], e_, int(round(2.0 / e_)))
    H = [m.hamiltonian(E4, zs[k], rs[k]) for k in range(len(zs))]
    errs.append(max(abs(np.array(H) - H[0])))
sl_ = np.polyfit(np.log(epss), np.log(errs), 1)[0]
check("leapfrog energy error over a fixed time is O(eps^2)", abs(sl_ - 2) < 0.15, "errors %s, slope %.3f" % (", ".join("%.2e" % x for x in errs), sl_))
rng = np.random.default_rng(24)
worst = 0.0
for _ in range(200):
    z = rng.standard_normal(2)
    r_ = rng.standard_normal(2)
    zs, rs = m.leapfrog(g4, z, r_, 0.3, 7)
    H0, H1 = m.hamiltonian(E4, z, r_), m.hamiltonian(E4, zs[-1], rs[-1])
    lhs = math.exp(-H0) * min(1.0, math.exp(H0 - H1))   # flow R -> R' (11.68), with the erratum's sign
    rhs = math.exp(-H1) * min(1.0, math.exp(H1 - H0))   # flow R' -> R (11.69)
    worst = max(worst, abs(lhs - rhs) / lhs)
check("Exercise 11.17: exp(-H) min(1, exp(H - H*)) = exp(-H*) min(1, exp(H* - H)) for 200 random leapfrog moves; with volume "
      "preservation and reversibility this is detailed balance for HMC", worst < 1e-14, "max relative difference %.1e" % worst)
mu2, S2 = np.zeros(2), m.rotated_gaussian(1.0, 0.1)
P2 = np.linalg.inv(S2)
Eg = lambda z: 0.5 * float(z @ P2 @ z)
gg = lambda z: P2 @ z
hm = m.hmc(np.random.default_rng(25), Eg, gg, np.zeros(2), 8000, 0.15, 15, jitter=0.1)
Xc = hm["chain"][200:]
check("HMC on a correlated Gaussian (sigma 1 and 0.1): sample covariance matches", np.allclose(np.cov(Xc.T), S2, atol=0.06),
      "cov %s, acceptance %.3f" % (np.round(np.cov(Xc.T), 3).tolist(), hm["acc"].mean()))
hb = m.hmc(np.random.default_rng(26), lambda z: 0.5 * float(z @ z), lambda z: z, np.zeros(1), 40000, 1.6, 3)
nob = []
z = np.zeros(1)
rr = np.random.default_rng(26)
for _ in range(40000):
    zs, rs = m.leapfrog(lambda q: q, z, rr.standard_normal(1), 1.6, 3)
    z = zs[-1]
    nob.append(z[0])
check("HMC removes the leapfrog bias: N(0, 1) with eps = 1.6 has variance 1 with the accept step; without it the variance is wrong",
      abs(hb["chain"][100:, 0].var() - 1) < 0.05 and abs(np.var(nob[100:]) - 1) > 0.3,
      "with: %.3f (acceptance %.3f); without: %.3f" % (hb["chain"][100:, 0].var(), hb["acc"].mean(), np.var(nob[100:])))
check("for a Gaussian, HMC's mean acceptance is E[min(1, exp(-dH))] and E[exp(-dH)] = 1 (Jarzynski-type identity of volume-preserving reversible maps)",
      abs(np.mean(np.exp(-hb["dH"])) - 1) < 0.05, "mean exp(-dH) = %.4f" % np.mean(np.exp(-hb["dH"])))

# ----------------------------------------------------------------------------- 11.6
print("== Estimating the partition function (Section 11.6)")
E1 = lambda z: 0.25 * z ** 4 - z ** 2 + 0.3 * z      # a double well
ZE = integrate.quad(lambda z: math.exp(-E1(z)), -12, 12, limit=200)[0]
G = lambda z: 0.5 * (z / 1.6) ** 2
ZG = 1.6 * math.sqrt(2 * math.pi)
zg = 1.6 * np.random.default_rng(27).standard_normal(400000)
est, se = m.partition_ratio(E1, G, zg)
check("(11.72) with the factor 1/L (erratum): Z_E / Z_G from samples of p_G matches quadrature", within(est, ZE / ZG, se),
      "%.5f +- %.5f versus %.5f" % (est, se, ZE / ZG))
def sampler(alpha, L, seed=[30]):
    seed[0] += 1
    Ea = lambda z: (1 - alpha) * G(z) + alpha * E1(z)
    return m.slice_sample(np.random.default_rng(seed[0]), lambda z: -Ea(z), 0.0, L, 1.0)["chain"][1:]
alphas = np.linspace(0, 1, 6)
est_c, fac = m.chained_partition_ratio(G, E1, sampler, alphas, 20000)
exact_f = [integrate.quad(lambda z, a=b_: math.exp(-((1 - a) * G(z) + a * E1(z))), -12, 12, limit=200)[0] /
           integrate.quad(lambda z, a=a_: math.exp(-((1 - a) * G(z) + a * E1(z))), -12, 12, limit=200)[0] for a_, b_ in zip(alphas[:-1], alphas[1:])]
check("chaining (11.74)-(11.75) with 5 intermediate ratios, each from slice-sampler draws, reproduces Z_E / Z_G",
      abs(est_c / (ZE / ZG) - 1) < 0.02, "%.5f versus %.5f; factors %s" % (est_c, ZE / ZG, ", ".join("%.3f" % f for f in fac)))

# ----------------------------------------------------------------------------- the demos' generator
print("== The demos' random numbers (JSRandom = DK.rng)")
js = m.JSRandom(1)
first = [js.u() for _ in range(3)]
js2 = m.JSRandom(42)
big = np.array([js2.u() for _ in range(100000)])
check("JSRandom (mulberry32) is uniform on [0, 1): mean, variance, KS", abs(big.mean() - 0.5) < 0.004 and abs(big.var() - 1 / 12) < 0.002 and stats.kstest(big, "uniform").pvalue > 0.01,
      "first values for seed 1: %s" % ", ".join("%.10f" % v for v in first))

print("\n%d checks: %d PASS, %d FAIL" % (len(RESULTS), sum(RESULTS), len(RESULTS) - sum(RESULTS)))