import os
import time
from pathlib import Path

os.environ.setdefault("QLROM_DATA", str(Path.home() / ".cache" / "qlrom"))

import matplotlib.pyplot as plt
import numpy as np
import torch

from qlroms.intrusive_qlroms import ks1d, ks2d
from qlroms.intrusive_qlroms.build_ks import build_fom, get_full_trajectory
from qlroms.utils.plots import close_pdf, field_image, save_figure, start_pdf

t_start = time.time()
CASES_2D = list(ks2d.TEST_CASES)
CASES_1D = list(ks1d.TEST_CASES)
print(f"ks2d cases: {CASES_2D}\nks1d cases: {CASES_1D}")

# every figure below also becomes a page of figs/00_test_cases_for_tutorials.pdf
start_pdf("00_test_cases_for_tutorials")

ks2d cases: ['periodic', 'travelling', 'quasi-periodic', 'chaotic', 'chaotic_B']
ks1d cases: ['quasi-periodic', 'chaotic']

PosixPath('../outputs/00_test_cases_for_tutorials.pdf')

def case_row(module, name):
    cfg = dict(module.TEST_CASES[name])
    fom, merged = build_fom(module, case=name)
    grid = f"{cfg.get('Nx')}" + (f"x{cfg['Ny']}" if "Ny" in cfg else "")
    visc = (f"nu1={cfg['nu1']}, nu2={cfg['nu2']}" if "nu1" in cfg
            else f"visc={cfg.get('visc', cfg.get('nu'))}, Lx={cfg.get('Lx', 0):.1f}")
    return dict(case=name, grid=grid, dt=cfg["dt"], visc=visc,
                lamb1=cfg.get("lamb1", float("nan")),
                K=cfg.get("K", None), r=cfg.get("r", None),
                Ntrain=merged["Ntrain"], Ntest=merged["Ntest"], i0=merged["i0"],
                dof=int(cfg.get("Nx", 0) * cfg.get("Ny", 1)))


rows = ([("ks2d", case_row(ks2d, c)) for c in CASES_2D]
        + [("ks1d", case_row(ks1d, c)) for c in CASES_1D])
print(f"{'module':>6s} {'case':>15s} {'grid':>9s} {'dofs':>6s} {'dt':>6s} {'lamb1':>6s} "
      f"{'K':>4s} {'r':>4s} {'Ntrain':>7s} {'viscosity':>26s}")
for mod, R in rows:
    print(f"{mod:>6s} {R['case']:>15s} {R['grid']:>9s} {R['dof']:>6d} {R['dt']:>6.3g} "
          f"{R['lamb1']:>6.3g} {str(R['K'] or '-'):>4s} {str(R['r'] or '-'):>4s} "
          f"{R['Ntrain']:>7d} {R['visc']:>26s}")

module            case      grid   dofs     dt  lamb1    K    r  Ntrain                  viscosity
  ks2d        periodic     32x32   1024   0.01      0    -    -   20000           nu1=0.5, nu2=0.2
  ks2d      travelling     32x32   1024    0.1      0    5   25   20000          nu1=0.5, nu2=0.35
  ks2d  quasi-periodic     32x32   1024   0.01      0    -    -   20000           nu1=0.5, nu2=0.1
  ks2d         chaotic     64x64   4096   0.01    1.8    -    -   20000           nu1=0.1, nu2=0.1
  ks2d       chaotic_B     64x64   4096   0.01    1.8   40   50   20000           nu1=0.3, nu2=0.1
  ks1d  quasi-periodic       128    128    0.1      0    6   30  100000 visc=0.22535211267605634, Lx=6.3
  ks1d         chaotic       128    128   0.05  0.062   10   30  800000          visc=1.0, Lx=62.8

/home/anovoama/.conda/envs/qlrom/lib/python3.11/site-packages/torch/cuda/__init__.py:188: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 12080). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /__w/pytorch/pytorch/c10/cuda/CUDAFunctions.cpp:119.)
  return torch._C._cuda_getDeviceCount() > 0

def short_run(module, case, n_keep=1500):
    """Fresh integration past the case's own transient -- no cache, a few seconds."""
    fom, cfg = build_fom(module, case=case)
    psi, t = fom.time_integrate(Nt=int(fom.t_transient / fom.dt))
    fom.update_history(psi, t)
    psi, t = fom.time_integrate(Nt=n_keep)
    fom.update_history(psi, t)
    u = fom.get_observable_hist(Nt=n_keep, loc="all")
    return fom, cfg, torch.as_tensor(u[:, :, 0].T, dtype=torch.float64)


runs = {}
fig, axs = plt.subplots(len(CASES_2D), 4, figsize=(9.5, 2.1 * len(CASES_2D)),
                        layout="constrained")
for row, case in enumerate(CASES_2D):
    fom, cfg, Xc = short_run(ks2d, case)
    runs[case] = (fom, cfg, Xc)
    shape = (int(fom.Nx), int(fom.Ny))
    idx = np.linspace(0, Xc.shape[1] - 1, 4).astype(int)
    vmax = float(Xc[:, idx].abs().max())
    for col, n in enumerate(idx):
        field_image(axs[row, col], Xc[:, n], shape, vmax=vmax)
        if row == 0:
            axs[row, col].set_title(f"$t_{col}$", fontsize=9)
    axs[row, 0].set_ylabel(f"{case}\n{shape[0]}x{shape[1]}", fontsize=8)
fig.suptitle("ks2d test cases, past the transient", fontsize=11)
save_figure(fig)
plt.show()

case_1d = "chaotic" if "chaotic" in CASES_1D else CASES_1D[0]
fom1, cfg1, X1 = short_run(ks1d, case_1d, n_keep=1500)
fig, ax = plt.subplots(figsize=(9.5, 2.6), layout="constrained")
im = ax.imshow(X1.numpy(), aspect="auto", origin="lower", cmap="RdBu_r",
               extent=[0, X1.shape[1] * float(fom1.dt), 0, float(cfg1.get("Lx", 1.0))])
ax.set(xlabel="t", ylabel="x", title=f"ks1d {case_1d}: the usual space-time picture")
fig.colorbar(im, ax=ax, shrink=0.9)
save_figure(fig)
plt.show()

print(f"{'case':>15s} {'r for 99%':>10s} {'r for 99.9%':>12s} {'r for 99.99%':>13s} "
      f"{'canonical r':>12s}")
spectra = {}
for case, (fom, cfg, Xc) in runs.items():
    Xc = Xc - Xc.mean(dim=1, keepdim=True)
    s = torch.linalg.svdvals(Xc[:, ::3])
    energy = torch.cumsum(s ** 2, 0) / (s ** 2).sum()
    spectra[case] = energy.numpy()
    need = [int(torch.searchsorted(energy, torch.tensor(q)) + 1) for q in (0.99, 0.999, 0.9999)]
    print(f"{case:>15s} {need[0]:>10d} {need[1]:>12d} {need[2]:>13d} "
          f"{str(ks2d.TEST_CASES[case].get('r', '-')):>12s}")

fig, ax = plt.subplots(figsize=(6.4, 3.4), layout="constrained")
for case, e in spectra.items():
    ax.semilogy(np.arange(1, len(e) + 1), 1.0 - e, lw=1.4, label=case)
ax.set(xlabel="POD modes (global basis)", ylabel="residual energy", xlim=(0, 60), ylim=(1e-12, 1))
ax.grid(alpha=0.25, lw=0.5)
fig.legend(loc="outside right upper", frameon=False, fontsize=9)
save_figure(fig)
plt.show()

           case  r for 99%  r for 99.9%  r for 99.99%  canonical r
       periodic          3            4             5            -
     travelling         12           19            27           25
 quasi-periodic          8           10            13            -

        chaotic         35           54            75            -
      chaotic_B         12           18            24           50

fom, cfg = build_fom(ks2d, case="travelling", overrides={"Ntrain": 20_000, "Ntest": 3_000})
X = get_full_trajectory(Ntot=cfg["i0"] + cfg["Ntrain"] + cfg["Ntest"], model=fom, i0=cfg["i0"])
Xtrain, Xtest = X[:, :cfg["Ntrain"]], X[:, cfg["Ntrain"]:]

root = Path(os.environ["QLROM_DATA"])
cached = sorted(root.glob("**/full_trajectory.pth"))
print(f"{'cached trajectory':>62s} {'size':>9s} {'columns':>9s}")
for f in cached:
    n_cols = tuple(torch.load(f, map_location="meta", weights_only=True).shape)[1]
    print(f"{str(f.relative_to(root)):>62s} {f.stat().st_size / 2**30:>8.2f}G {n_cols:>9d}")
print(f"\nnot in this list = generated on first use, at the cost of one FOM run.")

                                             cached trajectory      size   columns
    ks1d/nu_0.2254_Lx6.2832_Nx128/dt0.1000/full_trajectory.pth     0.03G     27000
   ks1d/nu_1.0000_Lx62.8319_Nx128/dt0.0500/full_trajectory.pth     0.22G    235000
    ks1d/nu_1.0000_Lx62.8319_Nx64/dt0.1000/full_trajectory.pth     0.02G     34600
 ks2d/nu1_0.30_nu2_0.10_Nx64_Ny64/dt0.0100/full_trajectory.pth     7.78G    255000
 ks2d/nu1_0.50_nu2_0.20_Nx32_Ny32/dt0.0100/full_trajectory.pth     0.25G     33000
 ks2d/nu1_0.50_nu2_0.35_Nx32_Ny32/dt0.1000/full_trajectory.pth     0.56G     73000

not in this list = generated on first use, at the cost of one FOM run.

import contextlib
import io
import os
import time
from itertools import product
from pathlib import Path

os.environ.setdefault("QLROM_DATA", str(Path.home() / ".cache" / "qlrom"))

import matplotlib.pyplot as plt
import numpy as np
import torch
from dynamodels.physical import KS, KS2D
from ntsa import lyapunov as lyap
from ntsa import tools as nt
from ntsa.characterize import characterize
from scipy.signal import find_peaks

from qlroms import free_run, metrics
from qlroms.intrusive_qlroms import ks1d, ks2d
from qlroms.intrusive_qlroms.build_ks import build_fom, get_full_trajectory
from qlroms.model import QLModel
from qlroms.utils.config import get_case
from qlroms.utils.plots import field_image, outputs_dir, save_figure

plt.rcParams["figure.dpi"] = 160


C_FOM, C_ROM, C_MARK = "#2a78d6", "#eb6834", "#1baf7a"
CASE_COLORS = ["#2a78d6", "#1baf7a", "#b05fd6", "#d64a4a"]
SEED = 0

N_SIG = 12_000        # samples fed to every signal-level diagnostic
N_REC = 1_200         # samples in a recurrence window (recurrence_matrix is O(T^2))
N_FC = 2_000          # closed-loop steps used only for the prediction-horizon contrast
EXP_WINDOW_MIN = 0.15  # a fitted lambda_1 counts only if its window spans this much of the record

# Each case brings its own (K, r) and its own characterization timescales from
# qlroms.utils.config -- t_lam / t_lam_rom are the Lyapunov horizons, which must cover the
# growth phase and stop near saturation (the ks2d chaos ROM saturates sooner than its FOM,
# so it gets a shorter one). Only the WINDOWS are this notebook's own: the cases are
# generated far longer than a ~15 min notebook can characterize, and the text says so
# wherever a number depends on it.
WINDOWS = {"ks1d qp":    dict(Ntrain=10_000, Ntest=2_000),
           "ks1d chaos": dict(Ntrain=20_000, Ntest=2_000),
           "ks2d tw":    dict(Ntrain=15_000, Ntest=2_000),
           "ks2d chaos": dict(Ntrain=40_000, Ntest=2_000)}
CASES = {name: {**get_case(name), **window} for name, window in WINDOWS.items()}

_OUT = outputs_dir()          # notebooks/outputs, from the repo root or from tutorials/

def load_case(name):
    """Cached FOM trajectory plus the metadata every later section needs."""
    P = CASES[name]
    mod = P["module"]
    is2d = mod is ks2d
    fom, cfg = build_fom(mod, case=P["case"], overrides={"Ntrain": P["Ntrain"], "Ntest": P["Ntest"]})
    traj = get_full_trajectory(Ntot=cfg["i0"] + P["Ntrain"] + P["Ntest"], model=fom, i0=cfg["i0"])
    dt = float(fom.dt)
    obs_idx = np.asarray(fom.sensor_locations, dtype=int)
    # every diagnostic below reads ONE scalar sensor: the same physical point for FOM and ROM
    i_obs = 1
    # ntsa re-integrates the FOM (respawn + run_long) instead of replaying the cache, so seed it
    # from the LAST cached state: the KS2D state IS the physical field, the 1-D KS state is its rfft
    u_end = np.asarray(traj[:, -1], dtype=float)
    fom_run = nt.respawn(fom, psi0=(u_end if is2d else np.fft.rfft(u_end)))
    fom_run.t_transient, fom_run.t_lyap = P["t_transient"], P["t_lyap"]
    return dict(name=name, P=P, fom=fom, fom_run=fom_run, cfg=cfg, dt=dt, obs_idx=obs_idx, i_obs=i_obs,
                traj=traj, Xtrain=traj[:, :P["Ntrain"]], Xtest=traj[:, P["Ntrain"]:],
                x_fom=np.asarray(traj[obs_idx[i_obs], :N_SIG], dtype=float),
                is2d=is2d, lamb1_ref=float(fom.lamb1))


cases = {name: load_case(name) for name in CASES}
for c in cases.values():
    n_h = c["traj"].shape[0]
    print(f"{c['name']:11s} ({c['P']['case']:14s}): N_h = {n_h:5d}, dt = {c['dt']:.3f}, "
          f"{c['traj'].shape[1]} cached snapshots | sensor row {c['obs_idx'][c['i_obs']]} "
          f"| analysis window {N_SIG} samples = {N_SIG * c['dt']:.0f} t.u. | recorded lamb1 = {c['lamb1_ref']}")

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks1d/nu_0.2254_Lx6.2832_Nx128/dt0.1000/full_trajectory.pth.

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks1d/nu_1.0000_Lx62.8319_Nx128/dt0.0500/full_trajectory.pth.

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks2d/nu1_0.50_nu2_0.35_Nx32_Ny32/dt0.1000/full_trajectory.pth.

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks2d/nu1_0.30_nu2_0.10_Nx64_Ny64/dt0.0100/full_trajectory.pth.

ks1d qp     (quasi-periodic): N_h =   128, dt = 0.100, 27000 cached snapshots | sensor row 42 | analysis window 12000 samples = 1200 t.u. | recorded lamb1 = 0.0
ks1d chaos  (chaotic       ): N_h =   128, dt = 0.050, 52000 cached snapshots | sensor row 42 | analysis window 12000 samples = 600 t.u. | recorded lamb1 = 0.062
ks2d tw     (travelling    ): N_h =  1024, dt = 0.100, 27000 cached snapshots | sensor row 341 | analysis window 12000 samples = 1200 t.u. | recorded lamb1 = 0.0
ks2d chaos  (chaotic_B     ): N_h =  4096, dt = 0.010, 52000 cached snapshots | sensor row 1365 | analysis window 12000 samples = 120 t.u. | recorded lamb1 = 1.8

ACF_TAIL_DISCRETE = 0.3   # |rho| still this large at the end of the lag window => discrete spectrum


def spectral_lines(psd_f, psd, prominence=0.05, n_max=8, m_max=16, gen_max=4):
    """Prominent PSD lines, and the fewest generators whose integer combinations explain them.

    `find_peaks` sees LINES where `argsort` sees BINS: the two strongest bins of one sharp
    line are that line and its leakage neighbour. A combination of order sum(m) inherits
    half a bin of quantization from each multiple, hence the order-dependent tolerance.
    Multiples above `m_max` do not count as explanations, and a set of lines needing
    `gen_max` generators is not a lattice -- both caps also keep the search finite.
    """
    df = float(psd_f[1] - psd_f[0])
    idx, props = find_peaks(psd, prominence=prominence * psd.max())
    strongest = idx[np.argsort(props["prominences"])[::-1][:n_max]]
    lines, gens = np.sort(psd_f[strongest]), []
    for fk in lines:
        combos = product(*(range(min(int(fk / g) + 2, m_max)) for g in gens))
        if any(abs(fk - float(np.dot(ms, gens))) <= 0.5 * df * (1 + sum(ms))
               for ms in combos if any(ms)):
            continue
        gens.append(float(fk))
        if len(gens) >= gen_max:
            break
    return lines, gens, len(idx)


def signal_report(x, dt, tag):
    """Every scalar diagnostic ntsa computes from one signal, in one dict."""
    out = dict(tag=tag, x=x, dt=dt)
    f, psd = nt.fun_PSD(dt, x[None, :])
    out["f"], out["psd"] = f[1:], psd[0][1:]                      # bin 0 is the mean, always dropped
    out["f_top"] = float(out["f"][np.argmax(out["psd"])])         # tallest bin: a line only on a torus
    out["lines"], out["gen"], out["n_lines"] = spectral_lines(out["f"], out["psd"])
    # the centroid is the broadband statistic that stays meaningful on a continuum
    out["f_c"] = float(np.sum(out["f"] * out["psd"]) / np.sum(out["psd"]))
    n_lags = min(len(x) // 4, 4000)
    out["acf"] = nt.autocorrelation(x, n_lags)
    zc = np.flatnonzero(out["acf"] < 0)
    out["t_dec"] = float(zc[0] * dt) if zc.size else float(n_lags * dt)
    out["acf_tail"] = float(np.abs(out["acf"][3 * len(out["acf"]) // 4:]).max())
    out["discrete"] = bool(out["acf_tail"] > ACF_TAIL_DISCRETE)   # gates the lattice reading
    max_lag = min(len(x) // 10, 800)
    out["ami"] = nt.average_mutual_information(x, max_lag)
    out["zeta"] = int(nt.optimal_lag(x, max_lag=max_lag))
    out["dim"], out["fnn"] = nt.false_nearest_neighbours(x, out["zeta"])
    out["Y"] = nt.delay_embed(x, max(out["dim"], 3), out["zeta"])
    out["D2"], out["cd"] = nt.correlation_dimension(out["Y"])
    # D2 is fitted over the 2nd-50th percentile of the pair distances: a LARGE-scale slope.
    # Split it to see whether the object keeps that dimension as r shrinks.
    log_r, log_C = out["cd"]
    out["D2_lo"] = float(np.polyfit(log_r[:8], log_C[:8], 1)[0])    # small-r end of that window
    out["D2_hi"] = float(np.polyfit(log_r[-8:], log_C[-8:], 1)[0])  # large-r end
    return out


def spectrum_text(s):
    """PSD verdict: a line lattice only where the autocorrelation says there are lines."""
    if not s["discrete"]:
        return f"broadband, {s['n_lines']:2d} periodogram maxima, no line lattice"
    return (f"{s['n_lines']:2d} prominent lines, the strongest {len(s['lines'])} on "
            f"{len(s['gen'])} generator(s) {np.round(s['gen'], 4)}")


t0 = time.time()
for c in cases.values():
    c["sig"] = signal_report(c["x_fom"], c["dt"], f"{c['name']} FOM")
    s = c["sig"]
    print(f"{c['name']:11s}: PSD {spectrum_text(s)}; tallest f = {s['f_top']:.4f}, centroid "
          f"{s['f_c']:.4f} | ACF first zero {s['t_dec']:.1f} t.u., tail |rho| {s['acf_tail']:.2f} "
          f"| optimal lag zeta = {s['zeta']} samples = {s['zeta'] * c['dt']:.2f} t.u. | FNN "
          f"embedding dim = {s['dim']} | D2 = {s['D2']:.2f} (small-r {s['D2_lo']:.2f}, "
          f"large-r {s['D2_hi']:.2f})")
print(f"  (all four signal reports in {time.time() - t0:.1f} s)")

ks1d qp    : PSD  6 prominent lines, the strongest 6 on 2 generator(s) [0.0208 0.0658]; tallest f = 0.0208, centroid 0.3295 | ACF first zero 12.1 t.u., tail |rho| 0.69 | optimal lag zeta = 119 samples = 11.90 t.u. | FNN embedding dim = 10 | D2 = 0.50 (small-r 0.72, large-r 0.51)

ks1d chaos : PSD broadband, 18 periodogram maxima, no line lattice; tallest f = 0.0067, centroid 0.3996 | ACF first zero 18.2 t.u., tail |rho| 0.23 | optimal lag zeta = 77 samples = 3.85 t.u. | FNN embedding dim = 3 | D2 = 2.01 (small-r 2.11, large-r 1.71)

ks2d tw    : PSD  7 prominent lines, the strongest 7 on 2 generator(s) [0.0175 0.03  ]; tallest f = 0.0300, centroid 0.3299 | ACF first zero 6.6 t.u., tail |rho| 0.81 | optimal lag zeta = 50 samples = 5.00 t.u. | FNN embedding dim = 4 | D2 = 2.28 (small-r 2.11, large-r 2.37)

ks2d chaos : PSD broadband, 18 periodogram maxima, no line lattice; tallest f = 0.0750, centroid 3.0566 | ACF first zero 2.9 t.u., tail |rho| 0.16 | optimal lag zeta = 85 samples = 0.85 t.u. | FNN embedding dim = 3 | D2 = 2.02 (small-r 2.28, large-r 1.77)
  (all four signal reports in 5.5 s)

fig, axs = plt.subplots(1, 4, figsize=(13, 2.9), layout="constrained")
for ax, (c, col) in zip(axs, zip(cases.values(), CASE_COLORS, strict=True), strict=True):
    s = c["sig"]
    ax.semilogy(s["f"], s["psd"], lw=0.8, color=col)
    for fp in s["lines"]:
        ax.axvline(fp, color="0.4", ls=":", lw=0.8)
    for fg in (s["gen"] if s["discrete"] else []):
        ax.axvline(fg, color=C_MARK, ls="-", lw=1.0, alpha=0.8)
    verdict = f"{len(s['gen'])} generator(s)" if s["discrete"] else "broadband"
    ax.set(xlim=(0, min(8 * s["f_top"], s["f"][-1])), xlabel="$f$",
           title=f"{c['name']}\n{verdict}, tallest $f={s['f_top']:.4f}$")
    ax.grid(alpha=0.25, lw=0.5)
axs[0].set_ylabel("PSD")
fig.suptitle("(a) Power spectra: prominent lines (dotted) and the generators they sit on (green); "
             "the lattice is only read where the ACF tail says the spectrum is discrete", fontsize=11)
plt.show()

fig, axs = plt.subplots(2, 4, figsize=(13, 5.0), layout="constrained")
for j, (c, col) in enumerate(zip(cases.values(), CASE_COLORS, strict=True)):
    s = c["sig"]
    lags = np.arange(len(s["acf"])) * c["dt"]
    axs[0, j].plot(lags, s["acf"], lw=0.9, color=col)
    axs[0, j].axhline(0, color="0.4", lw=0.6)
    axs[0, j].set(xlim=(0, 40 * s["zeta"] * c["dt"]), xlabel=r"lag $\tau$",
                  title=f"{c['name']}: ACF (tail {s['acf_tail']:.2f})")
    ami_lags = np.arange(1, len(s["ami"]) + 1) * c["dt"]
    axs[1, j].plot(ami_lags, s["ami"], lw=0.9, color=col)
    axs[1, j].axvline(s["zeta"] * c["dt"], color=C_MARK, lw=1.4,
                      label=fr"$\zeta={s['zeta']}$ ({s['zeta'] * c['dt']:.2f} t.u.)")
    axs[1, j].set(xlim=(0, 8 * s["zeta"] * c["dt"]), xlabel=r"lag $\tau$", title="AMI $I(\\tau)$")
    axs[1, j].legend(frameon=False, fontsize=8)
    for ax in (axs[0, j], axs[1, j]):
        ax.grid(alpha=0.25, lw=0.5)
fig.suptitle("(b) Autocorrelation and average mutual information; the green line is the delay ntsa picks",
             fontsize=11)
plt.show()

fig, ax = plt.subplots(figsize=(5.6, 3.2), layout="constrained")
for (c, col) in zip(cases.values(), CASE_COLORS, strict=True):
    s = c["sig"]
    d = np.arange(1, len(s["fnn"]) + 1)
    ax.plot(d, s["fnn"], "o-", ms=4, lw=1.3, color=col, label=f"{c['name']} (d = {s['dim']})")
ax.axhline(0.01, color="0.3", ls="--", lw=0.9)
ax.set(xlabel="embedding dimension $d$", ylabel="false nearest neighbours", yscale="log",
       title="(c) FNN: where the attractor stops self-intersecting")
ax.legend(frameon=False, fontsize=8)
ax.grid(alpha=0.25, lw=0.5)
plt.show()

fig = plt.figure(figsize=(13, 3.4), layout="constrained")
for j, (c, col) in enumerate(zip(cases.values(), CASE_COLORS, strict=True)):
    s = c["sig"]
    Y = s["Y"][::2]
    ax = fig.add_subplot(1, 4, j + 1, projection="3d")
    ax.plot(Y[:, 0], Y[:, 1], Y[:, 2], lw=0.6, color=col, alpha=0.8)
    ax.set(xlabel="$x(t)$", ylabel=r"$x(t+\zeta)$", zlabel=r"$x(t+2\zeta)$")
    ax.set_title(f"{c['name']}\n$\\zeta={s['zeta']}$, $d={s['dim']}$", fontsize=9)
fig.suptitle("(a) Delay embedding: the attractor reconstructed from one sensor", fontsize=11)
plt.show()

fig, axs = plt.subplots(1, 4, figsize=(13, 3.2), layout="constrained")
for ax, (c, col) in zip(axs, zip(cases.values(), CASE_COLORS, strict=True), strict=True):
    s = c["sig"]
    xm, xn, _ = nt.first_return_map(s["x"])
    ax.plot(xm, xn, ".", ms=2.5, color=col, alpha=0.6)
    lim = [min(xm.min(), xn.min()), max(xm.max(), xn.max())]
    ax.plot(lim, lim, "-", color="0.4", lw=0.8)
    n_cl = nt.count_peak_clusters(np.concatenate([xm, xn[-1:]]), float(np.ptp(s["x"])))
    ax.set(xlabel="$x_m$", ylabel="$x_{m+1}$",
           title=f"{c['name']}\n{len(xm) + 1} maxima, {n_cl} clusters", aspect="equal")
    ax.grid(alpha=0.25, lw=0.5)
fig.suptitle("(b) First return map of the successive maxima", fontsize=11)
plt.show()

fig, axs = plt.subplots(1, 4, figsize=(13, 3.2), layout="constrained")
for ax, (c, col) in zip(axs, zip(cases.values(), CASE_COLORS, strict=True), strict=True):
    s = c["sig"]
    P = nt.poincare_section(s["x"], s["zeta"])
    ax.plot(P[:, 0], P[:, 1], ".", ms=2.5, color=col, alpha=0.6)
    ax.set(xlabel="$x(t)$", ylabel=r"$x(t+\zeta)$", title=f"{c['name']}\n{len(P)} crossings")
    ax.grid(alpha=0.25, lw=0.5)
fig.suptitle(r"(c) Poincare section on the plane $x(t+2\zeta)=\mathrm{median}$, upward crossings",
             fontsize=11)
plt.show()

fig, axs = plt.subplots(1, 4, figsize=(13, 3.3), layout="constrained")
for ax, c in zip(axs, cases.values(), strict=True):
    s = c["sig"]
    R = nt.recurrence_matrix(s["Y"][:N_REC], eps_frac=0.10)
    T = R.shape[0] * c["dt"]
    ax.imshow(R, cmap="binary", origin="lower", extent=[0, T, 0, T], interpolation="nearest")
    ax.set(xlabel="$t$", ylabel="$t'$", title=f"{c['name']}\nrecurrence rate {R.mean():.2f}")
fig.suptitle("(d) Recurrence plots: unbroken diagonals = periodic, broken = chaotic", fontsize=11)
plt.show()

fig, ax = plt.subplots(figsize=(5.6, 3.4), layout="constrained")
for (c, col) in zip(cases.values(), CASE_COLORS, strict=True):
    log_r, log_C = c["sig"]["cd"]
    ax.plot(log_r, log_C, lw=1.4, color=col, label=f"{c['name']}: $D_2$ = {c['sig']['D2']:.2f}")
ax.set(xlabel=r"$\log r$", ylabel=r"$\log C(r)$",
       title="(e) Grassberger-Procaccia correlation sum")
ax.legend(frameon=False, fontsize=8)
ax.grid(alpha=0.25, lw=0.5)
plt.show()

def lyapunov_verdict(lam1, res):
    """(lam1 to trust or None, printable string) -- see the non-normality caveat above."""
    t = res["t"]
    frac = float((t[res["i2"]] - t[res["i1"]]) / (t[-1] - t[0]))
    if not np.isfinite(lam1):
        return None, "rejected (no growth)", frac
    if frac < EXP_WINDOW_MIN:
        return None, f"transient ({100 * frac:.0f}% win)", frac
    return float(lam1), f"{lam1:.3f}", frac


def saturation_rate(res):
    """(nats climbed, time taken, mean rate) from the start of the record to saturation.

    A record that spends `span` time units climbing `climb` nats is described by an
    exponential of rate climb/span. A claimed exponent much LARGER than that is
    contradicted by the record itself: it would have exhausted the same climb sooner.
    """
    m, t = res["mean_log_sep"], res["t"]
    hit = np.flatnonzero(m >= np.log(res["sat"]))
    i_sat = int(hit[0]) if hit.size else len(m) - 1
    climb, span = float(m[i_sat] - m[0]), float(t[i_sat] - t[0])
    return climb, span, (climb / span if span > 0 else np.nan)


t0 = time.time()
for c in cases.values():
    P = c["P"]
    lam1, lam1_std, res = lyap.leading_lyapunov(c["fom_run"], n_pert=3, t_run=P["t_lam"])
    lam_ok, lam_str, frac = lyapunov_verdict(lam1, res)
    climb, span, rate = saturation_rate(res)
    c["lyap_fom"] = dict(lam1=lam1, std=lam1_std, res=res, ok=lam_ok, text=lam_str, frac=frac,
                         climb=climb, span=span, rate=rate)
    ref = c["lamb1_ref"]
    sat_txt = (f"climbs {climb:4.1f} nats in {span:5.1f} t.u., mean rate {rate:.3f}/t.u."
               if lam_ok is not None else
               f"plateaus at {100 * res['sat'] / res['diam']:5.2f}% of the attractor diameter")
    # the recorded exponent as a TIME: how long exp(lamb1 t) needs to cover the same climb
    ref_txt = (f" (at the recorded {ref}, that climb takes {climb / ref:.1f} t.u.)"
               if lam_ok is not None and ref > 0 else "")
    print(f"{c['name']:11s}: lam1 = {lam_str:22s} (raw fit {lam1:.3f}, R2 = {res['r2']:.2f}, "
          f"window {100 * frac:4.1f}% of {P['t_lam']:.0f} t.u.) | {sat_txt}{ref_txt} | "
          f"config records {ref}")
print(f"  (four leading_lyapunov runs in {time.time() - t0:.0f} s)")

ks1d qp    : lam1 = transient (9% win)     (raw fit 0.299, R2 = 0.72, window  9.1% of 200 t.u.) | plateaus at 13.40% of the attractor diameter | config records 0.0

ks1d chaos : lam1 = 0.067                  (raw fit 0.067, R2 = 0.98, window 45.6% of 250 t.u.) | climbs 13.0 nats in 197.5 t.u., mean rate 0.066/t.u. (at the recorded 0.062, that climb takes 210.1 t.u.) | config records 0.062

ks2d tw    : lam1 = rejected (no growth)   (raw fit nan, R2 = 0.00, window  1.7% of 120 t.u.) | plateaus at  0.00% of the attractor diameter | config records 0.0

ks2d chaos : lam1 = 0.303                  (raw fit 0.303, R2 = 0.95, window 39.2% of 90 t.u.) | climbs 15.1 nats in  53.0 t.u., mean rate 0.286/t.u. (at the recorded 1.8, that climb takes 8.4 t.u.) | config records 1.8
  (four leading_lyapunov runs in 87 s)

fig, axs = plt.subplots(1, 4, figsize=(13, 3.0), layout="constrained")
for ax, (c, col) in zip(axs, zip(cases.values(), CASE_COLORS, strict=True), strict=True):
    L = c["lyap_fom"]
    res = L["res"]
    ax.plot(res["t"], res["mean_log_sep"], lw=1.1, color=col)
    i1, i2 = res["i1"], res["i2"]
    ax.plot(res["t"][i1:i2], res["mean_log_sep"][i1:i2], lw=2.2, color=C_ROM,
            label=f"fit: {L['text']}")
    ax.axhline(np.log(res["diam"]), color="0.4", ls="--", lw=0.8)
    ax.set(xlabel="$t$", ylabel=r"$\langle\log\|\delta\|\rangle$", title=f"{c['name']}")
    ax.legend(frameon=False, fontsize=8, loc="lower right")
    ax.grid(alpha=0.25, lw=0.5)
fig.suptitle("Perturbation growth: the dashed line is the attractor diameter. Only a curve that keeps "
             "climbing is chaos", fontsize=10)
plt.show()

rows = []
for c in cases.values():
    L = c["lyap_fom"]
    label, ev = nt.classify_regime(c["x_fom"], c["dt"], lam1=L["ok"],
                                   lam1_std=(L["std"] if L["ok"] is not None else 0.0),
                                   t_total=N_SIG * c["dt"])
    c["regime_fom"], c["evidence_fom"] = label, ev
    ratio = f"{ev['f2'] / ev['f1']:.3f}" if (ev["f1"] and ev["f2"]) else "-"
    rows.append((c["name"], L["text"], c["lamb1_ref"], label, ev["n_clusters"], c["sig"]["dim"],
                 c["sig"]["D2"], len(c["sig"]["gen"]) if c["sig"]["discrete"] else 0,
                 ratio, ev["rational_match"] or "-"))

# the last three columns are the two spectral verdicts side by side: OUR generator count
# (0 = broadband, no lattice quoted) and the f2/f1 ratio classify_regime forms internally
hdr = (f"{'case':<12} {'lam1 measured':>22} {'lamb1 cfg':>10} {'classify_regime':>18} "
       f"{'clusters':>9} {'d_FNN':>6} {'D2':>6} {'n_gen':>6} {'its f2/f1':>10} {'rational':>9}")
print(hdr)
print("-" * len(hdr))
for name, lam_s, lam_ref, label, ncl, dim, D2, ngen, ratio, rat in rows:
    print(f"{name:<12} {lam_s:>22} {lam_ref:>10.3f} {label:>18} {ncl:>9d} {dim:>6d} {D2:>6.2f} "
          f"{ngen:>6d} {ratio:>10} {rat:>9}")

case                  lam1 measured  lamb1 cfg    classify_regime  clusters  d_FNN     D2  n_gen  its f2/f1  rational
---------------------------------------------------------------------------------------------------------------------
ks1d qp          transient (9% win)      0.000 limit_cycle_period_3         3     10   0.50      2      3.160         -
ks1d chaos                    0.067      0.062            chaotic        35      3   2.01      0      2.500       5/2
ks2d tw        rejected (no growth)      0.000      quasiperiodic        59      4   2.28      2      1.583         -
ks2d chaos                    0.303      1.800            chaotic        32      3   2.02      0      1.333       4/3

class KSL(KS):
    """1-D KS with L exposed as a sweepable parameter (KS.params is empty by design)."""

    params = ["L"]


def cold_sweep(model, param, values, **kwargs):
    """`bifurcation_sweep` one value at a time, each restarted from the SAME seed state.

    A one-value branch is a branch of one: `continuation=True` still buys us the fresh
    `respawn` per run that the ETDRK4 operators need, but nothing is carried across
    parameter values, so no column can inherit the previous column's attractor.
    """
    peaks = None
    for v in values:
        with contextlib.redirect_stderr(io.StringIO()):        # one tqdm bar per value, silenced
            _, pk = nt.bifurcation_sweep(model, param, [float(v)], continuation=True, **kwargs)
        peaks = peaks or {ext: [[] for _ in per_obs] for ext, per_obs in pk.items()}
        for ext, per_obs in pk.items():
            for iq, col in enumerate(per_obs):
                peaks[ext][iq].append(col[0])
    return np.asarray(values, dtype=float), peaks


t0 = time.time()
L_vals = np.linspace(8.0, 72.0, 33)
ks_seed = KSL(Nx=128, dt=0.25, nu=1.0, L=float(L_vals[0]), Nq=1)
L_sweep, L_peaks = cold_sweep(ks_seed, "L", L_vals, t_transient=200.0, t_sample=1000.0,
                              extrema=("max", "min"), Nq=1, seed=SEED)
n_pts = sum(len(p) for p in L_peaks["max"][0])
print(f"1-D KS sweep: {len(L_sweep)} values of L at nu=1, {n_pts} maxima collected "
      f"({time.time() - t0:.0f} s)")


def sweep_summary(name, values, peaks, every=1):
    """Columns of a bifurcation diagram as numbers: how many extrema, how spread, how clustered."""
    # count_peak_clusters needs ONE reference range for the whole sweep, else every column is
    # rescaled to its own spread and a fixed point reads as a continuum
    span = float(np.ptp(np.concatenate([p for p in peaks["max"][0] if len(p)])))
    for v, pk in list(zip(values, peaks["max"][0], strict=True))[::every]:
        n_cl = nt.count_peak_clusters(pk, span) if len(pk) > 1 else len(pk)
        print(f"  {name} = {v:6.3f}: {len(pk):4d} maxima, spread "
              f"{np.ptp(pk) if len(pk) else 0.0:7.3f}, {n_cl:4d} clusters")


sweep_summary("L", L_sweep, L_peaks, every=2)

L_CASE = {"ks1d qp": float(cases["ks1d qp"]["fom"].L / np.sqrt(cases["ks1d qp"]["fom"].nu)),
          "ks1d chaos": float(cases["ks1d chaos"]["fom"].L / np.sqrt(cases["ks1d chaos"]["fom"].nu))}
fig, axs = nt.plot_bifurcation(L_sweep, L_peaks, r"$\tilde L = L/\sqrt{\nu}$", [r"$u(x_0)$"])
plt.setp(axs[0].lines, ms=4)   # readable markers: plot_bifurcation's default ms=1.5 is too fine here
for name, Lc in L_CASE.items():
    axs[0].axvline(Lc, color=C_MARK, lw=1.4, ls="--")
    axs[0].text(Lc, axs[0].get_ylim()[1], f" {name}\n $\\tilde L$={Lc:.1f}", color=C_MARK,
                fontsize=8, va="top")
axs[0].set_title("1-D KS: extrema of one sensor after transients (coarse sweep)", fontsize=10)
save_figure(fig)
plt.show()

[run_long] trimmed 340 t.u. of residual transient drift (34% of the record) — raise t_transient to keep the full horizon

1-D KS sweep: 33 values of L at nu=1, 1602 maxima collected (31 s)
  L =  8.000:    0 maxima, spread   0.000,    0 clusters
  L = 12.000:   60 maxima, spread   2.621,    2 clusters
  L = 16.000:    0 maxima, spread   0.000,    0 clusters
  L = 20.000:   51 maxima, spread   2.388,   51 clusters
  L = 24.000:   14 maxima, spread   2.451,   14 clusters
  L = 28.000:   71 maxima, spread   0.143,   71 clusters
  L = 32.000:    7 maxima, spread   0.004,    1 clusters
  L = 36.000:   57 maxima, spread   4.676,   57 clusters
  L = 40.000:   55 maxima, spread   4.444,   55 clusters
  L = 44.000:   59 maxima, spread   4.531,   59 clusters
  L = 48.000:   63 maxima, spread   4.677,   63 clusters
  L = 52.000:   60 maxima, spread   4.313,   60 clusters
  L = 56.000:   36 maxima, spread   3.431,   36 clusters
  L = 60.000:   58 maxima, spread   4.605,   58 clusters
  L = 64.000:   54 maxima, spread   3.991,   54 clusters
  L = 68.000:   57 maxima, spread   4.358,   57 clusters
  L = 72.000:   64 maxima, spread   4.510,   64 clusters

class KS2Dnu2(KS2D):
    """2-D KS with nu2 exposed as a sweepable parameter."""

    params = ["nu2"]


class KS2Dnu1(KS2D):
    """2-D KS with nu1 exposed as a sweepable parameter."""

    params = ["nu1"]


t0 = time.time()
nu2_vals = np.linspace(0.05, 0.50, 10)          # step 0.05: hits the 0.1 / 0.2 / 0.35 cases exactly
seed2 = KS2Dnu2(Nx=32, Ny=32, nu1=0.5, nu2=float(nu2_vals[0]), dt=0.1, Nq=1)
nu2_sweep, nu2_peaks = nt.bifurcation_sweep(seed2, "nu2", nu2_vals, continuation=True,
                                            t_transient=300.0, t_sample=250.0, extrema=("max",),
                                            Nq=1, seed=SEED)
print(f"2-D KS anisotropy sweep (nu1 = 0.5): {len(nu2_sweep)} values ({time.time() - t0:.0f} s)")

t0 = time.time()
nu1_vals = np.linspace(0.05, 0.60, 12)          # step 0.05: hits the 0.1 / 0.3 / 0.5 cases exactly
seed1 = KS2Dnu1(Nx=32, Ny=32, nu1=float(nu1_vals[0]), nu2=0.1, dt=0.05, Nq=1)
nu1_sweep, nu1_peaks = nt.bifurcation_sweep(seed1, "nu1", nu1_vals, continuation=True,
                                            t_transient=150.0, t_sample=150.0, extrema=("max",),
                                            Nq=1, seed=SEED)
print(f"2-D KS dissipation sweep (nu2 = 0.1): {len(nu1_sweep)} values ({time.time() - t0:.0f} s)")

nu2 sweep:   0%|          | 0/10 [00:00<?, ?it/s]

nu2 sweep:  10%|█         | 1/10 [00:04<00:44,  4.90s/it]

nu2 sweep:  20%|██        | 2/10 [00:09<00:38,  4.87s/it]

nu2 sweep:  30%|███       | 3/10 [00:14<00:32,  4.69s/it]

nu2 sweep:  40%|████      | 4/10 [00:18<00:27,  4.66s/it]

nu2 sweep:  50%|█████     | 5/10 [00:22<00:22,  4.47s/it]

nu2 sweep:  60%|██████    | 6/10 [00:27<00:17,  4.41s/it]

nu2 sweep:  70%|███████   | 7/10 [00:31<00:13,  4.45s/it]

nu2 sweep:  80%|████████  | 8/10 [00:36<00:09,  4.50s/it]

nu2 sweep:  90%|█████████ | 9/10 [00:40<00:04,  4.41s/it]

for sweep, peaks, lab, marks, fname in [
        (nu2_sweep, nu2_peaks, r"$\nu_2$ (at $\nu_1=0.5$)",
         {"quasi-periodic": 0.10, "periodic": 0.20, "ks2d tw": 0.35}, "04_bifurcation_ks2d_nu2.png"),
        (nu1_sweep, nu1_peaks, r"$\nu_1$ (at $\nu_2=0.1$)",
         {"chaotic": 0.10, "ks2d chaos": 0.30, "quasi-periodic": 0.50}, "04_bifurcation_ks2d_nu1.png")]:
    fig, axs = nt.plot_bifurcation(sweep, peaks, lab, [r"$u(x_0)$"])
    plt.setp(axs[0].lines, ms=4)
    for name, v in marks.items():
        axs[0].axvline(v, color=C_MARK, lw=1.4, ls="--")
        axs[0].text(v, axs[0].get_ylim()[1], f" {name}", color=C_MARK, fontsize=8, va="top",
                    rotation=90)
    axs[0].set_title(f"2-D KS: sensor extrema vs {lab} (coarse sweep, 32x32)", fontsize=10)
    save_figure(fig)
    plt.show()

sweep_summary("nu2", nu2_sweep, nu2_peaks)
sweep_summary("nu1", nu1_sweep, nu1_peaks)

  nu2 =  0.050:    0 maxima, spread   0.000,    0 clusters
  nu2 =  0.100:   35 maxima, spread  14.184,   35 clusters
  nu2 =  0.150:   36 maxima, spread   7.940,    6 clusters
  nu2 =  0.200:   50 maxima, spread   0.005,    1 clusters
  nu2 =  0.250:   18 maxima, spread   0.002,    1 clusters
  nu2 =  0.300:   17 maxima, spread   6.163,   11 clusters
  nu2 =  0.350:   12 maxima, spread   5.457,   12 clusters
  nu2 =  0.400:    6 maxima, spread   4.919,    6 clusters
  nu2 =  0.450:  417 maxima, spread   0.000,    1 clusters
  nu2 =  0.500:    9 maxima, spread   0.000,    1 clusters
  nu1 =  0.050:  251 maxima, spread  17.759,  251 clusters
  nu1 =  0.100:  131 maxima, spread  13.978,  131 clusters
  nu1 =  0.150:   90 maxima, spread  14.514,   90 clusters
  nu1 =  0.200:   53 maxima, spread  12.534,   53 clusters
  nu1 =  0.250:   44 maxima, spread   4.556,   44 clusters
  nu1 =  0.300:   37 maxima, spread   9.628,   37 clusters
  nu1 =  0.350:   17 maxima, spread   3.632,    2 clusters
  nu1 =  0.400:   16 maxima, spread   3.773,    2 clusters
  nu1 =  0.450:   18 maxima, spread  10.230,   18 clusters
  nu1 =  0.500:   15 maxima, spread   6.558,   15 clusters
  nu1 =  0.550:   14 maxima, spread   3.800,    5 clusters
  nu1 =  0.600:   14 maxima, spread   2.997,   14 clusters

print(f"SUMMARY: {len(CASES_2D)} ks2d cases + {len(CASES_1D)} ks1d case, "
      f"{len(cached)} trajectories cached under {root}\n"
      f"  canonical settings: " + ", ".join(
          f"{c} K={ks2d.TEST_CASES[c].get('K', '-')} r={ks2d.TEST_CASES[c].get('r', '-')}"
          for c in CASES_2D if "K" in ks2d.TEST_CASES[c]) + "\n"
      f"  global POD modes for 99.9% energy: " + ", ".join(
          f"{c}: {int(np.searchsorted(spectra[c], 0.999) + 1)}" for c in CASES_2D) + "\n"
      f"  runtime {time.time() - t_start:.0f}s")

SUMMARY: 5 ks2d cases + 2 ks1d case, 6 trajectories cached under /home/anovoama/.cache/qlrom
  canonical settings: travelling K=5 r=25, chaotic_B K=40 r=50
  global POD modes for 99.9% energy: periodic: 4, travelling: 19, quasi-periodic: 10, chaotic: 54, chaotic_B: 18
  runtime 331s

close_pdf()

15 figures written to ../outputs/00_test_cases_for_tutorials.pdf

PosixPath('../outputs/00_test_cases_for_tutorials.pdf')

column	meaning
no dots / one dot	fixed point
$k$ dots	period-$k$ cycle
a few tight clusters	frequency-locked or quasi-periodic
a dense band	chaos

Tutorial 0: the test cases¶

1. The cases¶

2. What they look like¶

3. How hard is each one to reduce?¶

4. Getting a trajectory (and the cache)¶

5. What the tutorials assume¶

6. Characterizing the cases¶

7. Signal-level analysis: what one scalar knows¶

8. Attractor geometry: five pictures of the same object¶

9. Regime classification and the leading Lyapunov exponent¶

10. Bifurcation diagrams -- a full-order-model object¶