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=(13, 4.2), log_scale=True)

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

Panels: 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=(13, 4.2), log_scale: bool = True):
    """Heatmaps over the (K, r) grid, with the K=1 global ROM as the bottom row.

    Panels: 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, 3 if has_kld else 2, 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)

    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, 0], M, norm_m, "viridis", "projection MSE")

    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, 1], 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, 2], 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_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

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