L02 — Noise in Gene Circuits

"""
L02 - Noise in Gene Circuits: deterministic ODEs vs exact stochastic simulation (Gillespie SSA).

Two classic synthetic-biology circuits from 2000:
  * the genetic toggle switch (Gardner, Cantor & Collins) - bistable; noise makes it flip,
    the ODE never can;
  * the repressilator (Elowitz & Leibler) - a limit-cycle clock; noise blurs its period.

Run from the repo root:   uv run python L02_gene_circuit_noise/sim.py
"""
import os

os.environ.setdefault("OMP_NUM_THREADS", "1")  # parallelism comes from 4 worker processes instead

import time
from multiprocessing import Pool
from pathlib import Path

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib.collections import EllipseCollection, LineCollection
from matplotlib.colors import LinearSegmentedColormap
from scipy import fft, ndimage, sparse, stats
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
from scipy.signal import find_peaks
from scipy.sparse.linalg import spsolve

OUT = Path(__file__).parent / "output"
SEED = 2000  # the year both circuits were published
N_PROC = 4  # CPU etiquette: never more than 4 workers

C_U, C_V, C_ODE = "#4cc9f0", "#f72585", "#ffffff"  # toggle repressor 1, repressor 2, deterministic
C_GENES = ["#4cc9f0", "#f72585", "#ffd166"]  # lacI, tetR, cI
GENES = ["LacI", "TetR", "cI"]


# =============================================================================================
# 1. The algorithm: Gillespie's direct method
# =============================================================================================
def gillespie(x0, stoich, propensity, t_end, dt, rng):
    """Gillespie's direct method (1977): exact sample paths of the chemical master equation.

    It runs many independent cells at once. Every array has one row per cell, and every loop
    iteration fires exactly ONE reaction in EVERY cell (each cell keeps its own clock). The five
    numbered steps are the textbook algorithm, vectorised across cells.

    x0         (cells, species) initial molecule counts
    stoich     (reactions, species) change in the counts when each reaction fires
    propensity f(x) -> (cells, reactions): rate a_j(x) of every reaction [1/time]
    Returns t_grid (G,) and X (cells, G, species): the state sampled every dt. The path is
    piecewise constant, so X at time s is the state held between the jumps around s.
    """
    x = np.array(x0, dtype=np.int64)  # copy, we mutate it
    cells = np.arange(len(x))
    t = np.zeros(len(x))
    t_grid = np.arange(0.0, t_end + dt / 2, dt)
    t_sample = np.append(t_grid, np.inf)  # padded, so a finished cell never samples again
    k = np.zeros(len(x), dtype=np.int64)  # index of each cell's next sample
    X = np.empty((len(x), len(t_grid), x.shape[1]), dtype=np.int32)
    while k.min() < len(t_grid):
        a = propensity(x)  # 1. propensities a_j(x)
        cum = np.cumsum(a, axis=1)
        a0 = cum[:, -1]  # 2. total rate a0 = sum_j a_j
        tau = -np.log(1.0 - rng.random(len(x))) / a0  # 3. waiting time tau ~ Exp(a0)
        due = t_sample[k] < t + tau  # (bookkeeping: store the state for the grid
        while due.any():  #  times that pass before this jump)
            X[cells[due], k[due]] = x[due]
            k[due] += 1
            due = t_sample[k] < t + tau
        j = (cum < (rng.random(len(x)) * a0)[:, None]).sum(axis=1)  # 4. pick j w.p. a_j / a0
        x += stoich[j]  # 5. fire reaction j,
        t += tau  #    advance the clock
    return t_grid, X


# =============================================================================================
# 2. Models: each is an ODE right-hand side + a reaction list (stoichiometry, propensities)
# =============================================================================================
# --- birth-death: 0 -> X at rate k, X -> 0 at rate gamma*x. Stationary law: Poisson(k/gamma).
BD_K, BD_GAMMA = 10.0, 1.0
BD_STOICH = np.array([[1], [-1]])


def bd_propensity(x):
    return np.column_stack([np.full(len(x), BD_K), BD_GAMMA * x[:, 0]])


# --- toggle switch (Gardner 2000), dimensionless: du/dt = a/(1+v^n) - u, dv/dt = a/(1+u^n) - v
# Time is in protein lifetimes. The SSA uses counts U = omega*u, V = omega*v (omega = system size).
TOG_ALPHA, TOG_N = 5.0, 2
TOG_STOICH = np.array([[1, 0], [-1, 0], [0, 1], [0, -1]])


def toggle_rhs(t, y):
    u, v = y
    return [TOG_ALPHA / (1 + v**TOG_N) - u, TOG_ALPHA / (1 + u**TOG_N) - v]


def toggle_propensity(omega):
    def a(x):
        U, V = x[:, 0], x[:, 1]
        return np.column_stack([omega * TOG_ALPHA / (1 + (V / omega) ** TOG_N), U,  # make U, lose U
                                omega * TOG_ALPHA / (1 + (U / omega) ** TOG_N), V])  # make V, lose V
    return a


def toggle_fixed_points():
    """All ODE fixed points: roots of g(u) = u - a/(1+v(u)^n) with v(u) = a/(1+u^n) (bracket + brentq)."""
    v_of = lambda u: TOG_ALPHA / (1 + u**TOG_N)
    g = lambda u: u - TOG_ALPHA / (1 + v_of(u) ** TOG_N)
    grid = np.linspace(0, TOG_ALPHA + 1, 4001)
    s = np.sign(g(grid))
    pts = []
    for i in np.flatnonzero(s[:-1] * s[1:] < 0):
        u = brentq(g, grid[i], grid[i + 1], xtol=1e-12)
        v = v_of(u)
        dfdv = -TOG_ALPHA * TOG_N * v ** (TOG_N - 1) / (1 + v**TOG_N) ** 2
        dgdu = -TOG_ALPHA * TOG_N * u ** (TOG_N - 1) / (1 + u**TOG_N) ** 2
        eig = np.linalg.eigvals([[-1, dfdv], [dgdu, -1]])
        pts.append((u, v, bool(eig.real.max() < 0)))
    return pts  # sorted by u: [B (v high), saddle, A (u high)]


# --- repressilator (Elowitz & Leibler 2000, Box 1), dimensionless, time in mRNA lifetimes:
#   dm_i/dt = -m_i + alpha/(1 + p_{i-1}^n) + alpha0,   dp_i/dt = -beta (p_i - m_i)
# LacI represses tetR, TetR represses cI, cI represses lacI.
REP_ALPHA, REP_ALPHA0, REP_BETA, REP_N = 216.0, 0.216, 0.2, 2
# Copy-number scale from the paper's Fig. 1 numbers: K_M = 40 monomers/cell, 20 proteins per
# transcript, protein/mRNA half-lives 10/2 min. So p = 1 is 40 proteins and m = 1 is 0.4 mRNAs.
MRNA_PER_UNIT, PROT_PER_UNIT = 0.4, 40.0
T_UNIT_MIN = 2 / np.log(2)  # time unit = mRNA lifetime = half-life / ln 2 [min]
I3, Z3 = np.eye(3, dtype=int), np.zeros((3, 3), dtype=int)
REP_STOICH = np.block([[I3, Z3], [-I3, Z3], [Z3, I3], [Z3, -I3]])  # transcribe, decay, translate, decay


def repressilator_rhs(t, y, beta=REP_BETA):
    m, p = y[:3], y[3:]
    return np.concatenate([REP_ALPHA / (1 + np.roll(p, 1) ** REP_N) + REP_ALPHA0 - m, beta * (m - p)])


def repressilator_propensity(s):
    """s scales the copy numbers (s = 1 is the paper's cell); noise shrinks like 1/sqrt(s)."""
    om, op = s * MRNA_PER_UNIT, s * PROT_PER_UNIT

    def a(x):
        M, P = x[:, :3], x[:, 3:]
        return np.hstack([om * (REP_ALPHA / (1 + (P[:, [2, 0, 1]] / op) ** REP_N) + REP_ALPHA0),
                          M, REP_BETA * op / om * M, REP_BETA * P])
    return a


MODELS = {"toggle": (TOG_STOICH, toggle_propensity), "repressilator": (REP_STOICH, repressilator_propensity)}


# =============================================================================================
# 3. The exact alternative: solve the chemical master equation (CME) directly (toggle only)
# =============================================================================================
def toggle_cme(omega):
    """Sparse CME generator Q (Q[x, y] = rate x -> y) on the lattice 0..N for U and V.

    Finite state projection: births past N are dropped, which loses a negligible amount of mass.
    """
    N = int(TOG_ALPHA * omega + 8 * np.sqrt(TOG_ALPHA * omega) + 10)
    U, V = [g.ravel() for g in np.meshgrid(np.arange(N + 1), np.arange(N + 1), indexing="ij")]
    a = toggle_propensity(omega)(np.column_stack([U, V]))  # the very same rates the SSA uses
    rows, cols, vals = [], [], []
    for j, (dU, dV) in enumerate(TOG_STOICH):
        U2, V2 = U + dU, V + dV
        ok = (U2 >= 0) & (U2 <= N) & (V2 >= 0) & (V2 <= N) & (a[:, j] > 0)
        rows.append((U * (N + 1) + V)[ok])
        cols.append((U2 * (N + 1) + V2)[ok])
        vals.append(a[ok, j])
    Q = sparse.csr_matrix((np.concatenate(vals), (np.concatenate(rows), np.concatenate(cols))),
                          shape=(len(U), len(U)))
    return Q - sparse.diags(np.asarray(Q.sum(axis=1)).ravel()), U, V, N


def cme_mfpt(Q, target):
    """Mean first-passage time into `target` from every state: Q h = -1 off target, h = 0 on it."""
    free = ~target
    h = np.zeros(Q.shape[0])
    h[free] = spsolve(Q[free][:, free].tocsc(), -np.ones(free.sum()))
    return h


def cme_stationary(Q):
    """Stationary distribution: Q^T pi = 0 with one balance equation swapped for sum(pi) = 1."""
    A = sparse.vstack([Q.T.tocsr()[:-1], sparse.csr_matrix(np.ones((1, Q.shape[0])))]).tocsc()
    b = np.zeros(Q.shape[0])
    b[-1] = 1
    return spsolve(A, b)


# =============================================================================================
# 4. Worker-process jobs (top-level so they pickle on Windows)
# =============================================================================================
def ssa_job(model, param, x0, t_end, dt, seed):
    stoich, prop = MODELS[model]
    return gillespie(x0, stoich, prop(param), t_end, dt, np.random.default_rng(seed))


def toggle_fpt_job(omega, x_A, delta, t_max, seed, cells=512, dt=0.25, chunk=1000.0):
    """First-passage time from attractor A into B = {V - U >= delta*omega}, censored at t_max.

    Runs in chunks to bound memory. Restarting from the last state is exact (Markov property).
    """
    rng = np.random.default_rng(seed)
    x = np.tile(x_A, (cells, 1))
    t_hit = np.full(cells, np.inf)
    t0 = 0.0
    while t0 < t_max and np.isinf(t_hit).any():
        t, X = gillespie(x, TOG_STOICH, toggle_propensity(omega), chunk, dt, rng)
        in_B = X[:, :, 1] - X[:, :, 0] >= delta * omega
        # ponytail: hit time known to +-dt/2 (grid sampling), exact event-level hit test if it matters
        first = np.where(in_B.any(axis=1), t0 + t[in_B.argmax(axis=1)] - dt / 2, np.inf)
        t_hit = np.minimum(t_hit, first)
        x, t0 = X[:, -1], t0 + chunk
    return omega, np.minimum(t_hit, t0), np.isfinite(t_hit)


# =============================================================================================
# 5. Analyses and figures
# =============================================================================================
def birth_death():
    """Verification: SSA of 0 <-> X must reproduce the Poisson law exactly (not just its mean)."""
    rng = np.random.default_rng([SEED, 0])
    t, X = gillespie(np.zeros((20000, 1), int), BD_STOICH, bd_propensity, 10.0, 0.05, rng)
    lam_t = BD_K / BD_GAMMA * (1 - np.exp(-BD_GAMMA * t))  # from x(0)=0 the law is Poisson(lam(t)) at all t
    n, lam = X[:, -1, 0], lam_t[-1]
    lo, hi = (int(q) for q in stats.poisson.ppf([1e-3, 1 - 1e-3], lam))
    obs = np.bincount(np.clip(n, lo, hi) - lo, minlength=hi - lo + 1)
    exp = stats.poisson.pmf(np.arange(lo, hi + 1), lam)
    exp[0], exp[-1] = stats.poisson.cdf(lo, lam), stats.poisson.sf(hi - 1, lam)  # pooled tails
    chi2, p_chi2 = stats.chisquare(obs, exp * len(n))
    kk = np.arange(n.max() + 1)
    tvd = 0.5 * np.abs(np.bincount(n) / len(n) - stats.poisson.pmf(kk, lam)).sum()
    res = dict(mean=n.mean(), var=n.var(ddof=1), lam=lam, se=np.sqrt(lam / len(n)), fano=n.var(ddof=1) / n.mean(),
               chi2=chi2, dof=len(obs) - 1, p=p_chi2, tvd=tvd)

    fig, (a1, a2) = plt.subplots(1, 2, figsize=(14, 5.2))
    for i in range(40):
        a1.step(t, X[i, :, 0], where="post", lw=0.6, alpha=0.35, color=C_U)
    a1.plot(t, X[:, :, 0].mean(0), color="#ffd166", lw=2.2, label="SSA ensemble mean (20 000 cells)")
    a1.plot(t, lam_t, "--", color=C_ODE, lw=1.6, label=r"ODE  $\dot x = k - \gamma x$")
    a1.fill_between(t, lam_t - np.sqrt(lam_t), lam_t + np.sqrt(lam_t), color=C_ODE, alpha=0.12,
                    label=r"Poisson $\pm\sqrt{\lambda(t)}$")
    a1.set(xlabel=r"time [$1/\gamma$]", ylabel="molecules $x$", title="Birth-death: 40 exact sample paths")
    a1.legend(loc="lower right", fontsize=9)
    for ti, c in zip([0.25, 1.0, 10.0], ["#f72585", "#ffd166", C_U]):
        i = int(round(ti / 0.05))
        cnt = np.bincount(X[:, i, 0], minlength=30)[:30] / X.shape[0]
        a2.bar(np.arange(30), cnt, width=0.85, color=c, alpha=0.45)
        a2.plot(np.arange(30), stats.poisson.pmf(np.arange(30), lam_t[i]), "o-", ms=4, lw=1, color=c,
                label=fr"$t={ti:g}$: SSA (bars) vs Poisson({lam_t[i]:.2f})")
    a2.set(xlabel="molecules $x$", ylabel="probability", xlim=(-0.7, 25),
           title=fr"Stationary: TVD = {tvd:.4f}, Fano = {res['fano']:.3f}, $\chi^2$ p = {p_chi2:.2f}")
    a2.legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(OUT / "birth_death_verification.png", dpi=150)
    plt.close(fig)
    return res


def rep_ode_reference():
    """Settle onto the limit cycle, return (state at a LacI peak, ODE period in time units)."""
    sol = solve_ivp(repressilator_rhs, (0, 1500), [1, 0, 0, 2, 1, 3.0], method="LSODA",
                    rtol=1e-9, atol=1e-11, dense_output=True)
    tt = np.linspace(900, 1500, 600001)
    p1 = sol.sol(tt)[3]
    pk, _ = find_peaks(p1)
    return sol.sol(tt[pk[-1]]), np.diff(tt[pk]).mean()


def rep_periods(t, P, T_ode, amp):
    """Peak-to-peak intervals of every cell's protein trace (peaks >= half an ODE period apart)."""
    dt = t[1] - t[0]
    out = [np.diff(t[find_peaks(p, distance=int(0.5 * T_ode / dt), prominence=0.3 * amp)[0]]) for p in P]
    return np.concatenate(out)


def spectrum(t, P):
    """Cell-averaged power spectrum of mean-removed traces (Hann window, 8x zero padding)."""
    y = (P - P.mean(axis=1, keepdims=True)) * np.hanning(P.shape[1])
    n = 8 * P.shape[1]
    S = (np.abs(fft.rfft(y, n=n, axis=1)) ** 2).mean(0)
    f = fft.rfftfreq(n, d=t[1] - t[0])
    return f, S / S.max()


def solver_benchmark():
    """Explicit (RK45, DOP853) vs implicit/stiff (BDF, Radau) vs auto-switching (LSODA).

    Stiffness knob: protein lifetime / mRNA lifetime = 1/beta. Real cells sit at 5..100+.
    """
    methods = ["RK45", "DOP853", "LSODA", "BDF", "Radau"]
    ratios = [5, 50, 500, 5000]
    y0 = [1, 0, 0, 2, 1, 3.0]
    rows = []
    for r in ratios:
        beta, t_end = 1 / r, 20.0 * r  # 20 protein lifetimes
        ref = solve_ivp(repressilator_rhs, (0, t_end), y0, method="Radau", rtol=1e-11, atol=1e-12, args=(beta,)).y[:, -1]
        for m in methods:
            best = np.inf
            for rep in range(3):
                t0 = time.perf_counter()
                sol = solve_ivp(repressilator_rhs, (0, t_end), y0, method=m, rtol=1e-6, atol=1e-9, args=(beta,))
                el = time.perf_counter() - t0
                best = min(best, el)
                if el > 1.0:  # slow runs: once is enough
                    break
            err = np.abs(sol.y[:, -1] - ref).max() / np.abs(ref).max()
            rows.append((r, m, best, sol.nfev, sol.njev, err))
    return methods, ratios, rows


# =============================================================================================
# 6. Main
# =============================================================================================
def main():
    T0 = time.perf_counter()
    OUT.mkdir(exist_ok=True)
    plt.style.use("dark_background")
    plt.rcParams.update({"font.size": 11, "axes.grid": True, "grid.alpha": 0.15, "axes.titlesize": 12,
                         "figure.facecolor": "#0b0c10", "axes.facecolor": "#0b0c10", "savefig.facecolor": "#0b0c10"})

    # ---------------- (a) verification of the SSA against an exact law ----------------
    bd = birth_death()
    print(f"[birth-death] mean {bd['mean']:.4f} (exact {bd['lam']:.4f} +- {bd['se']:.4f}), "
          f"Fano {bd['fano']:.4f} (exact 1), chi2 {bd['chi2']:.1f}/{bd['dof']} dof p={bd['p']:.3f}, TVD {bd['tvd']:.4f}")

    # ---------------- (b) toggle: deterministic picture + exact CME ----------------
    fps = toggle_fixed_points()
    (uB, vB, sB), (us, vs, ss), (uA, vA, sA) = fps
    delta = (uA - vA) / 2  # B = {v - u >= delta}: halfway between saddle and attractor
    print(f"[toggle ODE] fixed points: " + ", ".join(f"({u:.3f},{v:.3f}) {'stable' if s else 'saddle'}" for u, v, s in fps))
    omegas_cme = np.arange(2, 21)
    x_A = {om: np.array([round(uA * om), round(vA * om)]) for om in omegas_cme}
    tau_cme, pi = {}, {}
    for om in omegas_cme:
        Q, U, V, N = toggle_cme(om)
        h = cme_mfpt(Q, V - U >= delta * om)
        tau_cme[om] = h[x_A[om][0] * (N + 1) + x_A[om][1]]
        if om in (4, 6):
            pi[om] = (cme_stationary(Q).reshape(N + 1, N + 1), N)
    print("[toggle CME] exact MFPT A->B: " + ", ".join(f"W={om}:{tau_cme[om]:.3g}" for om in omegas_cme))

    # ---------------- (c) repressilator ODE ----------------
    y_lc, T_ode = rep_ode_reference()
    print(f"[repressilator ODE] period {T_ode:.3f} mRNA lifetimes = {T_ode * T_UNIT_MIN:.1f} min")

    # ---------------- (d) all stochastic simulations, 4 worker processes ----------------
    t_sim = time.perf_counter()
    omegas_ssa = list(range(2, 11))
    scales = [0.25, 0.5, 1.0, 2.0, 4.0]
    with Pool(N_PROC) as pool:
        jobs = {}
        # longest first so the 4 workers stay balanced; SSA cost ~ s * t_end, so s=4 runs half as long
        for i, s in enumerate(scales[::-1]):
            cells, t_end = (200, 200.0) if s > 2 else (100, 400.0)
            x0 = np.tile(np.round(np.r_[y_lc[:3] * s * MRNA_PER_UNIT, y_lc[3:] * s * PROT_PER_UNIT]), (cells, 1))
            jobs[("rep", s)] = pool.apply_async(ssa_job, ("repressilator", s, x0, t_end, 0.5, [SEED, 10 + i]))
        for om in omegas_ssa[::-1]:
            jobs[("fpt", om)] = pool.apply_async(toggle_fpt_job, (om, x_A[om], delta, 3 * tau_cme[om], [SEED, 100 + om]))
        for om in (3, 6, 9):
            jobs[("trace", om)] = pool.apply_async(ssa_job, ("toggle", om, x_A[om][None], 4000.0, 0.5, [SEED, 200 + om]))
        jobs[("pop", 4)] = pool.apply_async(ssa_job, ("toggle", 4, np.tile(x_A[4], (400, 1)), 2000.0, 1.0, [SEED, 300]))
        res = {key: job.get() for key, job in jobs.items()}
    print(f"[SSA] all stochastic runs done in {time.perf_counter() - t_sim:.1f} s (4 processes)")

    # ---------------- (e) toggle analyses ----------------
    fpt = {}
    for om in omegas_ssa:
        _, T, hit = res[("fpt", om)]
        tau_hat = T.sum() / hit.sum()  # exponential MLE with right-censoring
        half = 1.96 / np.sqrt(hit.sum())
        fpt[om] = (tau_hat, tau_hat * np.exp(-half), tau_hat * np.exp(half), hit.sum(), len(T), T, hit)
    z = {om: np.log(fpt[om][0] / tau_cme[om]) * np.sqrt(fpt[om][3]) for om in omegas_ssa}
    big = omegas_cme >= 10
    slope, icpt = np.polyfit(omegas_cme[big], np.log([tau_cme[o] for o in omegas_cme[big]]), 1)
    r2 = np.corrcoef(omegas_cme[big], np.log([tau_cme[o] for o in omegas_cme[big]]))[0, 1] ** 2
    for om in omegas_ssa:
        print(f"  W={om:2d}: SSA tau = {fpt[om][0]:9.1f}  [{fpt[om][1]:.1f}, {fpt[om][2]:.1f}]  "
              f"({fpt[om][3]}/{fpt[om][4]} hits)   CME exact = {tau_cme[om]:9.1f}   z = {z[om]:+.2f}")
    print(f"[toggle] ln tau = {icpt:.3f} + {slope:.4f} W  (CME, W>=10), R^2 = {r2:.5f}")

    t_pop, X_pop = res[("pop", 4)]
    P4, N4 = pi[4]
    U_ssa = X_pop[:, t_pop >= 500, 0].ravel()
    p_ssa = np.bincount(U_ssa, minlength=N4 + 1)[: N4 + 1] / U_ssa.size
    p_cme = P4.sum(axis=1)
    tvd_toggle = 0.5 * np.abs(p_ssa - p_cme).sum()
    print(f"[toggle W=4] stationary P(U): SSA vs CME TVD = {tvd_toggle:.4f}")

    fig_toggle(fps, pi, res, tvd_toggle, p_ssa, p_cme)
    fig_switching(omegas_cme, tau_cme, fpt, slope, icpt, r2)
    gif_toggle_population(fps, t_pop, X_pop, tau_cme[4])
    make_hero(pi[6], res[("trace", 6)], fps, omegas_cme, tau_cme, fpt)

    # ---------------- (f) repressilator analyses ----------------
    t_rep = res[("rep", 1.0)][0]
    ode = solve_ivp(repressilator_rhs, (0, t_rep[-1]), y_lc, method="LSODA", rtol=1e-9, atol=1e-11, t_eval=t_rep).y
    amp = ode[3].max() - ode[3].min()
    rep = {}
    print(f"[repressilator ODE] LacI peak {ode[3].max():.2f} K_M")
    for s in scales:
        t_s = res[("rep", s)][0]
        P = res[("rep", s)][1][:, :, 3] / (s * PROT_PER_UNIT)  # LacI in units of K_M
        dt = t_s[1] - t_s[0]
        pk = [p[find_peaks(p, distance=int(0.5 * T_ode / dt), prominence=0.3 * amp)[0]] for p in P]
        per = rep_periods(t_s, P, T_ode, amp) * T_UNIT_MIN
        f, S = spectrum(t_s, P)
        f_pk = f[1:][np.argmax(S[1:])]
        rep[s] = dict(per=per, mean=per.mean(), sd=per.std(ddof=1), cv=per.std(ddof=1) / per.mean(),
                      f=f, S=S, T_fft=T_UNIT_MIN / f_pk, peak=np.concatenate(pk).mean())
        print(f"[repressilator s={s}] period {rep[s]['mean']:.1f} +- {rep[s]['sd']:.1f} min "
              f"(CV {rep[s]['cv']:.3f}, n={len(per)}), FFT peak period {rep[s]['T_fft']:.1f} min, "
              f"mean LacI peak {rep[s]['peak']:.1f} K_M")
    cv_all = np.polyfit(np.log(scales), np.log([rep[s]["cv"] for s in scales]), 1)[0]
    lna = [s for s in scales if s >= 1]  # linear-noise theory is a large-copy-number law
    cv_slope = np.polyfit(np.log(lna), np.log([rep[s]["cv"] for s in lna]), 1)[0]
    print(f"[repressilator] CV ~ s^{cv_slope:.3f} for s>=1, s^{cv_all:.3f} over all s  (linear-noise theory: -0.5)")
    fig_repressilator(t_rep, ode, res, rep, T_ode, cv_slope, scales)
    gif_repressilator_cells(t_rep, ode, res[("rep", 1.0)][1])

    # ---------------- (g) ODE solver benchmark (run alone for clean timings) ----------------
    methods, ratios, rows = solver_benchmark()
    print("[solvers]  ratio  method   time[s]    nfev   njev   rel.err")
    for r, m, el, nfev, njev, err in rows:
        print(f"           {r:5d}  {m:7s} {el:8.4f} {nfev:7d} {njev:6d}   {err:.1e}")
    fig_solvers(methods, ratios, rows)

    # ---------------- (h) verification summary ----------------
    zs = np.array(list(z.values()))
    checks = [
        ("birth-death mean = k/gamma (|z|<4)", f"{bd['lam']:.3f}", f"{bd['mean']:.3f}", abs(bd["mean"] - bd["lam"]) < 4 * bd["se"]),
        ("birth-death Fano factor = 1 (+-0.03)", "1.000", f"{bd['fano']:.3f}", abs(bd["fano"] - 1) < 0.03),
        ("birth-death chi^2 GOF vs Poisson (p>0.001)", "p>0.001", f"p={bd['p']:.3f}", bd["p"] > 1e-3),
        ("birth-death total-variation distance (<0.02)", "0", f"{bd['tvd']:.4f}", bd["tvd"] < 0.02),
        ("toggle ODE: 2 stable + 1 saddle", "2+1", f"{sA + sB:d}+{(not ss):d}", sA and sB and not ss),
        ("toggle SSA vs CME stationary P(U), W=4 (TVD<0.03)", "0", f"{tvd_toggle:.4f}", tvd_toggle < 0.03),
        ("toggle SSA MFPT vs exact CME, all W (|z|<3)", "|z|<3", f"max|z|={np.abs(zs).max():.2f}", np.abs(zs).max() < 3),
        ("toggle ln(tau) linear in W, CME W>=10 (R^2>0.999)", ">0.999", f"{r2:.5f}", r2 > 0.999),
        ("repressilator ODE period vs Elowitz 160+-40 min", "120-200", f"{T_ode * T_UNIT_MIN:.0f}", 120 <= T_ode * T_UNIT_MIN <= 200),
        ("repressilator SSA -> ODE (Kurtz): s=4 period within 5%", f"{T_ode * T_UNIT_MIN:.1f}", f"{rep[4.0]['mean']:.1f}",
         abs(rep[4.0]["mean"] / (T_ode * T_UNIT_MIN) - 1) < 0.05),
        ("repressilator period CV ~ s^-1/2 for s>=1 (+-0.15)", "-0.50", f"{cv_slope:.3f}", abs(cv_slope + 0.5) < 0.15),
    ]
    print("\n=== VERIFICATION ===")
    for name, exp_, got, ok in checks:
        print(f"  [{'PASS' if ok else 'FAIL'}] {name:55s} expected {exp_:>10s}  got {got}")
    print(f"\nTotal runtime {time.perf_counter() - T0:.1f} s")
    assert all(c[3] for c in checks), "verification failed"


# ---------------------------------------------------------------------------------------------
# figures
# ---------------------------------------------------------------------------------------------
def phase_plane(ax, fps, lim=6.2):
    """Nullclines, vector field and fixed points of the toggle ODE."""
    g = np.linspace(0, lim, 60)
    uu, vv = np.meshgrid(g, g)
    du, dv = toggle_rhs(0, [uu, vv])
    ax.streamplot(uu, vv, du, dv, color="#555a66", density=1.1, linewidth=0.6, arrowsize=0.7)
    s = np.linspace(0, lim, 400)
    ax.plot(TOG_ALPHA / (1 + s**TOG_N), s, color=C_U, lw=2, label=r"$\dot u=0$ nullcline")
    ax.plot(s, TOG_ALPHA / (1 + s**TOG_N), color=C_V, lw=2, label=r"$\dot v=0$ nullcline")
    for u, v, stable in fps:
        ax.plot(u, v, "o", ms=10, mfc="white" if stable else "none", mec="white", mew=2, zorder=5)
    ax.set(xlim=(0, lim), ylim=(0, lim), xlabel="repressor 1  $u = U/\\Omega$", ylabel="repressor 2  $v = V/\\Omega$")
    ax.set_aspect("equal")


def fig_toggle(fps, pi, res, tvd, p_ssa, p_cme):
    fig = plt.figure(figsize=(17, 8.6))
    gs = fig.add_gridspec(4, 2, width_ratios=[1, 1.45], hspace=0.55, wspace=0.24)
    ax = fig.add_subplot(gs[:, 0])
    P6, N6 = pi[6]
    g6 = np.arange(N6 + 1) / 6
    ax.contourf(g6, g6, np.log10(np.maximum(P6.T, 1e-12)), levels=np.linspace(-6, np.log10(P6.max()), 14),
                cmap="magma", vmin=-6, vmax=np.log10(P6.max()) + 1.5, zorder=0)
    phase_plane(ax, fps)
    t6, X6 = res[("trace", 6)]
    ax.plot(X6[0, :4000, 0] / 6, X6[0, :4000, 1] / 6, color="white", lw=0.4, alpha=0.3, label=r"SSA path, $\Omega=6$")
    ax.plot([], [], "o", mfc="white", mec="white", label="stable (ODE attractor)")
    ax.plot([], [], "o", mfc="none", mec="white", label="saddle")
    ax.legend(loc="upper right", fontsize=9, framealpha=0.8)
    ax.set_title("Phase plane + exact stationary density $\\log_{10}P$ (CME, $\\Omega=6$)")

    ode = solve_ivp(toggle_rhs, (0, 4000), [fps[2][0], fps[2][1]], t_eval=res[("trace", 3)][0], rtol=1e-9)
    for row, om in enumerate((3, 6, 9)):
        a = fig.add_subplot(gs[row, 1])
        t, X = res[("trace", om)]
        a.plot(t, X[0, :, 0] / om, color=C_U, lw=0.6, label="$U/\\Omega$ (SSA)")
        a.plot(t, X[0, :, 1] / om, color=C_V, lw=0.6, alpha=0.8, label="$V/\\Omega$ (SSA)")
        a.plot(ode.t, ode.y[0], "--", color=C_ODE, lw=1.4, label="$u$ (ODE)")
        a.set(xlim=(0, 4000), ylim=(-0.3, 8.5), ylabel="conc.")
        a.set_title(f"$\\Omega = {om}$ molecules per unit: the ODE sits in state A forever, the cell flips", fontsize=10.5, loc="left")
        if row == 0:
            a.legend(ncol=3, fontsize=8.5, loc="upper right", framealpha=0.8)
    a.set_xlabel("time [protein lifetimes]")
    a = fig.add_subplot(gs[3, 1])
    k = np.arange(len(p_cme))
    a.bar(k, p_ssa, width=1.0, color=C_U, alpha=0.5, label="SSA, 400 cells, $t\\geq500$")
    a.plot(k, p_cme, color="#ffd166", lw=1.8, label="exact CME (sparse solve)")
    for u, v, st in fps:
        if st:
            a.axvline(u * 4, color=C_ODE, ls="--", lw=1.2)
    a.plot([], [], "--", color=C_ODE, label="ODE: point attractors")
    a.set(xlabel="$U$ molecules ($\\Omega=4$)", ylabel="$P(U)$", xlim=(-0.5, 36))
    a.set_title(f"Stationary law is bimodal. SSA vs CME: TVD = {tvd:.4f}", fontsize=10.5, loc="left")
    a.legend(fontsize=8.5, loc="upper center", framealpha=0.8)
    fig.suptitle("Genetic toggle switch (Gardner et al. 2000): deterministic ODE vs exact stochastic simulation",
                 fontsize=15, y=0.99)
    fig.savefig(OUT / "toggle_ode_vs_ssa.png", dpi=140, bbox_inches="tight")
    plt.close(fig)


def fig_switching(omegas_cme, tau_cme, fpt, slope, icpt, r2):
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(15, 5.8), gridspec_kw={"width_ratios": [1.25, 1]})
    om = np.array(omegas_cme)
    a1.semilogy(om, [tau_cme[o] for o in om], "-", color="#ffd166", lw=2.2, label="exact CME mean first-passage time")
    ws = sorted(fpt)
    tau = np.array([fpt[w][0] for w in ws])
    err = np.array([[fpt[w][0] - fpt[w][1] for w in ws], [fpt[w][2] - fpt[w][0] for w in ws]])
    a1.errorbar(ws, tau, yerr=err, fmt="o", color=C_U, ms=7, capsize=4, label="Gillespie SSA, 512 cells (95% CI)")
    xf = np.linspace(2, 20, 50)
    a1.semilogy(xf, np.exp(icpt + slope * xf), ":", color=C_V, lw=1.8,
                label=fr"fit $\tau \propto e^{{\Delta\Phi\,\Omega}}$, $\Delta\Phi$ = {slope:.3f}  ($R^2$={r2:.4f})")
    a1.plot([], [], "--", color=C_ODE, label=r"ODE: $\tau = \infty$ for every $\Omega$")
    a1.set(xlabel=r"system size $\Omega$ (molecules per concentration unit)", ylabel="mean switching time [protein lifetimes]",
           title="Noise-induced switching slows exponentially with system size", ylim=(10, None))
    a1.legend(fontsize=9, loc="upper left")
    a1.text(0.98, 0.04, "cell cycle ~ 1 protein lifetime (dilution)\nso $\\Omega=20$: ~ once per " +
            f"{tau_cme[20]:.1e} generations", transform=a1.transAxes, ha="right", fontsize=9, color="#aaaaaa")
    for w, c in zip((3, 6, 9), ("#ffd166", C_U, C_V)):
        _, _, _, _, _, T, hit = fpt[w]
        ts = np.sort(T[hit])
        surv = 1 - np.arange(1, len(ts) + 1) / len(T)
        a2.semilogy(ts / tau_cme[w], surv, color=c, lw=1.8, label=fr"SSA, $\Omega={w}$")
    xx = np.linspace(0, 4, 50)
    a2.semilogy(xx, np.exp(-xx), "--", color=C_ODE, lw=1.4, label=r"memoryless: $e^{-t/\tau_{\rm CME}}$")
    a2.set(xlabel=r"$t\,/\,\tau_{\rm CME}$", ylabel="P(not yet switched)", xlim=(0, 4), ylim=(1e-2, 1.05),
           title="Waiting times are exponential (a Poisson process of rare flips)")
    a2.legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(OUT / "toggle_switching_time.png", dpi=150)
    plt.close(fig)


def gif_toggle_population(fps, t, X, tau):
    """400 identical cells start in state A. ODE: they stay. SSA: noise splits the population."""
    frames = np.flatnonzero(t <= 300)[::2]
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(11, 5), dpi=80, gridspec_kw={"width_ratios": [1, 1.15]})
    s = np.linspace(0, 6.2, 300)
    a1.plot(TOG_ALPHA / (1 + s**TOG_N), s, color=C_U, lw=1.2, alpha=0.6)
    a1.plot(s, TOG_ALPHA / (1 + s**TOG_N), color=C_V, lw=1.2, alpha=0.6)
    a1.plot(fps[1][0], fps[1][1], "o", ms=8, mfc="none", mec="white")
    jitter = np.random.default_rng(1).uniform(-0.1, 0.1, (X.shape[0], 2))  # only so equal counts don't overlap
    sc = a1.scatter(X[:, 0, 0] / 4 + jitter[:, 0], X[:, 0, 1] / 4 + jitter[:, 1], s=9, c=np.zeros(X.shape[0]),
                    cmap=LinearSegmentedColormap.from_list("uv", [C_U, C_V]), vmin=0, vmax=1)
    a1.plot(fps[2][0], fps[2][1], "*", ms=20, color="white", mec="k", label="ODE cell (all 400 identical)")
    a1.set(xlim=(0, 6.2), ylim=(0, 6.2), xlabel="$u = U/\\Omega$", ylabel="$v = V/\\Omega$")
    a1.set_aspect("equal")
    a1.legend(loc="upper right", fontsize=9)
    title = a1.set_title("")
    frac = (X[:, :, 1] > X[:, :, 0]).mean(0)
    a2.plot(t[frames], 0.5 * (1 - np.exp(-2 * t[frames] / tau)), "--", color="#ffd166", lw=1.6,
            label=r"two-state Markov: $\frac{1}{2}(1-e^{-2t/\tau_{\rm CME}})$")
    a2.axhline(0, color=C_ODE, lw=2, label="ODE: nobody ever switches")
    (line,) = a2.plot([], [], color=C_V, lw=2.2, label="SSA: fraction with $V>U$")
    a2.set(xlim=(0, t[frames[-1]]), ylim=(-0.03, 0.62), xlabel="time [protein lifetimes]", ylabel="fraction of cells in state B")
    a2.legend(loc="lower right", fontsize=9)
    fig.suptitle("Toggle switch, $\\Omega=4$: identical cells, identical genes, different fates", fontsize=13)
    fig.tight_layout()

    def update(i):
        k = frames[i]
        sc.set_offsets(np.column_stack([X[:, k, 0] / 4 + jitter[:, 0], X[:, k, 1] / 4 + jitter[:, 1]]))
        sc.set_array((X[:, k, 1] > X[:, k, 0]).astype(float))
        title.set_text(f"t = {t[k]:5.0f}   in state B: {frac[k] * 100:4.1f}%")
        line.set_data(t[frames[: i + 1]], frac[frames[: i + 1]])
        return sc, title, line

    anim = animation.FuncAnimation(fig, update, frames=len(frames), blit=False)
    anim.save(OUT / "toggle_population.gif", writer=animation.PillowWriter(fps=15))
    plt.close(fig)


def make_hero(pi6, trace6, fps, omegas_cme, tau_cme, fpt):
    """Gallery still: exact probability landscape + one Gillespie path + the exponential law."""
    P, N = pi6
    om = 6
    n = int(6.4 * om)
    L = np.log10(np.maximum(P[:n, :n], 1e-12))
    ext = (0, (n - 1) / om, 0, (n - 1) / om)
    fig = plt.figure(figsize=(16, 9), facecolor="black")
    ax = fig.add_axes([0.035, 0.07, 0.46, 0.86], facecolor="black")
    ax.imshow(ndimage.zoom(L, 6, order=3).T, origin="lower", extent=ext, cmap="magma", vmin=-7.5, vmax=L.max() + 1.0,
              interpolation="bicubic")
    ax.contour(ndimage.zoom(L, 3, order=3).T, levels=np.arange(-7, -1, 1), extent=ext, colors="white", linewidths=0.4, alpha=0.3)
    g = np.linspace(0, ext[1], 40)
    uu, vv = np.meshgrid(g, g)
    du, dv = toggle_rhs(0, [uu, vv])
    ax.streamplot(uu, vv, du, dv, color=(1, 1, 1, 0.18), density=1.0, linewidth=0.6, arrowsize=0.6)
    t, X = trace6
    xy = X[0, :6001] / om + np.random.default_rng(3).normal(0, 0.03, (6001, 2))  # tiny jitter: lattice -> readable line
    seg = np.stack([xy[:-1], xy[1:]], axis=1)
    u_wins = X[0, 1:6001, 0] > X[0, 1:6001, 1]
    colors = np.where(u_wins[:, None], [0.3, 0.79, 0.94, 0.22], [0.97, 0.15, 0.52, 0.3])
    ax.add_collection(LineCollection(seg, colors=colors, linewidths=0.5))
    for u, v, st in fps:
        ax.plot(u, v, "o", ms=13 if st else 10, mfc="white" if st else "none", mec="white", mew=2.2, zorder=5)
    ax.text(fps[2][0], fps[2][1] + 0.35, "ODE attractor A", color="white", ha="center", fontsize=11)
    ax.text(fps[0][0] + 0.3, fps[0][1] + 0.1, "ODE attractor B", color="white", fontsize=11)
    ax.text(fps[1][0] + 0.25, fps[1][1] + 0.2, "saddle", color="white", fontsize=10)
    ax.set(xlim=ext[:2], ylim=ext[2:])
    ax.set_xlabel("repressor 1  (molecules / $\\Omega$)", color="#bbbbbb")
    ax.set_ylabel("repressor 2  (molecules / $\\Omega$)", color="#bbbbbb")
    ax.tick_params(colors="#888888")
    ax.grid(False)

    fig.text(0.53, 0.9, "Noise flips a genetic switch", fontsize=32, color="white", weight="bold")
    fig.text(0.53, 0.855, "Gardner-Collins toggle switch: deterministic ODE vs exact Gillespie simulation.\n"
             "Background: exact stationary probability (chemical master equation, $\\Omega=6$).\n"
             "Line: one simulated cell, coloured by which repressor is winning.",
             fontsize=13, color="#cccccc", va="top")
    a = fig.add_axes([0.56, 0.45, 0.415, 0.25], facecolor="black")
    a.plot(t, X[0, :, 0] / om, color=C_U, lw=0.5)
    a.plot(t, X[0, :, 1] / om, color=C_V, lw=0.5, alpha=0.85)
    a.axhline(fps[2][0], color="white", ls="--", lw=1.6)
    a.set_title("One simulated cell keeps flipping. The ODE (dashed) stays in A forever.", color="white", fontsize=12, loc="left")
    a.set(xlim=(0, t[-1]), ylim=(-0.2, 8.8), ylabel="concentration")
    a.set_xlabel("time [protein lifetimes]", color="#bbbbbb")
    a.grid(False)
    a = fig.add_axes([0.56, 0.08, 0.415, 0.27], facecolor="black")
    o = np.array(omegas_cme)
    a.semilogy(o, [tau_cme[w] for w in o], color="#ffd166", lw=2.5, label="exact (master equation)")
    a.semilogy(sorted(fpt), [fpt[w][0] for w in sorted(fpt)], "o", color=C_U, ms=7, label="Gillespie SSA")
    a.set(xlabel="system size $\\Omega$ (molecules per concentration unit)", ylabel="mean switch time")
    a.legend(loc="upper left", fontsize=10, frameon=False)
    a.text(0.97, 0.08, "grows like $e^{\\,\\Delta\\Phi\\,\\Omega}$\nODE: never", transform=a.transAxes, ha="right",
           color="white", fontsize=13)
    a.grid(alpha=0.15)
    fig.savefig(OUT / "hero.png", dpi=120, facecolor="black")
    plt.close(fig)


def fig_repressilator(t, ode, res, rep, T_ode, cv_slope, scales):
    hrs = t * T_UNIT_MIN / 60
    X1 = res[("rep", 1.0)][1]
    fig = plt.figure(figsize=(17, 10))
    gs = fig.add_gridspec(3, 3, height_ratios=[1, 1, 1.25], hspace=0.5, wspace=0.32)
    s_col = dict(zip(scales, ["#f72585", "#ffd166", C_U, "#80ffdb", "#c77dff"]))
    a = fig.add_subplot(gs[0, :2])
    for i in range(3):
        a.plot(hrs, ode[3 + i] * PROT_PER_UNIT, color=C_GENES[i], lw=1.6, label=GENES[i])
    a.set(xlim=(0, 16), ylabel="proteins / cell", title="ODE: a perfect clock (limit cycle)")
    a.legend(ncol=3, fontsize=9, loc="upper right")
    a = fig.add_subplot(gs[1, :2])
    for i in range(3):
        a.plot(hrs, X1[0, :, 3 + i], color=C_GENES[i], lw=1.1)
    a2 = a.twinx()
    a2.fill_between(hrs, X1[0, :, 0], step="post", color="#aaaaaa", alpha=0.35, lw=0)
    a2.set_ylim(0, 4 * max(X1[0, :, 0].max(), 1))
    a2.tick_params(axis="y", colors="#aaaaaa")
    a2.grid(False)
    a.set(xlim=(0, 16), xlabel="time [h]", ylabel="proteins / cell",
          title="Gillespie SSA, one cell at the paper's copy numbers (grey: LacI mRNA, right axis)")
    a = fig.add_subplot(gs[:2, 2])
    a.plot(X1[0, :, 3], X1[0, :, 4], color=C_U, lw=0.5, alpha=0.7, label=f"SSA (one cell, {hrs[-1]:.0f} h)")
    a.plot(ode[3] * PROT_PER_UNIT, ode[4] * PROT_PER_UNIT, color=C_ODE, lw=2, label="ODE limit cycle")
    a.set(xlabel="LacI proteins", ylabel="TetR proteins", title="Phase space")
    a.legend(fontsize=9, loc="upper right")
    # bottom row: period distribution, spectra, CV scaling
    a = fig.add_subplot(gs[2, 0])
    bins = np.linspace(60, 220, 50)
    for s, c in s_col.items():
        a.hist(rep[s]["per"], bins=bins, histtype="step", lw=1.8, color=c, density=True,
               label=f"s={s:g}: {rep[s]['mean']:.0f}$\\pm${rep[s]['sd']:.0f} min")
    a.axvline(T_ode * T_UNIT_MIN, color=C_ODE, ls="--", lw=1.6, label=f"ODE {T_ode * T_UNIT_MIN:.0f} min")
    a.axvspan(120, 200, color="#80ffdb", alpha=0.08, label="observed 160$\\pm$40 min")
    a.set(xlabel="peak-to-peak period [min]", ylabel="density", title="Period distribution (find_peaks)")
    a.legend(fontsize=8)
    a = fig.add_subplot(gs[2, 1])
    for s, c in s_col.items():
        a.semilogy(rep[s]["f"] / T_UNIT_MIN * 60, rep[s]["S"], color=c, lw=1.4, label=f"SSA s={s:g}")
    fo, So = spectrum(t, ode[3][None])
    a.semilogy(fo / T_UNIT_MIN * 60, So, "--", color=C_ODE, lw=1.2, label="ODE")
    a.set(xlim=(0, 2.0), ylim=(1e-5, 1.5), xlabel="frequency [cycles / h]", ylabel="power (normalised)",
          title="Power spectrum (FFT): noise broadens the line")
    a.legend(fontsize=8)
    a = fig.add_subplot(gs[2, 2])
    cv = [rep[s]["cv"] for s in scales]
    a.loglog(scales, cv, "o-", color=C_U, ms=8, lw=2, label=f"SSA (fit $s\\geq1$: CV $\\propto s^{{{cv_slope:.2f}}}$)")
    a.loglog(scales, cv[2] * (np.array(scales) / scales[2]) ** -0.5, "--", color="#ffd166", label="$s^{-1/2}$ (linear noise)")
    a.set_xticks(scales, [f"{s:g}" for s in scales])
    a.minorticks_off()
    a.axhline(40 / 160, color="#80ffdb", ls=":", label="observed CV (Elowitz 2000)")
    a.set(xlabel="copy-number scale $s$ (s=1: paper's cell)", ylabel="period CV", title="Clock precision vs molecule count")
    a.legend(fontsize=8)
    fig.suptitle("Repressilator (Elowitz & Leibler 2000): ODE limit cycle vs exact stochastic simulation", fontsize=15, y=0.99)
    fig.savefig(OUT / "repressilator.png", dpi=140, bbox_inches="tight")
    plt.close(fig)


def gif_repressilator_cells(t, ode, X):
    """Elowitz-style time-lapse: identical cells blinking. ODE: in lock-step. SSA: they drift apart."""
    rng = np.random.default_rng(7)
    n = X.shape[0]
    side = int(np.sqrt(n))
    gx, gy = np.meshgrid(np.arange(side), np.arange(side))
    xy = np.column_stack([gx.ravel(), gy.ravel()]) + rng.normal(0, 0.08, (n, 2))
    ang = rng.uniform(0, 180, n)
    ode_p = ode[3] * PROT_PER_UNIT
    vmax = ode_p.max()
    gfp = LinearSegmentedColormap.from_list("gfp", ["#02140a", "#0c5c2a", "#39ff14", "#eaffd0"])
    frames = np.arange(0, len(t), 6)
    hrs = t * T_UNIT_MIN / 60
    fig = plt.figure(figsize=(10, 6.4), dpi=64)
    gs = fig.add_gridspec(2, 2, height_ratios=[1.35, 1])
    cols = []
    for c, name in enumerate(["ODE: 100 identical deterministic cells", "SSA: 100 cells with molecular noise"]):
        a = fig.add_subplot(gs[0, c])
        ec = EllipseCollection(np.full(n, 0.85), np.full(n, 0.38), ang, offsets=xy, units="xy",
                               offset_transform=a.transData, cmap=gfp, clim=(0, vmax))
        a.add_collection(ec)
        a.set(xlim=(-0.8, side - 0.2), ylim=(-0.8, side - 0.2), title=name)
        a.set_aspect("equal")
        a.axis("off")
        cols.append(ec)
    a = fig.add_subplot(gs[1, :])
    for i in range(6):
        a.plot(hrs, X[i, :, 3], color="#39ff14", lw=0.5, alpha=0.35)
    a.plot(hrs, X[:, :, 3].mean(0), color="#39ff14", lw=2.2, label="SSA population mean (dephases)")
    a.plot(hrs, ode_p, color=C_ODE, lw=1.2, alpha=0.8, label="ODE (every cell)")
    cursor = a.axvline(0, color="#ffd166", lw=1.5)
    a.set(xlim=(0, hrs[-1]), xlabel="time [h]", ylabel="LacI reporter [proteins]")
    a.legend(loc="upper right", fontsize=9, ncol=2)
    sup = fig.suptitle("")
    fig.tight_layout(rect=(0, 0, 1, 0.95))

    def update(i):
        k = frames[i]
        cols[0].set_array(np.full(n, ode_p[k]))
        cols[1].set_array(X[:, k, 3].astype(float))
        cursor.set_xdata([hrs[k]])
        sup.set_text(f"Repressilator time-lapse   t = {hrs[k]:5.1f} h")
        return cols + [cursor, sup]

    anim = animation.FuncAnimation(fig, update, frames=len(frames), blit=False)
    anim.save(OUT / "repressilator_cells.gif", writer=animation.PillowWriter(fps=15))
    plt.close(fig)


def fig_solvers(methods, ratios, rows):
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(14, 5.2))
    colors = dict(zip(methods, ["#f72585", "#ff8fab", "#ffd166", C_U, "#80ffdb"]))
    for m in methods:
        r = [x for x in rows if x[1] == m]
        ls = "--" if m in ("RK45", "DOP853") else "-"
        a1.loglog(ratios, [x[3] for x in r], ls, marker="o", color=colors[m], lw=2, label=m)
        a2.loglog(ratios, [x[2] for x in r], ls, marker="o", color=colors[m], lw=2, label=m)
    a1.set(xlabel="stiffness ratio  protein lifetime / mRNA lifetime  ($1/\\beta$)", ylabel="RHS evaluations",
           title="Cost to integrate 20 protein lifetimes (rtol 1e-6)")
    a2.set(xlabel="stiffness ratio  $1/\\beta$", ylabel="wall time [s] (best of 3)",
           title="Explicit solvers (dashed) are stability-limited")
    a1.legend(fontsize=9)
    a2.legend(fontsize=9)
    a1.axvline(5, color="#888", ls=":", lw=1)
    a1.text(5.3, a1.get_ylim()[0] * 3, "paper's\nrepressilator", fontsize=8, color="#aaa")
    fig.suptitle("scipy.integrate.solve_ivp on the repressilator: explicit vs stiff solvers", fontsize=13)
    fig.tight_layout()
    fig.savefig(OUT / "solver_benchmark.png", dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    main()