qlroms.utils.diagnosis_plots

Figures for the ql-ROM diagnostics in qlroms.utils.diagnosis.

Every function here takes the plain result dicts of sweep_qlrom_diagnosis / compare_local_global_reconstruction (or bare snapshot arrays) and returns a matplotlib figure. Nothing in this module knows which model family produced the trajectories -- the family only ever enters through diagnosis.forecast_fn.

Used by both the offline scripts (scripts/ks/run_sweep.py, scripts/ks/run_diagnosis.py, which write multi-page PDFs) and the diagnosis tutorial notebook.

plot_kr_error_map(diag, figsize=(17, 4.2), log_scale=True)

Heatmaps over the (K, r) grid, with the K=1 global ROM as the bottom row.

Panels: the a posteriori forecast error, the a priori projection MSE, the global/local MSE ratio (>1 = local better), and the cluster-occupancy KLD of the closed-loop run.

Source code in qlroms/utils/diagnosis_plots.py

def plot_kr_error_map(diag: dict, figsize=(17, 4.2), log_scale: bool = True):
    """Heatmaps over the (K, r) grid, with the K=1 global ROM as the bottom row.

    Panels: the a posteriori forecast error, the a priori projection MSE, the global/local
    MSE ratio (>1 = local better), and the cluster-occupancy KLD of the closed-loop run.
    """
    K_list = np.asarray(diag["K_list"])
    r_list = np.asarray(diag["r_list"])
    K_all = np.concatenate([[1], K_list])
    has_kld = diag.get("local_kld") is not None

    fig, axes = plt.subplots(1, 4 if has_kld else 3, figsize=figsize,
                             squeeze=False, layout="constrained")

    def heat(ax, data, norm, cmap_name, title):
        cmap = plt.get_cmap(cmap_name).copy()
        cmap.set_bad("#bfbfbf")
        data = np.where(np.isfinite(data), np.clip(data, norm.vmin, norm.vmax), np.nan)
        im = ax.imshow(data, aspect="auto", origin="lower", norm=norm, cmap=cmap,
                       extent=[r_list[0] - 0.5, r_list[-1] + 0.5, -0.5, len(K_all) - 0.5])
        ax.set_xticks(r_list, [str(r) for r in r_list])
        ax.set_yticks(np.arange(len(K_all)), [str(k) for k in K_all])
        ax.set(xlabel="$r$ (modes)", ylabel="$K$ (clusters)")
        ax.set_title(title, fontsize=10)
        ax.axhline(0.5, color="white", lw=3)          # K=1 sits below the line
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    E = np.vstack([np.asarray(diag["global_errors"])[None, :], np.asarray(diag["local_errors"])])
    fin = E[np.isfinite(E)]
    lo = max(float(np.min(fin)), 1e-12) if fin.size else 1e-12
    hi = max(float(np.max(fin)), lo * 1.01) if fin.size else 1.0
    heat(axes[0, 0], E, colors.Normalize(lo, hi), "viridis_r",
         "forecast error (a posteriori)")
    best = np.unravel_index(np.nanargmin(E), E.shape)
    axes[0, 0].plot(r_list[best[1]], best[0], marker="*", ms=14, mfc="none",
                    mec="white", mew=1.6, ls="none")

    M_loc = np.asarray(diag["local_proj_mse"])
    M_glo = np.asarray(diag["global_proj_mse"])
    M = np.vstack([M_glo[None, :], M_loc])
    vmin = max(float(np.nanmin(M)), 1e-6) if log_scale else float(np.nanmin(M))
    vmax = max(float(np.nanmax(M)), vmin * 10)
    norm_m = colors.LogNorm(vmin, vmax) if log_scale else colors.Normalize(vmin, vmax)
    heat(axes[0, 1], M, norm_m, "viridis", "projection MSE (a priori)")

    ratio = np.vstack([np.ones((1, len(r_list))), M_glo[None, :] / np.clip(M_loc, 1e-12, None)])
    fin = ratio[np.isfinite(ratio)]
    norm_r = colors.TwoSlopeNorm(vmin=min(float(np.min(fin)), 0.95), vcenter=1.0,
                                 vmax=max(float(np.percentile(fin, 95)), 1.05))
    heat(axes[0, 2], ratio, norm_r, "RdYlGn", "global / local MSE  (>1: local better)")

    if has_kld:
        L = np.vstack([np.asarray(diag["global_kld"])[None, :], np.asarray(diag["local_kld"])])
        fin = L[np.isfinite(L)]
        lo = max(float(np.min(fin)), 1e-6) if fin.size else 1e-6
        hi = max(float(np.max(fin)), lo * 10) if fin.size else 1.0
        heat(axes[0, 3], L, colors.LogNorm(lo, hi), "plasma",
             r"occupancy KLD $D_{KL}(P_{ref}\,\|\,P_{ROM})$")
    return fig, axes

plot_parameter_selection(diag, figsize=(12, 3.6))

How K and r are picked: the BIC elbow, MSE vs K at the chosen r, MSE vs r.

Source code in qlroms/utils/diagnosis_plots.py

def plot_parameter_selection(diag: dict, figsize=(12, 3.6)):
    """How K and r are picked: the BIC elbow, MSE vs K at the chosen r, MSE vs r."""
    K_list = np.asarray(diag["K_list"])
    r_list = np.asarray(diag["r_list"])
    K_opt, r_check = int(diag["K_opt"]), int(diag["r_check"])
    loc = np.asarray(diag["local_proj_mse"])
    glo = np.asarray(diag["global_proj_mse"])
    loc_tr = np.asarray(diag.get("local_proj_mse_train", np.full_like(loc, np.nan)))
    glo_tr = np.asarray(diag.get("global_proj_mse_train", np.full_like(glo, np.nan)))

    k_vals = np.asarray(diag["bic_k_values"], dtype=int)
    bic = np.asarray(diag["bic_values"], dtype=float)
    order = np.argsort(k_vals)
    k_vals, bic = k_vals[order], bic[order]
    dbic = np.diff(bic) / np.diff(k_vals).astype(float)

    iK = int(np.argmin(np.abs(K_list - K_opt)))
    ir = int(np.argmin(np.abs(r_list - r_check)))
    floor = 1e-16

    fig, (a, b, c) = plt.subplots(1, 3, figsize=figsize, layout="constrained")
    a.plot(k_vals[:-1], dbic, "o-", color="k", ms=4, lw=1.2)
    a.axvline(K_opt, color="k", ls="--", lw=1.2)
    a.set(xlabel="$K$", ylabel=r"$\Delta BIC/\Delta K$")
    a.set_title(f"BIC elbow: $K_{{opt}}={K_opt}$", fontsize=10)

    b.plot(K_list, np.maximum(loc[:, ir], floor), "o-", color="tab:red", ms=4, lw=1.2,
           label=f"ql-ROM test ($r={r_check}$)")
    if np.isfinite(loc_tr).any():
        b.scatter(K_list, np.maximum(loc_tr[:, ir], floor), marker="x", color="tab:red",
                  s=45, lw=1.5, zorder=6, label=f"ql-ROM train ($r={r_check}$)")
    b.set(xlabel="$K$", ylabel="projection MSE", yscale="log")
    b.legend(frameon=False, fontsize=8)

    c.scatter(r_list, np.maximum(glo, floor), color="tab:blue", s=45, zorder=5, label="g-ROM test")
    c.scatter(r_list, np.maximum(loc[iK], floor), facecolor="none", edgecolor="tab:red", lw=1.8,
              s=55, zorder=5, label=f"ql-ROM test ($K={K_opt}$)")
    if np.isfinite(glo_tr).any():
        c.scatter(r_list, np.maximum(glo_tr, floor), marker="x", color="tab:blue", s=40, lw=1.4,
                  zorder=6, label="g-ROM train")
    if np.isfinite(loc_tr).any():
        c.scatter(r_list, np.maximum(loc_tr[iK], floor), marker="x", color="tab:red", s=40, lw=1.4,
                  zorder=6, label=f"ql-ROM train ($K={K_opt}$)")
    c.axvline(r_check, color="k", ls="--", lw=1.2)
    c.set(xlabel="$r$", ylabel="projection MSE", yscale="log")
    c.legend(frameon=False, fontsize=8)
    for ax in (a, b, c):
        ax.grid(alpha=0.25, lw=0.5)
    return fig, (a, b, c)

plot_cluster_occupancy(diag, figsize=(12, 3.6))

P(c_k) of the data vs of the ROMs, a priori (representation) and a posteriori.

Source code in qlroms/utils/diagnosis_plots.py

def plot_cluster_occupancy(diag: dict, figsize=(12, 3.6)):
    """P(c_k) of the data vs of the ROMs, a priori (representation) and a posteriori."""
    K_opt, r_check = int(diag["K_opt"]), int(diag["r_check"])
    k = np.arange(K_opt)
    w = 0.25

    fig, (ax_r, ax_f) = plt.subplots(1, 2, figsize=figsize, sharey=True, layout="constrained")

    def draw(ax, P_ref, P_ql, P_gr, ref_label, kl_ql, kl_gr, title):
        ax.bar(k - w, P_ref, width=w, color="tab:gray", alpha=0.85, label=ref_label)
        ax.bar(k, P_ql, width=w, color="tab:red", alpha=0.85, label=f"ql-ROM  KL={kl_ql:.2e}")
        ax.bar(k + w, P_gr, width=w, color="tab:orange", alpha=0.85, label=f"g-ROM  KL={kl_gr:.2e}")
        ax.set_xticks(k, [f"$c_{{{i}}}$" for i in k])
        ax.set_xlabel("cluster $k$")
        ax.set_title(title, fontsize=10)
        ax.legend(frameon=False, fontsize=8)

    draw(ax_r, np.asarray(diag["P_test"]), np.asarray(diag["P_qlrom_repr"]),
         np.asarray(diag["P_grom_repr"]), "test data",
         float(diag["kl_ql_repr"]), float(diag["kl_gr_repr"]),
         rf"representation ($K={K_opt}$, $r={r_check}$)")
    draw(ax_f, np.asarray(diag["P_test"]), np.asarray(diag["P_qlrom_te"]),
         np.asarray(diag["P_grom_te"]), "test data",
         float(diag["kl_qlrom"]), float(diag["kl_grom"]),
         rf"forecast ($K={K_opt}$, $r={r_check}$)")
    ax_r.set_ylabel(r"$P(\mathbf{c}_k)$")
    return fig, (ax_r, ax_f)

plot_transition_matrix(diag, figsize=(12, 7))

Markov transition matrices T[i, j] = P(c_j | c_i): data vs g-ROM vs ql-ROM.

Source code in qlroms/utils/diagnosis_plots.py

def plot_transition_matrix(diag: dict, figsize=(12, 7)):
    """Markov transition matrices T[i, j] = P(c_j | c_i): data vs g-ROM vs ql-ROM."""
    K_opt, r_check = int(diag["K_opt"]), int(diag["r_check"])
    ticks = np.arange(0, K_opt, max(1, K_opt // 8))
    panels = [
        ("data (test)", diag["T_test"], "Blues"),
        (f"g-ROM representation ($r={r_check}$)", diag["T_grom_repr"], "Blues"),
        (f"ql-ROM representation ($K={K_opt}$, $r={r_check}$)", diag["T_qlrom_repr"], "Blues"),
        ("data (test)", diag["T_test"], "Greens"),
        (f"g-ROM forecast ($r={r_check}$)", diag["T_grom_te"], "Greens"),
        (f"ql-ROM forecast ($K={K_opt}$, $r={r_check}$)", diag["T_qlrom_te"], "Greens"),
    ]
    fig, axes = plt.subplots(2, 3, figsize=figsize, layout="constrained")
    for ax, (title, T, cmap) in zip(axes.flat, panels, strict=True):
        T = np.asarray(T)
        im = ax.imshow(T, origin="upper", cmap=cmap, vmin=0, vmax=float(T.max()) or 1.0)
        ax.set_title(title, fontsize=9)
        ax.set(xlabel="to cluster", ylabel="from cluster")
        ax.set_xticks(ticks, ticks + 1)
        ax.set_yticks(ticks, ticks + 1)
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    return fig, axes

plot_chart_affiliation(labels_true, labels_rom, dt=1.0, K=None, t0=0.0, figsize=(12, 4.0), cmap_name=None, labels=('data', 'ql-ROM'))

Which chart holds the trajectory, step by step: data against closed-loop ROM.

The occupancy histogram says how often each chart is visited, this says when -- so a ROM that keeps the right occupancy while running out of phase, or one that sticks in a chart it should have left, is visible here and nowhere else. Returns (fig, axes, stats) with stats = agreement fraction, per-series switch counts, and the step index where the two first part company.

Source code in qlroms/utils/diagnosis_plots.py

def plot_chart_affiliation(labels_true, labels_rom, dt: float = 1.0, K: int | None = None,
                           t0: float = 0.0, figsize=(12, 4.0), cmap_name: str | None = None,
                           labels=("data", "ql-ROM")):
    """Which chart holds the trajectory, step by step: data against closed-loop ROM.

    The occupancy histogram says how often each chart is visited, this says *when* --
    so a ROM that keeps the right occupancy while running out of phase, or one that
    sticks in a chart it should have left, is visible here and nowhere else.
    Returns (fig, axes, stats) with stats = agreement fraction, per-series switch
    counts, and the step index where the two first part company.
    """
    y_true = np.asarray(labels_true, dtype=int).reshape(-1)
    y_rom = np.asarray(labels_rom, dtype=int).reshape(-1)
    n = min(y_true.size, y_rom.size)
    y_true, y_rom = y_true[:n], y_rom[:n]
    K = int(K if K is not None else max(y_true.max(), y_rom.max()) + 1)
    t = t0 + np.arange(n) * dt

    same = y_true == y_rom
    agree = float(same.mean())
    first_split = int(np.argmax(~same)) if (~same).any() else -1
    stats = {"agreement": agree, "first_divergence_step": first_split,
             "switches_true": int((np.diff(y_true) != 0).sum()),
             "switches_rom": int((np.diff(y_rom) != 0).sum())}

    from .plots import chart_cmap
    cmap = plt.get_cmap(cmap_name) if cmap_name else chart_cmap(K)[0]
    norm = colors.BoundaryNorm(np.arange(K + 1) - 0.5, K)

    fig, (ax_s, ax_l) = plt.subplots(2, 1, figsize=figsize, sharex=True,
                                     gridspec_kw={"height_ratios": [1, 2.2]},
                                     layout="constrained")
    ax_s.imshow(np.vstack([y_true, y_rom]), aspect="auto", interpolation="nearest",
                cmap=cmap, norm=norm, extent=[t[0], t[-1], 1.5, -0.5])
    ax_s.set_yticks([0, 1], list(labels))
    ax_s.set_title(f"chart affiliation vs time -- agreement {100 * agree:.0f}%, "
                   f"switches {stats['switches_true']} (data) / {stats['switches_rom']} (ROM)",
                   fontsize=10)

    ax_l.step(t, y_true, where="post", color="0.35", lw=1.4, label=labels[0])
    ax_l.step(t, y_rom, where="post", color="tab:red", lw=1.2, ls="--", label=labels[1])
    ax_l.fill_between(t, -0.5, K - 0.5, where=~same, step="post",
                      color="tab:red", alpha=0.12, lw=0)
    if first_split >= 0:
        ax_l.axvline(t[first_split], color="tab:red", lw=1.0, ls=":")
    ax_l.set(xlabel="$t$ [s]", ylabel="chart $k$", ylim=(-0.5, K - 0.5),
             yticks=np.arange(K))
    ax_l.grid(alpha=0.25, lw=0.5)
    ax_l.legend(frameon=False, fontsize=8, ncol=2, loc="upper right")
    return fig, (ax_s, ax_l), stats

plot_error_timeseries(diag, dt, figsize=(11, 3.2), y_cap=200.0)

Per-timestep relative error (%) of the global ROM and of every local K, at r_check.

Source code in qlroms/utils/diagnosis_plots.py

def plot_error_timeseries(diag: dict, dt: float, figsize=(11, 3.2), y_cap: float = 200.0):
    """Per-timestep relative error (%) of the global ROM and of every local K, at r_check."""
    K_all = [1] + list(diag["K_list"])
    err = np.asarray(diag["err_t"])
    t = np.asarray(diag["t"]) * dt

    fig, axs = plt.subplots(1, 2, figsize=figsize, sharex=True, sharey=True, layout="constrained")
    axs[0].plot(t, err[0], color="tab:blue", lw=1.5, label="g-ROM ($K=1$)")
    cmap = plt.get_cmap("YlOrRd")
    for i, K in enumerate(K_all[1:]):
        axs[1].plot(t, err[i + 1], lw=1.0, alpha=0.9,
                    color=cmap(0.3 + 0.7 * i / max(len(K_all) - 2, 1)), label=f"ql-ROM $K={K}$")

    fin = np.where(np.isfinite(err) & (err >= 0), err, np.nan)
    q95 = np.array([np.nanpercentile(row, 95) if np.isfinite(row).any() else np.nan for row in fin])
    ok = q95[np.isfinite(q95)]
    y_max = min(max(1.35 * float(np.percentile(ok, 90)), 1e-2), y_cap) if ok.size else 1.0
    for ax in axs:
        ax.set(ylim=(0, y_max), xlim=(t[0], t[-1]), xlabel="$t$")
        ax.legend(frameon=False, fontsize=8, ncols=max(1, (len(K_all) - 1) // 8 + 1))
        ax.grid(alpha=0.25, lw=0.5)
    axs[0].set_ylabel(r"$\varepsilon(t)$  [%]")
    axs[0].set_title(f"$r={diag['r_check']}$", fontsize=10)
    return fig, axs

plot_fom_vs_roms(Xtest, X_local, X_global, fom, K, r, n_snapshots=4, n_probes=4, figsize=(13, 14))

One page comparing FOM, global (K=1) ROM and local ql-ROM trajectories.

Rows: snapshots at n_snapshots instants, spatiotemporal maps (1-D only), probe time series, and the pointwise PDF p(u). Works for a 1-D case (fom.Nx) and a 2-D one (fom.Nx, fom.Ny).

Source code in qlroms/utils/diagnosis_plots.py

def plot_fom_vs_roms(Xtest, X_local, X_global, fom, K: int, r: int,
                     n_snapshots: int = 4, n_probes: int = 4, figsize=(13, 14)):
    """One page comparing FOM, global (K=1) ROM and local ql-ROM trajectories.

    Rows: snapshots at n_snapshots instants, spatiotemporal maps (1-D only), probe
    time series, and the pointwise PDF p(u). Works for a 1-D case (fom.Nx) and a
    2-D one (fom.Nx, fom.Ny).
    """
    X, XL, XG = (_to_numpy(a).real for a in (Xtest, X_local, X_global))
    Ndof, Nt = X.shape
    dt = float(fom.dt)
    is_2d = getattr(fom, "Ny", None) is not None
    snaps = np.linspace(0, Nt - 1, n_snapshots + 1, dtype=int)[:-1]
    vmin, vmax = float(X.min()), float(X.max())

    rows = [("FOM", X, "k", "-"), (f"g-ROM\n$r={r}$", XG, "tab:blue", "--"),
            (f"ql-ROM\n$({r},{K})$", XL, "tab:red", ":")]

    if is_2d:
        probes = [ix * fom.Ny + iy for ix in (fom.Nx // 4, 3 * fom.Nx // 4)
                  for iy in (fom.Ny // 4, 3 * fom.Ny // 4)]
        probe_labels = [f"$u(i={p})$" for p in probes]
    else:
        probes = np.linspace(0, Ndof - 1, n_probes + 2, dtype=int)[1:-1].tolist()
        xs = _to_numpy(fom.x).ravel() if hasattr(fom, "x") else np.arange(Ndof)
        probe_labels = [f"$u(x={xs[p]:.1f})$" for p in probes]

    n_map_rows = 0 if is_2d else 3
    n_top = 3 + n_map_rows + len(probes)
    n_cols = max(n_snapshots, 2)
    fig = plt.figure(figsize=(figsize[0], figsize[1] + 2.0), layout="constrained")
    # p(u) gets its own band at the bottom: a legend hung outside an axes is charged to
    # every column that axes spans, so keeping it in the main grid would eat the width of
    # the snapshot panels above it
    top, bottom = fig.subfigures(2, 1, height_ratios=[n_top, 2.5])
    # the full-width rows keep their legend in one narrow reserved column (and out of the
    # layout, so it is not paid for twice)
    gs = top.add_gridspec(n_top, n_cols + 1, width_ratios=[1] * n_cols + [0.62])

    # --- snapshots ---------------------------------------------------------
    ax_share = None
    for i, (lab, data, color, _) in enumerate(rows):
        for j, n in enumerate(snaps):
            # the 1-D line panels all live on the same (vmin, vmax): share that axis so
            # only the first column spends width on tick labels
            ax = top.add_subplot(gs[i, j], sharey=None if is_2d else ax_share)
            if ax_share is None:
                ax_share = ax
            col = np.clip(data[:, n], vmin, vmax)
            if is_2d:
                ax.imshow(col.reshape(fom.Nx, fom.Ny), cmap="RdBu_r", vmin=vmin, vmax=vmax,
                          origin="lower", aspect="auto")
                ax.set(xticks=[], yticks=[])
            else:
                ax.plot(col, color=color, lw=1.2)
                ax.set(ylim=(vmin, vmax), xticks=[])
                ax.tick_params(labelleft=(j == 0))
            if j == 0:
                ax.set_ylabel(lab, fontsize=8)
            if i == 0:
                ax.set_title(f"$t={n * dt:.1f}$", fontsize=9)

    # --- spatiotemporal maps (1-D only) ------------------------------------
    if not is_2d:
        for i, (lab, data, _, _) in enumerate(rows):
            ax = top.add_subplot(gs[3 + i, :n_cols])
            ax.imshow(np.clip(data, vmin, vmax), aspect="auto", origin="lower", cmap="RdBu_r",
                      vmin=vmin, vmax=vmax, extent=[0, Nt * dt, 0, Ndof])
            ax.set_ylabel(lab, fontsize=8)
            for n in snaps:
                ax.axvline(n * dt, color="w", ls=":", lw=0.8)
            if i == 2:
                ax.set_xlabel("$t$", fontsize=9)

    # --- probe time series -------------------------------------------------
    t = np.arange(Nt) * dt
    for i, (p, lab) in enumerate(zip(probes, probe_labels, strict=True)):
        ax = top.add_subplot(gs[3 + n_map_rows + i, :n_cols])
        names = ["FOM", f"global $r={r}$", f"local $(r,K)=({r},{K})$"]
        for name, (_, data, color, ls) in zip(names, rows, strict=True):
            ax.plot(t, data[p], color=color, ls=ls, lw=1.6 if color == "k" else 1.2,
                    label=name, zorder=1 if color == "k" else 2)
        m = 0.15 * (X[p].max() - X[p].min())
        ax.set(ylim=(X[p].min() - m, X[p].max() + m), xlim=(t[0], t[-1]), ylabel=lab)
        ax.grid(alpha=0.2)
        if i == 0:
            leg = ax.legend(frameon=False, fontsize=8, loc="upper left", bbox_to_anchor=(1.01, 1.0))
            leg.set_in_layout(False)      # it lives in the reserved column, not in the row
        if i == len(probes) - 1:
            ax.set_xlabel("$t$", fontsize=9)

    # --- pointwise PDF -----------------------------------------------------
    bins = np.linspace(vmin - 0.5, vmax + 0.5, 60)
    pdf_names = [f"global $r={r}$", f"local $(r,K)=({r},{K})$"]
    gs_pdf = bottom.add_gridspec(1, 2)
    for j, ((_, data, color, _), name) in enumerate(zip(rows[1:], pdf_names, strict=True)):
        ax = bottom.add_subplot(gs_pdf[0, j])
        ax.hist(X.ravel(), bins=bins, density=True, color="0.4", alpha=0.6, label="FOM")
        # no clipping: values outside the FOM range are simply not counted, so a
        # diverged run shows as a deficit rather than as a spike at the bin edge
        ax.hist(data.ravel(), bins=bins, density=True,
                histtype="step", color=color, lw=1.4, label=name)
        ax.set(xlabel="$u$", ylabel="$p(u)$" if j == 0 else None)
        ax.legend(frameon=False, fontsize=8, loc="upper left", bbox_to_anchor=(1.01, 1.0))
    return fig

plot_attractor_charts(A_train, A_test, labels_train, labels_test, centroids, dt, figsize=(13.5, 4.0))

The attractor in three leading global POD coordinates, sliced by the charts.

Panels (a1, a2) and (a1, a3) colour the snapshots by chart (dots train, crosses test, stars the chart centroids). Panel (a1, a2) on the right colours them by time in dominant periods, t f0, with f0 the spectral peak of the training a1, so the colour bar runs from 0 to the record length in periods for any flow and time step.

Parameters:

Name	Type	Description	Default
A_train, A_test		(3, T) coordinates; the test window follows the training one.	required
labels_train, labels_test		(T,) chart ids.	required
centroids		(3, K) chart centroids in the same coordinates.	required
dt	float	snapshot time step.	required

Source code in qlroms/utils/diagnosis_plots.py

def plot_attractor_charts(A_train, A_test, labels_train, labels_test, centroids, dt: float,
                          figsize=(13.5, 4.0)):
    """The attractor in three leading global POD coordinates, sliced by the charts.

    Panels (a1, a2) and (a1, a3) colour the snapshots by chart (dots train, crosses test,
    stars the chart centroids). Panel (a1, a2) on the right colours them by time in
    dominant periods, t f0, with f0 the spectral peak of the training a1, so the colour
    bar runs from 0 to the record length in periods for any flow and time step.

    Args:
        A_train, A_test:            (3, T) coordinates; the test window follows the training one.
        labels_train, labels_test:  (T,) chart ids.
        centroids:                  (3, K) chart centroids in the same coordinates.
        dt:                         snapshot time step.
    """
    A_train, A_test, centroids = (np.asarray(a) for a in (A_train, A_test, centroids))
    labels_train, labels_test = np.asarray(labels_train), np.asarray(labels_test)
    K = centroids.shape[1]
    cmap, _ = chart_cmap(K)
    fig, axes = plt.subplots(1, 3, figsize=figsize, layout="constrained")
    for ax, (i, j) in zip(axes[:2], ((0, 1), (0, 2))):
        for k in range(K):
            m_tr, m_te = labels_train == k, labels_test == k
            ax.plot(A_train[i, m_tr], A_train[j, m_tr], ".", color=cmap(k), ms=3, alpha=0.7,
                    label=f"chart {k} (train)")
            ax.plot(A_test[i, m_te], A_test[j, m_te], "x", color=cmap(k), ms=4, mew=0.8, alpha=0.9)
        ax.plot(centroids[i], centroids[j], "*", color="k", ms=16, zorder=5)
        ax.set(xlabel=f"$a_{i + 1}$", ylabel=f"$a_{j + 1}$")
    axes[0].set_title(f"$K = {K}$ charts (dots train, crosses test)", fontsize=10)
    axes[1].set_title("stars: chart centroids", fontsize=10)
    axes[0].legend(frameon=False, fontsize=7, ncol=2)

    ax, n_tr = axes[2], A_train.shape[1]
    f0 = dominant_frequency(A_train[0], dt)
    tau = np.arange(n_tr + A_test.shape[1]) * dt * f0
    norm = plt.Normalize(0, tau[-1])
    s = ax.scatter(A_train[0], A_train[1], c=tau[:n_tr], norm=norm, s=4, cmap="viridis")
    ax.scatter(A_test[0], A_test[1], c=tau[n_tr:], norm=norm, s=10, cmap="viridis",
               marker="x", linewidths=0.7)
    fig.colorbar(s, ax=ax, label=f"$t f_0$ ($f_0 = {f0:.4f}$)", shrink=0.9)
    ax.set(xlabel="$a_1$", ylabel="$a_2$", title="coloured by time, not by chart")
    return fig

plot_forecast_error_curves(curves, dt, title=None, figsize=(7, 3.2), cmaps=('Blues', 'Oranges', 'Greens', 'Purples'))

Relative forecast error against time for the curves of forecast_error_curves.

Each model takes one colour map, in order (blues, oranges, ...). Where the active chart is known, the curve is shaded by chart, light for chart 0 to dark for the last one, so a colour change marks a chart switch. The legend carries each model's mean error.

Source code in qlroms/utils/diagnosis_plots.py

def plot_forecast_error_curves(curves: dict, dt: float, title: str | None = None, figsize=(7, 3.2),
                               cmaps=("Blues", "Oranges", "Greens", "Purples")):
    """Relative forecast error against time for the curves of `forecast_error_curves`.

    Each model takes one colour map, in order (blues, oranges, ...). Where the active chart
    is known, the curve is shaded by chart, light for chart 0 to dark for the last one, so
    a colour change marks a chart switch. The legend carries each model's mean error.
    """
    from matplotlib.collections import LineCollection

    fig, ax = plt.subplots(figsize=figsize, layout="constrained")
    for (name, c), cmap_name in zip(curves.items(), cmaps):
        err, chart = np.asarray(c["error"]), c["chart"]
        t = np.arange(1, len(err) + 1) * dt
        cmap = plt.get_cmap(cmap_name)
        label = f"{name} (mean {np.nanmean(err):.3f})"
        if chart is None:
            ax.plot(t, err, lw=1.3, color=cmap(0.7), label=label)
            continue
        chart = np.asarray(chart)
        shade = 0.35 + 0.6 * chart / max(chart.max(), 1)          # chart id -> colour-map level
        pts = np.column_stack([t, err])
        ax.add_collection(LineCollection(np.stack([pts[:-1], pts[1:]], axis=1), colors=cmap(shade[1:]), lw=1.5))
        ax.plot([], [], lw=1.5, color=cmap(0.65), label=label)   # legend proxy
    ax.autoscale_view()
    ax.set(xlabel="$t$ [s]", ylabel="relative error", yscale="log", title=title)
    ax.grid(alpha=0.25, lw=0.5)
    ax.legend(frameon=False, fontsize=8, title="shade: chart $k$ (light to dark)", title_fontsize=7)
    return fig

save_figures_pdf(figures, save_dir, filename='diagnosis.pdf')

Write figures (bare or (fig, title) pairs) to one multi-page PDF; return its path.

Source code in qlroms/utils/diagnosis_plots.py

def save_figures_pdf(figures, save_dir: str, filename: str = "diagnosis.pdf") -> str:
    """Write figures (bare or (fig, title) pairs) to one multi-page PDF; return its path."""
    from matplotlib.backends.backend_pdf import PdfPages

    os.makedirs(save_dir, exist_ok=True)
    path = os.path.join(save_dir, filename)
    with PdfPages(path) as pdf:
        for entry in figures:
            fig = entry[0] if isinstance(entry, tuple) else entry
            if fig is not None:
                pdf.savefig(fig, bbox_inches="tight")
    print(f"Saved figures -> {path}")
    return path