"""Bishop (2006), Chapter 5: Neural Networks.

Reproduces the numerical figures of Chapter 5 of C. M. Bishop, Pattern Recognition and Machine
Learning (Springer, 2006), pp. 225-290, written from scratch with numpy, scipy and matplotlib.
Every algorithm is a small function meant to be read next to the lecture notes
(ch05-neural-networks.html):

* Section 5.1, feed-forward network functions: Net (any number of layers of tanh hidden units;
  output units linear, logistic sigmoid or softmax (5.1)-(5.9); the weights and biases are one flat
  vector w, so w = (hidden-unit biases and weights, then output-unit biases and weights)).
* Sections 5.2-5.3, error functions and backpropagation: the sum-of-squares, cross-entropy and
  multiclass cross-entropy errors (5.11), (5.21), (5.24) all give the output delta y - t for the
  canonical link (5.18), (5.66), (5.69); Net.backward is (5.43)-(5.56); a mixture density
  network has its own error (5.153) and derivatives (5.155)-(5.157) (with the errata).
* Section 5.3.4, the Jacobian: Net.jacobian is (5.72)-(5.76).
* Section 5.4, the Hessian: hessian_diagonal ((5.79)-(5.81), the exact diagonal of Ricotti et al.
  and the Becker-Le Cun approximation), hessian_outer ((5.83)-(5.85), outer product / Levenberg-
  Marquardt), inverse_hessian_sequential ((5.86)-(5.89)), hessian_fd ((5.91)), hessian_exact
  (the two-layer formulas (5.92)-(5.95)), hessian_vector (the R-operator (5.96)-(5.108), for any
  depth) and hessian_rop (the full Hessian from W Hessian-vector products).
* Section 5.5, regularisation: weight decay (5.112), early stopping in a quadratic error
  (Figure 5.13, early_stopping_vs_weight_decay), warped digits (Figure 5.14, warp_digit) and
  tangent vectors (Figure 5.16, rotate_tangent).
* Section 5.6, mixture density networks: mdn_error, mdn_delta, mdn_parameters, mdn_mode.
* Section 5.7, Bayesian neural networks (Laplace approximation, MacKay's evidence framework):
  laplace_log_evidence_regression (5.175), evidence_regression (alternating a fit of w with
  (5.178)-(5.182)), predictive_regression ((5.172)-(5.173)), evidence_classification
  ((5.183)-(5.186)), predictive_classification ((5.188)-(5.190)).
* JSRandom (imported from ch11_sampling.py): the seeded generator DK.rng of the browser demos.
  data_sine, init_weights and adam_train below use only that generator's numbers and plain
  floating-point arithmetic, so backprop-demo.html, mixture-density-demo.html and
  bayesian-nn-demo.html reproduce them number for number (see the checks in the demos' pages).

Figures: 5.3, 5.4, 5.9, 5.10, 5.11, 5.12, 5.13, 5.14, 5.16, 5.19, 5.21, 5.22, 5.23 (the book's
other figures in this chapter are diagrams). Run `python3 ch05_neural_networks.py --list`,
`python3 ch05_neural_networks.py` (all figures, a few minutes) or
`python3 ch05_neural_networks.py --fig 5.3 5.21`. Files go to ../img/ch05/.

What differs from the book, and the settings chosen where the book gives none:
* Optimiser: the book uses scaled conjugate gradients; we use L-BFGS (scipy) for the figures, from
  several random starts when the error has bad local minima (we keep the best), and Adam or plain
  gradient descent in the demos. Any local minimiser reaches the same kind of solution.
* Figure 5.3: N = 50 points uniform on (-1, 1) (we use an equally spaced grid), no noise, three tanh hidden units, linear
  output. The caption says f(x) = sin(x), but on (-1, 1) that is almost a straight line; the book's panel (b) shows one
  full hump, so we use sin(2x), whose maximum is at x = 0.785;
  the dashed curves are the three hidden-unit outputs z_j(x) (the book draws them rescaled to fit).
* Figure 5.4: the synthetic classification data of Appendix A are not published, so we draw
  100 points (50 per class; JSRandom(3)) from our own version of its generative process: blue class
  N((-0.2, 0.0), diag(1.0^2, 1.2^2)), red class an equal mixture of N((1.0, 1.1), diag(0.6^2, 0.8^2))
  and N((0.9, -1.1), diag(0.6^2, 0.7^2)); see classification_data. The green curve is the 0.5
  contour of the true posterior. The dashed blue lines are a_j = 0 for each hidden unit, i.e.
  where z_j = tanh(a_j) = 0 (the book's caption says 0.5, which is not the midpoint of tanh).
* Figures 5.9 and 5.10: data from the sinusoid of Appendix A (x uniform on (0, 1), t = sin(2 pi x)
  plus noise with sd 0.3). Figure 5.9 uses the ten points x_n = (n - 1)/9 so that all three panels
  see the same data (as in the book), and the best of 8 starts of L-BFGS on the sum of squares plus a
  tiny weight decay (lambda = 1e-3), which only stops the weights of the over-flexible networks running
  off to huge values (without it the M = 3 and M = 10 fits contain near-vertical spikes between the points). Figure 5.10 uses 60 training points and a test set of 600 (L-BFGS, at most 1500 iterations, a few starts end far above the plotted range),
  reports the sum of squares sum (y - t)^2 on the test set (floor 600 x 0.09 = 54), and
  initialises each of the 30 starts per M with weights N(0, 10), as stated in the caption.
* Figure 5.11: the book's panel titles give the precisions with an erratum: every number on the
  right of '=' should be raised to the power -2. So panels are (a) all alphas 1, (b) alpha2w = 10^-2,
  (c) alpha1w = 10^-6, alpha1b = 10^-4, (d) alpha1w = alpha1b = 10^-6. 12 tanh hidden units, five
  prior draws per panel.
* Figure 5.12: 12 training and 12 validation points of the sinusoid (numpy seed 25), a network with 10 hidden
  units, plain full-batch gradient descent (step 0.004) from a small random start (N(0, 0.5^2));
  the curves show the RMS error sqrt(2 E / N) after every 400 gradient steps (50 values), and the dashed
  line marks the minimum of the validation error. (The book's units are not stated; RMS error has the right scale.)
* Figure 5.13: a quadratic error with Hessian eigenvalues 1 and 0.05 (read off the book's
  ellipse); the curve is the gradient-descent path from the origin, and the dotted curve the ridge
  path w(alpha) for weight decay. Both are computed in closed form.
* Figure 5.14: the digit is MNIST test image 0 of class 6 (MNIST digits are the book's 'off-line digits').
  The book smooths i.i.d. random displacements with Gaussians of width 0.01, 30 and 60 on a larger
  image. On our 28 x 28 images we use widths 0.01, 2.5 and 5 pixels, and rescale each field to
  an rms displacement of 1.2 pixels (0.6 pixels for the unsmoothed field), displacements drawn
  uniformly in (-1, 1) before smoothing; the warp is a cubic-spline resampling.
* Figure 5.16: tangent vector tau = d x(xi)/d xi at xi = 0 for rotation by xi (central finite
  difference of the cubic-spline rotated image), x + epsilon tau with epsilon = 15 degrees = 0.262
  rad, and the truly rotated image (clockwise, as in the book).
* Figures 5.19 and 5.21: x uniform on (0, 1), t = x + 0.3 sin(2 pi x) + noise uniform on
  (-0.1, 0.1), N = 250 points; the inverse problem swaps x and t. The network of 5.19 has five
  tanh hidden units and the sum-of-squares error; the mixture density network has three Gaussian
  components, five hidden units and 9 outputs (the book's own numbers). In panel (d) the conditional
  mode is approximated, as in the book, by the mean of the most probable component.
* Figures 5.22 and 5.23: our version of the classification data of Figure 5.4 with N = 100,
  a network with 8 tanh hidden units, a single hyperparameter alpha started at 0 as in the
  caption (we start at alpha = 0.01 to keep the Hessian well-conditioned), re-estimated by
  the evidence procedure (5.178)-(5.179) alternating with the fit of w, as in
  evidence_classification.
"""
import math

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from scipy import ndimage, optimize

import common
from ch11_sampling import JSRandom  # noqa: F401  (re-exported: the demos' generator)

LOG2PI = math.log(2 * math.pi)
OLIVE, MAG, LIME = "#b5a600", "#c73bc7", "#27d227"      # the dashed hidden-unit colours of Figure 5.3
SEED_DEMO = 5


def _u(rng):
    return rng.u() if isinstance(rng, JSRandom) else float(rng.random())


def _n(rng):
    return rng.norm() if isinstance(rng, JSRandom) else float(rng.standard_normal())


# =========================================================================== the network
class Net:
    """A feed-forward network with tanh hidden layers (Section 5.1).

    sizes = [D, M1, ..., K]: D inputs, hidden layers of M1, ... units, K output units.
    kind selects the error function and the output activation:
      'sse'  linear outputs, sum-of-squares (5.11)           delta_k = y_k - t_k          (5.18)
      'bce'  one logistic-sigmoid output, cross-entropy (5.21)   delta = y - t              (5.66)
      'ce'   softmax outputs, multiclass cross-entropy (5.24)    delta_k = y_k - t_k       (5.69)
      'mdn'  3 K raw outputs of a mixture density network, error (5.153)         (Section 5.6)
    The parameters are one flat vector w; layer l holds an array of shape (n_out, n_in + 1) whose
    column 0 is the bias, stored row by row: w = (W1.ravel(), W2.ravel(), ...). The bias is the
    weight w_j0 attached to an extra input x_0 = 1 (5.2).
    """

    def __init__(self, sizes, kind="sse", K_mix=None):
        self.sizes = list(sizes)
        self.kind = kind
        self.K_mix = K_mix
        self.L = len(sizes) - 1
        self.shapes = [(sizes[l + 1], sizes[l] + 1) for l in range(self.L)]
        self.offsets = np.concatenate([[0], np.cumsum([a * b for a, b in self.shapes])]).astype(int)
        self.W = int(self.offsets[-1])

    # ------------------------------------------------------------------ parameters
    def unpack(self, w):
        return [w[self.offsets[l]:self.offsets[l + 1]].reshape(self.shapes[l]) for l in range(self.L)]

    def pack(self, Ws):
        return np.concatenate([a.ravel() for a in Ws])

    # ------------------------------------------------------------------ forward propagation (5.48)-(5.49)
    def forward(self, w, X):
        """Forward propagation: returns a dict with the activations a_l, the unit outputs z_l
        (z_0 = x) and the output activations 'a' (N x K)."""
        X = np.atleast_2d(np.asarray(X, dtype=float))
        Ws = self.unpack(w)
        Zs, As = [X], []
        z = X
        for l, W in enumerate(Ws):
            a = z @ W[:, 1:].T + W[:, 0]
            As.append(a)
            z = np.tanh(a) if l < self.L - 1 else a
            Zs.append(z)
        return dict(X=X, A=As, Z=Zs, a=As[-1], Ws=Ws)

    def output(self, w, X):
        """The network outputs y(x, w): linear, logistic sigmoid (5.19) or softmax (5.25);
        for a mixture density network the raw activations."""
        a = self.forward(w, X)["a"]
        return self.act(a)

    def act(self, a):
        if self.kind == "bce":
            return common.sigmoid(a)
        if self.kind == "ce":
            e = np.exp(a - a.max(1, keepdims=True))
            return e / e.sum(1, keepdims=True)
        return a

    # ------------------------------------------------------------------ error and output delta
    def error(self, w, X, T):
        """The error E(w) = sum_n E_n(w): (5.11), (5.21), (5.24) or (5.153)."""
        a = self.forward(w, X)["a"]
        return self.error_from_a(a, T)

    def error_from_a(self, a, T):
        T = np.asarray(T, dtype=float).reshape(a.shape[0], -1)
        if self.kind == "sse":
            return 0.5 * np.sum((a - T) ** 2)
        if self.kind == "bce":
            # -t ln y - (1 - t) ln(1 - y) = softplus(a) - t a, written stably
            return float(np.sum(np.logaddexp(0.0, a) - T * a))
        if self.kind == "ce":
            m = a.max(1, keepdims=True)
            lse = np.log(np.exp(a - m).sum(1, keepdims=True)) + m
            return float(-np.sum(T * (a - lse)))
        if self.kind == "mdn":
            return mdn_error(a, T[:, 0], self.K_mix)
        raise ValueError(self.kind)

    def delta_out(self, a, T):
        """dE_n/da_k for every pattern, (N x K): y - t for the canonical links (5.18), (5.66), (5.69);
        (5.155)-(5.157) for the mixture density network."""
        T = np.asarray(T, dtype=float).reshape(a.shape[0], -1)
        if self.kind == "mdn":
            return mdn_delta(a, T[:, 0], self.K_mix)
        return self.act(a) - T

    def second_out(self, a):
        """M_nkk' = d2 E_n / da_k da_k' of (5.92), shape (N, K, K), for the canonical errors."""
        y = self.act(a)
        N, K = y.shape
        if self.kind == "sse":
            return np.tile(np.eye(K), (N, 1, 1))
        if self.kind == "bce":
            return (y * (1 - y))[:, :, None]
        if self.kind == "ce":
            return np.einsum("nk,kl->nkl", y, np.eye(K)) - y[:, :, None] * y[:, None, :]
        raise ValueError("no closed form for kind " + self.kind)

    # ------------------------------------------------------------------ backpropagation (5.55)-(5.56)
    def backward(self, c, delta, per_sample=False):
        """Backpropagate the output deltas delta (N x K) through the cache c of forward().
        Returns dE/dw (summed over patterns) or, with per_sample=True, the N x W matrix whose
        row n is dE_n/dw.   delta_j = h'(a_j) sum_k w_kj delta_k  (5.56),   dE_n/dw_ji = delta_j z_i (5.53)."""
        Ws, Zs = c["Ws"], c["Z"]
        N = delta.shape[0]
        parts = [None] * self.L
        d = delta
        for l in range(self.L - 1, -1, -1):
            zin = np.hstack([np.ones((N, 1)), Zs[l]])
            if per_sample:
                parts[l] = np.einsum("nj,ni->nji", d, zin).reshape(N, -1)
            else:
                parts[l] = (d.T @ zin).ravel()
            if l > 0:
                d = (d @ Ws[l][:, 1:]) * (1.0 - Zs[l] ** 2)          # h' = 1 - z^2 for tanh
        return np.hstack(parts) if per_sample else np.concatenate(parts)

    def grad(self, w, X, T):
        """dE/dw by backpropagation."""
        c = self.forward(w, X)
        return self.backward(c, self.delta_out(c["a"], T))

    def error_grad(self, w, X, T):
        c = self.forward(w, X)
        return self.error_from_a(c["a"], T), self.backward(c, self.delta_out(c["a"], T))

    def grad_activation(self, w, X):
        """b_n = d a_n / dw (N x K x W): the gradient of every output activation, from K backward passes."""
        c = self.forward(w, X)
        N, K = c["a"].shape
        out = np.zeros((N, K, self.W))
        for k in range(K):
            d = np.zeros((N, K))
            d[:, k] = 1.0
            out[:, k, :] = self.backward(c, d, per_sample=True)
        return out

    # ------------------------------------------------------------------ the Jacobian (Section 5.3.4)
    def jacobian(self, w, X):
        """J_nki = d y_k / d x_i (N x K x D) by the backpropagation (5.73)-(5.76); the output
        derivatives d y_k / d a_l are (5.75) for sigmoid and (5.76) for softmax outputs."""
        c = self.forward(w, X)
        N, K = c["a"].shape
        y = self.act(c["a"])
        if self.kind == "bce":
            G = np.zeros((N, K, K))
            G[:, 0, 0] = (y * (1 - y))[:, 0]
        elif self.kind == "ce":
            G = np.einsum("nk,kl->nkl", y, np.eye(K)) - y[:, :, None] * y[:, None, :]
        else:
            G = np.tile(np.eye(K), (N, 1, 1))
        Ws = c["Ws"]
        for l in range(self.L - 1, -1, -1):
            G = G @ Ws[l][:, 1:]                                   # d y / d z_l
            if l > 0:
                G = G * (1.0 - c["Z"][l] ** 2)[:, None, :]         # times h'(a), (5.74)
        return G

    # ------------------------------------------------------------------ fast Hessian-vector product (Section 5.4.6)
    def hessian_vector(self, w, X, T, v):
        """H v for E = sum_n E_n with a canonical link (delta = y - t), by the R-operator
        R{.} = v^T grad (5.96)-(5.97) applied to the forward and backward passes, for any depth.
        Cost O(W) per pattern. R{a_j} = sum v_ji z_i + sum w_ji R{z_i}; R{z} = h'(a) R{a} (5.102);
        R{delta_k} = R{y_k} (5.106); R{delta_j} = h'' R{a_j} sum_k w_kj delta_k + h' sum_k v_kj delta_k
        + h' sum_k w_kj R{delta_k} (5.107)-(5.108)."""
        c = self.forward(w, X)
        Ws, Zs, As = c["Ws"], c["Z"], c["A"]
        Vs = self.unpack(v)
        N = c["X"].shape[0]
        # forward R pass
        Ra, Rz = [], [np.zeros_like(Zs[0])]
        for l in range(self.L):
            r = Rz[l] @ Ws[l][:, 1:].T + Zs[l] @ Vs[l][:, 1:].T + Vs[l][:, 0]
            Ra.append(r)
            Rz.append((1 - Zs[l + 1] ** 2) * r if l < self.L - 1 else r)
        # output deltas and their R
        delta = self.delta_out(c["a"], T)
        y = self.act(c["a"])
        if self.kind == "sse":
            Rdelta = Ra[-1]
        elif self.kind == "bce":
            Rdelta = (y * (1 - y)) * Ra[-1]
        elif self.kind == "ce":
            Rdelta = y * Ra[-1] - y * np.sum(y * Ra[-1], axis=1, keepdims=True)
        else:
            raise ValueError("canonical links only")
        out = [None] * self.L
        d, Rd = delta, Rdelta
        for l in range(self.L - 1, -1, -1):
            zin = np.hstack([np.ones((N, 1)), Zs[l]])
            Rzin = np.hstack([np.zeros((N, 1)), Rz[l]])
            out[l] = (Rd.T @ zin + d.T @ Rzin).ravel()
            if l > 0:
                h1 = 1 - Zs[l] ** 2
                h2 = -2 * Zs[l] * h1
                Rd_new = h2 * Ra[l - 1] * (d @ Ws[l][:, 1:]) + h1 * (d @ Vs[l][:, 1:]) + h1 * (Rd @ Ws[l][:, 1:])
                d = h1 * (d @ Ws[l][:, 1:])
                Rd = Rd_new
        return np.concatenate(out)

    def hessian_rop(self, w, X, T):
        """The exact Hessian from W Hessian-vector products H e_i, (5.96)."""
        H = np.zeros((self.W, self.W))
        for i in range(self.W):
            e = np.zeros(self.W)
            e[i] = 1.0
            H[:, i] = self.hessian_vector(w, X, T, e)
        return 0.5 * (H + H.T)

    # ------------------------------------------------------------------ exact Hessian, two layers (5.92)-(5.95)
    def hessian_exact(self, w, X, T):
        """The exact Hessian of a two-layer network (L = 2) from the three blocks (5.93)-(5.95), summed
        over patterns, with M_kk' = d2E_n/da_k da_k' and delta_k = dE_n/da_k of (5.92)."""
        assert self.L == 2
        c = self.forward(w, X)
        W1, W2 = c["Ws"]
        N = c["X"].shape[0]
        M, D1 = W1.shape
        K = W2.shape[0]
        Xt = np.hstack([np.ones((N, 1)), c["X"]])
        Z = c["Z"][1]
        Zt = np.hstack([np.ones((N, 1)), Z])
        h1 = 1 - Z ** 2
        h2 = -2 * Z * h1
        a = c["a"]
        delta = self.delta_out(a, T)
        Mk = self.second_out(a)
        w2 = W2[:, 1:]                                       # (K, M)
        V = delta @ w2                                       # sum_k w_kj delta_k, (N, M)
        U = np.einsum("kj,nkm,ml->njl", w2, Mk, w2)          # sum_kk' w_kj M_kk' w_k'j', (N, M, M)
        H11 = np.zeros((M, D1, M, D1))
        # first term of (5.94): diagonal in j, j'
        t1 = np.einsum("ni,nl,nj->jil", Xt, Xt, h2 * V)
        for j in range(M):
            H11[j, :, j, :] += t1[j]
        # second term of (5.94)
        H11 += np.einsum("ni,nl,nj,nm,njm->jiml", Xt, Xt, h1, h1, U)
        H22 = np.einsum("nj,nl,nkm->kjml", Zt, Zt, Mk)       # (5.93): (K, M+1, K, M+1)
        # (5.95): one weight in each layer, first-layer (j, i), second-layer (k, j')
        H12 = np.einsum("ni,nj,nk->jik", Xt, h1, delta)[:, :, :, None] * np.eye(M + 1)[1:, None, None, :]
        B = np.einsum("kj,nkm->njm", w2, Mk)                 # sum_k' w_k'j M_k'k, (N, M, K)
        H12 = H12 + np.einsum("ni,nj,nl,njk->jikl", Xt, h1, Zt, B)
        n1, n2 = M * D1, K * (M + 1)
        H = np.zeros((n1 + n2, n1 + n2))
        H[:n1, :n1] = H11.reshape(n1, n1)
        H[n1:, n1:] = H22.reshape(n2, n2)
        H[:n1, n1:] = H12.reshape(n1, n2)
        H[n1:, :n1] = H[:n1, n1:].T
        return H

    # ------------------------------------------------------------------ diagonal approximation (5.79)-(5.81)
    def hessian_diagonal(self, w, X, T, becker_lecun=False):
        """Diagonal of the Hessian of a two-layer network. dE_n/dw_ji^2 = (d2E_n/da_j^2) z_i^2 (5.79) with
        d2E_n/da_j^2 = h'(a_j)^2 sum_kk' w_kj w_k'j M_kk' + h''(a_j) sum_k w_kj delta_k  (5.80) (exact diagonal,
        Ricotti et al.); with becker_lecun=True the off-diagonal M_kk' (k != k') are dropped (5.81)."""
        assert self.L == 2
        c = self.forward(w, X)
        W1, W2 = c["Ws"]
        N = c["X"].shape[0]
        Xt = np.hstack([np.ones((N, 1)), c["X"]])
        Z = c["Z"][1]
        Zt = np.hstack([np.ones((N, 1)), Z])
        h1 = 1 - Z ** 2
        h2 = -2 * Z * h1
        a = c["a"]
        delta = self.delta_out(a, T)
        Mk = self.second_out(a)
        w2 = W2[:, 1:]
        if becker_lecun:
            Mk = np.einsum("nkl,kl->nkl", Mk, np.eye(Mk.shape[1]))
        d2a = h1 ** 2 * np.einsum("kj,nkm,mj->nj", w2, Mk, w2) + h2 * (delta @ w2)     # (N, M)
        d1 = np.einsum("nj,ni->ji", d2a, Xt ** 2)                                      # (M, D+1)
        d2 = np.einsum("nkk,nj->kj", Mk, Zt ** 2)                                      # (K, M+1)
        return np.concatenate([d1.ravel(), d2.ravel()])

    # ------------------------------------------------------------------ outer product approximation (5.83)-(5.85)
    def hessian_outer(self, w, X, T=None):
        """H ~ sum_n B_n^T M_n B_n with B_n = d a_n / dw: (5.84) for the sum-of-squares error (M = I),
        (5.85) y (1 - y) b b^T for the sigmoid cross-entropy, and the softmax analogue."""
        c = self.forward(w, X)
        B = self.grad_activation(w, X)
        Mk = self.second_out(c["a"]) if self.kind != "mdn" else None
        return np.einsum("nkw,nkl,nlv->wv", B, Mk, B)

    def hessian_fd(self, w, X, T, eps=1e-5):
        """The Hessian by central differences of the backpropagated gradient (5.91)."""
        H = np.zeros((self.W, self.W))
        for i in range(self.W):
            e = np.zeros(self.W)
            e[i] = eps
            H[:, i] = (self.grad(w + e, X, T) - self.grad(w - e, X, T)) / (2 * eps)
        return 0.5 * (H + H.T)

    def grad_fd(self, w, X, T, eps=1e-6):
        """The gradient by central differences, dE/dw_i = (E(w_i + eps) - E(w_i - eps)) / (2 eps) + O(eps^2)
        (5.69), the check of Section 5.3.3."""
        g = np.zeros(self.W)
        for i in range(self.W):
            e = np.zeros(self.W)
            e[i] = eps
            g[i] = (self.error(w + e, X, T) - self.error(w - e, X, T)) / (2 * eps)
        return g

    def jacobian_fd(self, w, X, eps=1e-6):
        """The Jacobian by central differences (5.77)."""
        X = np.atleast_2d(X)
        N, D = X.shape
        K = self.sizes[-1]
        J = np.zeros((N, K, D))
        for i in range(D):
            e = np.zeros(D)
            e[i] = eps
            J[:, :, i] = (self.output(w, X + e) - self.output(w, X - e)) / (2 * eps)
        return J


def inverse_hessian_sequential(B, alpha0=1e-4, weights=None):
    """The inverse of H + alpha0 I for H = sum_n weights_n b_n b_n^T, one pattern at a time:
    H_{L+1}^{-1} = H_L^{-1} - (H_L^{-1} b)(b^T H_L^{-1}) / (1 + b^T H_L^{-1} b), (5.87)-(5.89), with H_0 = alpha0 I.
    B is N x W (rows b_n); weights_n (default 1) multiply b_n b_n^T as in (5.85)."""
    N, W = B.shape
    P = np.eye(W) / alpha0
    for n in range(N):
        b = B[n] * (1.0 if weights is None else math.sqrt(weights[n]))
        Pb = P @ b
        P = P - np.outer(Pb, Pb) / (1.0 + b @ Pb)
    return P


# =========================================================================== training
def init_weights(net, rng, sd1=1.0, sd2=None):
    """Draw the initial weights in flat order: first layer N(0, sd1^2), the other layers
    N(0, sd2^2) with sd2 = 1/sqrt(M) by default (M the number of units feeding the layer).
    Using JSRandom(seed) reproduces the demos' starts."""
    w = np.zeros(net.W)
    for l in range(net.L):
        n_in = net.sizes[l]
        sd = sd1 if l == 0 else (sd2 if sd2 is not None else 1.0 / math.sqrt(n_in))
        a, b = net.offsets[l], net.offsets[l + 1]
        for i in range(a, b):
            w[i] = sd * _n(rng)
    return w


def data_sine(N, noise, rng, lo=0.0, hi=1.0):
    """The sinusoidal data set of Appendix A: x uniform on (lo, hi), t = sin(2 pi x) + noise * N(0,1).
    Draw order (as in the demos): x_1, ..., x_N from the uniform stream, then t_1, ..., t_N."""
    x = np.array([lo + (hi - lo) * _u(rng) for _ in range(N)])
    t = np.sin(2 * np.pi * x) + noise * np.array([_n(rng) for _ in range(N)])
    return x[:, None], t[:, None]


def adam_train(net, w0, X, T, lr=0.01, steps=1000, wd=0.0, b1=0.9, b2=0.999, eps=1e-8, record=None):
    """Full-batch Adam on E(w)/N + (wd/2) w^T w, the optimiser of the demos (plain arithmetic, so JavaScript
    follows the same trajectory). record(step, w) is called after every step if given."""
    N = len(X)
    w = w0.copy()
    m = np.zeros_like(w)
    v = np.zeros_like(w)
    for t in range(1, steps + 1):
        g = net.grad(w, X, T) / N + wd * w
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g * g
        mh = m / (1 - b1 ** t)
        vh = v / (1 - b2 ** t)
        w = w - lr * mh / (np.sqrt(vh) + eps)
        if record is not None:
            record(t, w)
    return w


def gd_train(net, w0, X, T, lr, steps, wd=0.0, record=None):
    """Plain full-batch gradient descent on E(w) + (wd/2) w^T w (5.112); record(step, w) after every step."""
    w = w0.copy()
    for t in range(1, steps + 1):
        w = w - lr * (net.grad(w, X, T) + wd * w)
        if record is not None:
            record(t, w)
    return w


def fit_lbfgs(net, X, T, w0, wd=0.0, maxiter=3000):
    """Minimise E(w) + (wd/2) w^T w by L-BFGS with backpropagated gradients."""
    def f(w):
        e, g = net.error_grad(w, X, T)
        return e + 0.5 * np.sum(wd * w * w), g + wd * w
    r = optimize.minimize(f, w0, jac=True, method="L-BFGS-B", options=dict(maxiter=maxiter, maxfun=maxiter * 2, ftol=1e-14, gtol=1e-9))
    return r.x


def fit_best_of(net, X, T, rng, starts=6, sd1=1.0, wd=0.0, maxiter=3000, sd2=None):
    """The best of several random starts (the error surface has local minima, Figure 5.5)."""
    best, bw = np.inf, None
    for _ in range(starts):
        w = fit_lbfgs(net, X, T, init_weights(net, rng, sd1, sd2), wd, maxiter)
        e = net.error(w, X, T) + 0.5 * np.sum(wd * w * w)
        if e < best:
            best, bw = e, w
    return bw


def newton_polish(net, w, X, T, alpha, iters=30, beta=1.0):
    """Minimise beta E_D(w) + (alpha/2) w^T w to high accuracy with damped Newton steps on the exact
    Hessian (used before every Laplace approximation, and mirrored in the demo)."""
    def f(w):
        return beta * net.error(w, X, T) + 0.5 * np.sum(alpha * w * w)
    mu = 1e-3
    for _ in range(iters):
        g = beta * net.grad(w, X, T) + alpha * w
        if np.linalg.norm(g) < 1e-10:
            break
        H = beta * net.hessian_exact(w, X, T) + np.diag(alpha * np.ones(net.W))
        f0 = f(w)
        while True:
            try:
                step = np.linalg.solve(H + mu * np.eye(net.W), g)
            except np.linalg.LinAlgError:
                mu *= 10
                continue
            if f(w - step) <= f0 + 1e-14:
                w = w - step
                mu = max(mu / 10, 1e-12)
                break
            mu *= 10
            if mu > 1e12:
                return w
    return w


# =========================================================================== mixture density networks (Section 5.6)
def mdn_parameters(a, K):
    """Split the 3K output activations (order: a^pi, a^sigma, a^mu) into mixing coefficients (softmax (5.150)),
    standard deviations sigma = exp(a^sigma) (5.151) and means mu = a^mu (5.152). Shapes (N, K)."""
    a_pi, a_s, mu = a[:, :K], a[:, K:2 * K], a[:, 2 * K:3 * K]
    e = np.exp(a_pi - a_pi.max(1, keepdims=True))
    return e / e.sum(1, keepdims=True), np.exp(a_s), mu


def mdn_log_components(a, t, K):
    pi, s, mu = mdn_parameters(a, K)
    lp = np.log(pi) - 0.5 * LOG2PI - np.log(s) - 0.5 * ((t[:, None] - mu) / s) ** 2
    return lp


def mdn_error(a, t, K):
    """E = - sum_n ln sum_k pi_k N(t_n | mu_k, sigma_k^2)  (5.153), for scalar targets (L = 1)."""
    return float(-np.sum(common.logsumexp(mdn_log_components(a, t, K), axis=1)))


def mdn_delta(a, t, K):
    """dE_n/da for the three groups of outputs: dE/da^pi_k = pi_k - gamma_k (5.155), dE/da^mu_k = gamma_k (mu_k - t)/sigma_k^2
    (5.156), dE/da^sigma_k = gamma_k (L - (t - mu_k)^2/sigma_k^2) (5.157, with the erratum), L = 1, where gamma_k = (5.154)."""
    pi, s, mu = mdn_parameters(a, K)
    lp = mdn_log_components(a, t, K)
    gam = np.exp(lp - common.logsumexp(lp, axis=1)[:, None])
    d = np.zeros_like(a)
    d[:, :K] = pi - gam
    d[:, K:2 * K] = gam * (1.0 - ((t[:, None] - mu) / s) ** 2)
    d[:, 2 * K:] = gam * (mu - t[:, None]) / s ** 2
    return d


def mdn_density(a, tgrid, K):
    """p(t | x) of (5.148) on a grid of t for every row of a: array (len(a), len(tgrid))."""
    pi, s, mu = mdn_parameters(a, K)
    z = (tgrid[None, None, :] - mu[:, :, None]) / s[:, :, None]
    return np.einsum("nk,nkg->ng", pi, np.exp(-0.5 * z ** 2) / (math.sqrt(2 * math.pi) * s[:, :, None]))


def mdn_mode(a, K):
    """The approximate conditional mode used for Figure 5.21(d): the centre of the most probable component."""
    pi, s, mu = mdn_parameters(a, K)
    return mu[np.arange(len(a)), pi.argmax(1)]


# =========================================================================== Bayesian neural networks (Section 5.7)
def laplace_log_evidence(E_map, H_data, alpha, beta, N, W, regression=True):
    """ln p(D | alpha, beta) of the Laplace approximation, (5.175) for regression:
         -E(w_MAP) - 1/2 ln|A| + W/2 ln alpha + N/2 ln beta - N/2 ln 2 pi,   A = beta H_D + alpha I  (5.166),
    with E(w) = beta E_D + alpha/2 w^T w; H_D is the Hessian of the sum-of-squares part. For classification
    (5.183): -E(w_MAP) - 1/2 ln|A| + W/2 ln alpha, E the regularised cross-entropy, A = H_D + alpha I (beta = 1)."""
    A = (beta if regression else 1.0) * H_data + alpha * np.eye(W)
    s, ld = np.linalg.slogdet(A)
    out = -E_map - 0.5 * ld + 0.5 * W * math.log(alpha)
    if regression:
        out += 0.5 * N * math.log(beta) - 0.5 * N * LOG2PI
    return out


def _gamma_alpha(lam, alpha):
    """gamma = sum lam_i / (alpha + lam_i) (5.179), lam the eigenvalues of beta H (or of H)."""
    lam = np.maximum(np.asarray(lam, dtype=float), 0.0)   # the exact H of a poorly fitted network can have small negative eigenvalues; treat them as 0
    return float(np.sum(lam / (alpha + lam)))


def evidence_regression(net, X, T, w0, alpha0=1.0, beta0=1.0, outer=12, history=False):
    """MacKay's evidence procedure (Section 5.7.2): alternate (i) a fit of w by minimising
    beta E_D + alpha/2 w^T w, (ii) the eigenvalues lam_i of beta H and the updates
    gamma = sum lam_i/(alpha + lam_i) (5.179), alpha = gamma / w^T w (5.178),
    1/beta = 2 E_D/(N - gamma) (5.180-5.182). Returns a dict with w, alpha, beta, gamma, A (the posterior
    precision), ln evidence and, with history=True, the trace of (alpha, beta, gamma, ln evidence)."""
    N = len(X)
    w, alpha, beta = w0.copy(), alpha0, beta0
    hist = []
    for it in range(outer):
        if it == 0:
            w = fit_lbfgs(net, X, T, w, wd=alpha / beta, maxiter=500)
        w = newton_polish(net, w, X, T, alpha, beta=beta)
        H = net.hessian_exact(w, X, T)
        lam = np.linalg.eigvalsh(beta * H)
        gam = _gamma_alpha(lam, alpha)
        ED = net.error(w, X, T)
        lev = laplace_log_evidence(beta * ED + 0.5 * alpha * w @ w, H, alpha, beta, N, net.W)
        hist.append((alpha, beta, gam, lev))
        a_new = gam / (w @ w)
        b_new = (N - min(gam, N - 1.0)) / (2 * ED)       # (5.180-5.182); gamma < N - 1 keeps beta positive early on
        if abs(a_new - alpha) < 1e-6 * alpha and abs(b_new - beta) < 1e-6 * beta:
            alpha, beta = a_new, b_new
            break
        alpha, beta = a_new, b_new
    w = newton_polish(net, w, X, T, alpha, beta=beta)
    H = net.hessian_exact(w, X, T)
    A = beta * H + alpha * np.eye(net.W)
    lam = np.linalg.eigvalsh(beta * H)
    ED = net.error(w, X, T)
    res = dict(w=w, alpha=alpha, beta=beta, gamma=_gamma_alpha(lam, alpha), A=A, H=H, eig=lam,
               log_evidence=laplace_log_evidence(beta * ED + 0.5 * alpha * w @ w, H, alpha, beta, N, net.W))
    if history:
        res["history"] = hist
    return res


def predictive_regression(net, w, A, beta, Xg):
    """Mean y(x, w_MAP) and variance sigma^2(x) = 1/beta + g^T A^{-1} g of (5.172)-(5.173), g = d y / d w."""
    y = net.output(w, Xg)[:, 0]
    G = net.grad_activation(w, Xg)[:, 0, :]
    var = 1.0 / beta + np.einsum("nw,nw->n", G @ np.linalg.inv(A), G)
    return y, var


def weight_groups(net, split_biases=False):
    """Hyperparameter groups in the spirit of Section 5.5.1: one alpha per layer of weights (biases included),
    or with split_biases=True separate groups 2l (weights) and 2l + 1 (biases) of layer l."""
    g = np.zeros(net.W, dtype=int)
    for l in range(net.L):
        blk = np.full(net.shapes[l], l, dtype=int)
        if split_biases:
            blk[:, 1:] = 2 * l
            blk[:, 0] = 2 * l + 1
        g[net.offsets[l]:net.offsets[l + 1]] = blk.ravel()
    return g


def evidence_classification(net, X, T, w0, alpha0=0.01, outer=30, groups=None):
    """Evidence procedure for a two-class network (Section 5.7.3): fit w with the regularised cross-entropy
    E + 1/2 sum_i alpha_g(i) w_i^2, then alpha_c = gamma_c / sum_{i in c} w_i^2 with
    gamma_c = sum_{i in c} (1 - alpha_c [A^{-1}]_ii) (MacKay 1992; with a single group this is gamma =
    sum lam_i/(alpha + lam_i), (5.178)-(5.179) with beta = 1), A = H + diag(alpha); and repeat. groups
    (one integer per weight) defines the hyperparameter groups; the default is one alpha for all weights."""
    grp = np.zeros(net.W, dtype=int) if groups is None else np.asarray(groups)
    ng = grp.max() + 1
    alphas = np.full(ng, float(alpha0))
    w = w0.copy()

    def update(w, alphas):
        H = net.hessian_exact(w, X, T)
        Ainv = np.linalg.inv(H + np.diag(alphas[grp]))
        gam = np.array([np.sum(1.0 - alphas[c] * np.diag(Ainv)[grp == c]) for c in range(ng)])
        sq = np.array([np.sum(w[grp == c] ** 2) for c in range(ng)])
        return gam, sq, H
    for it in range(outer):
        w = fit_lbfgs(net, X, T, w, wd=alphas[grp], maxiter=800)
        w = newton_polish(net, w, X, T, alphas[grp])
        gam, sq, H = update(w, alphas)
        new = np.clip(np.maximum(gam, 1e-6) / np.maximum(sq, 1e-12), 1e-6, 1e6)
        done = np.all(np.abs(new - alphas) < 1e-4 * alphas)
        alphas = new
        if done:
            break
    w = newton_polish(net, w, X, T, alphas[grp])
    gam, sq, H = update(w, alphas)
    A = H + np.diag(alphas[grp])
    E = net.error(w, X, T) + 0.5 * np.sum(alphas[grp] * w * w)
    ld = np.linalg.slogdet(A)[1]
    lev = -E - 0.5 * ld + 0.5 * sum(np.sum(grp == c) * math.log(alphas[c]) for c in range(ng))
    lam = np.linalg.eigvalsh(H)
    return dict(w=w, alpha=alphas if ng > 1 else float(alphas[0]), alphas=alphas, groups=grp, A=A, H=H, eig=lam,
                gamma=float(gam.sum()), gammas=gam, log_evidence=float(lev))


def kappa(s2):
    """kappa(sigma^2) = (1 + pi sigma^2 / 8)^(-1/2), (4.154)."""
    return 1.0 / np.sqrt(1.0 + np.pi * s2 / 8.0)


def predictive_classification(net, w, A, Xg):
    """p(t = 1 | x, D) ~ sigma(kappa(sigma_a^2) a_MAP) (5.190, with the erratum a_MAP), where
    sigma_a^2 = g^T A^{-1} g (5.188) and g = d a / d w (5.186). Returns (marginalised, plain point estimate)."""
    a = net.forward(w, Xg)["a"][:, 0]
    G = net.grad_activation(w, Xg)[:, 0, :]
    s2 = np.einsum("nw,nw->n", G @ np.linalg.inv(A), G)
    return common.sigmoid(kappa(s2) * a), common.sigmoid(a), s2


# =========================================================================== regularisation (Section 5.5)
def early_stopping_vs_weight_decay(lams, w_ml, eta, tau, alpha):
    """Quadratic error E = 1/2 sum_i lam_i (w_i - w_ML,i)^2 in the Hessian eigenbasis. Gradient descent
    from w = 0 with step eta for tau steps gives w_i = (1 - (1 - eta lam_i)^tau) w_ML,i; weight decay gives
    w_i = lam_i / (lam_i + alpha) w_ML,i (the ridge solution (3.28)). Returns both (Section 5.5.2, Figure 5.13)."""
    lams, w_ml = np.asarray(lams, float), np.asarray(w_ml, float)
    return (1.0 - (1.0 - eta * lams) ** tau) * w_ml, lams / (lams + alpha) * w_ml


def warp_digit(img, sigma, rms, rng):
    """Synthetic warping of Figure 5.14: i.i.d. uniform displacements (dx, dy) on every pixel, smoothed by a
    Gaussian of width sigma (pixels), rescaled to the given rms displacement, and the image resampled by cubic
    splines at the displaced coordinates. Returns the warped image and the two fields."""
    H, Wd = img.shape
    dx = rng.uniform(-1, 1, (H, Wd))
    dy = rng.uniform(-1, 1, (H, Wd))
    dx, dy = ndimage.gaussian_filter(dx, sigma, mode="reflect"), ndimage.gaussian_filter(dy, sigma, mode="reflect")
    s = rms / math.sqrt(np.mean(dx ** 2 + dy ** 2))
    dx, dy = dx * s, dy * s
    yy, xx = np.mgrid[0:H, 0:Wd].astype(float)
    out = ndimage.map_coordinates(img, [yy + dy, xx + dx], order=3, mode="constant", cval=0.0)
    return np.clip(out, 0, 1), dx, dy


def rotate_image(img, deg):
    """Rotate a digit clockwise by deg degrees about its centre (cubic-spline resampling)."""
    return ndimage.rotate(img, -deg, reshape=False, order=3, mode="constant", cval=0.0)


def rotate_tangent(img, h_deg=0.5):
    """The tangent vector tau = dx(xi)/dxi at xi = 0 of the rotation (Section 5.5.4, (5.126)), per radian,
    by a central difference of the rotated image."""
    sm = ndimage.gaussian_filter(img, 0.6)          # a light smoothing keeps spline ringing out of the derivative
    return (rotate_image(sm, h_deg) - rotate_image(sm, -h_deg)) / (2 * math.radians(h_deg))


# =========================================================================== data sets
def classification_data(n_per_class=50, seed=3):
    """Our version of the synthetic two-class data of Appendix A (Figure A.7): equal priors, the blue class
    (t = 0) one Gaussian and the red class (t = 1) a mixture of two Gaussians (parameters in the module docstring).
    Drawn with JSRandom(seed), point by point (blue class first, then red: for a red point one uniform picks the
    component, then two normals), so bayesian-nn-demo.html draws exactly the same points. Returns X (N x 2), t (N,)."""
    rng = JSRandom(seed)
    X = []
    for _ in range(n_per_class):
        X.append([-0.2 + 1.0 * rng.norm(), 0.0 + 1.2 * rng.norm()])
    for _ in range(n_per_class):
        if rng.u() < 0.5:
            X.append([1.0 + 0.6 * rng.norm(), 1.1 + 0.8 * rng.norm()])
        else:
            X.append([0.9 + 0.6 * rng.norm(), -1.1 + 0.7 * rng.norm()])
    return np.array(X), np.r_[np.zeros(n_per_class), np.ones(n_per_class)]


def true_posterior_red(X):
    """p(red | x) of the generating mixture (equal priors)."""
    def g(X, m, s):
        z = (X - m) / s
        return np.exp(-0.5 * (z ** 2).sum(-1)) / (2 * np.pi * np.prod(s))
    pb = g(X, np.array([-0.2, 0.0]), np.array([1.0, 1.2]))
    pr = 0.5 * g(X, np.array([1.0, 1.1]), np.array([0.6, 0.8])) + 0.5 * g(X, np.array([0.9, -1.1]), np.array([0.6, 0.7]))
    return pr / (pr + pb)


def inverse_problem_data(N=250, seed=7, rng=None):
    """The forward-problem data of Figure 5.19: x uniform on (0, 1), t = x + 0.3 sin(2 pi x) + U(-0.1, 0.1).
    With rng=JSRandom(seed) the order of draws is x_1..x_N, then the noises (as in mixture-density-demo.html)."""
    if rng is None:
        g = np.random.default_rng(seed)
        x = g.random(N)
        return x, x + 0.3 * np.sin(2 * np.pi * x) + g.uniform(-0.1, 0.1, N)
    x = np.array([rng.u() for _ in range(N)])
    t = x + 0.3 * np.sin(2 * np.pi * x) + np.array([0.2 * rng.u() - 0.1 for _ in range(N)])
    return x, t


def evidence_step_regression(net, X, T, w, alpha, beta):
    """One round of the demo's evidence procedure (Section 5.7.2): damped-Newton fit of w for the current
    (alpha, beta), then gamma (5.179), alpha = gamma/w^T w (5.178), 1/beta = 2 E_D/(N - gamma) (5.180-5.182).
    Returns w, the new alpha, beta, gamma, ln evidence (5.175) at the OLD (alpha, beta), and the eigenvalues of beta H."""
    N = len(X)
    w = newton_polish(net, w, X, T, alpha, beta=beta)
    H = net.hessian_exact(w, X, T)
    lam = np.linalg.eigvalsh(beta * H)
    gam = _gamma_alpha(lam, alpha)
    ED = net.error(w, X, T)
    lev = laplace_log_evidence(beta * ED + 0.5 * alpha * w @ w, H, alpha, beta, N, net.W)
    return w, gam / (w @ w), (N - min(gam, N - 1.0)) / (2 * ED), gam, lev, lam


# =========================================================================== figures
def _axes_box(ax, xlim=None, ylim=None, ticks=False):
    ax.set_xlim(xlim) if xlim else None
    ax.set_ylim(ylim) if ylim else None
    if not ticks:
        ax.set_xticks([])
        ax.set_yticks([])


def fig5_3(out_dir=None):
    """Figure 5.3: a two-layer network with three tanh hidden units fits x^2, sin(x), |x| and the step."""
    fns = [("x^2", lambda x: x ** 2), ("sin(2x)", lambda x: np.sin(2 * x)), ("|x|", np.abs), ("H(x)", lambda x: (x > 0).astype(float))]
    rng = np.random.default_rng(1)
    x = np.linspace(-1, 1, 50)
    xs = np.linspace(-1, 1, 400)
    fig, axes = common.panel_grid(2, 2, size=(3.4, 2.7))
    for i, (name, f) in enumerate(fns):
        t = f(x)
        net = Net([1, 3, 1])
        w = fit_best_of(net, x[:, None], t[:, None], rng, starts=12 if i > 1 else 6, sd1=2.0 if i == 3 else 1.0, maxiter=6000)
        c = net.forward(w, xs[:, None])
        ax = axes[i // 2, i % 2]
        # scale each hidden-unit output to the data range, as the book does, only for display
        for j, col in enumerate([OLIVE, MAG, LIME]):
            z = c["Z"][1][:, j]
            lo, hi = t.min(), t.max()
            zz = lo + (z - z.min()) / max(z.max() - z.min(), 1e-9) * (hi - lo) if np.ptp(z) > 1e-9 else z
            ax.plot(xs, zz, "--", color=col, lw=1.2)
        ax.plot(xs, c["a"][:, 0], color=common.RED, lw=1.8)
        ax.plot(x, t, ".", color=common.BLUE, ms=4)
        ax.set_title("(%s)  f(x) = %s,  E = %.2g" % ("abcd"[i], name, net.error(w, x[:, None], t[:, None])), fontsize=9.5)
        _axes_box(ax, (-1.05, 1.05))
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-3", out_dir)


def fig5_4(out_dir=None):
    """Figure 5.4: two tanh hidden units and a logistic output on the two-class synthetic data."""
    X, t = classification_data()
    rng = np.random.default_rng(2)
    net = Net([2, 2, 1], "bce")
    w = fit_best_of(net, X, t[:, None], rng, starts=30, sd1=1.0, maxiter=4000)
    g = np.linspace(-2.5, 2.5, 300), np.linspace(-3, 3, 300)
    GX, GY = np.meshgrid(*g)
    P = np.c_[GX.ravel(), GY.ravel()]
    y = net.output(w, P)[:, 0].reshape(GX.shape)
    W1, _ = net.unpack(w)
    fig, ax = plt.subplots(figsize=(4.6, 4.4))
    ax.plot(X[t == 0, 0], X[t == 0, 1], "o", mfc="none", mec=common.BLUE, mew=1.4, ms=6)
    ax.plot(X[t == 1, 0], X[t == 1, 1], "x", color=common.RED, mew=1.4, ms=6)
    for j in range(2):
        a = (P @ W1[j, 1:] + W1[j, 0]).reshape(GX.shape)
        ax.contour(GX, GY, a, [0], colors=[common.BLUE], linestyles="--", linewidths=0.9)
    ax.contour(GX, GY, y, [0.5], colors=[common.RED], linewidths=1.5)
    ax.contour(GX, GY, true_posterior_red(P).reshape(GX.shape), [0.5], colors=[common.GREEN], linewidths=1.5)
    ax.set_xlim(-2.5, 2.5)
    ax.set_ylim(-3, 3)
    ax.set_title("2 tanh hidden units: E = %.2f, errors %d / %d" % (net.error(w, X, t[:, None]), int(((y_at(net, w, X) > 0.5) != (t == 1)).sum()), len(t)), fontsize=9.5)
    return common.savefig(fig, 5, "fig5-4", out_dir)


def y_at(net, w, X):
    return net.output(w, X)[:, 0]


def fig5_9(out_dir=None):
    """Figure 5.9: networks with M = 1, 3 and 10 hidden units fitted to ten sinusoidal points."""
    rng = np.random.default_rng(4)
    x = np.linspace(0, 1, 10)
    t = np.sin(2 * np.pi * x) + 0.3 * rng.standard_normal(10)
    rng = np.random.default_rng(1)
    xs = np.linspace(0, 1, 300)
    fig, axes = common.panel_grid(1, 3, size=(3.3, 2.7))
    for ax, M in zip(axes[0], (1, 3, 10)):
        net = Net([1, M, 1])
        w = fit_best_of(net, x[:, None], t[:, None], rng, starts=8, sd1=1.0, wd=1e-3, maxiter=3000)
        ax.plot(xs, net.output(w, xs[:, None])[:, 0], color=common.RED, lw=1.8)
        ax.plot(x, t, "x", color=common.BLUE, mew=1.8, ms=6)
        ax.plot(xs, np.sin(2 * np.pi * xs), color=common.GREEN, lw=0.8, alpha=0.7)
        ax.set_ylim(-1.6, 1.6)
        ax.set_title("M = %d,  E = %.3f" % (M, net.error(w, x[:, None], t[:, None])), fontsize=9.5)
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-9", out_dir)


def fig5_10(out_dir=None):
    """Figure 5.10: test-set error against the number of hidden units, 30 random starts per M."""
    rng = np.random.default_rng(10)
    xtr = rng.random(60)
    ttr = np.sin(2 * np.pi * xtr) + 0.3 * rng.standard_normal(60)
    xte = rng.random(600)
    tte = np.sin(2 * np.pi * xte) + 0.3 * rng.standard_normal(600)
    Ms = list(range(1, 11))
    res = {}
    for M in Ms:
        net = Net([1, M, 1])
        vals = []
        for s in range(30):
            w0 = math.sqrt(10.0) * rng.standard_normal(net.W)
            w = fit_lbfgs(net, xtr[:, None], ttr[:, None], w0, maxiter=1500)
            vals.append(2 * net.error(w, xte[:, None], tte[:, None]))
        res[M] = vals
    fig, ax = plt.subplots(figsize=(4.6, 3.4))
    for M in Ms:
        ax.plot([M] * 30, res[M], "+", color=common.BLUE, ms=5)
    best = min(Ms, key=lambda m: min(res[m]))
    ax.plot([best], [min(res[best])], "o", mfc="none", mec=common.RED, ms=9)
    ax.set_xlabel("number of hidden units M")
    ax.set_ylabel("test-set sum of squares")
    ax.set_xlim(0, 11)
    ax.set_ylim(50, 200)
    ax.set_title("best test error at M = %d (%.1f); noise floor %.0f; larger values off scale" % (best, min(res[best]), 600 * 0.09), fontsize=8.5)
    return common.savefig(fig, 5, "fig5-10", out_dir)


def fig5_11(out_dir=None):
    """Figure 5.11: functions drawn from the prior of a 12-unit tanh network, four hyperparameter settings."""
    panels = [dict(a1w=1, a1b=1, a2w=1, a2b=1), dict(a1w=1, a1b=1, a2w=1e-2, a2b=1),
              dict(a1w=1e-6, a1b=1e-4, a2w=1, a2b=1), dict(a1w=1e-6, a1b=1e-6, a2w=1, a2b=1)]
    net = Net([1, 12, 1])
    xs = np.linspace(-1, 1, 400)[:, None]
    rng = np.random.default_rng(11)
    cols = [common.RED, LIME, common.BLUE, MAG, OLIVE]
    fig, axes = common.panel_grid(2, 2, size=(3.6, 2.6))
    for p, ax in zip(panels, axes.ravel()):
        sd = np.concatenate([np.full(12, 1 / math.sqrt(p["a1b"])), np.full(12, 1 / math.sqrt(p["a1w"])),
                             [1 / math.sqrt(p["a2b"])], np.full(12, 1 / math.sqrt(p["a2w"]))])
        # layout of w: hidden unit j: (bias, weight); output: (bias, 12 weights)
        sd_w = np.concatenate([np.ravel(np.c_[np.full(12, 1 / math.sqrt(p["a1b"])), np.full(12, 1 / math.sqrt(p["a1w"]))]),
                               [1 / math.sqrt(p["a2b"])], np.full(12, 1 / math.sqrt(p["a2w"]))])
        for col in cols:
            w = sd_w * rng.standard_normal(net.W)
            ax.plot(xs[:, 0], net.output(w, xs)[:, 0], color=col, lw=1.5)
        ax.set_xlim(-1, 1)
        f = lambda v: "1" if v == 1 else "10^{%d}" % round(math.log10(v))
        ax.set_title(r"$\alpha_1^w=%s,\ \alpha_1^b=%s,\ \alpha_2^w=%s,\ \alpha_2^b=%s$" % (f(p["a1w"]), f(p["a1b"]), f(p["a2w"]), f(p["a2b"])), fontsize=8.5)
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-11", out_dir)


def early_stopping_run(seed=25, M=10, N=12, lr=0.004, steps=20000, every=400):
    """Gradient-descent run behind Figure 5.12; returns iterations, training RMS error and validation RMS error."""
    rng = np.random.default_rng(seed)
    x = rng.random(2 * N)
    t = np.sin(2 * np.pi * x) + 0.3 * rng.standard_normal(2 * N)
    Xtr, Ttr, Xva, Tva = x[:N, None], t[:N, None], x[N:, None], t[N:, None]
    net = Net([1, M, 1])
    w0 = init_weights(net, rng, sd1=0.5, sd2=0.5)
    its, etr, eva = [], [], []

    def rec(step, w):
        if step % every == 0:
            its.append(step)
            etr.append(math.sqrt(2 * net.error(w, Xtr, Ttr) / N))
            eva.append(math.sqrt(2 * net.error(w, Xva, Tva) / N))
    gd_train(net, w0, Xtr, Ttr, lr, steps, record=rec)
    return np.array(its), np.array(etr), np.array(eva)


def fig5_12(out_dir=None):
    """Figure 5.12: training and validation error during a gradient-descent training run."""
    its, etr, eva = early_stopping_run()
    k = int(np.argmin(eva))
    fig, axes = common.panel_grid(1, 2, size=(3.8, 2.8))
    for ax, y, col in zip(axes[0], (etr, eva), (common.RED, common.BLUE)):
        ax.plot(range(1, len(y) + 1), y, "o-", color=col, ms=3.5, mfc="none")
        ax.axvline(k + 1, color="k", ls="--", lw=1)
        ax.set_xlabel("iteration step (x 400 gradient steps)")
    axes[0, 0].set_title("training error (RMS)", fontsize=9.5)
    axes[0, 1].set_title("validation error (RMS), minimum at step %d" % (k + 1), fontsize=9.5)
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-12", out_dir)


def fig5_13(out_dir=None):
    """Figure 5.13: early stopping versus weight decay for a quadratic error."""
    lams = np.array([0.1, 1.0])           # eigenvalues of the Hessian; the axes are its eigenvectors
    w_ml = np.array([2.0, 1.0])
    eta = 0.2
    fig, ax = plt.subplots(figsize=(5.2, 3.4))
    th = np.linspace(0, 2 * np.pi, 300)
    lev = 0.03                            # contour E = lev, semi-axes sqrt(2 lev / lambda_i)
    ax.plot(w_ml[0] + math.sqrt(2 * lev / lams[0]) * np.cos(th), w_ml[1] + math.sqrt(2 * lev / lams[1]) * np.sin(th), color=common.RED, lw=1.4)
    taus = np.linspace(0, 150, 600)
    path = np.array([early_stopping_vs_weight_decay(lams, w_ml, eta, tau, 1.0)[0] for tau in taus])
    alphas = np.logspace(2, -3, 400)
    ridge = np.array([early_stopping_vs_weight_decay(lams, w_ml, eta, 1.0, a)[1] for a in alphas])
    ax.plot(path[:, 0], path[:, 1], color=common.BLUE, lw=1.8, label=r"gradient descent from $\mathbf{0}$ (steps $\tau$)")
    ax.plot(ridge[:, 0], ridge[:, 1], ":", color=common.GREEN, lw=2.0, label=r"weight decay, $\alpha$ from $\infty$ to 0")
    tau_s = 12
    es = early_stopping_vs_weight_decay(lams, w_ml, eta, tau_s, 1.0 / (eta * tau_s))
    ax.plot(*es[0], "o", color=common.BLUE, ms=6)
    ax.annotate(r"early stop $\widetilde{\mathbf{w}}$, $\tau=%d$" % tau_s, es[0], xytext=(es[0][0] - 0.1, es[0][1] + 0.38), fontsize=8.5, color=common.BLUE)
    ax.plot(*es[1], "s", color=common.GREEN, ms=7, mfc="none")
    ax.annotate(r"weight decay, $\alpha=1/(\eta\tau)=%.2f$" % (1.0 / (eta * tau_s)), es[1], xytext=(es[1][0] + 0.15, es[1][1] - 0.3), fontsize=8.5, color=common.GREEN)
    ax.plot(*w_ml, "o", color=common.RED, ms=6)
    ax.annotate(r"$\mathbf{w}_{ML}$", w_ml, xytext=(w_ml[0] + 0.12, w_ml[1] + 0.05), fontsize=9)
    ax.set_xlabel(r"$w_1$ (small eigenvalue $\lambda_1 = 0.1$)")
    ax.set_ylabel(r"$w_2$ ($\lambda_2 = 1$)")
    ax.set_xlim(-0.2, 3.0)
    ax.set_ylim(-0.2, 1.7)
    ax.set_aspect("equal")
    ax.legend(fontsize=8, loc="lower right")
    return common.savefig(fig, 5, "fig5-13", out_dir)


def fig5_14(out_dir=None):
    """Figure 5.14: synthetic warping of a handwritten digit with smoothed random displacement fields."""
    X, y = common.mnist(train=False, digits=(6,))
    img = X[0].reshape(28, 28)
    rng = np.random.default_rng(14)
    fig = plt.figure(figsize=(7.0, 3.0))
    gs = fig.add_gridspec(2, 4, width_ratios=[1.9, 1, 1, 1], wspace=0.08, hspace=0.08)
    a0 = fig.add_subplot(gs[:, 0])
    a0.imshow(img, cmap="Blues", vmin=0, vmax=1.2)
    a0.set_xticks([])
    a0.set_yticks([])
    a0.set_title("original", fontsize=9)
    for c, (sig, rms) in enumerate([(0.01, 0.6), (2.5, 1.2), (5.0, 1.2)]):
        w, dx, dy = warp_digit(img, sig, rms, rng)
        a = fig.add_subplot(gs[0, c + 1])
        a.imshow(w, cmap="Blues", vmin=0, vmax=1.2)
        a.set_title("width %g px" % sig, fontsize=8.5)
        b = fig.add_subplot(gs[1, c + 1])
        yy, xx = np.mgrid[0:28:2, 0:28:2]
        b.quiver(xx, yy, dx[::2, ::2], -dy[::2, ::2], color=common.BLUE, angles="xy", scale_units="xy", scale=1.0, width=0.012)
        b.set_xlim(-1, 28)
        b.set_ylim(28, -1)
        for ax in (a, b):
            ax.set_xticks([])
            ax.set_yticks([])
    return common.savefig(fig, 5, "fig5-14", out_dir)


def fig5_16(out_dir=None):
    """Figure 5.16: a digit, the tangent vector of a rotation, x + epsilon tau, and the truly rotated digit."""
    X, y = common.mnist(train=False, digits=(6,))
    img = X[0].reshape(28, 28)
    tau = rotate_tangent(img)
    eps = math.radians(15)
    xs = [img, tau, img + eps * tau, rotate_image(img, 15.0)]
    titles = ["(a) original x", r"(b) tangent vector $\tau$", r"(c) $x+\epsilon\tau$, $\epsilon=15^\circ$", "(d) true rotation by 15 degrees"]
    fig, axes = common.panel_grid(1, 4, size=(2.3, 2.5))
    for k, (ax, im, ti) in enumerate(zip(axes[0], xs, titles)):
        if k == 1:
            v = np.abs(im).max()
            ax.imshow(im, cmap=LinearSegmentedColormap.from_list("yb", ["#e6b800", "#ffffff", "#1f3fbf"]), vmin=-v, vmax=v)   # blue positive, yellow negative
        else:
            ax.imshow(im, cmap="Blues", vmin=0, vmax=1.2)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_title(ti, fontsize=8.5)
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-16", out_dir)


def train_mdn(x, t, K=3, M=5, seed=0, steps=6000, lr=0.01):
    """A mixture density network by Adam (full batch), the same optimiser as the demo. Returns net and weights."""
    net = Net([1, M, 3 * K], "mdn", K_mix=K)
    rng = JSRandom(seed)
    w0 = init_weights(net, rng, sd1=1.0, sd2=0.5)
    # bias the sigma outputs towards small widths: mixed-in start, ln 0.1 (unused when the demo does the same)
    w = adam_train(net, w0, x[:, None], t[:, None], lr=lr, steps=steps)
    return net, w


def fig5_19(out_dir=None):
    """Figure 5.19: the forward problem fitted by a sum-of-squares network and the inverse problem."""
    x, t = inverse_problem_data()
    rng = np.random.default_rng(19)
    fig, axes = common.panel_grid(1, 2, size=(3.8, 3.2))
    xs = np.linspace(0, 1, 300)
    for ax, (a, b), nm in zip(axes[0], ((x, t), (t, x)), ("forward problem", "inverse problem")):
        net = Net([1, 5, 1])
        w = fit_best_of(net, a[:, None], b[:, None], rng, starts=8, sd1=1.0, wd=1e-3, maxiter=4000)
        ax.plot(a, b, "o", mfc="none", mec=LIME, ms=4)
        ax.plot(xs, net.output(w, xs[:, None])[:, 0], color=common.RED, lw=2)
        ax.set_xlim(-0.05, 1.1)
        ax.set_ylim(-0.1, 1.1)
        ax.set_title("%s, 5 hidden units, E = %.2f" % (nm, net.error(w, a[:, None], b[:, None])), fontsize=9)
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-19", out_dir)


def fig5_21(out_dir=None):
    """Figure 5.21: a mixture density network on the inverse problem: pi_k(x), mu_k(x), p(t|x) and the mode."""
    x, t = inverse_problem_data()
    xi, ti = t, x                          # the inverse problem: input t_forward, target x_forward
    K = 3
    seed = 2                               # of seeds 1-8 (Adam, 15000 steps, step 0.01) this has among the lowest errors
    net, w = train_mdn(xi, ti, K=K, M=5, seed=seed, steps=15000, lr=0.01)
    e = net.error(w, xi[:, None], ti[:, None])
    xs = np.linspace(0, 1, 300)
    a = net.forward(w, xs[:, None])["a"]
    pi, s, mu = mdn_parameters(a, K)
    tg = np.linspace(0, 1, 250)
    dens = mdn_density(a, tg, K)
    cols = [common.BLUE, common.RED, "#1b7a1b"]
    fig, axes = common.panel_grid(2, 2, size=(3.9, 3.0))
    for k in range(K):
        axes[0, 0].plot(xs, pi[:, k], color=cols[k], lw=2)
        axes[0, 1].plot(xs, mu[:, k], color=cols[k], lw=2)
    axes[0, 1].plot(xi, ti, ".", color=LIME, ms=2.5, alpha=0.5)
    axes[1, 0].contour(xs, tg, dens.T, levels=dens.max() * np.geomspace(0.01, 0.95, 14), cmap="jet", linewidths=0.9)
    axes[1, 1].plot(xi, ti, "o", mfc="none", mec=LIME, ms=3.5)
    axes[1, 1].plot(xs, mdn_mode(a, K), ".", color=common.RED, ms=3)
    for ax, lab in zip(axes.ravel(), "abcd"):
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        common.panel_label(ax, lab, "upper left")
    axes[0, 0].set_title(r"mixing coefficients $\pi_k(x)$", fontsize=9)
    axes[0, 1].set_title(r"means $\mu_k(x)$", fontsize=9)
    axes[1, 0].set_title(r"density $p(t|x)$", fontsize=9)
    axes[1, 1].set_title("approximate conditional mode", fontsize=9)
    axes[1, 0].set_xlabel("x")
    axes[1, 1].set_xlabel("x")
    fig.suptitle("E = %.1f, N = 250, K = 3, 5 hidden units (seed %d)" % (e, seed), fontsize=9, y=1.0)
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-21", out_dir)


def classification_evidence_fit(seed=22, M=8, two_groups=False):
    """Maximum likelihood versus evidence-regularised networks on the classification data (Figures 5.22, 5.23)."""
    X, t = classification_data()
    T = t[:, None]
    net = Net([2, M, 1], "bce")
    rng = np.random.default_rng(seed)
    w_ml = fit_best_of(net, X, T, rng, starts=6, sd1=1.5, maxiter=6000)
    ev = evidence_classification(net, X, T, w_ml, alpha0=0.01, groups=weight_groups(net) if two_groups else None)
    return net, X, t, w_ml, ev


def fig5_22(out_dir=None):
    """Figure 5.22: the evidence framework on the synthetic two-class data (8 hidden units)."""
    net, X, t, w_ml, ev = classification_evidence_fit()
    GX, GY = np.meshgrid(np.linspace(-2.5, 2.5, 300), np.linspace(-3, 3, 300))
    P = np.c_[GX.ravel(), GY.ravel()]
    fig, ax = plt.subplots(figsize=(4.6, 4.4))
    ax.plot(X[t == 0, 0], X[t == 0, 1], "o", mfc="none", mec=common.BLUE, mew=1.4, ms=6)
    ax.plot(X[t == 1, 0], X[t == 1, 1], "x", color=common.RED, mew=1.4, ms=6)
    ax.contour(GX, GY, true_posterior_red(P).reshape(GX.shape), [0.5], colors=[common.GREEN], linewidths=1.5)
    ax.contour(GX, GY, net.output(w_ml, P)[:, 0].reshape(GX.shape), [0.5], colors=["k"], linewidths=1.3)
    ax.contour(GX, GY, net.output(ev["w"], P)[:, 0].reshape(GX.shape), [0.5], colors=[common.RED], linewidths=1.6)
    ax.set_xlim(-2.5, 2.5)
    ax.set_ylim(-3, 3)
    err = lambda w: int(((net.output(w, X)[:, 0] > 0.5) != (t == 1)).sum())
    ax.set_title("black: max likelihood (%d err.); red: evidence, gamma = %.1f (%d err.)" % (err(w_ml), ev["gamma"], err(ev["w"])), fontsize=7.5)
    return common.savefig(fig, 5, "fig5-22", out_dir)


def fig5_23(out_dir=None):
    """Figure 5.23: the Laplace approximation for a classification network: p(t=1|x) contours,
    the point estimate (5.185) on the left and the marginalised predictive (5.190) on the right."""
    net, X, t, w_ml, ev = classification_evidence_fit()
    GX, GY = np.meshgrid(np.linspace(-2.5, 2.5, 120), np.linspace(-3, 3, 120))
    P = np.c_[GX.ravel(), GY.ravel()]
    pm, pp, s2 = predictive_classification(net, ev["w"], ev["A"], P)
    fig, axes = common.panel_grid(1, 2, size=(4.2, 4.0))
    for ax, p, ti in zip(axes[0], (pp, pm), ("(a) point estimate (5.185)", "(b) marginalised (5.190)")):
        ax.plot(X[t == 0, 0], X[t == 0, 1], "o", mfc="none", mec=common.BLUE, mew=1.2, ms=5)
        ax.plot(X[t == 1, 0], X[t == 1, 1], "x", color=common.RED, mew=1.2, ms=5)
        Z = p.reshape(GX.shape)
        ax.contour(GX, GY, Z, [0.5], colors=[common.GREEN], linewidths=1.6)
        ax.contour(GX, GY, Z, [0.1, 0.3, 0.7, 0.9], colors=["#08306b", "#4292c6", "#e08214", "#7f2704"], linewidths=1.0)
        ax.set_xlim(-2.5, 2.5)
        ax.set_ylim(-3, 3)
        ax.set_title(ti, fontsize=9)
    fig.tight_layout()
    return common.savefig(fig, 5, "fig5-23", out_dir)


FIGURES = {"5.3": fig5_3, "5.4": fig5_4, "5.9": fig5_9, "5.10": fig5_10, "5.11": fig5_11, "5.12": fig5_12,
           "5.13": fig5_13, "5.14": fig5_14, "5.16": fig5_16, "5.19": fig5_19, "5.21": fig5_21,
           "5.22": fig5_22, "5.23": fig5_23}

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