"""Bishop (2006), Chapter 11: Sampling Methods.

Reproduces the numerical figures of Chapter 11 of C. M. Bishop, Pattern Recognition and
Machine Learning (Springer, 2006), pp. 523-558, written from scratch with numpy, scipy and
matplotlib. Every sampler is a small function meant to be read next to the lecture notes
(ch11-sampling-methods.html):

* Section 11.1.1, transformations of uniform numbers: inverse_cdf (11.5)-(11.6),
  exponential_from_uniform (11.7), cauchy_from_uniform (11.8, Exercises 11.3 and 11.7),
  box_muller (11.10)-(11.12, polar form, with the erratum ln r^2), gaussian_cholesky (Exercise 11.5).
* Section 11.1.2, rejection sampling: rejection_sample (11.13)-(11.14), with the gamma density
  (11.15) and the scaled Cauchy envelope (11.16) in gamma_cauchy_envelope.
* Section 11.1.3, adaptive rejection sampling (Gilks and Wild 1992): ars_hull builds the
  piecewise exponential envelope (11.17) from tangent lines, adaptive_rejection_sample draws
  from it (Exercises 11.8 and 11.9).
* Sections 11.1.4 and 11.1.5: importance_sampling (11.19)-(11.23) with the effective sample
  size, sir (sampling-importance-resampling, (11.25)-(11.27)).
* Section 11.2: metropolis (11.33), metropolis_hastings (11.44), mh_transition_matrix (the
  kernel on a finite space, for the detailed-balance check (11.40), (11.45)),
  random_walk_second_moment ((11.34)-(11.36), Exercise 11.10).
* Section 11.3: gibbs_gaussian (11.46)-(11.48), overrelaxed_gibbs (Adler's over-relaxation
  (11.50)), ordered_overrelaxed_gibbs (Neal 1999).
* Section 11.4: slice_sample (11.51)-(11.52), with stepping out and shrinking (Neal 2003).
* Section 11.5: leapfrog (11.64)-(11.66), hamiltonian (11.57), hmc (hybrid Monte Carlo,
  acceptance (11.67), random direction of integration).
* Section 11.6: partition_ratio (11.72, with the erratum: the factor 1/L) and
  chained_partition_ratio (11.74)-(11.75).
* Diagnostics: autocorrelation, integrated_autocorr_time, effective_sample_size.
* JSRandom: the seeded generator of the browser demos (DK.rng in courses/demo-kit.js:
  mulberry32 uniforms and Box-Muller normals). Passing JSRandom(seed) instead of a numpy
  Generator makes metropolis, gibbs_gaussian, slice_sample, hmc, ... use exactly the numbers
  of the demos, which is how the demos are cross-checked against this file.

Figures: 11.2 to 11.14. Run `python3 ch11_sampling.py --list`, `python3 ch11_sampling.py`
(all figures, about a minute) or `python3 ch11_sampling.py --fig 11.9 11.13`. The PNGs are
written to ../img/ch11/.

What differs from the book, and the settings chosen where the book gives none:
* The book's Figures 11.2, 11.3, 11.4, 11.6, 11.8, 11.10, 11.11, 11.13 and 11.14 are sketches.
  We compute each of them from a real sampler and add panels with the numbers behind the
  sketch. The bimodal density used in Figures 11.2, 11.4, 11.8 and 11.13 (and in the demos) is
  p(z) = 0.55 N(z|-2, 0.75^2) + 0.45 N(z|2, 0.95^2); its unnormalised version p_tilde drops the
  factors 1/(sqrt(2 pi) s_k), so Z_p = sqrt(2 pi) (0.55 x 0.75 + 0.45 x 0.95) = 2.1055.
* Figure 11.4: proposal q(z) = N(z|0, 2.2^2), k the smallest constant with k q >= p_tilde
  (found numerically): k = 4.841, so the acceptance (1/k) int p_tilde (11.14) is 0.435.
* Figure 11.5: the curves are Gam(z|10, 1) (peak 0.132 at z = 9, as in the book) and the
  Cauchy envelope with c = a - 1 = 9, b^2 = 2a - 1 = 19 and the smallest k; we find k = 1.804,
  so the acceptance rate is 1/k = 0.554.
* Figure 11.6: target Gam(z|10, 1) again (log-concave); starting abscissae 4, 8.5, 14, whose
  envelope accepts with probability 0.863. Following the book, a point is added to the
  envelope only when it is rejected (Gilks and Wild also use a lower squeeze function; we do
  not, to keep the code short).
* Figure 11.7: sigma_p = 1, sigma_q = 1.1 (read off the figure); the right panel adds the
  acceptance (sigma_p/sigma_q)^D against D (1/20959 for sigma_q/sigma_p = 1.01, D = 1000, the
  book's "1/20,000"). Its Monte Carlo points use the fact that the acceptance probability
  p(z)/(k q(z)) depends on z only through |z|^2, which under q is sigma_q^2 times a chi-squared
  variable with D degrees of freedom; this is exactly equivalent to drawing the D coordinates
  and is much cheaper for D = 1000.
* Figure 11.8: q(z) = N(z|0, 2^2), f(z) = tanh(0.6 (z + 1)); a badly placed q = N(2, 0.8^2) for
  comparison.
* Figure 11.9: the book gives the proposal sd (0.2) and the number of candidates (150) but
  not the target. The ellipse in the figure has centre (1.5, 1.5), marginal sd 1.06 and
  correlation 0.76 if read as a one-sd contour, as the caption says. With that target a
  proposal sd of 0.2 is rejected only about 20 times in 150, and the chain would leave the
  ellipse most of the time; read as the two-sd contour, the target (covariance
  [[0.28, 0.2125], [0.2125, 0.28]], principal sds 0.70 and 0.26) gives about 40 rejections in 150
  and a chain that stays mostly inside, as in the book. We use the second reading, draw both
  contours, start at (1.0, 1.0) and use numpy seed 16, which gives 43 rejections as in the book.
* Figures 11.10 and 11.12: random-walk Metropolis with proposal sd rho = sigma_min on Gaussians
  with sigma_max / sigma_min = 2, 4, 8, 16 and 32 (integrated autocorrelation time by Sokal's
  automatic window, c = 5, 16 chains), and Gibbs sampling with correlation 0.5 ... 0.995
  (exact IAT from the linear sweep map, measured IAT by batch means over 100 chains). File
  fig11-12 is our companion to the book's Figure 11.11 (the book's Figure 11.12 is a
  graphical-model sketch of Markov blankets, not a numerical figure).
* Figure 11.13: w = 0.5 for the stepping-out panel, numpy seed 0 (four steps out, three
  candidates rejected before one is accepted).
* Figure 11.14: panels (a) and (b) shear a rectangle of phase space with the potential
  E(z) = 1 - cos(2 z) and eps = 0.6 (the book's sketch has no numbers); panels (c) and (d) use
  the standard Gaussian E(z) = z^2 / 2.
"""
import math

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, Polygon
from scipy import special, stats

import common

LOG2PI = math.log(2 * math.pi)
M32 = 0xFFFFFFFF

# The bimodal density of Figures 11.2, 11.4, 11.8 and 11.13 (and of the demos).
BIMODAL_W = np.array([0.55, 0.45])
BIMODAL_MU = np.array([-2.0, 2.0])
BIMODAL_SD = np.array([0.75, 0.95])
BIMODAL_Z = float(np.sqrt(2 * np.pi) * (BIMODAL_W * BIMODAL_SD).sum())  # int p_tilde dz

# The target of Figure 11.9 (read off the book's ellipse).
FIG9_MU = np.array([1.5, 1.5])
FIG9_SIGMA = np.array([[0.28, 0.2125], [0.2125, 0.28]])  # the book's ellipse is its 2-sd contour


# =========================================================================== the demos' generator
def _imul(a, b):
    return (a * b) & M32


class JSRandom:
    """The seeded random numbers of the browser demos (DK.rng in courses/demo-kit.js).

    Uniforms come from the mulberry32 generator and normals from the basic Box-Muller transform,
    the second normal of each pair being cached, exactly as in JavaScript. random() and
    standard_normal() mimic the numpy Generator methods used by the samplers in this file.
    """

    def __init__(self, seed):
        self.a = int(seed) & M32
        self.spare = None

    def u(self):
        a = (self.a + 0x6D2B79F5) & M32
        self.a = a
        t = _imul(a ^ (a >> 15), 1 | a)
        t = ((t + _imul(t ^ (t >> 7), 61 | t)) & M32) ^ t
        return ((t ^ (t >> 14)) & M32) / 4294967296.0

    def _pos(self):
        x = 0.0
        while x == 0.0:
            x = self.u()
        return x

    def norm(self):
        if self.spare is not None:
            s, self.spare = self.spare, None
            return s
        r = math.sqrt(-2.0 * math.log(self._pos()))
        th = 2.0 * math.pi * self.u()
        self.spare = r * math.sin(th)
        return r * math.cos(th)

    def random(self, size=None):
        if size is None:
            return self.u()
        return np.array([self.u() for _ in range(int(np.prod(size)))]).reshape(size)

    def standard_normal(self, size=None):
        if size is None:
            return self.norm()
        return np.array([self.norm() for _ in range(int(np.prod(size)))]).reshape(size)


def _rng(rng):
    return np.random.default_rng(rng) if rng is None or isinstance(rng, (int, np.integer)) else rng


# =========================================================================== densities used below
def bimodal_p_tilde(z):
    """Unnormalised bimodal density: sum_k w_k exp(-(z - mu_k)^2 / (2 s_k^2)); its integral is BIMODAL_Z."""
    z = np.asarray(z, dtype=float)[..., None]
    return (BIMODAL_W * np.exp(-0.5 * ((z - BIMODAL_MU) / BIMODAL_SD) ** 2)).sum(-1)


def bimodal_pdf(z):
    return bimodal_p_tilde(z) / BIMODAL_Z


def bimodal_cdf(z):
    z = np.asarray(z, dtype=float)[..., None]
    return (BIMODAL_W * BIMODAL_SD * np.sqrt(2 * np.pi) * special.ndtr((z - BIMODAL_MU) / BIMODAL_SD)).sum(-1) / BIMODAL_Z


def gamma_logpdf(z, a, b=1.0):
    """ln Gam(z|a, b) of (11.15); -inf for z <= 0."""
    z = np.asarray(z, dtype=float)
    with np.errstate(divide="ignore", invalid="ignore"):
        out = a * np.log(b) + (a - 1) * np.log(z) - b * z - special.gammaln(a)
    return np.where(z > 0, out, -np.inf)


def gauss_pdf(z, mu, sd):
    return np.exp(-0.5 * ((np.asarray(z, dtype=float) - mu) / sd) ** 2) / (sd * np.sqrt(2 * np.pi))


# =========================================================================== Section 11.1.1
def inverse_cdf(u, cdf, lo, hi, tol=1e-12):
    """Transformation method (11.5)-(11.6): y = h^{-1}(u) with h the CDF, by bisection on [lo, hi].

    If u is uniform on (0, 1) then y has density h'(y) = p(y) (Exercise 11.2). Vectorised over u.
    """
    u = np.asarray(u, dtype=float)
    a, b = np.full(u.shape, float(lo)), np.full(u.shape, float(hi))
    while np.max(b - a) > tol:
        m = 0.5 * (a + b)
        below = cdf(m) < u
        a, b = np.where(below, m, a), np.where(below, b, m)
    return 0.5 * (a + b)


def exponential_from_uniform(u, lam=1.0):
    """(11.7): y = -ln(1 - u) / lam is exponential with rate lam when u is uniform on (0, 1)."""
    return -np.log1p(-np.asarray(u, dtype=float)) / lam


def cauchy_from_uniform(u, c=0.0, b=1.0):
    """Exercises 11.3 and 11.7: y = c + b tan(pi (u - 1/2)) has density (1/(pi b)) / (1 + (y - c)^2 / b^2)."""
    return c + b * np.tan(np.pi * (np.asarray(u, dtype=float) - 0.5))


def box_muller(rng, n):
    """Polar Box-Muller method (11.10)-(11.11): returns n standard normals and the number of pairs tried.

    Each try draws z1 = 2u - 1 and z2 = 2u - 1 (in that order); pairs outside the unit disc (or at
    its centre) are discarded, so a fraction pi/4 is kept. A kept pair gives two independent
    N(0, 1) values y_i = z_i (-2 ln r^2 / r^2)^{1/2} with r^2 = z1^2 + z2^2 (the book prints ln z_i
    in the numerators; the erratum corrects it to ln r^2).
    """
    rng = _rng(rng)
    out, tries = [], 0
    while len(out) < n:
        z1, z2 = 2 * rng.random() - 1, 2 * rng.random() - 1
        tries += 1
        r2 = z1 * z1 + z2 * z2
        if r2 >= 1.0 or r2 == 0.0:
            continue
        f = math.sqrt(-2.0 * math.log(r2) / r2)
        out.extend([z1 * f, z2 * f])
    return np.array(out[:n]), tries


def gaussian_cholesky(rng, mu, Sigma, n):
    """Exercise 11.5: y = mu + L z with Sigma = L L^T and z ~ N(0, I) has mean mu and covariance Sigma."""
    rng = _rng(rng)
    L = np.linalg.cholesky(Sigma)
    return mu + rng.standard_normal((n, len(mu))) @ L.T


# =========================================================================== Section 11.1.2
def rejection_sample(rng, n, log_p_tilde, q_sample, log_q, log_k, max_tries=10 ** 7):
    """Rejection sampling (Section 11.1.2): draw n samples from p = p_tilde / Z_p.

    Each try draws z0 ~ q (q_sample(rng)), then u0 uniform on [0, k q(z0)] and keeps z0 if
    u0 <= p_tilde(z0). This needs k q(z) >= p_tilde(z) for all z; the acceptance probability is
    (1/k) int p_tilde dz (11.14). Returns (samples, number of tries, record of (z0, u0, kept)).
    """
    rng = _rng(rng)
    out, rec, tries = [], [], 0
    while len(out) < n and tries < max_tries:
        z0 = q_sample(rng)
        u0 = rng.random() * math.exp(log_k + log_q(z0))
        tries += 1
        keep = u0 <= math.exp(log_p_tilde(z0))
        rec.append((z0, u0, keep))
        if keep:
            out.append(z0)
    return np.array(out), tries, rec


def max_log_ratio(log_p, log_q, lo, hi, n=20001):
    """ln k for the smallest k with k q >= p on [lo, hi]: a fine grid, then golden-section refinement."""
    z = np.linspace(lo, hi, n)
    d = log_p(z) - log_q(z)
    i = int(np.nanargmax(d))
    a, b = z[max(i - 1, 0)], z[min(i + 1, n - 1)]
    g = (math.sqrt(5) - 1) / 2
    f = lambda t: float(log_p(np.array(t)) - log_q(np.array(t)))
    for _ in range(80):
        c, e = b - g * (b - a), a + g * (b - a)
        if f(c) > f(e):
            b = e
        else:
            a = c
    return max(f(0.5 * (a + b)), float(np.nanmax(d)))


def gamma_cauchy_envelope(a, b=1.0):
    """The comparison function of Figure 11.5 for Gam(z|a, b), a > 1.

    The Cauchy proposal has location c = (a - 1)/b and scale sqrt(2a - 1)/b (the book's c = a - 1,
    b^2 = 2a - 1 for b = 1); k is the smallest constant with k q(z) >= Gam(z|a, b).
    Returns dict(c, s, log_k, k, log_q, sample).
    """
    c, s = (a - 1) / b, math.sqrt(2 * a - 1) / b
    log_q = lambda z: -np.log(np.pi * s) - np.log1p(((np.asarray(z, dtype=float) - c) / s) ** 2)
    log_k = max_log_ratio(lambda z: gamma_logpdf(z, a, b), log_q, 1e-6, c + 60 * s)
    sample = lambda rng: c + s * math.tan(math.pi * (rng.random() - 0.5))
    return dict(c=c, s=s, log_k=log_k, k=math.exp(log_k), log_q=log_q, sample=sample)


# =========================================================================== Section 11.1.3
def ars_hull(x, h, dh, lo=-np.inf, hi=np.inf):
    """Upper hull of a concave ln p_tilde from tangents at the sorted abscissae x (Gilks and Wild 1992).

    Piece i is the tangent at x_i, u_i(z) = h_i + s_i (z - x_i), used on [zb_i, zb_{i+1}], where
    zb_i is where the tangents at x_{i-1} and x_i cross (zb_0 = lo, zb_K = hi). In density terms
    q(z) is proportional to exp(u_i(z)) = k_i exp(s_i (z - x_i)) with k_i = exp(h_i): the piecewise
    exponential envelope (11.17) (with the erratum's sign, lambda_i = -s_i). Exercise 11.8 asks for
    the normalising constants: piece i has mass exp(h_i) int exp(s_i (z - x_i)) dz over its interval.
    Returns dict(x, hx, s, zb, logm) with logm the log masses of the pieces.
    """
    x = np.sort(np.asarray(x, dtype=float))
    hx, s = h(x), dh(x)
    zb = np.empty(len(x) + 1)
    zb[0], zb[-1] = lo, hi
    zb[1:-1] = (hx[1:] - hx[:-1] - x[1:] * s[1:] + x[:-1] * s[:-1]) / (s[:-1] - s[1:])
    logm = np.empty(len(x))
    for i in range(len(x)):
        a, b, si = zb[i] - x[i], zb[i + 1] - x[i], s[i]
        if abs(si) < 1e-300:
            logm[i] = hx[i] + math.log(b - a)
        elif si > 0:  # int_a^b e^{s t} dt = e^{s b} (1 - e^{-s (b - a)}) / s
            logm[i] = hx[i] + si * b + math.log(-math.expm1(-si * (b - a))) - math.log(si)
        else:
            logm[i] = hx[i] + si * a + math.log(-math.expm1(si * (b - a))) - math.log(-si)
    return dict(x=x, hx=hx, s=s, zb=zb, logm=logm)


def ars_upper(hull, z):
    """The envelope ln q_tilde(z) = u_i(z) on the piece containing z."""
    z = np.asarray(z, dtype=float)
    i = np.clip(np.searchsorted(hull["zb"], z, side="right") - 1, 0, len(hull["x"]) - 1)
    return hull["hx"][i] + hull["s"][i] * (z - hull["x"][i])


def ars_draw(hull, u1, u2):
    """Exercise 11.9: one draw from the piecewise exponential envelope from two uniforms.

    u1 picks the piece with probability proportional to its mass; within piece i (slope s on
    [a, b]) the inverse CDF of the truncated exponential gives z = a + ln(1 + u2 (e^{s (b - a)} - 1)) / s.
    """
    m = np.exp(hull["logm"] - hull["logm"].max())
    i = min(int(np.searchsorted(np.cumsum(m) / m.sum(), u1, side="right")), len(m) - 1)
    a, b, s = hull["zb"][i], hull["zb"][i + 1], hull["s"][i]
    if abs(s) < 1e-300:
        return a + u2 * (b - a)
    if np.isinf(b):
        return a + math.log1p(-u2) / s
    if np.isinf(a):
        return b + math.log(u2) / s
    return a + math.log1p(u2 * math.expm1(s * (b - a))) / s


def adaptive_rejection_sample(rng, n, h, dh, x_init, lo=-np.inf, hi=np.inf):
    """Adaptive rejection sampling (Section 11.1.3, Gilks and Wild 1992) for a log-concave p_tilde.

    h = ln p_tilde and dh its derivative; x_init are the starting abscissae (the leftmost must have
    positive slope if lo = -inf, the rightmost negative slope if hi = inf). Each try uses three
    uniforms (piece, position, acceptance). As in the book, a rejected point is added to the
    abscissae, which tightens the envelope. Returns (samples, tries, list of abscissa sets).
    """
    rng = _rng(rng)
    xs = list(np.sort(x_init))
    hull = ars_hull(xs, h, dh, lo, hi)
    out, tries, history = [], 0, [list(xs)]
    while len(out) < n:
        z = ars_draw(hull, rng.random(), rng.random())
        tries += 1
        if math.log(max(rng.random(), 1e-300)) <= float(h(np.array(z)) - ars_upper(hull, z)):
            out.append(z)
        else:
            xs = sorted(xs + [z])
            hull = ars_hull(xs, h, dh, lo, hi)
            history.append(list(xs))
    return np.array(out), tries, history


# =========================================================================== Sections 11.1.4 and 11.1.5
def importance_sampling(z, log_p_tilde, log_q_tilde, f):
    """Importance sampling with samples z ~ q (Section 11.1.4).

    r_l = p_tilde(z_l) / q_tilde(z_l); normalised weights w_l = r_l / sum_m r_m (11.23);
    E[f] is estimated by sum_l w_l f(z_l) (11.22) and Z_p / Z_q by the mean of r_l (11.21).
    The effective sample size 1 / sum_l w_l^2 (Kong 1992) equals L for equal weights and 1 when
    one weight dominates. Returns dict(estimate, w, ratio, ess).
    """
    lr = log_p_tilde(z) - log_q_tilde(z)
    m = lr.max()
    r = np.exp(lr - m)
    w = r / r.sum()
    return dict(estimate=float(np.dot(w, f(z))), w=w, ratio=float(np.exp(m) * r.mean()), ess=float(1.0 / np.sum(w ** 2)))


def sir(rng, z, w, L=None):
    """Sampling-importance-resampling (Section 11.1.5): L draws from the discrete distribution on z with
    probabilities w (one uniform each, inverse CDF of the cumulative weights)."""
    rng = _rng(rng)
    L = len(z) if L is None else L
    c = np.cumsum(w)
    c[-1] = 1.0
    idx = np.array([int(np.searchsorted(c, rng.random(), side="right")) for _ in range(L)])
    return np.asarray(z)[np.minimum(idx, len(z) - 1)]


# =========================================================================== Section 11.2
def metropolis(rng, log_p, z0, n, step):
    """Metropolis algorithm (11.33) with an isotropic Gaussian proposal N(z*|z, step^2 I).

    Per iteration: d normals for the proposal, then one uniform u; the candidate is accepted if
    u < p_tilde(z*) / p_tilde(z) (always when the ratio is >= 1). On rejection the current state is
    repeated. Returns dict(chain (n+1 x d), cand (n x d), acc (n booleans)).
    """
    rng = _rng(rng)
    z = np.array(z0, dtype=float)
    lp = log_p(z)
    chain, cand, acc = [z.copy()], [], []
    for _ in range(n):
        zs = z + step * rng.standard_normal(len(z))
        lps = log_p(zs)
        ok = math.log(max(rng.random(), 1e-300)) < lps - lp
        cand.append(zs)
        acc.append(ok)
        if ok:
            z, lp = zs, lps
        chain.append(z.copy())
    return dict(chain=np.array(chain), cand=np.array(cand), acc=np.array(acc))


def metropolis_parallel(rng, log_p, Z0, n, step):
    """M independent Metropolis chains at once (rows of Z0); returns an (n+1, M, d) array and the acceptance rate."""
    rng = _rng(rng)
    Z = np.array(Z0, dtype=float)
    lp = log_p(Z)
    out, nacc = [Z.copy()], 0
    for _ in range(n):
        Zs = Z + step * rng.standard_normal(Z.shape)
        lps = log_p(Zs)
        ok = np.log(rng.random(len(Z))) < lps - lp
        Z[ok], lp[ok] = Zs[ok], lps[ok]
        nacc += ok.sum()
        out.append(Z.copy())
    return np.array(out), nacc / (n * len(Z))


def metropolis_hastings(rng, log_p, z0, n, propose, log_q):
    """Metropolis-Hastings (11.44) with a general proposal.

    propose(rng, z) draws z* ~ q(.|z) and log_q(a, b) = ln q(a|b). The candidate is accepted with
    probability min(1, p_tilde(z*) q(z|z*) / (p_tilde(z) q(z*|z))). Returns dict(chain, acc).
    """
    rng = _rng(rng)
    z = np.array(z0, dtype=float)
    chain, acc = [z.copy()], []
    for _ in range(n):
        zs = np.asarray(propose(rng, z), dtype=float)
        la = log_p(zs) + log_q(z, zs) - log_p(z) - log_q(zs, z)
        ok = math.log(max(rng.random(), 1e-300)) < la
        acc.append(ok)
        if ok:
            z = zs
        chain.append(z.copy())
    return dict(chain=np.array(chain), acc=np.array(acc))


def mh_transition_matrix(p, Q, hastings=True):
    """The Metropolis-Hastings kernel on a finite space: T[a, b] = Q[a, b] A(b, a) for b != a, the rest on T[a, a].

    p is the (unnormalised) target and Q[a, b] = q(b|a) the proposal. With hastings=False the plain
    Metropolis ratio p(b)/p(a) is used even if Q is not symmetric (which breaks detailed balance).
    """
    p, Q = np.asarray(p, dtype=float), np.asarray(Q, dtype=float)
    n = len(p)
    T = np.zeros((n, n))
    for a in range(n):
        for b in range(n):
            if a != b and Q[a, b] > 0:
                ratio = p[b] * Q[b, a] / (p[a] * Q[a, b]) if hastings else p[b] / p[a]
                T[a, b] = Q[a, b] * min(1.0, ratio)
        T[a, a] = 1.0 - T[a].sum()
    return T


def random_walk_second_moment(tau):
    """E[(z^(tau))^2] for the lazy random walk (11.34)-(11.36) started at 0, by propagating its distribution exactly.

    Exercise 11.10: the answer is tau / 2.
    """
    p = np.zeros(2 * tau + 1)
    p[tau] = 1.0
    out = [0.0]
    k = np.arange(-tau, tau + 1)
    for _ in range(tau):
        p = 0.5 * p + 0.25 * np.roll(p, 1) + 0.25 * np.roll(p, -1)
        out.append(float((p * k ** 2).sum()))
    return np.array(out)


# =========================================================================== Section 11.3
def gaussian_conditional(mu, Lam, z, i):
    """Mean and sd of z_i given the other coordinates under N(mu, Lam^{-1}) (Section 2.3.1):
    mean mu_i - (1/Lam_ii) sum_{j != i} Lam_ij (z_j - mu_j), variance 1/Lam_ii."""
    d = z - mu
    m = mu[i] - (Lam[i] @ d - Lam[i, i] * d[i]) / Lam[i, i]
    return m, 1.0 / math.sqrt(Lam[i, i])


def overrelaxed_gibbs(rng, mu, Sigma, z0, n_sweeps, alpha):
    """Over-relaxed Gibbs sampling for a Gaussian (Adler 1981; (11.50), with the erratum alpha^2).

    Coordinates are updated in turn (11.46)-(11.48); with conditional mean m_i and sd s_i,
    z_i' = m_i + alpha (z_i - m_i) + s_i (1 - alpha^2)^{1/2} nu, nu ~ N(0, 1) (one normal per update).
    alpha = 0 is ordinary Gibbs sampling. Returns dict(chain: state after every sweep,
    path: state after every single-coordinate update, for drawing the staircase).
    """
    rng = _rng(rng)
    mu = np.asarray(mu, dtype=float)
    Lam = np.linalg.inv(Sigma)
    z = np.array(z0, dtype=float)
    chain, path = [z.copy()], [z.copy()]
    for _ in range(n_sweeps):
        for i in range(len(z)):
            m, s = gaussian_conditional(mu, Lam, z, i)
            z[i] = m + alpha * (z[i] - m) + s * math.sqrt(1 - alpha * alpha) * rng.standard_normal()
            path.append(z.copy())
        chain.append(z.copy())
    return dict(chain=np.array(chain), path=np.array(path))


def gibbs_gaussian(rng, mu, Sigma, z0, n_sweeps):
    """Gibbs sampling (Section 11.3) for N(mu, Sigma): each z_i in turn is redrawn from p(z_i | z_\\i)."""
    return overrelaxed_gibbs(rng, mu, Sigma, z0, n_sweeps, 0.0)


def ordered_overrelaxed_gibbs(rng, mu, Sigma, z0, n_sweeps, K):
    """Ordered over-relaxation (Neal 1999) for a Gaussian: K draws from p(z_i | z_\\i) are pooled with z_i,
    sorted, and z_i moves to the value whose rank mirrors its own (rank r -> rank K - r).
    K normals per update; K = 1 behaves like Gibbs sampling with a reflection, large K like alpha -> -1."""
    rng = _rng(rng)
    mu = np.asarray(mu, dtype=float)
    Lam = np.linalg.inv(Sigma)
    z = np.array(z0, dtype=float)
    chain, path = [z.copy()], [z.copy()]
    for _ in range(n_sweeps):
        for i in range(len(z)):
            m, s = gaussian_conditional(mu, Lam, z, i)
            pool = np.sort(np.append(m + s * rng.standard_normal(K), z[i]))
            r = int(np.searchsorted(pool, z[i]))
            z[i] = pool[K - r]
            path.append(z.copy())
        chain.append(z.copy())
    return dict(chain=np.array(chain), path=np.array(path))


# =========================================================================== Section 11.4
def slice_sample(rng, log_p_tilde, z0, n, w, max_steps=200, record=False):
    """Univariate slice sampling (Section 11.4; Neal 2003) with stepping out and shrinking.

    Per iteration: u = U1 p_tilde(z) defines the slice {z: p_tilde(z) > u} (11.51)-(11.52); an
    interval of width w is placed at random around z (U2); each end is moved out by w until it
    lies outside the slice; candidates are drawn uniformly from the interval (one uniform each)
    and the interval is shrunk to the candidate's side whenever a candidate falls outside the slice.
    Returns dict(chain, evals per iteration, trace (if record): per iteration the level, the
    initial interval, the stepped-out interval and the list of (candidate, accepted)).
    """
    rng = _rng(rng)
    z = float(z0)
    lz = log_p_tilde(z)
    chain, evals, trace = [z], [], []
    for _ in range(n):
        lu = lz + math.log(max(rng.random(), 1e-300))
        a = z - w * rng.random()
        b = a + w
        ne = 2
        init = (a, b)
        k = 0
        while log_p_tilde(a) > lu and k < max_steps:
            a -= w
            ne += 1
            k += 1
        k = 0
        while log_p_tilde(b) > lu and k < max_steps:
            b += w
            ne += 1
            k += 1
        out = (a, b)
        cands = []
        while True:
            zc = a + (b - a) * rng.random()
            lc = log_p_tilde(zc)
            ne += 1
            if lc > lu:
                cands.append((zc, True))
                z, lz = zc, lc
                break
            cands.append((zc, False))
            if zc < z:
                a = zc
            else:
                b = zc
        chain.append(z)
        evals.append(ne)
        if record:
            trace.append(dict(level=math.exp(lu), init=init, out=out, cands=cands))
    return dict(chain=np.array(chain), evals=np.array(evals), trace=trace)


# =========================================================================== Section 11.5
def hamiltonian(E, z, r):
    """H(z, r) = E(z) + K(r) with K(r) = |r|^2 / 2 (11.56)-(11.57)."""
    return E(z) + 0.5 * float(np.dot(r, r))


def leapfrog(grad_E, z, r, eps, L):
    """L leapfrog steps (11.64)-(11.66), each a half step of r, a full step of z and a half step of r.

    Written step by step (not merging consecutive half steps) so that the trajectory, including the
    momentum at the integer times, can be drawn. Returns arrays zs, rs of shape (L+1, d).
    """
    z, r = np.array(z, dtype=float), np.array(r, dtype=float)
    zs, rs = [z.copy()], [r.copy()]
    for _ in range(L):
        r = r - 0.5 * eps * grad_E(z)   # (11.64)
        z = z + eps * r                 # (11.65)
        r = r - 0.5 * eps * grad_E(z)   # (11.66)
        zs.append(z.copy())
        rs.append(r.copy())
    return np.array(zs), np.array(rs)


def hmc(rng, E, grad_E, z0, n, eps, L, jitter=0.0, keep_paths=False):
    """Hybrid Monte Carlo (Section 11.5.2; Duane et al. 1987).

    Per iteration: d normals for the momentum r ~ N(0, I) (Exercise 11.16), one uniform for the
    direction (step +eps if u < 1/2, else -eps; this makes the proposal symmetric), one uniform for
    a random step size eps (1 + jitter (2u - 1)) if jitter > 0, then L leapfrog steps, then one
    uniform u: the end point is accepted if u < exp{H(z, r) - H(z*, r*)} (11.67).
    Returns dict(chain, acc, dH (= H* - H), paths (if keep_paths)).
    """
    rng = _rng(rng)
    z = np.array(z0, dtype=float)
    chain, acc, dH, paths = [z.copy()], [], [], []
    for _ in range(n):
        r = rng.standard_normal(len(z))
        sgn = 1.0 if rng.random() < 0.5 else -1.0
        e = eps * (1 + jitter * (2 * rng.random() - 1)) if jitter > 0 else eps
        zs, rs = leapfrog(grad_E, z, r, sgn * e, L)
        d = hamiltonian(E, zs[-1], rs[-1]) - hamiltonian(E, z, r)
        ok = math.log(max(rng.random(), 1e-300)) < -d
        acc.append(ok)
        dH.append(d)
        if keep_paths:
            paths.append(zs)
        if ok:
            z = zs[-1].copy()
        chain.append(z.copy())
    return dict(chain=np.array(chain), acc=np.array(acc), dH=np.array(dH), paths=paths)


# =========================================================================== Section 11.6
def partition_ratio(E, G, samples_G):
    """Importance-sampling estimate of Z_E / Z_G (11.72): the mean of exp(-E(z_l) + G(z_l)) over z_l ~ p_G.

    The book omits the factor 1/L on the last line (erratum). Returns (estimate, standard error).
    """
    t = np.exp(-E(samples_G) + G(samples_G))
    return float(t.mean()), float(t.std(ddof=1) / math.sqrt(len(t)))


def chained_partition_ratio(E1, EM, sampler, alphas, L):
    """Chaining (11.74)-(11.75): Z_M / Z_1 as a product of ratios Z_{j+1} / Z_j between the
    distributions with energies E_alpha = (1 - alpha) E1 + alpha EM, each estimated by (11.72)
    from L samples of p_{alpha_j} drawn by sampler(alpha, L). Returns (estimate, list of factors)."""
    factors = []
    for a0, a1 in zip(alphas[:-1], alphas[1:]):
        z = sampler(a0, L)
        Ea = lambda x, a=a0: (1 - a) * E1(x) + a * EM(x)
        Eb = lambda x, a=a1: (1 - a) * E1(x) + a * EM(x)
        factors.append(partition_ratio(Eb, Ea, z)[0])
    return float(np.prod(factors)), factors


# =========================================================================== diagnostics
def autocorrelation(x):
    """Sample autocorrelation function of a 1-D series (FFT, biased estimator)."""
    x = np.asarray(x, dtype=float) - np.mean(x)
    n = len(x)
    f = np.fft.rfft(x, 2 * n)
    ac = np.fft.irfft(f * np.conj(f))[:n]
    return ac / ac[0]


def integrated_autocorr_time(x, c=5.0):
    """tau = 1 + 2 sum_{t >= 1} rho_t, summed up to Sokal's automatic window (the first M >= c tau(M)).

    For a 2-D array (chains in columns) the autocorrelations are averaged over chains first.
    """
    x = np.asarray(x, dtype=float)
    rho = autocorrelation(x) if x.ndim == 1 else np.mean([autocorrelation(col) for col in x.T], axis=0)
    tau = 2 * np.cumsum(rho) - 1
    for M in range(1, len(tau)):
        if M >= c * tau[M]:
            return float(tau[M])
    return float(tau[-1])


def effective_sample_size(x):
    x = np.asarray(x, dtype=float)
    return x.shape[0] * (1 if x.ndim == 1 else x.shape[1]) / integrated_autocorr_time(x)


# =========================================================================== plotting helpers
GREEN_BOOK = "#22b822"


def _arrow(ax, p, q, color="k", lw=1.6, both=True):
    ax.add_patch(FancyArrowPatch(p, q, arrowstyle="<|-|>" if both else "-|>", mutation_scale=12, color=color, lw=lw))


def _gauss2_logpdf(mu, Sigma):
    P = np.linalg.inv(Sigma)
    c = -LOG2PI - 0.5 * np.log(np.linalg.det(Sigma))
    def f(Z):
        D = np.asarray(Z, dtype=float) - mu
        return c - 0.5 * np.einsum("...i,ij,...j->...", D, P, D)
    return f


def corr_gaussian(r, s=1.0, mu=(0.0, 0.0)):
    return np.asarray(mu, dtype=float), s * s * np.array([[1.0, r], [r, 1.0]])


# =========================================================================== figures
def fig11_2(out_dir=None):
    """Figure 11.2: the transformation method y = h^{-1}(z), with exponential and Cauchy examples (11.7), (11.8)."""
    rng = np.random.default_rng(2)
    fig, ax = common.panel_grid(1, 3, size=(3.6, 2.9))
    a = ax[0, 0]
    y = np.linspace(-5, 5.5, 600)
    p, h = bimodal_pdf(y), bimodal_cdf(y)
    scale = 0.95 / p.max() * 0.7
    a.plot(y, p * scale, color=common.RED, lw=2, label=r"$p(y)$ (scaled)")
    a.plot(y, h, color=common.BLUE, lw=2, label=r"$h(y)=\int_{-\infty}^y p$")
    a.axhline(1, color="k", ls="--", lw=1)
    for zu in [0.15, 0.4, 0.62, 0.9]:
        yy = float(inverse_cdf(zu, bimodal_cdf, -10, 10))
        a.plot([-5, yy, yy], [zu, zu, 0], color=common.GREY, lw=0.9, ls=":")
        a.plot([yy], [0], "v", color=common.GREY, ms=4)
    a.set_xlim(-5, 5.5)
    a.set_ylim(0, 1.08)
    a.set_xlabel("$y$")
    a.set_title("(a) $y = h^{-1}(z)$, $z$ uniform")
    a.legend(loc="center right", fontsize=8, frameon=False)
    u = rng.random(20000)
    yb = inverse_cdf(u, bimodal_cdf, -10, 10, tol=1e-9)
    a.hist(yb, bins=80, density=True, weights=np.full(len(yb), scale), color=common.RED, alpha=0.18)
    b = ax[0, 1]
    ye = exponential_from_uniform(u, 1.0)
    b.hist(ye, bins=np.linspace(0, 6, 61), density=True, color=common.BLUE, alpha=0.35, label="20 000 samples")
    t = np.linspace(0, 6, 300)
    b.plot(t, np.exp(-t), color=common.BLUE, lw=2, label=r"$\lambda e^{-\lambda y}$, $\lambda=1$")
    b.set_xlim(0, 6)
    b.set_xlabel("$y = -\\ln(1-z)/\\lambda$")
    b.set_title("(b) exponential (11.7)")
    b.legend(fontsize=8, frameon=False)
    c = ax[0, 2]
    yc = cauchy_from_uniform(u)
    c.hist(yc, bins=np.linspace(-8, 8, 81), density=True, color=common.GREEN, alpha=0.35,
           weights=np.full(len(yc), 1.0), label="20 000 samples")
    t = np.linspace(-8, 8, 400)
    c.plot(t, 1 / (np.pi * (1 + t * t)), color=common.GREEN, lw=2, label=r"$\frac{1}{\pi}\frac{1}{1+y^2}$")
    c.set_xlim(-8, 8)
    c.set_xlabel(r"$y = \tan(\pi(z - 1/2))$")
    c.set_title("(c) Cauchy (11.8)")
    c.legend(fontsize=8, frameon=False)
    # The Cauchy histogram is normalised over all draws, so the bars integrate to P(|y| < 8) = 0.92.
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-2", out_dir)


def fig11_3(out_dir=None):
    """Figure 11.3: the polar Box-Muller method: uniform points in the unit disc become pairs of Gaussians (11.10)-(11.12)."""
    rng = np.random.default_rng(3)
    z = 2 * rng.random((1500, 2)) - 1
    r2 = (z ** 2).sum(1)
    inside = r2 < 1
    y = z[inside] * np.sqrt(-2 * np.log(r2[inside]) / r2[inside])[:, None]
    fig, ax = common.panel_grid(1, 3, size=(3.2, 3.0))
    a = ax[0, 0]
    a.add_patch(plt.Circle((0, 0), 1, color=common.GREEN, alpha=0.18))
    a.add_patch(plt.Circle((0, 0), 1, fill=False, color=GREEN_BOOK, lw=2))
    a.plot([-1, 1, 1, -1, -1], [-1, -1, 1, 1, -1], color=common.BLUE, lw=2)
    a.plot(z[inside, 0], z[inside, 1], ".", color="#1a7f1a", ms=2.2, rasterized=True)
    a.plot(z[~inside, 0], z[~inside, 1], "x", color=common.GREY, ms=2.5, mew=0.7, rasterized=True)
    a.set_aspect("equal")
    a.set_xlim(-1.05, 1.05)
    a.set_ylim(-1.05, 1.05)
    a.set_xlabel("$z_1$")
    a.set_ylabel("$z_2$")
    a.set_title("(a) kept %d of %d pairs (%.3f; $\\pi/4$ = %.3f)" % (inside.sum(), len(z), inside.mean(), np.pi / 4), fontsize=9)
    b = ax[0, 1]
    b.plot(y[:, 0], y[:, 1], ".", color=common.BLUE, ms=2.2, rasterized=True)
    b.set_aspect("equal")
    b.set_xlim(-4, 4)
    b.set_ylim(-4, 4)
    b.set_xlabel("$y_1$")
    b.set_ylabel("$y_2$")
    b.set_title("(b) the transformed pairs (11.10), (11.11)", fontsize=9)
    c = ax[0, 2]
    big, _ = box_muller(np.random.default_rng(33), 40000)
    c.hist(big, bins=np.linspace(-4, 4, 65), density=True, color=common.BLUE, alpha=0.35, label="40 000 values")
    t = np.linspace(-4, 4, 300)
    c.plot(t, gauss_pdf(t, 0, 1), color=common.RED, lw=2, label="$N(y|0,1)$")
    c.set_xlim(-4, 4)
    c.set_title("(c) mean %.3f, variance %.3f" % (big.mean(), big.var()), fontsize=9)
    c.legend(fontsize=7.5, frameon=False, loc="upper left")
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-3", out_dir)


FIG4_Q = (0.0, 2.2)


def fig4_setup():
    """Proposal q = N(0, 2.2^2) for the bimodal p_tilde, and the smallest k (Figure 11.4)."""
    mq, sq = FIG4_Q
    log_q = lambda z: np.log(gauss_pdf(z, mq, sq))
    log_pt = lambda z: np.log(np.maximum(bimodal_p_tilde(z), 1e-300))
    log_k = max_log_ratio(log_pt, log_q, -12, 12)
    return dict(mq=mq, sq=sq, log_q=log_q, log_pt=log_pt, log_k=log_k, k=math.exp(log_k),
                sample=lambda rng: mq + sq * rng.standard_normal())


def fig11_4(out_dir=None):
    """Figure 11.4: rejection sampling of the bimodal p_tilde under the comparison function k q(z)."""
    S = fig4_setup()
    rng = np.random.default_rng(4)
    xs, tries, rec = rejection_sample(rng, 3000, lambda z: float(S["log_pt"](np.array(z))), S["sample"],
                                      lambda z: float(S["log_q"](np.array(z))), S["log_k"])
    z = np.linspace(-7, 7, 600)
    kq, pt = S["k"] * gauss_pdf(z, S["mq"], S["sq"]), bimodal_p_tilde(z)
    fig, ax = common.panel_grid(1, 2, size=(4.3, 3.1))
    a = ax[0, 0]
    a.fill_between(z, pt, kq, color="0.8")
    a.plot(z, kq, color=common.BLUE, lw=2, label="$kq(z)$")
    a.plot(z, pt, color=common.RED, lw=2, label=r"$\widetilde{p}(z)$")
    R = np.array([(r[0], r[1]) for r in rec[:400]])
    K = np.array([r[2] for r in rec[:400]])
    a.plot(R[K, 0], R[K, 1], ".", color=common.GREEN, ms=3, label="kept")
    a.plot(R[~K, 0], R[~K, 1], ".", color="k", ms=2.2, label="rejected")
    z0, kq0 = -0.9, S["k"] * gauss_pdf(-0.9, S["mq"], S["sq"])
    a.plot([z0, z0], [0, kq0], color="k", lw=1.3)
    a.text(z0, kq0 + 0.03, "$kq(z_0)$", ha="center", fontsize=9)
    a.set_xlim(-7, 7)
    a.set_ylim(0, None)
    a.set_xlabel("$z$")
    a.set_title("(a) pairs $(z_0, u_0)$, first 400 tries", fontsize=9.5)
    a.legend(fontsize=7.5, loc="upper right", frameon=False)
    b = ax[0, 1]
    b.hist(xs, bins=np.linspace(-6, 6, 61), density=True, color=common.GREEN, alpha=0.4, label="3000 kept samples")
    b.plot(z, bimodal_pdf(z), color=common.RED, lw=2, label="$p(z)$")
    b.set_xlim(-6, 6)
    b.set_xlabel("$z$")
    acc = len(xs) / tries
    b.set_title("(b) acceptance %.3f; (11.14): $\\frac{1}{k}\\int\\widetilde{p}$ = %.3f" % (acc, BIMODAL_Z / S["k"]), fontsize=9.5)
    b.legend(fontsize=8, frameon=False)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-4", out_dir)


def fig11_5(out_dir=None):
    """Figure 11.5: Gam(z|10, 1) (green) under the scaled Cauchy comparison function (red), and the accepted samples."""
    a_, env = 10.0, gamma_cauchy_envelope(10.0)
    rng = np.random.default_rng(5)
    xs, tries, _ = rejection_sample(rng, 20000, lambda z: float(gamma_logpdf(z, a_)), env["sample"],
                                    lambda z: float(env["log_q"](z)), env["log_k"])
    z = np.linspace(0, 30, 600)
    fig, ax = common.panel_grid(1, 2, size=(4.2, 3.0))
    a = ax[0, 0]
    a.plot(z, env["k"] * np.exp(env["log_q"](z)), color="#ff1a1a", lw=2.2, label=r"$kq(z)$, Cauchy (11.16)")
    a.plot(z, np.exp(gamma_logpdf(z, a_)), color="#22d022", lw=2.2, label=r"Gam$(z|10,1)$")
    a.set_xlim(0, 30)
    a.set_ylim(0, 0.15)
    a.set_xlabel("$z$")
    a.set_ylabel("$p(z)$")
    a.set_title("(a) $c = 9$, $b^2 = 19$, $k$ = %.3f" % env["k"], fontsize=9.5)
    a.legend(fontsize=8, frameon=False)
    b = ax[0, 1]
    b.hist(xs, bins=np.linspace(0, 30, 61), density=True, color="#22d022", alpha=0.35, label="20 000 accepted")
    b.plot(z, np.exp(gamma_logpdf(z, a_)), color="#1a9f1a", lw=2)
    b.set_xlim(0, 30)
    b.set_xlabel("$z$")
    b.set_title("(b) acceptance %.4f, $1/k$ = %.4f" % (len(xs) / tries, 1 / env["k"]), fontsize=9.5)
    b.legend(fontsize=8, frameon=False)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-5", out_dir)


def gamma_ars_functions(a=10.0):
    h = lambda z: (a - 1) * np.log(z) - z      # ln p_tilde for Gam(z|a, 1), up to a constant
    dh = lambda z: (a - 1) / np.asarray(z, dtype=float) - 1
    return h, dh


def fig11_6(out_dir=None):
    """Figure 11.6: adaptive rejection sampling of Gam(z|10, 1): the tangent envelope of ln p and how it tightens."""
    a_ = 10.0
    h, dh = gamma_ars_functions(a_)
    x0 = [4.0, 8.5, 14.0]
    rng = np.random.default_rng(6)
    xs, tries, hist = adaptive_rejection_sample(rng, 2000, h, dh, x0, lo=0.0)
    z = np.linspace(0.3, 26, 800)
    Zt = math.exp(special.gammaln(a_))  # int exp(h) dz = Gamma(a)
    fig, ax = common.panel_grid(1, 3, size=(3.5, 2.9))
    a = ax[0, 0]
    H0 = ars_hull(x0, h, dh, 0.0)
    a.plot(z, h(z), color=common.RED, lw=2, label=r"$\ln\widetilde{p}(z)$")
    a.plot(z, ars_upper(H0, z), color=common.BLUE, lw=2, label="tangent envelope")
    for xi in x0:
        a.plot([xi, xi], [h(z).min() - 3, h(xi)], color=common.BLUE, ls="--", lw=1.2)
    a.set_ylim(h(z).min() - 3, h(z).max() + 2)
    a.set_xticks(x0)
    a.set_xticklabels(["$z_1$", "$z_2$", "$z_3$"])
    a.set_title("(a) three tangents (11.17)", fontsize=9.5)
    a.legend(fontsize=7.5, frameon=False, loc="lower center")
    b = ax[0, 1]
    b.plot(z, np.exp(h(z)) / Zt, color=common.RED, lw=2, label="Gam$(z|10,1)$")
    for xsset, col, ls in [(x0, common.BLUE, "-"), (hist[min(3, len(hist) - 1)], common.PURPLE, "--"), (hist[-1], common.GREEN, ":")]:
        Hh = ars_hull(xsset, h, dh, 0.0)
        mass = np.exp(Hh["logm"]).sum()
        b.plot(z, np.exp(ars_upper(Hh, z)) / Zt, color=col, ls=ls, lw=1.6,
               label="%d points: area %.3f" % (len(xsset), mass / Zt))
    b.set_xlim(0, 26)
    b.set_xlabel("$z$")
    b.set_title("(b) envelope in density units", fontsize=9.5)
    b.legend(fontsize=7.5, frameon=False)
    c = ax[0, 2]
    npts = [len(s) for s in hist]
    acc = [Zt / np.exp(ars_hull(s, h, dh, 0.0)["logm"]).sum() for s in hist]
    c.step(npts, acc, where="post", color=common.BLUE, lw=2)
    c.plot(npts, acc, "o", color=common.BLUE, ms=3)
    c.set_ylim(0.5, 1.01)
    c.set_xlabel("number of abscissae")
    c.set_ylabel("acceptance probability")
    c.set_title("(c) %d draws: %d tries, %d rejections" % (len(xs), tries, tries - len(xs)), fontsize=9.5)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-6", out_dir)


def highdim_acceptance(rng, D, ratio, n):
    """Monte Carlo acceptance of rejection sampling N(0, I) by k N(0, ratio^2 I) in D dimensions.

    The acceptance probability of z is p(z) / (k q(z)) = exp(-|z|^2 (1 - 1/ratio^2) / 2), a function
    of |z|^2 only, and |z|^2 / ratio^2 is chi-squared with D degrees of freedom under q.
    """
    R2 = ratio ** 2 * rng.chisquare(D, n)
    return float(np.mean(rng.random(n) < np.exp(-0.5 * R2 * (1 - 1 / ratio ** 2))))


def fig11_7(out_dir=None):
    """Figure 11.7: rejection sampling of N(0, 1) with k N(0, 1.1^2), and the acceptance (sigma_p/sigma_q)^D against D."""
    sp, sq = 1.0, 1.1
    z = np.linspace(-5, 5, 500)
    k = sq / sp
    fig, ax = common.panel_grid(1, 2, size=(4.2, 3.0))
    a = ax[0, 0]
    a.plot(z, gauss_pdf(z, 0, sp), color="#22d022", lw=2.2, label=r"$p(z)$, $\sigma_p = 1$")
    a.plot(z, k * gauss_pdf(z, 0, sq), color="#ff1a1a", lw=2.2, label=r"$kq(z)$, $\sigma_q = 1.1$")
    a.set_xlim(-5, 5)
    a.set_ylim(0, 0.5)
    a.set_yticks([0, 0.25, 0.5])
    a.set_xlabel("$z$")
    a.set_ylabel("$p(z)$")
    a.set_title("(a) $D = 1$: acceptance $1/k$ = %.3f" % (1 / k), fontsize=9.5)
    a.legend(fontsize=8, frameon=False)
    b = ax[0, 1]
    rng = np.random.default_rng(7)
    Ds = np.array([1, 2, 5, 10, 20, 50, 100, 200, 500, 1000])
    for ratio, col in [(1.01, common.BLUE), (1.1, common.RED)]:
        Dl = np.linspace(1, 1000, 400)
        b.semilogy(Dl, ratio ** -Dl, color=col, lw=1.8, label=r"$(\sigma_p/\sigma_q)^D$, $\sigma_q/\sigma_p$ = %g" % ratio)
        mc = np.array([highdim_acceptance(rng, D, ratio, 2_000_000) for D in Ds])
        b.semilogy(Ds[mc > 0], mc[mc > 0], "o", color=col, ms=4)
    b.annotate("1/%d at $D$ = 1000" % round(1.01 ** 1000), (1000, 1.01 ** -1000), (330, 2e-6), fontsize=8,
               arrowprops=dict(arrowstyle="->", lw=0.8))
    b.set_ylim(1e-7, 2)
    b.set_xlabel("dimension $D$")
    b.set_ylabel("acceptance rate")
    b.set_title("(b) theory (lines) and simulation (dots)", fontsize=9.5)
    b.legend(fontsize=7.5, frameon=False, loc="upper right")
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-7", out_dir)


FIG8_Q = (0.0, 2.0)


def fig8_f(z):
    return np.tanh(0.6 * (np.asarray(z, dtype=float) + 1))


def fig11_8(out_dir=None):
    """Figure 11.8: importance sampling: p (red), q (green), f (blue), the weights, the estimate and the effective sample size."""
    from scipy import integrate
    truth = integrate.quad(lambda t: fig8_f(t) * bimodal_pdf(t), -15, 15, limit=200)[0]
    log_pt = lambda z: np.log(np.maximum(bimodal_p_tilde(z), 1e-300))
    z = np.linspace(-6.5, 6.5, 600)
    fig, ax = common.panel_grid(1, 3, size=(3.6, 3.0))
    a = ax[0, 0]
    mq, sq = FIG8_Q
    rng = np.random.default_rng(8)
    zs = mq + sq * rng.standard_normal(30)
    res = importance_sampling(zs, log_pt, lambda t: np.log(gauss_pdf(t, mq, sq)), fig8_f)
    a.plot(z, bimodal_pdf(z), color=common.RED, lw=2, label="$p(z)$")
    a.plot(z, gauss_pdf(z, mq, sq), color=GREEN_BOOK, lw=2, label="$q(z)$")
    a.plot(z, 0.3 * fig8_f(z), color=common.BLUE, lw=2, label="$f(z)$ (scaled)")
    a.vlines(zs, 0, res["w"] * 1.2, color="k", lw=1.2)
    a.plot(zs, res["w"] * 1.2, "o", color="k", ms=2.5, label="weights $w_l$ (30 draws)")
    a.axhline(0, color="k", lw=0.8)
    a.set_xlim(-6.5, 6.5)
    a.set_xlabel("$z$")
    a.set_title("(a) ESS = %.1f of 30" % res["ess"], fontsize=9.5)
    a.legend(fontsize=7, frameon=False, loc="lower right")
    b = ax[0, 1]
    Ls = np.unique(np.logspace(1, 4.3, 40).astype(int))
    for (m_, s_), col, lab in [((0.0, 2.0), GREEN_BOOK, "q = N(0, 2$^2$)"), ((2.0, 0.8), common.MAGENTA, "q = N(2, 0.8$^2$)")]:
        r2 = np.random.default_rng(81)
        big = m_ + s_ * r2.standard_normal(Ls[-1])
        est = [importance_sampling(big[:L], log_pt, lambda t: np.log(gauss_pdf(t, m_, s_)), fig8_f)["estimate"] for L in Ls]
        b.semilogx(Ls, est, color=col, lw=1.8, label=lab)
    b.axhline(truth, color="k", ls="--", lw=1.2, label="E[f] = %.3f" % truth)
    b.set_xlabel("number of samples $L$")
    b.set_title("(b) self-normalised estimate (11.22)", fontsize=9.5)
    b.legend(fontsize=7.5, frameon=False)
    c = ax[0, 2]
    sds = np.linspace(0.5, 6, 45)
    L = 2000
    for m_, col in [(0.0, GREEN_BOOK), (2.0, common.MAGENTA)]:
        r3 = np.random.default_rng(82)
        e = []
        for s_ in sds:
            zz = m_ + s_ * r3.standard_normal(L)
            e.append(importance_sampling(zz, log_pt, lambda t: np.log(gauss_pdf(t, m_, s_)), fig8_f)["ess"] / L)
        c.plot(sds, e, color=col, lw=1.8, label="mean of q = %g" % m_)
    c.set_ylim(0, 1)
    c.set_xlabel("sd of $q$")
    c.set_ylabel("ESS / L")
    c.set_title("(c) effective sample size, $L$ = %d" % L, fontsize=9.5)
    c.legend(fontsize=7.5, frameon=False)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-8", out_dir)


def fig9_run(seed=16, n=150, step=0.2, z0=(1.0, 1.0)):
    return metropolis(np.random.default_rng(seed), _gauss2_logpdf(FIG9_MU, FIG9_SIGMA), np.array(z0), n, step)


def fig11_9(out_dir=None):
    """Figure 11.9: Metropolis sampling of a correlated 2-D Gaussian with an isotropic proposal of sd 0.2; 150 candidates."""
    res = fig9_run()
    fig, ax = plt.subplots(figsize=(4.4, 4.2))
    ex, ey = common.ellipse_xy(FIG9_MU, FIG9_SIGMA, nsd=2.0)
    ax.plot(ex, ey, color="k", lw=2)
    ex, ey = common.ellipse_xy(FIG9_MU, FIG9_SIGMA, nsd=1.0)
    ax.plot(ex, ey, color="k", lw=0.9, ls="--")
    ch, cand, acc = res["chain"], res["cand"], res["acc"]
    for t in range(len(acc)):
        p, q = ch[t], cand[t]
        ax.plot([p[0], q[0]], [p[1], q[1]], color="#22d022" if acc[t] else "#ff1a1a", lw=1.8)
    ax.set_xlim(0, 3)
    ax.set_ylim(0, 3)
    ax.set_aspect("equal")
    ax.set_title("150 candidates: %d accepted (green), %d rejected (red)" % (acc.sum(), (~acc).sum()), fontsize=9.5)
    ax.text(0.05, 2.9, "ellipses: 1 sd (dashed) and 2 sd", fontsize=8, va="top")
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-9", out_dir)


def rotated_gaussian(smax, smin, angle=np.pi / 4):
    R = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
    return R @ np.diag([smax ** 2, smin ** 2]) @ R.T


def rwm_iat(ratio, n=20000, chains=20, seed=10):
    """Integrated autocorrelation time of the major-axis coordinate for random-walk Metropolis with rho = sigma_min = 1."""
    S = rotated_gaussian(ratio, 1.0, 0.0)
    rng = np.random.default_rng(seed)
    Z0 = gaussian_cholesky(rng, np.zeros(2), S, chains)
    out, acc = metropolis_parallel(rng, _gauss2_logpdf(np.zeros(2), S), Z0, n, 1.0)
    return integrated_autocorr_time(out[1:, :, 0]), acc


def fig11_10(out_dir=None):
    """Figure 11.10: an isotropic proposal of scale rho ~ sigma_min on an elongated Gaussian: random-walk behaviour."""
    fig, ax = common.panel_grid(1, 3, size=(3.6, 3.1))
    a = ax[0, 0]
    smax, smin = 1.0, 0.18
    ang = np.deg2rad(28)
    S = rotated_gaussian(smax, smin, ang)
    ex, ey = common.ellipse_xy(np.zeros(2), S, nsd=1.0)
    a.plot(ex, ey, color=common.RED, lw=2)
    rho = smin
    c0 = np.array([-0.55, -0.27])
    a.add_patch(plt.Circle(c0, rho, fill=False, color=common.BLUE, lw=2))
    _arrow(a, c0, c0 + rho * np.array([np.cos(np.pi / 4), np.sin(np.pi / 4)]), common.BLUE, both=False)
    a.text(*(c0 + 1.25 * rho * np.array([0.9, 0.9])), r"$\rho$", fontsize=11)
    u, v = np.array([np.cos(ang), np.sin(ang)]), np.array([-np.sin(ang), np.cos(ang)])
    off = 0.45 * v
    _arrow(a, -smax * u + off, smax * u + off)
    a.text(*(0.15 * u + off + 0.08 * v), r"$\sigma_{\max}$", fontsize=11, ha="center")
    p0 = -smax * u - 0.12 * u
    _arrow(a, p0 - smin * v, p0 + smin * v)
    a.text(*(p0 - 0.26 * v - 0.12 * u), r"$\sigma_{\min}$", fontsize=11, ha="center")
    a.set_aspect("equal")
    a.set_xlim(-1.25, 1.25)
    a.set_ylim(-1.0, 1.0)
    a.set_xticks([])
    a.set_yticks([])
    a.set_title(r"(a) $\sigma_{\max}/\sigma_{\min}$ = %.1f, $\rho = \sigma_{\min}$" % (smax / smin), fontsize=9.5)
    b = ax[0, 1]
    ratio = 10.0
    S10 = rotated_gaussian(ratio, 1.0, 0.0)
    M, T = 2000, 3000
    rng = np.random.default_rng(11)
    out, accr = metropolis_parallel(rng, _gauss2_logpdf(np.zeros(2), S10), np.zeros((M, 2)), T, 1.0)
    msd = ((out[:, :, 0] - out[0, :, 0]) ** 2).mean(1)
    tt = np.arange(T + 1)
    b.loglog(tt[1:], np.sqrt(msd[1:]), color=common.BLUE, lw=2, label="RMS distance along the long axis")
    b.loglog(tt[1:], np.sqrt(accr * tt[1:]), color="k", ls="--", lw=1.2, label=r"$\rho\sqrt{a\tau}$, $a$ = acceptance rate")
    b.axhline(ratio, color=common.RED, ls=":", lw=1.4, label=r"$\sigma_{\max}$ (start at the mean)")
    b.axvline(ratio ** 2, color=common.GREY, lw=1, ls="-.")
    b.text(ratio ** 2 * 1.15, 0.7, r"$\tau = (\sigma_{\max}/\sigma_{\min})^2$", fontsize=8)
    b.set_xlabel(r"step $\tau$")
    b.set_title(r"(b) $\sigma_{\max}/\sigma_{\min}$ = 10, %d chains, acc. %.2f" % (M, accr), fontsize=9.5)
    b.legend(fontsize=7, frameon=False, loc="upper left")
    c = ax[0, 2]
    ratios = np.array([2, 4, 8, 16, 32])
    iat = [rwm_iat(r_, n=max(4000, int(120 * r_ ** 2)), chains=16)[0] for r_ in ratios]
    c.loglog(ratios ** 2, iat, "o-", color=common.BLUE, lw=1.8, label="measured IAT")
    c.loglog(ratios ** 2, iat[-1] / ratios[-1] ** 2 * ratios ** 2, "k--", lw=1.1, label="slope 1")
    c.set_xlabel(r"$(\sigma_{\max}/\sigma_{\min})^2$")
    c.set_ylabel("integrated autocorrelation time")
    c.set_title(r"(c) steps per independent sample, $\rho=\sigma_{\min}$", fontsize=9.5)
    c.legend(fontsize=7.5, frameon=False)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-10", out_dir)


def fig11_11(out_dir=None):
    """Figure 11.11: Gibbs sampling of a correlated Gaussian: a few sweeps, the conditional (green), widths l and L, and many samples."""
    r = 0.9
    mu, S = corr_gaussian(r, 1.0, (0.0, 0.0))
    res = gibbs_gaussian(np.random.default_rng(12), mu, S, np.array([-0.9, -1.3]), 4)
    fig, ax = common.panel_grid(1, 2, size=(4.0, 3.9))
    a = ax[0, 0]
    ex, ey = common.ellipse_xy(mu, S, nsd=1.0)
    a.plot(ex, ey, color=common.RED, lw=2)
    P = res["path"]
    a.plot(P[:, 0], P[:, 1], color=common.BLUE, lw=2)
    a.plot(P[0, 0], P[0, 1], "o", color=common.BLUE, ms=4)
    l, L = math.sqrt(1 - r * r), 1.0
    z2 = 0.95
    m1 = r * z2
    a.axhline(z2, color="#22d022", ls="--", lw=1.4)
    a.plot([m1, m1], [-2.4, z2], color="#22d022", ls="--", lw=1.4)
    t = np.linspace(m1 - 4 * l, m1 + 4 * l, 200)
    a.plot(t, -2.4 + 0.45 * np.exp(-0.5 * ((t - m1) / l) ** 2), color="#22d022", lw=2)
    _arrow(a, (m1 - l, -1.65), (m1 + l, -1.65))
    a.text(m1, -1.55, "$l$", ha="center", fontsize=11)
    _arrow(a, (-L, 2.15), (L, 2.15))
    a.text(0, 2.25, "$L$", ha="center", fontsize=11)
    a.set_xlim(-2.4, 2.4)
    a.set_ylim(-2.4, 2.6)
    a.set_aspect("equal")
    a.set_xlabel("$z_1$")
    a.set_ylabel("$z_2$")
    a.set_title(r"(a) correlation 0.9: $l = \sqrt{1-0.9^2}$ = %.2f, $L$ = 1" % l, fontsize=9)
    b = ax[0, 1]
    res2 = gibbs_gaussian(np.random.default_rng(13), mu, S, np.zeros(2), 3000)
    C = res2["chain"][100:]
    b.plot(C[:, 0], C[:, 1], ".", color=common.BLUE, ms=1.6, alpha=0.6, rasterized=True)
    for nsd in (1, 2):
        ex, ey = common.ellipse_xy(mu, S, nsd=nsd)
        b.plot(ex, ey, color=common.RED, lw=1.6)
    Ce = np.cov(C.T)
    b.set_xlim(-3.3, 3.3)
    b.set_ylim(-3.3, 3.3)
    b.set_aspect("equal")
    b.set_xlabel("$z_1$")
    b.set_title("(b) 2900 sweeps: var %.2f, %.2f; corr %.3f" % (Ce[0, 0], Ce[1, 1], Ce[0, 1] / np.sqrt(Ce[0, 0] * Ce[1, 1])), fontsize=9)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-11", out_dir)


def overrelaxed_sweep_matrix(r, alpha):
    """The linear part A of one over-relaxed Gibbs sweep (z1 then z2) for a zero-mean, unit-variance
    bivariate Gaussian with correlation r: z' = A z + noise. alpha = 0 is Gibbs sampling."""
    A1 = np.array([[alpha, (1 - alpha) * r], [0.0, 1.0]])
    A2 = np.array([[1.0, 0.0], [(1 - alpha) * r, alpha]])
    return A2 @ A1


def exact_iat_z1(r, alpha=0.0):
    """Exact integrated autocorrelation time of z1 for the sweep z' = A z + noise in equilibrium.

    The lag-k autocovariance is A^k Sigma, so tau = 1 + 2 sum_k (A^k Sigma)_11 / Sigma_11
    = [(2 (I - A)^{-1} - I) Sigma]_11 / Sigma_11. For alpha = 0 this is (1 + r^2) / (1 - r^2).
    """
    A = overrelaxed_sweep_matrix(r, alpha)
    S = np.array([[1.0, r], [r, 1.0]])
    return float(((2 * np.linalg.inv(np.eye(2) - A) - np.eye(2)) @ S)[0, 0])


def gibbs2_many(rng, r, Z, n, alpha=0.0, K=None):
    """Many bivariate Gibbs chains at once (rows of Z; zero mean, unit variances, correlation r):
    over-relaxed with alpha, or ordered over-relaxation with K draws if K is given. Returns z1 per sweep."""
    s = math.sqrt(1 - r * r)
    Z = np.array(Z, dtype=float)
    out = np.empty((n, len(Z)))
    for t in range(n):
        for i in (0, 1):
            m = r * Z[:, 1 - i]
            if K is None:
                Z[:, i] = m + alpha * (Z[:, i] - m) + s * math.sqrt(1 - alpha * alpha) * rng.standard_normal(len(Z))
            else:
                pool = np.sort(np.c_[m[:, None] + s * rng.standard_normal((len(Z), K)), Z[:, i]], axis=1)
                rank = (pool < Z[:, i:i + 1]).sum(1)
                Z[:, i] = pool[np.arange(len(Z)), K - rank]
        out[t] = Z[:, 0]
    return out


def iat_batch_means(X, nb=40):
    """IAT from batch means: tau = b var(batch means) / var(x), batches of length b = n / nb, pooled over chains (columns)."""
    n = (X.shape[0] // nb) * nb
    b = n // nb
    means = X[:n].reshape(nb, b, -1).mean(1)
    return float(b * means.var(ddof=1) / X.var())


def gibbs_iat(r, alpha=0.0, K=None, chains=100, seed=14):
    """Measured IAT (in sweeps) of z1 at correlation r, from 100 chains started in equilibrium."""
    rng = np.random.default_rng(seed)
    n = int(max(4000, 400 / (1 - r * r)))
    Z0 = gaussian_cholesky(rng, np.zeros(2), np.array([[1.0, r], [r, 1.0]]), chains)
    return iat_batch_means(gibbs2_many(rng, r, Z0, n, alpha, K), nb=10)


def fig11_12(out_dir=None):
    """Our Figure 11.12 (companion to the book's Figure 11.11): Gibbs crawls when l << L; steps needed grow like (L/l)^2."""
    fig, ax = common.panel_grid(1, 3, size=(3.6, 3.2))
    r = 0.99
    mu, S = corr_gaussian(r)
    a = ax[0, 0]
    res = gibbs_gaussian(np.random.default_rng(15), mu, S, np.array([-1.6, -1.6]), 60)
    ores = overrelaxed_gibbs(np.random.default_rng(15), mu, S, np.array([-1.6, -1.6]), 60, -0.98)
    ex, ey = common.ellipse_xy(mu, S)
    a.plot(ex, ey, color=common.RED, lw=2)
    a.plot(res["path"][:, 0], res["path"][:, 1], color=common.BLUE, lw=1.2, label="Gibbs")
    a.plot(ores["path"][:, 0], ores["path"][:, 1], color=common.GREEN, lw=1.0, alpha=0.8, label=r"over-relaxed, $\alpha=-0.98$")
    a.set_aspect("equal")
    a.set_xlim(-1.8, 1.8)
    a.set_ylim(-1.8, 1.8)
    a.set_xlabel("$z_1$")
    a.set_ylabel("$z_2$")
    a.set_title(r"(a) correlation 0.99: $L/l$ = %.1f, 60 sweeps" % (1 / math.sqrt(1 - r * r)), fontsize=9)
    a.legend(fontsize=7, frameon=False, loc="upper left")
    b = ax[0, 1]
    T = 400
    rg = gibbs_gaussian(np.random.default_rng(16), mu, S, np.zeros(2), T)["chain"][:, 0]
    ro = overrelaxed_gibbs(np.random.default_rng(16), mu, S, np.zeros(2), T, -0.98)["chain"][:, 0]
    rd = mu[0] + np.random.default_rng(16).standard_normal(T + 1)
    b.plot(rd, color="0.75", lw=0.8, label="independent draws")
    b.plot(rg, color=common.BLUE, lw=1.3, label="Gibbs")
    b.plot(ro, color=common.GREEN, lw=1.1, label="over-relaxed")
    b.set_xlabel("sweep")
    b.set_ylabel("$z_1$")
    b.set_title("(b) traces of $z_1$", fontsize=9.5)
    b.legend(fontsize=7, frameon=False, loc="lower right", ncol=1)
    c = ax[0, 2]
    rs = np.array([0.5, 0.8, 0.9, 0.95, 0.98, 0.99, 0.995])
    LL = 1 / (1 - rs ** 2)
    rf = np.sqrt(1 - 1 / np.logspace(0.05, 2.35, 80))
    c.loglog(1 / (1 - rf ** 2), [exact_iat_z1(x) for x in rf], color=common.BLUE, lw=1.4, label="Gibbs: exact")
    c.loglog(1 / (1 - rf ** 2), [exact_iat_z1(x, -0.9) for x in rf], color=common.GREEN, lw=1.4, label=r"over-relaxed $\alpha=-0.9$: exact")
    c.loglog(LL, [gibbs_iat(x) for x in rs], "o", color=common.BLUE, ms=4.5, label="measured")
    c.loglog(LL, [gibbs_iat(x, alpha=-0.9) for x in rs], "s", color=common.GREEN, ms=4)
    c.loglog(LL, [gibbs_iat(x, K=20) for x in rs], "^-", color=common.ORANGE, ms=4, lw=1, label="ordered, $K$ = 20 (measured)")
    c.loglog(LL, 2 * LL, "k:", lw=1, label=r"$2(L/l)^2$")
    c.set_xlabel(r"$(L/l)^2 = 1/(1-r^2)$")
    c.set_ylabel("IAT of $z_1$ (sweeps)")
    c.set_ylim(0.05, 600)
    c.set_title("(c) sweeps per independent sample", fontsize=9.5)
    c.legend(fontsize=6.8, frameon=False, loc="upper left")
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-12", out_dir)


def fig11_13(out_dir=None):
    """Figure 11.13: slice sampling of the bimodal p_tilde: (a) the slice, (b) stepping out and shrinking, (c) many samples."""
    log_pt = lambda z: float(np.log(max(bimodal_p_tilde(z), 1e-300)))
    z = np.linspace(-5, 6, 800)
    pt = bimodal_p_tilde(z)
    fig, ax = common.panel_grid(1, 3, size=(3.6, 2.9))
    zt = -1.55
    u = 0.33 * BIMODAL_W[0]
    a = ax[0, 0]
    a.plot(z, pt, color=common.RED, lw=2)
    a.plot([zt, zt], [0, bimodal_p_tilde(zt)], color="k", lw=1.6)
    a.plot([-5, 6], [u, u], color=common.BLUE, ls="--", lw=1.6)
    on = pt > u
    zz = np.where(on, z, np.nan)
    a.plot(zz, np.full_like(z, u), color=common.BLUE, lw=3)
    a.text(zt - 0.15, u + 0.02, "$u$", ha="right", fontsize=11)
    a.set_xticks([zt])
    a.set_xticklabels([r"$z^{(\tau)}$"])
    a.set_yticks([])
    a.set_ylim(0, None)
    a.set_title(r"(a) the slice $\{z:\widetilde{p}(z) > u\}$", fontsize=9.5)
    b = ax[0, 1]
    w = 0.5
    res = slice_sample(np.random.default_rng(0), log_pt, zt, 1, w, record=True)
    tr = res["trace"][0]
    lev = tr["level"]
    b.plot(z, pt, color=common.RED, lw=2)
    b.plot([zt, zt], [0, bimodal_p_tilde(zt)], color="k", lw=1.6)
    i0, i1 = tr["init"]
    o0, o1 = tr["out"]
    b.plot([i0, i1], [lev + 0.06, lev + 0.06], color=common.GREY, lw=3, solid_capstyle="butt")
    b.text(i1 + 0.12, lev + 0.06, "first interval, width $w$", va="center", fontsize=7.5)
    b.plot([o0, o1], [lev, lev], color=common.BLUE, lw=2)
    b.plot(np.arange(o0, o1 + 1e-9, w), np.full(int(round((o1 - o0) / w)) + 1, lev), "|", color=common.BLUE, ms=7, mew=1.5)
    b.text(o0 - 0.1, lev, r"$z_{\min}$", ha="right", va="center", fontsize=10)
    b.text(o1 + 0.1, lev, r"$z_{\max}$", ha="left", va="center", fontsize=10)
    for k, (zc, ok) in enumerate(tr["cands"]):
        b.plot([zc], [lev - 0.05 - 0.035 * k], "o" if ok else "x", color=common.GREEN if ok else "k", ms=5)
    b.set_xticks([zt])
    b.set_xticklabels([r"$z^{(\tau)}$"])
    b.set_yticks([])
    b.set_ylim(0, None)
    b.set_title("(b) %d steps out, then %d candidates\n(x outside the slice, o accepted)" % (round((o1 - o0) / w) - 1, len(tr["cands"])), fontsize=9)
    c = ax[0, 2]
    big = slice_sample(np.random.default_rng(70), log_pt, 0.0, 20000, 1.0)
    c.hist(big["chain"][1:], bins=np.linspace(-5, 6, 89), density=True, color=common.BLUE, alpha=0.35, label="20 000 slice samples")
    c.plot(z, bimodal_pdf(z), color=common.RED, lw=2, label="$p(z)$")
    c.set_xlim(-5, 6)
    c.set_title("(c) %.1f density evaluations per sample" % big["evals"].mean(), fontsize=9.5)
    c.legend(fontsize=7.5, frameon=False)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-13", out_dir)


def _shoelace(P):
    x, y = P[:, 0], P[:, 1]
    return 0.5 * abs(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1)))


def fig11_14(out_dir=None):
    """Figure 11.14: leapfrog steps shear phase space without changing volume; a trajectory and its energy error."""
    gE = lambda z: 2 * np.sin(2 * z)        # E(z) = 1 - cos(2 z)
    eps = 0.6
    # boundary of a square in (z, r), densely sampled
    t = np.linspace(0, 1, 200, endpoint=False)
    z0, z1, r0, r1 = 0.1, 2.3, 0.3, 1.5
    B = np.concatenate([np.c_[z0 + (z1 - z0) * t, np.full_like(t, r0)], np.c_[np.full_like(t, z1), r0 + (r1 - r0) * t],
                        np.c_[z1 - (z1 - z0) * t, np.full_like(t, r1)], np.c_[np.full_like(t, z0), r1 - (r1 - r0) * t]])
    Sz, Sr = 0.55, 0.95
    St = np.concatenate([np.c_[Sz + 0.4 * t, np.full_like(t, r0)], np.c_[np.full_like(t, Sz + 0.4), r0 + (r1 - r0) * t],
                         np.c_[Sz + 0.4 - 0.4 * t, np.full_like(t, r1)], np.c_[np.full_like(t, Sz), r1 - (r1 - r0) * t]])
    half = lambda P: np.c_[P[:, 0], P[:, 1] - 0.5 * eps * gE(P[:, 0])]
    full = lambda P: np.c_[P[:, 0] + eps * P[:, 1], P[:, 1]]
    fig, ax = common.panel_grid(1, 4, size=(3.1, 2.9))
    for k, (P, Q, lab) in enumerate([(B, St, "(a) before"), (half(B), half(St), "(b) after the half step (11.64)")]):
        a = ax[0, k]
        a.add_patch(Polygon(P, closed=True, fill=False, ec=common.RED, lw=2))
        a.add_patch(Polygon(Q, closed=True, fc="#8080ff", ec=common.RED, lw=2))
        a.set_xlim(-0.1, 2.6)
        a.set_ylim(-0.6, 2.4)
        a.set_xlabel("$z_i$" if k == 0 else "$z_i'$")
        a.set_ylabel("$r_i$" if k == 0 else "$r_i'$")
        a.set_title("%s\narea %.4f" % (lab, _shoelace(P)), fontsize=9)
    one = half(full(half(B)))
    ax[0, 1].text(2.5, -0.5, "after a whole leapfrog\nstep: area %.4f" % _shoelace(one), fontsize=7.5, ha="right", va="bottom")
    c = ax[0, 2]
    E = lambda z: 0.5 * float(np.dot(z, z))
    g = lambda z: z
    th = np.linspace(0, 2 * np.pi, 300)
    H0 = 0.5 * (1.0 ** 2 + 1.0 ** 2)
    c.plot(np.sqrt(2 * H0) * np.cos(th), np.sqrt(2 * H0) * np.sin(th), color="k", lw=1.2, label="exact flow ($H$ constant)")
    for e_, col in [(0.3, common.BLUE), (1.2, common.ORANGE)]:
        zs, rs = leapfrog(g, np.array([1.0]), np.array([1.0]), e_, int(round(2 * np.pi / e_)))
        c.plot(zs[:, 0], rs[:, 0], "o-", color=col, ms=3, lw=1.1, label=r"leapfrog $\epsilon$ = %g" % e_)
    c.set_aspect("equal")
    c.set_xlim(-2.1, 2.1)
    c.set_ylim(-2.1, 3.4)
    c.set_xlabel("$z$")
    c.set_ylabel("$r$")
    c.set_title(r"(c) $E(z) = z^2/2$, one period", fontsize=9)
    c.legend(fontsize=6.5, frameon=False, loc="upper center", ncol=1)
    d = ax[0, 3]
    epss = np.logspace(-2.3, 0, 20)
    err = []
    for e_ in epss:
        zs, rs = leapfrog(g, np.array([1.0]), np.array([1.0]), e_, int(math.ceil(2 * np.pi / e_)))
        err.append(max(abs(0.5 * zs[k, 0] ** 2 + 0.5 * rs[k, 0] ** 2 - H0) for k in range(len(zs))))
    d.loglog(epss, err, "o-", color=common.BLUE, ms=3.5, lw=1.4, label=r"max $|H - H_0|$ over a period")
    d.loglog(epss, err[0] / epss[0] ** 2 * epss ** 2, "k--", lw=1, label=r"slope 2: $O(\epsilon^2)$")
    d.set_xlabel(r"step size $\epsilon$")
    d.set_title("(d) energy error of leapfrog", fontsize=9)
    d.legend(fontsize=7, frameon=False)
    fig.tight_layout()
    common.savefig(fig, 11, "fig11-14", out_dir)


FIGURES = {
    "11.2": fig11_2, "11.3": fig11_3, "11.4": fig11_4, "11.5": fig11_5, "11.6": fig11_6, "11.7": fig11_7,
    "11.8": fig11_8, "11.9": fig11_9, "11.10": fig11_10, "11.11": fig11_11, "11.12": fig11_12,
    "11.13": fig11_13, "11.14": fig11_14,
}

if __name__ == "__main__":
    common.run(FIGURES, 11)