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
from qlroms.utils.plots import close_pdf, save_figure, start_pdf
# the figures only this tutorial draws (tutorials/plots_00.py)
from plots_00 import plot_ks1d_spacetime, plot_ks2d_snapshots, plot_residual_energy

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('figs/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 = {case: short_run(ks2d, case) for case in CASES_2D}
save_figure(plot_ks2d_snapshots(runs))
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)
save_figure(plot_ks1d_spacetime(X1, float(fom1.dt), float(cfg1.get("Lx", 1.0)), case_1d))
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}")

save_figure(plot_residual_energy(spectra))
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.31G     40000
 ks2d/nu1_0.50_nu2_0.35_Nx32_Ny32/dt0.1000/full_trajectory.pth     0.86G    113000

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

import contextlib
import io
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
from dynamodels.physical import KS, KS2D
from ntsa import lyapunov as lyap
from ntsa import tools as nt

# the signal diagnostics that tutorial 3 shares (qlroms/utils/characterization.py)
from qlroms.utils.characterization import (load_characterization_case, lyapunov_verdict, saturation_rate,
                                           signal_report, spectrum_text)
from qlroms.utils.plots import save_figure
# the figures only this tutorial draws (tutorials/plots_00.py)
from plots_00 import (plot_acf_ami, plot_bifurcation_ks1d, plot_bifurcation_ks2d, plot_correlation_sum,
                      plot_delay_embedding, plot_fnn, plot_perturbation_growth, plot_poincare, plot_recurrence,
                      plot_return_maps, plot_spectra)

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

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))
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
# its module's TEST_CASES -- 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)}

# cached FOM trajectory, the sensor signal x_fom and the respawned FOM that ntsa integrates
cases = {name: load_characterization_case(name, n_sig=N_SIG, **window) for name, window in WINDOWS.items()}
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

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.1777 | ACF first zero 6.7 t.u., tail |rho| 0.81 | optimal lag zeta = 57 samples = 5.70 t.u. | FNN embedding dim = 4 | D2 = 2.30 (small-r 2.10, large-r 2.48)

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 3.6 s)

plot_spectra(cases)
plt.show()
plot_acf_ami(cases)
plt.show()
plot_fnn(cases)
plt.show()

plot_delay_embedding(cases)
plt.show()

plot_return_maps(cases)
plt.show()

plot_poincare(cases)
plt.show()

plot_recurrence(cases, N_REC)
plt.show()

plot_correlation_sum(cases)
plt.show()

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, min_window=EXP_WINDOW_MIN)
    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 72 s)

plot_perturbation_growth(cases)
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.30      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 = {name: float(cases[name]["fom"].L / np.sqrt(cases[name]["fom"].nu)) for name in ("ks1d qp", "ks1d chaos")}
save_figure(plot_bifurcation_ks1d(L_sweep, L_peaks, L_CASE))
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 (26 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 (KS2D labels both nu1 and
    nu2, and dynamodels requires the labelled parameters to be declared)."""

    params = ["nu1", "nu2"]


class KS2Dnu1(KS2D):
    """2-D KS with nu1 exposed as a sweepable parameter (KS2D labels both nu1 and
    nu2, and dynamodels requires the labelled parameters to be declared)."""

    params = ["nu1", "nu2"]


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.97s/it]

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

nu2 sweep:  30%|███       | 3/10 [00:13<00:31,  4.49s/it]

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

nu2 sweep:  50%|█████     | 5/10 [00:23<00:24,  4.92s/it]

nu2 sweep:  60%|██████    | 6/10 [00:29<00:21,  5.31s/it]

nu2 sweep:  70%|███████   | 7/10 [00:35<00:16,  5.53s/it]

nu2 sweep:  80%|████████  | 8/10 [00:42<00:11,  5.85s/it]

nu2 sweep:  90%|█████████ | 9/10 [00:48<00:05,  5.90s/it]

save_figure(plot_bifurcation_ks2d(nu2_sweep, nu2_peaks, r"$\nu_2$ (at $\nu_1=0.5$)",
                                  {"quasi-periodic": 0.10, "periodic": 0.20, "ks2d tw": 0.35}))
plt.show()
save_figure(plot_bifurcation_ks2d(nu1_sweep, nu1_peaks, r"$\nu_1$ (at $\nu_2=0.1$)",
                                  {"chaotic": 0.10, "ks2d chaos": 0.30, "quasi-periodic": 0.50}))
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 285s

close_pdf()

15 figures written to figs/00_test_cases_for_tutorials.pdf

PosixPath('figs/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. Loading 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¶