qlroms.utils.plots

Figure helpers shared by the tutorials and the scripts.

Only the pieces that were being copy-pasted between notebooks live here: the chart palette, symmetric field images, the truth/model/difference snapshot grid, the coordinate and relative-error panels, chart-path statistics and the outputs directory. Everything takes plain arrays (torch tensors are converted), so nothing here knows which model family produced them.

Figures that read a qlroms.utils.diagnosis result dict stay in qlroms.utils.diagnosis_plots.

model_color(label, fallback=0)

The colour label keeps everywhere; unknown labels fall back to the chart palette.

Source code in qlroms/utils/plots.py

def model_color(label, fallback=0):
    """The colour `label` keeps everywhere; unknown labels fall back to the chart palette."""
    key = str(label).lower()
    for name, color in MODEL_COLORS.items():
        if name in key:
            return color
    return CHART_COLORS[fallback % len(CHART_COLORS)]

chart_cmap(K)

Discrete colormap + norm for chart ids 0..K-1 (the palette cycles past 10 charts).

Source code in qlroms/utils/plots.py

def chart_cmap(K: int):
    """Discrete colormap + norm for chart ids 0..K-1 (the palette cycles past 10 charts)."""
    cmap = ListedColormap([CHART_COLORS[k % len(CHART_COLORS)] for k in range(K)])
    return cmap, BoundaryNorm(np.arange(K + 1) - 0.5, K)

field_image(ax, data, shape=None, vmax=None, cmap='RdBu_r', origin='lower', ticks=False, **kw)

imshow a field (or an operator matrix) on a symmetric scale, with the ticks off.

shape reshapes a flattened snapshot, e.g. (Nx, Ny); pass None for data that is already 2-D. vmax defaults to the largest magnitude in data, so pass the one of the reference field to put several panels on the same scale. ticks=True keeps the axes, for a map whose axes mean something (extent=[0, T, 0, L]).

Source code in qlroms/utils/plots.py

def field_image(ax, data, shape=None, vmax=None, cmap="RdBu_r", origin="lower", ticks=False, **kw):
    """imshow a field (or an operator matrix) on a symmetric scale, with the ticks off.

    `shape` reshapes a flattened snapshot, e.g. (Nx, Ny); pass None for data that is
    already 2-D. `vmax` defaults to the largest magnitude in `data`, so pass the one
    of the reference field to put several panels on the same scale. `ticks=True` keeps
    the axes, for a map whose axes mean something (`extent=[0, T, 0, L]`).
    """
    d = _to_numpy(data)
    if shape is not None:
        d = d.reshape(shape)
    vmax = float(np.abs(d).max()) if vmax is None else float(vmax)
    im = ax.imshow(d, cmap=cmap, vmin=-vmax, vmax=vmax, origin=origin, **kw)
    if not ticks:
        ax.set(xticks=[], yticks=[])
    return im

snapshot_grid(rows, idx, shape, dt=1.0, col_titles=None, vmax=None, suptitle=None, figsize=None)

Grid of fields: one row per (label, snapshots) pair, one column per index in idx.

All panels share the colour scale of the first row, so the difference row reads as an error and not as its own field.

Source code in qlroms/utils/plots.py

def snapshot_grid(rows, idx, shape, dt=1.0, col_titles=None, vmax=None, suptitle=None,
                  figsize=None):
    """Grid of fields: one row per (label, snapshots) pair, one column per index in `idx`.

    All panels share the colour scale of the first row, so the difference row reads as
    an error and not as its own field.
    """
    if vmax is None:
        vmax = float(np.abs(_to_numpy(rows[0][1])).max())
    if col_titles is None:
        col_titles = [f"$t = {n * dt:.2f}$" for n in idx]
    fig, axs = plt.subplots(len(rows), len(idx), squeeze=False, layout="constrained",
                            figsize=figsize or (2.3 * len(idx), 1.9 * len(rows)))
    for j, n in enumerate(idx):
        for i, (label, X) in enumerate(rows):
            im = field_image(axs[i, j], X[:, n], shape, vmax=vmax)
            if j == 0:
                axs[i, j].set_ylabel(label, fontsize=9)
        axs[0, j].set_title(col_titles[j], fontsize=9)
    fig.colorbar(im, ax=axs, shrink=0.7, pad=0.01)
    if suptitle:
        fig.suptitle(suptitle, fontsize=10)
    return fig, axs

snapshot_triptych(X_true, X_model, err, shape, dt, tag='model', title=None, threshold=0.5, figsize=(6.8, 6.0))

Fields before / at / after the valid time -- the first crossing of threshold.

Rows: truth, model, difference. X_true is the truth over the same window as X_model, and err its relative-error series (both start at the same snapshot).

Source code in qlroms/utils/plots.py

def snapshot_triptych(X_true, X_model, err, shape, dt, tag="model", title=None,
                      threshold=0.5, figsize=(6.8, 6.0)):
    """Fields before / at / after the valid time -- the first crossing of `threshold`.

    Rows: truth, model, difference. `X_true` is the truth over the same window as
    `X_model`, and `err` its relative-error series (both start at the same snapshot).
    """
    err = np.asarray(err, dtype=float)
    n_fc = X_model.shape[1]
    cross = np.flatnonzero(err > threshold)
    n_ph = int(cross[0]) if cross.size else n_fc - 1
    idx = [max(1, n_ph // 2), n_ph, min(n_fc - 1, max(int(1.6 * n_ph), n_ph + 2))]
    titles = [f"$t = {n * dt:.2f}$ ({when} $T_{{ph}}$)\nrel err {err[n]:.2f}"
              for n, when in zip(idx, ["before", "at", "after"], strict=True)]
    rows = [("truth", X_true), (tag, X_model), ("difference", X_true - X_model)]
    return snapshot_grid(rows, idx, shape, col_titles=titles, suptitle=title, figsize=figsize)

plot_coords(axes, z_true, dt, preds=None, t0=0.0, pad=0.6)

Truth vs predictions, one panel per row of z_true (r, T).

preds maps a label to the predicted coordinates, optionally with a colour and a linestyle. The y-range is set by the truth, so a diverging run leaves the panel as a gap instead of squashing everything else flat.

Source code in qlroms/utils/plots.py

def plot_coords(axes, z_true, dt, preds=None, t0=0.0, pad=0.6):
    """Truth vs predictions, one panel per row of `z_true` (r, T).

    `preds` maps a label to the predicted coordinates, optionally with a colour and a
    linestyle. The y-range is set by the truth, so a diverging run leaves the panel as
    a gap instead of squashing everything else flat.
    """
    z_true = _to_numpy(z_true)
    t = t0 + np.arange(z_true.shape[1]) * dt
    for j, ax in enumerate(axes):
        margin = pad * (z_true[j].max() - z_true[j].min())
        lo, hi = z_true[j].min() - margin, z_true[j].max() + margin
        # truth underneath, so a model that sits on top of it still shows
        ax.plot(t, z_true[j], color="0.25", lw=2.0, label="truth", zorder=1)
        for label, spec in (preds or {}).items():
            z, color, ls = _style(spec)
            z = _to_numpy(z)
            zj = np.where((z[j] < lo) | (z[j] > hi), np.nan, z[j])   # off-panel -> gap,
            ax.plot(t[:z.shape[1]], zj, color=color, lw=1.2, ls=ls,  # not a vertical streak
                    label=label, zorder=2)
        ax.set(ylabel=f"$z_{{{j + 1}}}$", ylim=(lo, hi))
        ax.grid(alpha=0.25, lw=0.5)
    axes[-1].set_xlabel("time [t.u.]")

plot_error_curves(ax, curves, dt, threshold=0.5, t0=0.0, xlabel='time [t.u.]', ylabel='relative error', lw=1.2, log=True)

Relative-error series, with the T_ph threshold drawn.

curves maps a label to the error series, optionally with a colour and a linestyle. log=False gives a linear axis, which reads better when every model stays bounded.

Source code in qlroms/utils/plots.py

def plot_error_curves(ax, curves, dt, threshold=0.5, t0=0.0, xlabel="time [t.u.]",
                      ylabel="relative error", lw=1.2, log=True):
    """Relative-error series, with the T_ph threshold drawn.

    `curves` maps a label to the error series, optionally with a colour and a linestyle.
    `log=False` gives a linear axis, which reads better when every model stays bounded.
    """
    plot = ax.semilogy if log else ax.plot
    for label, spec in curves.items():
        err, color, ls = _style(spec)
        err = _to_numpy(err)
        plot(t0 + np.arange(len(err)) * dt, err, color=color, ls=ls, lw=lw, label=label)
    if threshold:
        ax.axhline(threshold, color="0.5", ls=":", lw=1)
    ax.set(xlabel=xlabel, ylabel=ylabel)
    ax.grid(alpha=0.25, lw=0.5)

plot_evaluation(z_true, models, dt, ids_true=None, threshold=0.5, t0=0.0, title=None, n_coords=3, figsize=(11, 6), log_error=False)

Truth and every model on one figure: coordinates, relative error, active chart.

models maps a label to a dict holding any of "z" (coordinates, drawn vertically offset), "err" (relative-error series) and "ids" (active chart id per step). Adding a second model adds curves to the same three panels instead of a second figure: the truth is drawn once, each model keeps its model_color in every panel and every case, and the legend sits outside the axes.

Returns (fig, metrics). metrics[label] carries T_ph (time of the first crossing of threshold, None if it never crosses), err_mean, err_final and chart_match (fraction of steps in the same chart as ids_true).

Source code in qlroms/utils/plots.py

def plot_evaluation(z_true, models, dt, ids_true=None, threshold=0.5, t0=0.0, title=None,
                    n_coords=3, figsize=(11, 6), log_error=False):
    """Truth and every model on one figure: coordinates, relative error, active chart.

    `models` maps a label to a dict holding any of "z" (coordinates, drawn vertically
    offset), "err" (relative-error series) and "ids" (active chart id per step). Adding
    a second model adds curves to the same three panels instead of a second figure: the
    truth is drawn once, each model keeps its `model_color` in every panel and every
    case, and the legend sits outside the axes.

    Returns `(fig, metrics)`. `metrics[label]` carries `T_ph` (time of the first crossing
    of `threshold`, None if it never crosses), `err_mean`, `err_final` and `chart_match`
    (fraction of steps in the same chart as `ids_true`).
    """
    z_true = _to_numpy(z_true)
    ids_true = None if ids_true is None else np.asarray(_to_numpy(ids_true), dtype=int)
    show_ids = ids_true is not None or any(m.get("ids") is not None for m in models.values())

    fig = plt.figure(figsize=figsize, layout="constrained")
    gs = fig.add_gridspec(2, 2, height_ratios=[1.35, 1.0])
    ax_z = fig.add_subplot(gs[0, :])
    ax_e = fig.add_subplot(gs[1, 0] if show_ids else gs[1, :])
    ax_k = fig.add_subplot(gs[1, 1]) if show_ids else None

    # (a) leading coordinates, stacked with a common offset so one panel holds them all
    nz = min(n_coords, z_true.shape[0])
    span = float(np.ptp(z_true[:nz], axis=1).max()) or 1.0
    off = 1.25 * span * np.arange(nz)
    band = 0.62 * span            # half-height each coordinate may use before it overlaps
    t = t0 + np.arange(z_true.shape[1]) * dt
    for j in range(nz):
        ax_z.plot(t, z_true[j] + off[j], color=MODEL_COLORS["truth"], lw=1.8, zorder=2,
                  label="truth" if j == 0 else None)
    # the truth sets the scale: a diverged model leaves a gap instead of flattening the panel
    mid = np.array([0.5 * (z_true[j].max() + z_true[j].min()) for j in range(nz)])
    ax_z.set_ylim(off[0] + mid[0] - band, off[-1] + mid[-1] + band)
    ax_z.set(yticks=off, yticklabels=[f"$z_{{{j + 1}}}$" for j in range(nz)],
             xlabel="time [t.u.]")
    ax_z.set_title("(a) leading coordinates (vertically offset)", fontsize=10)
    ax_z.grid(alpha=0.25, lw=0.5)

    metrics, curves = {}, {}
    for i, (label, m) in enumerate(models.items()):
        color = model_color(label, fallback=i + 1)
        z = m.get("z")
        if z is not None:
            z = _to_numpy(z)
            for j in range(min(nz, z.shape[0])):
                zj = np.where(np.abs(z[j] - mid[j]) > band, np.nan, z[j])
                ax_z.plot(t[:z.shape[1]], zj + off[j], color=color, ls="--", lw=1.2,
                          zorder=3, label=label if j == 0 else None)
        err = None if m.get("err") is None else np.asarray(_to_numpy(m["err"]), dtype=float)
        if err is not None:
            # a diverged loop overflows: draw it up to the blow-up, then leave a gap
            fin = np.isfinite(err)
            curves[label] = (np.where(fin, err, np.nan), color, "--")
        ids = None if m.get("ids") is None else np.asarray(_to_numpy(m["ids"]), dtype=int)
        if ids is not None and ax_k is not None:
            ax_k.step(t0 + np.arange(len(ids)) * dt, ids, where="post", color=color,
                      ls="--", lw=1.2, zorder=3, label=label)
        cross = (np.flatnonzero(~np.isfinite(err) | (err > threshold)) if err is not None
                 else np.array([], dtype=int))
        # accuracy is only meaningful while the loop still tracks, so average over the
        # valid window; a diverged run is described by T_ph, not by its blow-up size
        valid = err[:cross[0]] if err is not None and cross.size else err
        n = 0 if ids_true is None or ids is None else min(len(ids), len(ids_true))
        metrics[label] = dict(
            T_ph=float(t0 + cross[0] * dt) if cross.size else None,
            err_mean=float(valid.mean()) if valid is not None and valid.size else None,
            chart_match=float((ids[:n] == ids_true[:n]).mean()) if n else None,
        )

    # (b) relative error, with the horizon of every model written where it can be read
    if curves:
        plot_error_curves(ax_e, curves, dt, threshold=threshold, t0=t0, log=log_error)
        if not log_error:      # a diverged curve must not flatten everyone else to zero
            top = max([m["err_mean"] or 0.0 for m in metrics.values()] + [threshold])
            ax_e.set_ylim(0.0, 1.6 * top)
        for i, (label, mt) in enumerate(metrics.items()):
            if mt["err_mean"] is None:
                continue
            horizon = (f"$T_{{ph}} = {mt['T_ph']:.1f}$" if mt["T_ph"] is not None
                       else "$T_{ph} >$ window")
            ax_e.text(0.98, 0.94 - 0.09 * i, f"{label}: {horizon}", transform=ax_e.transAxes,
                      ha="right", va="top", fontsize=8, color=model_color(label, fallback=i + 1))
    ax_e.set_title("(b) relative error of the closed loop", fontsize=10)

    # (c) which local model is in charge, truth against every closed loop
    if ax_k is not None:
        if ids_true is not None:
            ax_k.step(t0 + np.arange(len(ids_true)) * dt, ids_true, where="post",
                      color=MODEL_COLORS["truth"], lw=1.8, zorder=2, label="truth")
        ax_k.set(xlabel="time [t.u.]", ylabel="chart id")
        ax_k.set_title("(c) active chart", fontsize=10)
        ax_k.grid(alpha=0.25, lw=0.5)

    handles, labels = ax_z.get_legend_handles_labels()
    if not handles and ax_k is not None:
        handles, labels = ax_k.get_legend_handles_labels()
    fig.legend(handles, labels, loc="outside lower center", ncols=min(len(labels), 5),
               frameon=False, fontsize=9)
    if title:
        fig.suptitle(title, fontsize=11)
    return fig, metrics

chart_stats(ids, K=None)

Occupancy, segments and transition matrix P(j | i) of a chart-id path.

Returns occupancy (K,), the seg_len / seg_lab of every contiguous visit (in steps, so multiply by dt for times), and the row-normalized T (K, K).

Source code in qlroms/utils/plots.py

def chart_stats(ids, K=None):
    """Occupancy, segments and transition matrix P(j | i) of a chart-id path.

    Returns `occupancy` (K,), the `seg_len` / `seg_lab` of every contiguous visit
    (in steps, so multiply by dt for times), and the row-normalized `T` (K, K).
    """
    ids = np.asarray(ids, dtype=int)
    K = int(ids.max()) + 1 if K is None else int(K)
    starts = np.concatenate(([0], np.flatnonzero(np.diff(ids)) + 1))
    T = np.zeros((K, K))
    np.add.at(T, (ids[:-1], ids[1:]), 1.0)
    return dict(occupancy=np.bincount(ids, minlength=K) / len(ids),
                seg_len=np.diff(np.concatenate((starts, [len(ids)]))),
                seg_lab=ids[starts],
                T=T / np.maximum(T.sum(axis=1, keepdims=True), 1.0))

outputs_dir()

tutorials/figs, resolved from the repo root or from tutorials/ itself.

Source code in qlroms/utils/plots.py

def outputs_dir():
    """`tutorials/figs`, resolved from the repo root or from `tutorials/` itself."""
    out = Path("tutorials/figs") if Path("tutorials").is_dir() else Path("figs")
    out.mkdir(parents=True, exist_ok=True)
    return out

start_pdf(name)

Collect every figure of this notebook into outputs_dir()/<name>.pdf, one per page.

Call it once in the setup cell and close_pdf() in the last cell. Figures are picked up from save_figure and from every plt.show(), so no cell has to opt in, and nothing is written twice. Pages land as they come, so a forgotten close still leaves a valid file.

Source code in qlroms/utils/plots.py

def start_pdf(name):
    """Collect every figure of this notebook into `outputs_dir()/<name>.pdf`, one per page.

    Call it once in the setup cell and `close_pdf()` in the last cell. Figures are picked up
    from `save_figure` and from every `plt.show()`, so no cell has to opt in, and nothing is
    written twice. Pages land as they come, so a forgotten close still leaves a valid file.
    """
    close_pdf()
    path = outputs_dir() / f"{name}.pdf"
    _PDF.update(name=name, pages=PdfPages(path), n=0, show=plt.show)
    plt.show = _show_to_pdf
    atexit.register(close_pdf)
    return path

close_pdf()

Close the open figure PDF, if any, and say how many pages it got.

Source code in qlroms/utils/plots.py

def close_pdf():
    """Close the open figure PDF, if any, and say how many pages it got."""
    if _PDF["pages"] is None:
        return None
    _flush_open_figures()                         # anything still open at the end
    path = outputs_dir() / f"{_PDF['name']}.pdf"
    _PDF["pages"].close()
    if _PDF["show"] is not None:
        plt.show = _PDF["show"]
    print(f"{_PDF['n']} figures written to {path}")
    _PDF.update(name=None, pages=None, n=0, show=None)
    return path

save_figure(fig, name=None, dpi=150)

Add fig to the notebook's open PDF (see start_pdf) and, if name is given, to a PNG.

Source code in qlroms/utils/plots.py

def save_figure(fig, name=None, dpi=150):
    """Add `fig` to the notebook's open PDF (see `start_pdf`) and, if `name` is given, to a PNG."""
    if _PDF["pages"] is not None:
        _add_page(fig)
    if name is None:
        return None
    path = outputs_dir() / name
    fig.savefig(path, dpi=dpi)
    print(f"Figure saved to {path}")
    return path