torchtsetlin.viz

torchtsetlin.viz

Visualisation helpers (matplotlib is an optional dependency).

Every function returns the matplotlib Axes it drew on, so figures can be composed.

plot_history

plot_history(history, metrics: Optional[Sequence[str]] = None, ax=None)

Plot metric curves of a :class:~torchtsetlin.train.History.

Source code in src/torchtsetlin/viz.py

def plot_history(history, metrics: Optional[Sequence[str]] = None, ax=None):
    """Plot metric curves of a :class:`~torchtsetlin.train.History`."""
    plt = _plt()
    if ax is None:
        _, ax = plt.subplots(figsize=(6, 4))
    keys = metrics or [k for k in history.keys() if k.endswith("accuracy") or k.endswith("mae") or k.endswith("f1")]
    for k in keys:
        ys = history[k]
        xs = [i + 1 for i, v in enumerate(ys) if v is not None]
        ax.plot(xs, [v for v in ys if v is not None], marker=".", label=k)
    ax.set_xlabel("epoch")
    ax.grid(alpha=0.3)
    ax.legend()
    return ax

plot_confusion_matrix

plot_confusion_matrix(cm, class_names: Optional[Sequence[str]] = None, normalize: bool = False, ax=None, cmap='Blues')

Heat-map of a confusion matrix (rows = true, columns = predicted).

Source code in src/torchtsetlin/viz.py

def plot_confusion_matrix(cm, class_names: Optional[Sequence[str]] = None, normalize: bool = False, ax=None, cmap="Blues"):
    """Heat-map of a confusion matrix (rows = true, columns = predicted)."""
    plt = _plt()
    cm = _np(cm).astype(float)
    if normalize:
        cm = cm / np.clip(cm.sum(axis=1, keepdims=True), 1, None)
    if ax is None:
        _, ax = plt.subplots(figsize=(5, 4.5))
    im = ax.imshow(cm, cmap=cmap)
    K = cm.shape[0]
    names = list(class_names) if class_names is not None else [str(i) for i in range(K)]
    ax.set_xticks(range(K), names, rotation=45, ha="right")
    ax.set_yticks(range(K), names)
    ax.set_xlabel("predicted")
    ax.set_ylabel("true")
    thresh = cm.max() / 2 if cm.size else 0
    for i in range(K):
        for j in range(K):
            v = cm[i, j]
            ax.text(j, i, f"{v:.2f}" if normalize else f"{int(v)}", ha="center", va="center", color="white" if v > thresh else "black", fontsize=8)
    plt.colorbar(im, ax=ax, fraction=0.046)
    return ax

plot_clause_memory

plot_clause_memory(model, clause: int, feature_names: Optional[Sequence[str]] = None, ax=None)

Book-style memory plot of one clause: the automaton state of every literal with the include boundary drawn as a dashed line (states above it are memorized).

Source code in src/torchtsetlin/viz.py

def plot_clause_memory(model, clause: int, feature_names: Optional[Sequence[str]] = None, ax=None):
    """Book-style memory plot of one clause: the automaton state of every literal with the
    include boundary drawn as a dashed line (states above it are memorized)."""
    plt = _plt()
    names = model.literal_names(feature_names)
    states = _np(model.ta_state[clause]).astype(int) + 1  # 1-based like the book
    if ax is None:
        _, ax = plt.subplots(figsize=(max(6, len(names) * 0.35), 3.5))
    colors = ["tab:blue" if s > model.n_states else "lightgray" for s in states]
    ax.bar(range(len(names)), states, color=colors)
    ax.axhline(model.n_states + 0.5, color="red", linestyle="--", linewidth=1)
    ax.set_xticks(range(len(names)), names, rotation=90, fontsize=8)
    ax.set_ylim(0, 2 * model.n_states + 1)
    ax.set_ylabel("memory position")
    ax.set_title(f"clause {clause}: IF {model.clause_expression(clause, feature_names)}")
    return ax

plot_ta_states

plot_ta_states(model, clauses: Optional[Sequence[int]] = None, ax=None, cmap='RdBu_r')

Heat-map of all automata states (clauses x literals), centred on the include boundary.

Source code in src/torchtsetlin/viz.py

def plot_ta_states(model, clauses: Optional[Sequence[int]] = None, ax=None, cmap="RdBu_r"):
    """Heat-map of all automata states (clauses x literals), centred on the include boundary."""
    plt = _plt()
    st = model.ta_state
    if clauses is not None:
        st = st[list(clauses)]
    st = _np(st).astype(float) - (model.n_states - 0.5)
    if ax is None:
        _, ax = plt.subplots(figsize=(8, 4))
    im = ax.imshow(st, aspect="auto", cmap=cmap, vmin=-model.n_states, vmax=model.n_states, interpolation="nearest")
    ax.set_xlabel("literal (x_0..x_F-1, NOT x_0..NOT x_F-1)")
    ax.set_ylabel("clause")
    plt.colorbar(im, ax=ax, label="state − boundary")
    return ax

plot_literal_frequency

plot_literal_frequency(model, feature_names: Optional[Sequence[str]] = None, top_k: Optional[int] = 30, ax=None)

Bar plot of how many clauses include each literal.

Source code in src/torchtsetlin/viz.py

def plot_literal_frequency(model, feature_names: Optional[Sequence[str]] = None, top_k: Optional[int] = 30, ax=None):
    """Bar plot of how many clauses include each literal."""
    plt = _plt()
    freq = _np(model.literal_frequency())
    names = model.literal_names(feature_names)
    order = np.argsort(-freq)
    if top_k:
        order = order[:top_k]
    if ax is None:
        _, ax = plt.subplots(figsize=(max(6, len(order) * 0.3), 3.5))
    ax.bar(range(len(order)), freq[order], color=["tab:blue" if i < len(names) // 2 else "tab:orange" for i in order])
    ax.set_xticks(range(len(order)), [names[i] for i in order], rotation=90, fontsize=8)
    ax.set_ylabel("# clauses")
    return ax

plot_feature_importance

plot_feature_importance(importance, names: Optional[Sequence[str]] = None, top_k: Optional[int] = 20, ax=None, title: str = 'feature importance')

Horizontal bar plot of a (n,) importance vector.

Source code in src/torchtsetlin/viz.py

def plot_feature_importance(importance, names: Optional[Sequence[str]] = None, top_k: Optional[int] = 20, ax=None, title: str = "feature importance"):
    """Horizontal bar plot of a ``(n,)`` importance vector."""
    plt = _plt()
    imp = _np(importance).reshape(-1)
    names = list(names) if names is not None else [str(i) for i in range(len(imp))]
    order = np.argsort(-imp)
    if top_k:
        order = order[:top_k]
    if ax is None:
        _, ax = plt.subplots(figsize=(6, max(2.5, 0.3 * len(order))))
    ax.barh(range(len(order))[::-1], imp[order])
    ax.set_yticks(range(len(order))[::-1], [names[i] for i in order])
    ax.set_title(title)
    return ax

plot_conv_clause

plot_conv_clause(model, clause: int, plane: int = 0, ax=None, show_region: bool = True, depth: Optional[int] = None)

Draw the patch pattern of a convolutional clause: black = pixel must be on, white = must be off, gray = don't care. The title reports the allowed patch-position range.

For 3D models (:class:~torchtsetlin.models.Conv3dTsetlinMachine) the kd depth slices of the window are drawn side by side, separated by a blank column; pass depth to draw a single slice instead.

Source code in src/torchtsetlin/viz.py

def plot_conv_clause(model, clause: int, plane: int = 0, ax=None, show_region: bool = True, depth: Optional[int] = None):
    """Draw the patch pattern of a convolutional clause: black = pixel must be on, white =
    must be off, gray = don't care. The title reports the allowed patch-position range.

    For 3D models (:class:`~torchtsetlin.models.Conv3dTsetlinMachine`) the ``kd`` depth
    slices of the window are drawn side by side, separated by a blank column; pass
    ``depth`` to draw a single slice instead."""
    plt = _plt()
    patch = _np(model.clause_patch(clause))[plane].astype(float)
    slice_ticks: Sequence[float] = []
    slice_labels: Sequence[str] = []
    if patch.ndim == 3:  # (kd, kh, kw) — a 3D window
        if depth is not None:
            patch = patch[depth]
            slice_ticks, slice_labels = [(patch.shape[1] - 1) / 2], [f"d={depth}"]
        else:
            kd, kh, kw = patch.shape
            gap = np.full((kh, 1), np.nan)
            parts = []
            for i in range(kd):
                parts += [patch[i], gap] if i < kd - 1 else [patch[i]]
            patch = np.concatenate(parts, axis=1)
            slice_ticks = [i * (kw + 1) + (kw - 1) / 2 for i in range(kd)]
            slice_labels = [f"d={i}" for i in range(kd)]
    if ax is None:
        _, ax = plt.subplots(figsize=(3, 3))
    cmap = plt.get_cmap("gray").copy()
    cmap.set_bad("tab:blue", 0.35)  # the separator between depth slices
    ax.imshow(-patch, cmap=cmap, vmin=-1, vmax=1, interpolation="nearest")
    ax.set_xticks(slice_ticks, slice_labels, fontsize=7)
    ax.tick_params(length=0)
    ax.set_yticks([])
    title = f"clause {clause}"
    if show_region and model.position_encoding:
        r = model.clause_region(clause)
        box = " ".join(f"{a}∈[{r[a][0]},{r[a][1]}]" for a in ("d", "y", "x") if a in r)
        title += "\n" + box
    ax.set_title(title, fontsize=8)
    return ax

plot_conv_clauses

plot_conv_clauses(model, clauses: Sequence[int], ncols: int = 8, plane: int = 0, figsize_per: float = 1.6, depth: Optional[int] = None)

Grid of :func:plot_conv_clause panels. Returns the figure.

For 3D models each panel shows the kd depth slices side by side (or only depth).

Source code in src/torchtsetlin/viz.py

def plot_conv_clauses(model, clauses: Sequence[int], ncols: int = 8, plane: int = 0, figsize_per: float = 1.6, depth: Optional[int] = None):
    """Grid of :func:`plot_conv_clause` panels. Returns the figure.

    For 3D models each panel shows the ``kd`` depth slices side by side (or only ``depth``)."""
    plt = _plt()
    n = len(clauses)
    ncols = min(ncols, n)
    nrows = int(np.ceil(n / ncols))
    kd = model.patch_size[0] if len(model.patch_size) == 3 and depth is None else 1
    width = figsize_per * (1 + 0.8 * (kd - 1))
    fig, axes = plt.subplots(nrows, ncols, figsize=(width * ncols, figsize_per * 1.25 * nrows), squeeze=False)
    for ax in axes.flat:
        ax.axis("off")
    for ax, j in zip(axes.flat, clauses):
        ax.axis("on")
        plot_conv_clause(model, int(j), plane=plane, ax=ax, depth=depth)
    fig.tight_layout()
    return fig

plot_votes

plot_votes(votes, y=None, class_names: Optional[Sequence[str]] = None, ax=None, bins: int = 30)

Histogram of the winning-class vote sums, split by correct / incorrect if y given.

Source code in src/torchtsetlin/viz.py

def plot_votes(votes, y=None, class_names: Optional[Sequence[str]] = None, ax=None, bins: int = 30):
    """Histogram of the winning-class vote sums, split by correct / incorrect if ``y`` given."""
    plt = _plt()
    v = _np(votes)
    if ax is None:
        _, ax = plt.subplots(figsize=(6, 3.5))
    top = v.max(axis=1)
    if y is not None:
        yy = _np(y).reshape(-1)
        correct = v.argmax(axis=1) == yy
        ax.hist(top[correct], bins=bins, alpha=0.6, label="correct")
        ax.hist(top[~correct], bins=bins, alpha=0.6, label="incorrect")
        ax.legend()
    else:
        ax.hist(top, bins=bins)
    ax.set_xlabel("winning vote sum")
    ax.set_ylabel("count")
    return ax

plot_trustworthiness

plot_trustworthiness(curve, ax=None)

Plot accuracy vs confidence level (see :func:torchtsetlin.metrics.trustworthiness_curve).

Source code in src/torchtsetlin/viz.py

def plot_trustworthiness(curve, ax=None):
    """Plot accuracy vs confidence level (see :func:`torchtsetlin.metrics.trustworthiness_curve`)."""
    plt = _plt()
    if ax is None:
        _, ax = plt.subplots(figsize=(5, 3.5))
    ax.plot(_np(curve["confidence"]), _np(curve["accuracy"]), marker="o", label="accuracy")
    ax.plot(_np(curve["confidence"]), _np(curve["coverage"]), marker="s", label="coverage")
    ax.set_xlabel("confidence ≥")
    ax.legend()
    ax.grid(alpha=0.3)
    return ax

plot_clause_weights

plot_clause_weights(model, ax=None)

Weights of a weighted / coalesced model as an image (clauses x outputs) or bars.

Source code in src/torchtsetlin/viz.py

def plot_clause_weights(model, ax=None):
    """Weights of a weighted / coalesced model as an image (clauses x outputs) or bars."""
    plt = _plt()
    w = _np(model.weights)
    if ax is None:
        _, ax = plt.subplots(figsize=(6, 3.5))
    if w.ndim == 2:
        im = ax.imshow(w.T, aspect="auto", cmap="RdBu_r", vmin=-np.abs(w).max(), vmax=np.abs(w).max(), interpolation="nearest")
        ax.set_xlabel("clause")
        ax.set_ylabel("output")
        plt.colorbar(im, ax=ax)
    else:
        ax.bar(range(len(w)), w)
        ax.set_xlabel("clause")
        ax.set_ylabel("weight")
    return ax

segmentation_palette

segmentation_palette(n_classes: int, ignore_index: Optional[int] = None) -> np.ndarray

(n, 3) float RGB palette for label maps, indexable by class id.

Uses matplotlib's qualitative tab10/tab20 so neighbouring classes stay distinguishable, and paints ignore_index mid-grey. Index it with a label map: palette[labels] gives an (H, W, 3) image.

Source code in src/torchtsetlin/viz.py

def segmentation_palette(n_classes: int, ignore_index: Optional[int] = None) -> np.ndarray:
    """``(n, 3)`` float RGB palette for label maps, indexable by class id.

    Uses matplotlib's qualitative ``tab10``/``tab20`` so neighbouring classes stay
    distinguishable, and paints ``ignore_index`` mid-grey. Index it with a label map:
    ``palette[labels]`` gives an ``(H, W, 3)`` image.
    """
    plt = _plt()
    n = int(n_classes)
    size = max(n, (int(ignore_index) + 1) if ignore_index is not None else 0)
    cmap = plt.get_cmap("tab10" if n <= 10 else "tab20")
    pal = np.zeros((size, 3), dtype=float)
    for k in range(min(n, size)):
        pal[k] = cmap(k % cmap.N)[:3]
    if ignore_index is not None and 0 <= int(ignore_index) < size:
        pal[int(ignore_index)] = (0.5, 0.5, 0.5)
    return pal

plot_segmentation

plot_segmentation(image=None, target=None, pred=None, palette=None, class_names: Optional[Sequence[str]] = None, n_classes: Optional[int] = None, ignore_index: Optional[int] = None, titles: Optional[Sequence[str]] = None, figsize_per: float = 3.0)

Row of panels for one example: input image, ground truth, prediction and the error map.

Any of image / target / pred may be omitted. image may be a (3, H, W) or (H, W) array (floats in [0, 1] or ints in [0, 255]); target and pred are (H, W) label maps. When both are given a fourth panel marks the mispredicted pixels, which is usually where the interesting failures are.

Returns the figure.

Source code in src/torchtsetlin/viz.py

def plot_segmentation(
    image=None,
    target=None,
    pred=None,
    palette=None,
    class_names: Optional[Sequence[str]] = None,
    n_classes: Optional[int] = None,
    ignore_index: Optional[int] = None,
    titles: Optional[Sequence[str]] = None,
    figsize_per: float = 3.0,
):
    """Row of panels for one example: input image, ground truth, prediction and the error map.

    Any of ``image`` / ``target`` / ``pred`` may be omitted. ``image`` may be a ``(3, H, W)``
    or ``(H, W)`` array (floats in ``[0, 1]`` or ints in ``[0, 255]``); ``target`` and
    ``pred`` are ``(H, W)`` label maps. When both are given a fourth panel marks the
    mispredicted pixels, which is usually where the interesting failures are.

    Returns the figure.
    """
    plt = _plt()
    panels = []
    if image is not None:
        img = _np(image)
        if img.ndim == 3 and img.shape[0] in (1, 3):
            img = img.transpose(1, 2, 0)
        if img.ndim == 3 and img.shape[2] == 1:
            img = img[:, :, 0]
        if img.dtype.kind == "f" and img.max() > 1.5:
            img = img / 255.0
        panels.append(("image", img, None))
    maps = [m for m in (target, pred) if m is not None]
    if maps:
        K = int(n_classes) if n_classes is not None else int(max(_np(m).max() for m in maps)) + 1
        pal = np.asarray(palette) if palette is not None else segmentation_palette(K, ignore_index)
        if target is not None:
            panels.append(("ground truth", pal[np.clip(_np(target), 0, len(pal) - 1)], None))
        if pred is not None:
            panels.append(("prediction", pal[np.clip(_np(pred), 0, len(pal) - 1)], None))
    if target is not None and pred is not None:
        t, q = _np(target), _np(pred)
        err = (t != q)
        if ignore_index is not None:
            err = err & (t != int(ignore_index))
        panels.append((f"errors ({err.mean():.1%})", err.astype(float), "Reds"))

    names = list(titles) if titles is not None else [p[0] for p in panels]
    fig, axes = plt.subplots(1, len(panels), figsize=(figsize_per * len(panels), figsize_per * 1.1), squeeze=False)
    for ax, (_, data, cmap), name in zip(axes[0], panels, names):
        ax.imshow(data, cmap=cmap, interpolation="nearest", vmin=0 if cmap else None, vmax=1 if cmap else None)
        ax.set_title(name, fontsize=9)
        ax.set_xticks([])
        ax.set_yticks([])
    if class_names is not None and maps:
        import matplotlib.patches as mpatches

        pal = np.asarray(palette) if palette is not None else segmentation_palette(len(class_names), ignore_index)
        handles = [mpatches.Patch(color=pal[k], label=n) for k, n in enumerate(class_names)]
        fig.legend(handles=handles, loc="lower center", ncol=min(len(handles), 6), frameon=False, fontsize=8)
        fig.tight_layout(rect=(0, 0.08, 1, 1))
    else:
        fig.tight_layout()
    return fig

plot_vote_map

plot_vote_map(votes, class_index: int = 0, class_name: Optional[str] = None, ax=None, cmap='RdBu_r')

Heat map of one class's per-pixel vote sums (K, H, W) or (B, K, H, W).

Vote sums are the model's confidence surface: a class boundary sits where the winning margin collapses, so this shows where a label map is decided and where it is a coin flip.

Source code in src/torchtsetlin/viz.py

def plot_vote_map(votes, class_index: int = 0, class_name: Optional[str] = None, ax=None, cmap="RdBu_r"):
    """Heat map of one class's per-pixel vote sums ``(K, H, W)`` or ``(B, K, H, W)``.

    Vote sums are the model's *confidence surface*: a class boundary sits where the winning
    margin collapses, so this shows where a label map is decided and where it is a coin flip.
    """
    plt = _plt()
    v = _np(votes)
    if v.ndim == 4:
        v = v[0]
    m = v[int(class_index)]
    if ax is None:
        _, ax = plt.subplots(figsize=(4, 3.5))
    lim = float(np.abs(m).max()) or 1.0
    im = ax.imshow(m, cmap=cmap, vmin=-lim, vmax=lim, interpolation="nearest")
    ax.set_title(f"votes for {class_name or class_index}", fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.figure.colorbar(im, ax=ax, fraction=0.046)
    return ax