Taming chaos: inference, prediction and control of chaotic dynamics

import bz2
import gzip
import logging
import lzma
import struct
import time
from pathlib import Path

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import pytensor
import pytensor.tensor as pt
from matplotlib.colors import TwoSlopeNorm
from scipy import stats

from pymc_challenges import data

RANDOM_SEED = 46
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)  # several fits: no sampler banner per fit
BLUE, ORANGE, AQUA, GREY, PURPLE, RED, INK = (
    "#2a78d6", "#eb6834", "#1baf7a", "#8a8a86", "#8c5ac8", "#c8384e", "#222222")
CHAIN_COLOURS = [BLUE, ORANGE, AQUA, PURPLE]


def stacked(da):
    """(chain, draw, ...) -> (sample, ...) NumPy array."""
    return da.stack(sample=("chain", "draw")).transpose("sample", ...).to_numpy()


print(f"PyMC {pm.__version__}, ArviZ {az.__version__}, PyTensor {pytensor.__version__}")

def ricker_path(log_r, n0, steps):
    """Deterministic Ricker map on the log scale; log_r may be an array (vectorised)."""
    log_r = np.asarray(log_r, float)
    n = np.empty((steps,) + log_r.shape)
    n[0] = n0
    for t in range(steps - 1):
        n[t + 1] = log_r + n[t] - np.exp(n[t])
    return n


def ricker_lyapunov(log_r, steps=4000, burn=500):
    n = ricker_path(log_r, 0.1, steps + burn)[burn:]
    return np.log(np.abs(1 - np.exp(n)) + 1e-300).mean(axis=0)


log_r_grid = np.linspace(1.0, 4.2, 1600)
orbit = ricker_path(log_r_grid, 0.1, 700)[-200:]      # the long-run states for each r
lyap_grid = ricker_lyapunov(log_r_grid)
LOG_R_TRUE = 3.8                                      # Wood (2010)'s value
lam_true = ricker_lyapunov(np.array([LOG_R_TRUE]))[0]
print(f"Lyapunov exponent at log r = {LOG_R_TRUE}: {lam_true:.3f} per step "
      f"(an error grows e-fold every {1 / lam_true:.1f} steps)")

fig, axes = plt.subplots(3, 1, figsize=(10, 9), sharex=False,
                         gridspec_kw={"height_ratios": [2.2, 1, 1.6]})
ax = axes[0]
ax.plot(np.broadcast_to(log_r_grid, orbit.shape).ravel(), np.exp(orbit).ravel(), ",", color=INK,
        alpha=0.25)
ax.axvline(LOG_R_TRUE, color=RED, lw=1)
ax.set(ylabel="long-run population $N$", yscale="log", ylim=(1e-3, 30),
       title="Bifurcation diagram: where the Ricker population ends up, for each growth rate")
ax = axes[1]
ax.plot(log_r_grid, lyap_grid, color=INK, lw=0.7)
ax.fill_between(log_r_grid, 0, lyap_grid, where=lyap_grid > 0, color=RED, alpha=0.3, lw=0,
                label="chaos ($\\lambda > 0$)")
ax.axhline(0, color=GREY, lw=0.8)
ax.axvline(LOG_R_TRUE, color=RED, lw=1)
ax.set(xlabel="log growth rate  $\\log r$", ylabel="Lyapunov\nexponent $\\lambda$", ylim=(-1.5, 1.2))
ax.legend(loc="lower right")
# The butterfly effect: two populations one part in a million apart
ax = axes[2]
steps = 40
a = ricker_path(LOG_R_TRUE, 0.5, steps)
b = ricker_path(LOG_R_TRUE, 0.5 + 1e-6, steps)
ax.plot(np.exp(a), "o-", color=BLUE, ms=3, lw=1, label="start at $N_0$")
ax.plot(np.exp(b), "o--", color=ORANGE, ms=3, lw=1, label="start at $N_0$ + one part in a million")
diverge = np.argmax(np.abs(a - b) > 0.5)
ax.axvline(diverge, color=GREY, lw=0.8)
ax.text(diverge + 0.4, 0.9 * np.exp(a).max(), f"visibly different after {diverge} steps", fontsize=9)
ax.set(xlabel="time step", ylabel="population $N$",
       title=f"Sensitive dependence at log r = {LOG_R_TRUE}")
ax.legend(loc="upper left", fontsize=8);

def trajectory_loglik(log_r, n0=n_true[0], phi=PHI_TRUE, y=y_lab):
    path = ricker_path(log_r, n0, len(y))                       # (T, grid)
    return stats.poisson.logpmf(y[:, None], phi * np.exp(path)).sum(axis=0)


profile_grid = np.linspace(3.0, 4.5, 3000)
ll_traj = trajectory_loglik(profile_grid)
n_peaks = int(np.sum((ll_traj[1:-1] > ll_traj[:-2]) & (ll_traj[1:-1] > ll_traj[2:])))
print(f"local maxima of the trajectory-matching likelihood on [3, 4.5]: {n_peaks}; "
      f"global maximum at log r = {profile_grid[np.argmax(ll_traj)]:.3f} (truth {LOG_R_TRUE})")

# Why: the sensitivity of the path to log r grows like e^{lambda t}
def path_sensitivity(log_r, steps=60):
    """d n_t / d log r along the deterministic path (a tangent recursion)."""
    n, dn, out = 0.5, 0.0, []
    for _ in range(steps):
        out.append(abs(dn))
        n, dn = log_r + n - np.exp(n), 1 + (1 - np.exp(n)) * dn
    return np.array(out)


fig, axes = plt.subplots(1, 2, figsize=(12, 3.8), gridspec_kw={"width_ratios": [1.6, 1]})
ax = axes[0]
ax.plot(profile_grid, ll_traj, color=INK, lw=0.5)
ax.axvline(LOG_R_TRUE, color=RED, lw=1, label="truth")
ax.set(xlabel="log r", ylabel="log-likelihood", ylim=(ll_traj.max() - 4000, ll_traj.max() + 200),
       title=f"Trajectory matching: {n_peaks} local maxima")
ax.legend()
ax = axes[1]
for lr_, c, lab in [(1.5, BLUE, "log r = 1.5 (stable)"), (2.3, AQUA, "log r = 2.3 (cycle)"),
                    (LOG_R_TRUE, RED, f"log r = {LOG_R_TRUE} (chaos)")]:
    ax.semilogy(path_sensitivity(lr_) + 1e-12, color=c, label=lab)
ax.semilogy(np.arange(60), np.exp(lam_true * np.arange(60)), ":", color=RED, lw=1,
            label="$e^{\\lambda t}$")
ax.set(xlabel="time step t", ylabel="$|\\partial n_t / \\partial \\log r|$", ylim=(1e-2, 1e16),
       title="The gradient explodes")
ax.legend(fontsize=8);

post_ss = az.extract(idata_ss, var_names=["log_r", "sigma", "log_phi"])
log_r_d, sigma_d, log_phi_d = (post_ss[v].values for v in ["log_r", "sigma", "log_phi"])
n_last = n_draws[:, -1]
fc_rng = np.random.default_rng(RANDOM_SEED)
fc = np.empty((len(n_last), T_HOLD))
state = n_last.copy()
for h in range(T_HOLD):
    state = log_r_d + state - np.exp(state) + sigma_d * fc_rng.standard_normal(len(state))
    fc[:, h] = state
y_fc = fc_rng.poisson(np.exp(log_phi_d[:, None] + fc))

# The "climate": the long-run distribution of counts, from a long simulation at the posterior
n_long = np.empty(20000)
n_long[0] = 0.1
for t in range(len(n_long) - 1):
    n_long[t + 1] = log_r_d.mean() + n_long[t] - np.exp(n_long[t]) + sigma_d.mean() * fc_rng.standard_normal()
y_climate = fc_rng.poisson(np.exp(log_phi_d.mean() + n_long[500:]))

# Skill: how much better than the climate is the forecast at each lead time? (log score)
def log_score(samples, obs):
    pmf = np.bincount(np.minimum(samples, 999), minlength=1000) + 0.5
    return np.log(pmf[min(obs, 999)] / pmf.sum())


skill = np.array([log_score(y_fc[:, h], y_hold[h]) - log_score(y_climate, y_hold[h])
                  for h in range(T_HOLD)])
fc_q = np.quantile(y_fc, [0.05, 0.25, 0.5, 0.75, 0.95], axis=0)
clim_q = np.quantile(y_climate, [0.05, 0.95])

fig, axes = plt.subplots(1, 2, figsize=(12, 3.6), gridspec_kw={"width_ratios": [1.6, 1]})
ax = axes[0]
t_hold = np.arange(T_FIT, T_FIT + T_HOLD)
ax.plot(np.arange(T_FIT - 15, T_FIT), y_lab[-15:], "o-", color=INK, ms=3, lw=0.8)
ax.fill_between(t_hold, fc_q[0], fc_q[4], color=BLUE, alpha=0.25, lw=0, label="forecast 90%")
ax.fill_between(t_hold, fc_q[1], fc_q[3], color=BLUE, alpha=0.4, lw=0, label="forecast 50%")
ax.axhspan(*clim_q, color=GREY, alpha=0.12, lw=0, label="climate 90%")
ax.plot(t_hold, y_hold, "o-", color=RED, ms=3, lw=0.8, label="what happened")
ax.set(xlabel="generation", ylabel="count", yscale="symlog", title="Forecasts from the end of the data")
ax.legend(fontsize=8, ncols=2, loc="upper left")
ax = axes[1]
ax.bar(np.arange(1, T_HOLD + 1), skill, color=np.where(skill > 0, BLUE, GREY))
ax.axhline(0, color=INK, lw=0.8)
ax.set(xlabel="lead time (generations)", ylabel="log score gain over climate",
       title="Forecast skill vanishes within a few steps");
print("skill (log score gain) for leads 1-6:", skill[:6].round(2))

def ogy_run(log_r, u_max, sigma, steps=300, start=150, seed=0):
    r_ = np.random.default_rng(seed)
    target = np.log(log_r)
    n, us = np.empty(steps), np.zeros(steps)
    n[0] = 0.3
    for t in range(steps - 1):
        pred = log_r + n[t] - np.exp(n[t])
        if t >= start and abs(target - pred) < u_max:
            us[t] = target - pred
        n[t + 1] = pred + us[t] + sigma * r_.standard_normal()
    return n, us


fig, axes = plt.subplots(3, 1, figsize=(10, 7.5), sharex=True,
                         gridspec_kw={"height_ratios": [2, 1, 2]})
n_ctl, u_ctl = ogy_run(LOG_R_TRUE, u_max=0.1, sigma=0.0)
ax = axes[0]
ax.plot(np.exp(n_ctl), color=INK, lw=0.8)
ax.axvline(150, color=GREY, lw=0.8)
ax.axhline(LOG_R_TRUE, color=RED, lw=0.8, ls=":")
captured = 150 + np.argmax(np.abs(u_ctl[150:]) > 0)
ax.text(152, 2e-3, f"control on; captured after {captured - 150} steps of waiting", fontsize=9)
ax.set(ylabel="population", yscale="log", ylim=(1e-3, 40),
       title=f"OGY control of the chaotic Ricker map, |u| <= 0.1 (2.6% of log r), no noise")
ax = axes[1]
ax.bar(np.arange(len(u_ctl)), u_ctl, color=BLUE, width=1)
ax.set(ylabel="nudge $u_t$", ylim=(-0.12, 0.12))
n_ctl2, u_ctl2 = ogy_run(LOG_R_TRUE, u_max=1.0, sigma=0.05, seed=1)
ax = axes[2]
ax.plot(np.exp(n_ctl2), color=INK, lw=0.8)
ax.axvline(150, color=GREY, lw=0.8)
ax.axhline(LOG_R_TRUE, color=RED, lw=0.8, ls=":")
ax.set(xlabel="generation", ylabel="population", yscale="log", ylim=(1e-3, 40),
       title="With process noise sd 0.05 (and |u| <= 1): control must keep correcting every step");
print(f"noise-free: captured after {captured - 150} steps; "
      f"afterwards |u| <= {np.abs(u_ctl[captured + 5:]).max():.1e}")
print(f"with noise sd 0.05: mean |u| after capture = "
      f"{np.abs(u_ctl2[200:]).mean():.3f}, max = {np.abs(u_ctl2[200:]).max():.3f}")

def read_rdata(path):
    """Minimal reader for R's XDR serialisation (gzip, bzip2 or xz): named lists, vectors, strings."""
    raw = Path(path).read_bytes()
    buf = (bz2.decompress(raw) if raw[:2] == b"BZ" else
           lzma.decompress(raw) if raw[:6] == b"\xfd7zXZ\x00" else gzip.decompress(raw))
    assert buf[:7] in (b"RDX2\nX\n", b"RDX3\nX\n"), "not an XDR R data file"
    pos, refs = 7, []

    def i32():
        nonlocal pos
        pos += 4
        return struct.unpack(">i", buf[pos - 4:pos])[0]

    def item():
        nonlocal pos
        flags = i32()
        kind, has_attr, has_tag = flags & 0xFF, flags & (1 << 9), flags & (1 << 10)
        if kind == 254:                                     # NULL
            return None
        if kind == 255:                                     # reference to an earlier symbol
            return refs[(flags >> 8) - 1]
        if kind == 1:                                       # symbol
            refs.append(item())
            return refs[-1]
        if kind == 9:                                       # string
            n = i32()
            pos += max(n, 0)
            return None if n == -1 else buf[pos - n:pos].decode()
        if kind == 2:                                       # pairlist -> dict
            out = {}
            while kind == 2:
                _ = item() if has_attr else None
                tag = item() if has_tag else None
                out[tag] = item()
                flags = i32()
                kind, has_attr, has_tag = flags & 0xFF, flags & (1 << 9), flags & (1 << 10)
            return out
        n = i32()
        if kind in (10, 13):                                # logical / integer
            value, pos = np.frombuffer(buf, ">i4", n, pos).astype(int), pos + 4 * n
        elif kind == 14:                                    # double
            value, pos = np.frombuffer(buf, ">f8", n, pos).astype(float), pos + 8 * n
        elif kind in (16, 19):                              # character vector / list
            value = [item() for _ in range(n)]
        else:
            raise NotImplementedError(f"R type {kind}")
        attrs = item() if has_attr else None
        if kind == 19 and attrs and "names" in attrs:
            value = dict(zip(attrs["names"], value))
        return value

    version = i32()
    pos += 8
    if version == 3:
        pos += 4 + struct.unpack(">i", buf[pos:pos + 4])[0]
    return item()


data.describe("blowfly")
flies = read_rdata(data.path("blowfly"))["blowfly"]
N_obs = np.asarray(flies["pop"], float)
days = 2 * np.arange(len(N_obs))                     # one count every two days
print(f"{len(N_obs)} counts over {days[-1]} days; min {N_obs.min():.0f}, "
      f"median {np.median(N_obs):.0f}, max {N_obs.max():.0f} adults")
x_c = np.log(N_obs) - np.log(N_obs).mean()
acf = np.array([x_c[:-k] @ x_c[k:] / (x_c @ x_c) for k in range(1, 60)])
acf_peak = 1 + np.argmax(np.r_[0, (acf[1:-1] > acf[:-2]) & (acf[1:-1] > acf[2:]), 0] * acf)
print(f"autocorrelation of log counts peaks at lag {acf_peak} (= {2 * acf_peak} days)")

fig = plt.figure(figsize=(12, 6.5))
gs = fig.add_gridspec(2, 3)
ax = fig.add_subplot(gs[0, :])
ax.plot(days, N_obs, "o-", color=INK, ms=2.5, lw=0.8)
ax.set(xlabel="day", ylabel="adult flies", title="Nicholson's blowflies: one cage, one count every two days")
for i, lag in enumerate([1, 8, acf_peak]):
    ax = fig.add_subplot(gs[1, i])
    ax.plot(N_obs[:-lag], N_obs[lag:], "-", color=GREY, lw=0.5)
    ax.plot(N_obs[:-lag], N_obs[lag:], "o", color=BLUE, ms=2.5)
    ax.set(xlabel=f"$N_t$", ylabel=f"$N_{{t+{lag}}}$", xscale="log", yscale="log",
           title=f"lag {lag} ({2 * lag} days)")

TAUS = np.arange(4, 14)
T0 = TAUS.max()                                    # first step with every lag available
y_b = N_obs[T0:]
prev_b = N_obs[T0 - 1:-1]
lags_b = np.stack([N_obs[T0 - k:len(N_obs) - k] for k in TAUS])   # (delay, time)

with pm.Model(coords={"tau": TAUS}) as blowfly_model:
    log_P = pm.Normal("log_P", np.log(5), 1)
    log_N0 = pm.Normal("log_N0", np.log(1000), 1)
    s = pm.Beta("s", 4, 2)
    sigma = pm.HalfNormal("sigma", 0.5)
    m = pt.exp(log_P) * lags_b * pt.exp(-lags_b / pt.exp(log_N0)) + s * prev_b
    loglik_tau = pm.logp(pm.LogNormal.dist(pt.log(m), sigma), y_b[None, :]).sum(axis=1)
    pm.Potential("likelihood", pt.logsumexp(loglik_tau) - np.log(len(TAUS)))
    pm.Deterministic("p_tau", pt.exp(loglik_tau - pt.logsumexp(loglik_tau)), dims="tau")

# Prior predictive in NumPy (a Potential has no generative form): what do the priors allow?
def simulate_blowflies(P, N0, s, sigma, tau, steps, start, seed, control=None):
    """Simulate the delay model for many parameter draws at once (one row per draw).

    start: the first max(tau)+1 counts (oldest first). control(t, N_t) returns the number of
    adults added (+) or removed (-) right after the census at step t.
    """
    r_ = np.random.default_rng(seed)
    D = len(P)
    tau = np.broadcast_to(tau, (D,))
    width = tau.max() + 1
    hist = np.tile(np.asarray(start, float)[-width:][::-1], (D, 1))   # column k holds N_{t-k}
    out, us = np.empty((D, steps)), np.zeros((D, steps))
    rows = np.arange(D)
    for t in range(steps):
        if control is not None:
            us[:, t] = control(t, hist[:, 0])
            hist[:, 0] = hist[:, 0] + us[:, t]
        lag = hist[rows, tau]
        mean = P * lag * np.exp(-lag / N0) + s * hist[:, 0]
        new = mean * np.exp(sigma * r_.standard_normal(D)) if sigma is not None else mean
        hist = np.concatenate([new[:, None], hist[:, :-1]], axis=1)
        out[:, t] = new
    return out, us


pr = np.random.default_rng(RANDOM_SEED)
n_prior = 500
prior_draws = dict(P=np.exp(pr.normal(np.log(5), 1, n_prior)), N0=np.exp(pr.normal(np.log(1000), 1, n_prior)),
                   s=pr.beta(4, 2, n_prior), sigma=np.abs(pr.normal(0, 0.5, n_prior)),
                   tau=pr.choice(TAUS, n_prior))
with np.errstate(over="ignore", invalid="ignore"):
    prior_paths, _ = simulate_blowflies(**prior_draws, steps=180, start=N_obs[:T0 + 1], seed=1)
fig, axes = plt.subplots(1, 2, figsize=(12, 3.4), gridspec_kw={"width_ratios": [1.6, 1]})
ax = axes[0]
for i in range(12):
    ax.plot(2 * np.arange(180), prior_paths[i], lw=0.8, alpha=0.8)
ax.set(yscale="log", ylim=(1, 1e6), xlabel="day", ylabel="adults", title="Prior predictive: 12 simulated cages")
ax = axes[1]
prior_level = np.log10(np.median(prior_paths[:, 60:], axis=1))
ax.hist(prior_level[np.isfinite(prior_level)], bins=40, color=GREY)
ax.axvline(np.log10(np.median(N_obs)), color=RED, label="Nicholson's cage")
ax.set(xlabel="log10 median population (days 120-360)", ylabel="prior draws", title="Scale allowed by the prior")
ax.legend(fontsize=8);

sub = draw_rng.choice(len(P_d), 1000, replace=False)
# (1) one-step-ahead predictive intervals, conditional on the observed history
lag_obs = np.stack([N_obs[T0 - k:len(N_obs) - k] for k in tau_d[sub]])
mean_1 = P_d[sub, None] * lag_obs * np.exp(-lag_obs / N0_d[sub, None]) + s_d[sub, None] * prev_b
pred_1 = mean_1 * np.exp(sigma_b_d[sub, None] * draw_rng.standard_normal(mean_1.shape))
q1 = np.quantile(pred_1, [0.05, 0.5, 0.95], axis=0)
cover_1 = np.mean((y_b >= q1[0]) & (y_b <= q1[2]))
# (2) free-running simulations started from the first counts
sims, _ = simulate_blowflies(P_d[sub], N0_d[sub], s_d[sub], sigma_b_d[sub], tau_d[sub],
                             steps=len(N_obs) - T0 - 1, start=N_obs[:T0 + 1], seed=2)


def features(paths):
    """Phase-free summaries of (draws, time) series of counts."""
    logp = np.log(np.atleast_2d(paths))
    xc = logp - logp.mean(axis=1, keepdims=True)
    acfs = np.stack([(xc[:, :-k] * xc[:, k:]).sum(1) / (xc * xc).sum(1) for k in range(1, 41)], 1)
    lags = np.arange(1, 41)
    is_peak = (acfs[:, 1:-1] > acfs[:, :-2]) & (acfs[:, 1:-1] > acfs[:, 2:]) & (acfs[:, 1:-1] > 0.1)
    is_peak &= lags[1:-1] >= 8
    peak = np.where(is_peak.any(axis=1), lags[1:-1][np.argmax(is_peak, axis=1)], np.nan)
    return {"median log count": np.median(logp, axis=1),
            "sd of log count": logp.std(axis=1),
            "cycle length (steps)": peak,
            "trough: 5% quantile of log count": np.quantile(logp, 0.05, axis=1)}


f_sim, f_obs = features(sims), features(N_obs[T0 + 1:])
fig = plt.figure(figsize=(12, 7))
gs = fig.add_gridspec(2, 4)
ax = fig.add_subplot(gs[0, :2])
ax.fill_between(days[T0:], q1[0], q1[2], color=BLUE, alpha=0.3, lw=0, label="one-step 90% interval")
ax.plot(days[T0:], y_b, "o", color=INK, ms=2.5, label="observed")
ax.set(yscale="log", xlabel="day", ylabel="adults",
       title=f"(1) One step ahead: {100 * cover_1:.0f}% in the 90% band")
ax.title.set_fontsize(11)
ax.legend(fontsize=8)
ax = fig.add_subplot(gs[0, 2:])
ax.plot(days[T0 + 1:], N_obs[T0 + 1:], color=INK, lw=1.4, label="observed")
for i, c in zip(range(3), [BLUE, ORANGE, AQUA]):
    ax.plot(days[T0 + 1:], sims[i], color=c, lw=0.8, alpha=0.9)
ax.set(yscale="log", xlabel="day", title="(2) Free-running: 3 simulated cages vs the real one")
ax.title.set_fontsize(11)
ax.legend(fontsize=8)
for j, (name, vals) in enumerate(f_sim.items()):
    ax = fig.add_subplot(gs[1, j])
    vals = vals[np.isfinite(vals)]
    ax.hist(vals, bins=25, color=GREY)
    ax.axvline(f_obs[name][0], color=RED, lw=2)
    ax.set(title=name, yticks=[])
    ax.title.set_fontsize(10)
    tail = min(np.mean(vals <= f_obs[name][0]), np.mean(vals >= f_obs[name][0]))
    ax.text(0.03, 0.9, f"tail prob. {tail:.2f}", transform=ax.transAxes, fontsize=8)
print(f"one-step 90% interval coverage: {cover_1:.2f}")
print("simulated cycle length (steps), 5/25/50/75/95%:",
      np.nanquantile(f_sim["cycle length (steps)"], [0.05, 0.25, 0.5, 0.75, 0.95]),
      f"observed {f_obs['cycle length (steps)'][0]:.0f}")

def delay_lyapunov(P, N0, s, tau, sigma=None, steps=3000, burn=500, seed=0):
    """Largest Lyapunov exponent of the delay model (per step), vectorised over draws.

    sigma=None: the deterministic skeleton. With sigma: along noisy paths, using the same
    random multipliers for the path and its tangent (Ellner & Turchin's 'noisy' exponent).
    """
    r_ = np.random.default_rng(seed)
    D = len(P)
    tau = np.broadcast_to(tau, (D,))
    width = tau.max() + 1
    rows = np.arange(D)
    x = np.full((D, width), 1.0) * N0[:, None]
    v = r_.standard_normal((D, width))
    total = np.zeros(D)
    for t in range(steps + burn):
        lag, vlag = x[rows, tau], v[rows, tau]
        e = np.exp(sigma * r_.standard_normal(D)) if sigma is not None else 1.0
        births = P * lag * np.exp(-lag / N0) * e
        slope = P * np.exp(-lag / N0) * (1 - lag / N0) * e
        x = np.concatenate([(births + s * x[:, 0])[:, None], x[:, :-1]], axis=1)
        v = np.concatenate([(slope * vlag + s * v[:, 0])[:, None], v[:, :-1]], axis=1)
        norm = np.maximum(np.sqrt((v**2).sum(axis=1)), 1e-300)
        v /= norm[:, None]
        if t >= burn:
            total += np.log(norm)
    return total / steps


P_grid = np.exp(np.linspace(np.log(1), np.log(2000), 120))
s_grid = np.linspace(0.0, 0.97, 80)
PP, SS = np.meshgrid(P_grid, s_grid)
t0 = time.time()
lam_map = delay_lyapunov(PP.ravel(), np.full(PP.size, 1000.0), SS.ravel(), 8).reshape(PP.shape)
lam_skel = delay_lyapunov(P_d[sub], N0_d[sub], s_d[sub], tau_d[sub])
lam_noisy = delay_lyapunov(P_d[sub], N0_d[sub], s_d[sub], tau_d[sub], sigma=sigma_b_d[sub])
print(f"Lyapunov computations: {time.time() - t0:.0f} s")
print(f"skeleton exponent: median {np.median(lam_skel):.4f}, 90% interval "
      f"{np.quantile(lam_skel, 0.05):.4f} to {np.quantile(lam_skel, 0.95):.4f}; "
      f"P(lambda > 0.01) = {np.mean(lam_skel > 0.01):.3f}")
print(f"noisy exponent: median {np.median(lam_noisy):.4f}, 90% interval "
      f"{np.quantile(lam_noisy, 0.05):.4f} to {np.quantile(lam_noisy, 0.95):.4f}; "
      f"P(lambda > 0) = {np.mean(lam_noisy > 0):.3f}")
chaos_s = SS[lam_map > 0.01]
col = np.argmin(np.abs(P_grid - np.median(P_d)))
print(f"chaotic share of the map (lambda > 0.01): {np.mean(lam_map > 0.01):.2f}; highest survival "
      f"with chaos anywhere on the map: {chaos_s.max():.2f}; at P = {P_grid[col]:.0f} (the posterior "
      f"median) chaos needs s below {s_grid[lam_map[:, col] > 0.01].max():.2f}")

fig, axes = plt.subplots(1, 2, figsize=(12, 4.6), gridspec_kw={"width_ratios": [1.5, 1]})
ax = axes[0]
mesh = ax.pcolormesh(PP, SS, np.clip(lam_map, -0.3, 0.3), cmap="RdBu_r",
                     norm=TwoSlopeNorm(vcenter=0, vmin=-0.3, vmax=0.3), shading="auto")
fig.colorbar(mesh, ax=ax, label="Lyapunov exponent (per step)")
ax.plot(P_d[sub], s_d[sub], ".", color=INK, ms=1.5, alpha=0.4)
ax.set(xscale="log", xlabel="maximum production P (per adult per 2 days)", ylabel="adult 2-day survival s",
       title="Regime map of the delay model (tau = 8) with the posterior (black)")
ax.text(1.3, 0.05, "stable\nequilibrium", color="#08306b", fontsize=9)
ax.text(40, 0.9, "white = limit cycle ($\\lambda$ = 0)", fontsize=9)
ax.text(300, 0.15, "red = chaos", color="#67000d", fontsize=9)
ax = axes[1]
bins = np.linspace(-0.01, 0.01, 41)
ax.hist(np.clip(lam_skel, -0.01, 0.01), bins=bins, color=BLUE, alpha=0.7, label="skeleton (no noise)")
ax.hist(np.clip(lam_noisy, -0.01, 0.01), bins=bins, color=ORANGE, alpha=0.7, label="along noisy paths")
ax.axvline(0, color=INK, lw=0.8)
ax.set(xlabel="Lyapunov exponent (per step)", yticks=[], title="Posterior of the exponent")
ax.legend(fontsize=8);