Skip to content

xwm.plots

Figures, GIFs and JSON/LaTeX tables, with each palette's categorical limit enforced in code rather than documented. See Figures and tables.

Figures, animations and result tables.

Two palettes, picked by what the chart does: blue-orange (Wong's colourblind-safe family) for curves -- losses, training histories, error by horizon -- and viridis for magnitude and for many-way categorical comparisons. xwm.plots.style records the measured separation of each and enforces its series limit rather than documenting it.

Matplotlib and Pillow are optional dependencies (pip install xwm[plots]); importing this module without them succeeds, and the error surfaces only when a helper that needs them is actually called.

Modules:

Name Description
curves

Training-curve and spectrum plots.
export

Writing figures and animated GIFs to disk.
style

Plot palettes, with each one's categorical limit enforced in code.
tables

Result tables and metrics, as JSON and LaTeX.
visualize

Visualising frames, masks and latents.

Functions:

Name Description
plot_bars

A bar chart for comparing a single metric across a few configurations.
plot_history

Plot metrics from xwm.training.Trainer.fit's history.
plot_horizon

Error-versus-horizon curves -- the compounding-error picture.
plot_spectra

Overlay several embedding spectra, one line per named model.
plot_spectrum

Plot the normalised singular-value spectrum of an embedding matrix.
frames_to_uint8

Normalise a frame sequence to (T, H, W, 3) uint8.
save_figure

Save a matplotlib figure, creating parent directories.
save_gif

Write an animated GIF.
tile_frames

Lay several frame sequences side by side into one sequence.
upscale

Nearest-neighbour integer upscale, so pixels stay crisp rather than blurred.
palette_colors

n hex colours from a named palette, in a fixed order.
plot_style

Apply the xwm style for one block, leaving global rcParams untouched.
rc_params

Matplotlib rcParams for the xwm look: recessive axes, thin marks.
series_style

Colour, marker and line style for series index of n_series.
use_viridis

Apply the xwm style globally. Call once at the top of a script.
viridis_colors

n hex colours sampled evenly from viridis, in a fixed order.
viridis_style

Deprecated alias for plot_style, kept for callers that used it.
escape_latex

Escape LaTeX special characters. Metric names like feature_std need it.
format_cell

Render one cell: real numbers via float_format, everything else via str.
jsonable

Convert a value to something json can encode, without rounding.
latex_table

A tabular (optionally wrapped in table) using booktabs rules.
markdown_table

A GitHub-flavoured Markdown table, column-aligned for readability.
save_json

Write payload as pretty-printed JSON, creating parent directories.
save_metrics

Write a flat mapping of scalar metrics to <stem>.json.
save_table

Write <stem>.json (full precision) and <stem>.tex (formatted).
table_to_dict

The JSON payload for a table: full-precision values keyed by column.
plot_frames

Draw a strip of (T, C, H, W) frames.
plot_latent_pca

Scatter embeddings on their first two principal components.
plot_mask

Overlay a token mask on an image, dimming the masked patches.
plot_rollout

Compare a real trajectory against an imagined one, frame by frame.

plot_bars 

plot_bars(labels: Sequence[str], values: Sequence[float], *, ax=None, title: str | None = None, ylabel: str | None = None, annotate: bool = True, horizontal: bool = True, palette: str = MAGNITUDE_PALETTE)

A bar chart for comparing a single metric across a few configurations.

Horizontal by default: configuration names are long, and rotated tick labels are harder to read than a horizontal bar's left-aligned label.

Source code in xwm/plots/curves.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def plot_bars(
    labels: Sequence[str],
    values: Sequence[float],
    *,
    ax=None,
    title: str | None = None,
    ylabel: str | None = None,
    annotate: bool = True,
    horizontal: bool = True,
    palette: str = MAGNITUDE_PALETTE,
):
    """A bar chart for comparing a single metric across a few configurations.

    Horizontal by default: configuration names are long, and rotated tick labels
    are harder to read than a horizontal bar's left-aligned label.
    """
    plt = pyplot()
    n = len(labels)
    if n != len(values):
        raise ValueError(f"{n} labels but {len(values)} values")
    with plot_style(min(n, 5), palette):
        ax = ax or plt.subplots(figsize=(6, 0.5 * n + 1.6))[1]
        from .style import palette_colors

        colors = palette_colors(min(n, 5), palette)
        colors = [colors[i % len(colors)] for i in range(n)]
        positions = range(n)
        if horizontal:
            ax.barh(list(positions), list(values), color=colors, height=0.62)
            ax.set_yticks(list(positions), labels)
            ax.invert_yaxis()
            ax.grid(axis="x", visible=True)
            ax.grid(axis="y", visible=False)
            if ylabel:
                ax.set_xlabel(ylabel)
            if annotate:
                span = max(values) - min(min(values), 0.0) or 1.0
                for pos, value in zip(positions, values, strict=True):
                    ax.annotate(
                        f"{value:.4g}",
                        (value, pos),
                        textcoords="offset points",
                        xytext=(5, 0),
                        va="center",
                        fontsize=8,
                    )
                ax.set_xlim(right=max(values) + 0.16 * span)
        else:
            ax.bar(list(positions), list(values), color=colors, width=0.62)
            ax.set_xticks(list(positions), labels, rotation=20, ha="right")
            if ylabel:
                ax.set_ylabel(ylabel)
        if title:
            ax.set_title(title)
    return ax

plot_history 

plot_history(history: Sequence[dict[str, float]], *, keys: Sequence[str] | None = None, ax=None, logy: bool = False, title: str | None = None, ylabel: str = 'loss', label_last: bool = True, palette: str = CURVE_PALETTE)

Plot metrics from xwm.training.Trainer.fit's history.

Parameters:

Name Type Description Default
history Sequence[dict[str, float]]

the list of metric dicts fit returns.

 required
keys Sequence[str] | None

which metrics to draw; defaults to every loss* key.

 None
label_last bool

annotate each series' final value directly on the plot, so the reader does not have to match a colour to a legend entry to learn the number that matters.

 True
Source code in xwm/plots/curves.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def plot_history(
    history: Sequence[dict[str, float]],
    *,
    keys: Sequence[str] | None = None,
    ax=None,
    logy: bool = False,
    title: str | None = None,
    ylabel: str = "loss",
    label_last: bool = True,
    palette: str = CURVE_PALETTE,
):
    """Plot metrics from :meth:`xwm.training.Trainer.fit`'s history.

    Args:
        history: the list of metric dicts ``fit`` returns.
        keys: which metrics to draw; defaults to every ``loss*`` key.
        label_last: annotate each series' final value directly on the plot, so
            the reader does not have to match a colour to a legend entry to
            learn the number that matters.
    """
    plt = pyplot()
    if not history:
        raise ValueError("history is empty")
    if keys is None:
        keys = [k for k in history[0] if k.startswith("loss")]
    if not keys:
        raise ValueError("no metrics to plot")
    steps = [row["step"] for row in history]

    with plot_style(len(keys), palette):
        ax = ax or plt.subplots(figsize=(6, 3.8))[1]
        # Markers on every point become noise on long runs; show ~12 of them.
        every = max(1, len(steps) // 12)
        for i, key in enumerate(keys):
            values = [row[key] for row in history]
            style = series_style(i, len(keys), palette)
            ax.plot(steps, values, label=key, markevery=every, **style)
            if label_last:
                ax.annotate(
                    f"{values[-1]:.4g}",
                    (steps[-1], values[-1]),
                    textcoords="offset points",
                    xytext=(6, 0),
                    fontsize=8,
                    color=style["color"],
                    va="center",
                )
        ax.set_xlabel("step")
        ax.set_ylabel(ylabel)
        if logy:
            ax.set_yscale("log")
        if title:
            ax.set_title(title)
        # A legend is always present for two or more series.
        if len(keys) > 1:
            ax.legend(loc="upper right")
        ax.margins(x=0.12)
    return ax

plot_horizon 

plot_horizon(horizons: Sequence[int], series: dict[str, Sequence[float]], *, ax=None, title: str = 'Rollout error by horizon', ylabel: str = 'latent L1 error', palette: str = CURVE_PALETTE)

Error-versus-horizon curves -- the compounding-error picture.

Source code in xwm/plots/curves.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
def plot_horizon(
    horizons: Sequence[int],
    series: dict[str, Sequence[float]],
    *,
    ax=None,
    title: str = "Rollout error by horizon",
    ylabel: str = "latent L1 error",
    palette: str = CURVE_PALETTE,
):
    """Error-versus-horizon curves -- the compounding-error picture."""
    plt = pyplot()
    names = list(series)
    with plot_style(len(names), palette):
        ax = ax or plt.subplots(figsize=(6, 3.8))[1]
        for i, name in enumerate(names):
            ax.plot(
                list(horizons),
                list(series[name]),
                label=name,
                **series_style(i, len(names), palette),
            )
        ax.set_xlabel("prediction horizon (steps)")
        ax.set_ylabel(ylabel)
        ax.set_title(title)
        ax.set_xticks(list(horizons))
        if len(names) > 1:
            ax.legend(loc="lower right")
    return ax

plot_spectra 

plot_spectra(named: dict[str, Array], *, ax=None, title: str = 'Embedding spectra', palette: str = MAGNITUDE_PALETTE)

Overlay several embedding spectra, one line per named model.

The single most legible collapse diagnostic: a collapsed encoder's spectrum plunges after a few directions while a healthy one decays gently.

Source code in xwm/plots/curves.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def plot_spectra(
    named: dict[str, Array],
    *,
    ax=None,
    title: str = "Embedding spectra",
    palette: str = MAGNITUDE_PALETTE,
):
    """Overlay several embedding spectra, one line per named model.

    The single most legible collapse diagnostic: a collapsed encoder's spectrum
    plunges after a few directions while a healthy one decays gently.
    """
    plt = pyplot()
    from ..metrics.representation import rankme, singular_values

    names = list(named)
    with plot_style(len(names), palette):
        ax = ax or plt.subplots(figsize=(6, 3.8))[1]
        for i, name in enumerate(names):
            s = singular_values(named[name])
            style = series_style(i, len(names), palette)
            style["marker"] = "None"
            ax.plot(
                jnp.arange(1, s.shape[0] + 1),
                s / s[0],
                label=f"{name} (RankMe {float(rankme(named[name])):.0f})",
                **style,
            )
        ax.set_yscale("log")
        ax.set_xlabel("singular value index")
        ax.set_ylabel("magnitude (normalised)")
        ax.set_title(title)
        if len(names) > 1:
            ax.legend(loc="lower left")
    return ax

plot_spectrum 

plot_spectrum(z: Array, *, ax=None, title: str | None = None, label: str | None = None, palette: str = CURVE_PALETTE)

Plot the normalised singular-value spectrum of an embedding matrix.

A healthy representation decays gently; a collapsing one falls off a cliff after a handful of directions.

Source code in xwm/plots/curves.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def plot_spectrum(
    z: Array,
    *,
    ax=None,
    title: str | None = None,
    label: str | None = None,
    palette: str = CURVE_PALETTE,
):
    """Plot the normalised singular-value spectrum of an embedding matrix.

    A healthy representation decays gently; a collapsing one falls off a cliff
    after a handful of directions.
    """
    plt = pyplot()
    from ..metrics.representation import rankme, singular_values

    s = singular_values(z)
    with plot_style(1, palette):
        ax = ax or plt.subplots(figsize=(5.5, 3.8))[1]
        style = series_style(0, 1, palette)
        style.pop("marker")
        ax.plot(
            jnp.arange(1, s.shape[0] + 1),
            s / s[0],
            label=label or f"RankMe {float(rankme(z)):.1f} / {z.shape[-1]}",
            **style,
        )
        ax.set_yscale("log")
        ax.set_xlabel("singular value index")
        ax.set_ylabel("magnitude (normalised)")
        ax.set_title(title or f"RankMe = {float(rankme(z)):.1f} / {z.shape[-1]}")
    return ax

frames_to_uint8 

frames_to_uint8(frames: Array) -> ndarray

Normalise a frame sequence to (T, H, W, 3) uint8.

Accepts channels-first (T, C, H, W) or channels-last (T, H, W, C), grayscale or RGB, float in [0, 1] or already uint8.

Source code in xwm/plots/export.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def frames_to_uint8(frames: Array) -> np.ndarray:
    """Normalise a frame sequence to ``(T, H, W, 3)`` ``uint8``.

    Accepts channels-first ``(T, C, H, W)`` or channels-last ``(T, H, W, C)``,
    grayscale or RGB, float in ``[0, 1]`` or already ``uint8``.
    """
    arr = np.asarray(frames)
    if arr.ndim == 3:  # (T, H, W) grayscale
        arr = arr[:, None]
    if arr.ndim != 4:
        raise ValueError(f"expected a 3-D or 4-D frame sequence, got shape {arr.shape}")

    channel_counts = (1, 3, 4)
    first, last = arr.shape[1], arr.shape[-1]
    first_ok, last_ok = first in channel_counts, last in channel_counts
    if not (first_ok or last_ok):
        raise ValueError(
            f"cannot find a channel axis in shape {arr.shape}; expected 1, 3 or 4 "
            "channels at axis 1 (C, H, W) or axis 3 (H, W, C)"
        )
    # When both axes could be channels (e.g. a 3-pixel-wide image), prefer the
    # smaller one; xwm produces channels-first, so that is the right tiebreak.
    if first_ok and (not last_ok or first <= last):
        arr = arr.transpose(0, 2, 3, 1)

    if arr.dtype != np.uint8:
        arr = (np.clip(arr, 0.0, 1.0) * 255.0).round().astype(np.uint8)
    if arr.shape[-1] == 1:
        arr = np.repeat(arr, 3, axis=-1)
    return arr[..., :3]

save_figure 

save_figure(fig, path: str | Path, *, dpi: int = 150, close: bool = True) -> Path

Save a matplotlib figure, creating parent directories.

Source code in xwm/plots/export.py 
184
185
186
187
188
189
190
191
192
193
def save_figure(fig, path: str | Path, *, dpi: int = 150, close: bool = True) -> Path:
    """Save a matplotlib figure, creating parent directories."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    if close:
        import matplotlib.pyplot as plt

        plt.close(fig)
    return path

save_gif 

save_gif(path: str | Path, frames: Array | Sequence[Array], *, fps: int = 8, scale: int = 1, labels: Sequence[str] | None = None, loop: int = 0, colors: int = 256, dither: bool = False) -> Path

Write an animated GIF.

Parameters:

Name Type Description Default
frames Array | Sequence[Array]

one (T, C, H, W) sequence, or several to tile side by side.

 required
fps int

playback rate.

 8
scale int

integer upscale factor. Prefer rendering at the size you want -- nearest-neighbour upscaling turns every source pixel into a block and cannot add detail. Useful only for genuinely tiny sources such as the 32x32 sprite world.

 1
labels Sequence[str] | None

per-panel captions, drawn above each tile.

 None
loop int

0 loops forever.

 0
colors int

palette size, at most 256 (a GIF limit).

 256
dither bool

diffuse quantisation error. Off by default: on smooth renders it reads as grain, and with a 256-colour palette there is little error left to diffuse.

 False

A single palette is computed across all frames. Quantising each frame independently -- what Pillow does by default -- gives every frame its own palette, so colours shift frame to frame and the animation shimmers even when the underlying pixels barely change.

Source code in xwm/plots/export.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def save_gif(
    path: str | Path,
    frames: Array | Sequence[Array],
    *,
    fps: int = 8,
    scale: int = 1,
    labels: Sequence[str] | None = None,
    loop: int = 0,
    colors: int = 256,
    dither: bool = False,
) -> Path:
    """Write an animated GIF.

    Args:
        frames: one ``(T, C, H, W)`` sequence, or several to tile side by side.
        fps: playback rate.
        scale: integer upscale factor. Prefer rendering at the size you want --
            nearest-neighbour upscaling turns every source pixel into a block and
            cannot add detail. Useful only for genuinely tiny sources such as the
            32x32 sprite world.
        labels: per-panel captions, drawn above each tile.
        loop: ``0`` loops forever.
        colors: palette size, at most 256 (a GIF limit).
        dither: diffuse quantisation error. Off by default: on smooth renders it
            reads as grain, and with a 256-colour palette there is little error
            left to diffuse.

    A **single palette is computed across all frames**. Quantising each frame
    independently -- what Pillow does by default -- gives every frame its own
    palette, so colours shift frame to frame and the animation shimmers even
    when the underlying pixels barely change.
    """
    Image, _, _ = _pillow()
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    if isinstance(frames, (list, tuple)):
        stack = tile_frames(frames, labels=labels, scale=scale)
    else:
        stack = upscale(frames_to_uint8(frames), scale)
        if labels is not None:
            stack = tile_frames([frames], labels=labels, scale=scale)

    images = [Image.fromarray(f) for f in stack]
    quantized = _quantize_shared(images, colors=colors, dither=dither)
    quantized[0].save(
        path,
        save_all=True,
        append_images=quantized[1:],
        duration=max(int(round(1000 / fps)), 20),
        loop=loop,
        disposal=2,
        optimize=True,
    )
    return path

tile_frames 

tile_frames(sequences: Sequence[Array], *, labels: Sequence[str] | None = None, scale: int = 4, pad: int = 2, label_height: int = 14) -> ndarray

Lay several frame sequences side by side into one sequence.

Sequences of differing length are truncated to the shortest, so a comparison strip never silently pairs frame 5 of one rollout with frame 9 of another.

Source code in xwm/plots/export.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def tile_frames(
    sequences: Sequence[Array],
    *,
    labels: Sequence[str] | None = None,
    scale: int = 4,
    pad: int = 2,
    label_height: int = 14,
) -> np.ndarray:
    """Lay several frame sequences side by side into one sequence.

    Sequences of differing length are truncated to the shortest, so a comparison
    strip never silently pairs frame 5 of one rollout with frame 9 of another.
    """
    Image, ImageDraw, ImageFont = _pillow()
    panels = [upscale(frames_to_uint8(s), scale) for s in sequences]
    n = min(p.shape[0] for p in panels)
    height = max(p.shape[1] for p in panels)
    panels = [p[:n] for p in panels]

    header = label_height if labels is not None else 0
    widths = [p.shape[2] for p in panels]
    total_width = sum(widths) + pad * (len(panels) - 1)
    out = np.full((n, height + header, total_width, 3), 255, dtype=np.uint8)

    for t in range(n):
        x = 0
        for panel, width in zip(panels, widths, strict=True):
            out[t, header : header + panel.shape[1], x : x + width] = panel[t]
            x += width + pad

    if labels is not None:
        if len(labels) != len(panels):
            raise ValueError(f"{len(labels)} labels for {len(panels)} sequences")
        font = ImageFont.load_default()
        for t in range(n):
            img = Image.fromarray(out[t])
            draw = ImageDraw.Draw(img)
            x = 0
            for label, width in zip(labels, widths, strict=True):
                draw.text((x + 3, 2), str(label), fill=(20, 20, 20), font=font)
                x += width + pad
            out[t] = np.asarray(img)
    return out

upscale 

upscale(frames: ndarray, factor: int) -> ndarray

Nearest-neighbour integer upscale, so pixels stay crisp rather than blurred.

Source code in xwm/plots/export.py 
61
62
63
64
65
def upscale(frames: np.ndarray, factor: int) -> np.ndarray:
    """Nearest-neighbour integer upscale, so pixels stay crisp rather than blurred."""
    if factor <= 1:
        return frames
    return np.repeat(np.repeat(frames, factor, axis=1), factor, axis=2)

palette_colors 

palette_colors(n: int, palette: str = MAGNITUDE_PALETTE) -> list[str]

n hex colours from a named palette, in a fixed order.

Raises:

Type Description
ValueError

if n exceeds what the palette can separate legibly. Generating an extra hue past that point puts an adjacent pair below the legibility floor; split the chart into small multiples instead.
Source code in xwm/plots/style.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def palette_colors(n: int, palette: str = MAGNITUDE_PALETTE) -> list[str]:
    """``n`` hex colours from a named palette, in a fixed order.

    Raises:
        ValueError: if ``n`` exceeds what the palette can separate legibly.
            Generating an extra hue past that point puts an adjacent pair below
            the legibility floor; split the chart into small multiples instead.
    """
    if palette not in PALETTES:
        raise ValueError(f"unknown palette {palette!r}; choose from {sorted(PALETTES)}")
    limit = PALETTES[palette]
    if n < 1:
        raise ValueError(f"need at least one colour, got {n}")
    if n > limit:
        raise ValueError(
            f"{n} categorical series exceeds the {limit} that the {palette!r} palette "
            "can separate legibly (adjacent pairs fall below the normal-vision floor). "
            "Use small multiples, or group the smallest series into 'other'."
        )
    if palette == "blue-orange":
        return list(BLUE_ORANGE[:n])
    return viridis_colors(n)

plot_style 

plot_style(n_series: int = MAX_CATEGORICAL, palette: str = MAGNITUDE_PALETTE)

Apply the xwm style for one block, leaving global rcParams untouched.

Source code in xwm/plots/style.py 
192
193
194
195
196
197
198
@contextmanager
def plot_style(n_series: int = MAX_CATEGORICAL, palette: str = MAGNITUDE_PALETTE):
    """Apply the xwm style for one block, leaving global rcParams untouched."""
    import matplotlib.pyplot as plt

    with plt.rc_context(rc_params(n_series, palette)):
        yield

rc_params 

rc_params(n_series: int = MAX_CATEGORICAL, palette: str = MAGNITUDE_PALETTE) -> dict

Matplotlib rcParams for the xwm look: recessive axes, thin marks.

Source code in xwm/plots/style.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def rc_params(n_series: int = MAX_CATEGORICAL, palette: str = MAGNITUDE_PALETTE) -> dict:
    """Matplotlib rcParams for the xwm look: recessive axes, thin marks."""
    from cycler import cycler

    return {
        "axes.prop_cycle": cycler(color=palette_colors(n_series, palette)),
        "image.cmap": SEQUENTIAL_CMAP,
        "axes.grid": True,
        "axes.grid.axis": "y",
        "grid.color": _GRID,
        "grid.linewidth": 0.6,
        "grid.alpha": 0.8,
        "axes.edgecolor": _GRID,
        "axes.linewidth": 0.8,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.labelcolor": _INK,
        "axes.titlesize": 11,
        "text.color": _INK,
        "xtick.color": _INK,
        "ytick.color": _INK,
        "xtick.labelsize": 9,
        "ytick.labelsize": 9,
        "legend.frameon": False,
        "legend.fontsize": 9,
        "lines.linewidth": 2.0,
        "lines.markersize": 5.0,
        "figure.facecolor": "white",
        "savefig.facecolor": "white",
        "font.size": 10,
    }

series_style 

series_style(index: int, n_series: int, palette: str = MAGNITUDE_PALETTE) -> dict

Colour, marker and line style for series index of n_series.

Colour follows the series' identity (its index), never its rank, so filtering the chart never repaints the survivors.

Source code in xwm/plots/style.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def series_style(index: int, n_series: int, palette: str = MAGNITUDE_PALETTE) -> dict:
    """Colour, marker and line style for series ``index`` of ``n_series``.

    Colour follows the series' identity (its index), never its rank, so
    filtering the chart never repaints the survivors.
    """
    colors = palette_colors(n_series, palette)
    return {
        "color": colors[index],
        "marker": MARKERS[index % len(MARKERS)],
        "linestyle": LINESTYLES[index % len(LINESTYLES)],
        "linewidth": 2.0,
        "markersize": 5.0,
        "markeredgecolor": "white",
        "markeredgewidth": 0.6,
    }

use_viridis 

use_viridis(n_series: int = MAX_CATEGORICAL, palette: str = MAGNITUDE_PALETTE) -> None

Apply the xwm style globally. Call once at the top of a script.

Source code in xwm/plots/style.py 
185
186
187
188
189
def use_viridis(n_series: int = MAX_CATEGORICAL, palette: str = MAGNITUDE_PALETTE) -> None:
    """Apply the xwm style globally. Call once at the top of a script."""
    import matplotlib.pyplot as plt

    plt.rcParams.update(rc_params(n_series, palette))

viridis_colors 

viridis_colors(n: int, span: tuple[float, float] = VIRIDIS_SPAN) -> list[str]

n hex colours sampled evenly from viridis, in a fixed order.

Raises:

Type Description
ValueError

if n exceeds MAX_CATEGORICAL. Adding a sixth hue would put an adjacent pair below the legibility floor; split the chart into small multiples or group the tail into "other".
Source code in xwm/plots/style.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def viridis_colors(n: int, span: tuple[float, float] = VIRIDIS_SPAN) -> list[str]:
    """``n`` hex colours sampled evenly from viridis, in a fixed order.

    Raises:
        ValueError: if ``n`` exceeds :data:`MAX_CATEGORICAL`. Adding a sixth
            hue would put an adjacent pair below the legibility floor; split the
            chart into small multiples or group the tail into "other".
    """
    import matplotlib
    from matplotlib.colors import to_hex

    if n < 1:
        raise ValueError(f"need at least one colour, got {n}")
    if n > MAX_CATEGORICAL:
        raise ValueError(
            f"{n} categorical series exceeds the {MAX_CATEGORICAL} that viridis can "
            "separate legibly (adjacent pairs fall below the normal-vision floor). "
            "Use small multiples, or group the smallest series into 'other'."
        )
    cmap = matplotlib.colormaps[SEQUENTIAL_CMAP]
    lo, hi = span
    if n == 1:
        return [to_hex(cmap(0.5 * (lo + hi)))]
    step = (hi - lo) / (n - 1)
    return [to_hex(cmap(lo + i * step)) for i in range(n)]

viridis_style 

viridis_style(n_series: int = MAX_CATEGORICAL)

Deprecated alias for plot_style, kept for callers that used it.

Source code in xwm/plots/style.py 
201
202
203
204
205
@contextmanager
def viridis_style(n_series: int = MAX_CATEGORICAL):
    """Deprecated alias for :func:`plot_style`, kept for callers that used it."""
    with plot_style(n_series, MAGNITUDE_PALETTE):
        yield

escape_latex 

escape_latex(text: str) -> str

Escape LaTeX special characters. Metric names like feature_std need it.

Source code in xwm/plots/tables.py 
41
42
43
def escape_latex(text: str) -> str:
    """Escape LaTeX special characters. Metric names like ``feature_std`` need it."""
    return "".join(_LATEX_ESCAPES.get(ch, ch) for ch in text)

format_cell 

format_cell(value: Any, float_format: str = '{:.4f}') -> str

Render one cell: real numbers via float_format, everything else via str.

Unwraps 0-d NumPy and JAX scalars first. They are not Python float, so a bare isinstance check silently lets them through unformatted -- and they are exactly what a training loop hands to a results table.

Source code in xwm/plots/tables.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def format_cell(value: Any, float_format: str = "{:.4f}") -> str:
    """Render one cell: real numbers via ``float_format``, everything else via ``str``.

    Unwraps 0-d NumPy and JAX scalars first. They are not Python ``float``, so a
    bare ``isinstance`` check silently lets them through unformatted -- and they
    are exactly what a training loop hands to a results table.
    """
    if isinstance(value, bool):
        return "yes" if value else "no"
    if isinstance(value, int):
        return str(value)
    item = getattr(value, "item", None)
    if item is not None and getattr(value, "ndim", 0) == 0:
        value = item()
    if isinstance(value, bool):  # numpy bool unwraps to a Python bool
        return "yes" if value else "no"
    if isinstance(value, float):
        return float_format.format(value)
    if isinstance(value, int):
        return str(value)
    return str(value)

jsonable 

jsonable(value: Any) -> Any

Convert a value to something json can encode, without rounding.

NumPy and JAX scalars become Python numbers; non-finite floats become None, since JSON has no NaN or Infinity and emitting bare NaN produces a file that strict parsers reject.

Source code in xwm/plots/tables.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def jsonable(value: Any) -> Any:
    """Convert a value to something :mod:`json` can encode, without rounding.

    NumPy and JAX scalars become Python numbers; non-finite floats become
    ``None``, since JSON has no NaN or Infinity and emitting bare ``NaN``
    produces a file that strict parsers reject.
    """
    if isinstance(value, (str, bool)) or value is None:
        return value
    if isinstance(value, (int,)):
        return value
    item = getattr(value, "item", None)
    if item is not None and getattr(value, "ndim", 0) == 0:
        value = item()
    if isinstance(value, float):
        return value if math.isfinite(value) else None
    if isinstance(value, (list, tuple)):
        return [jsonable(v) for v in value]
    if isinstance(value, Mapping):
        return {str(k): jsonable(v) for k, v in value.items()}
    if hasattr(value, "tolist"):
        return jsonable(value.tolist())
    if isinstance(value, (int, float)):
        return value
    return str(value)

latex_table 

latex_table(headers: Sequence[str], rows: Sequence[Row], *, caption: str | None = None, label: str | None = None, align: str | None = None, float_format: str = '{:.4f}', escape: bool = True, booktabs: bool = True) -> str

A tabular (optionally wrapped in table) using booktabs rules.

Parameters:

Name Type Description Default
align str | None

column spec such as "lrrr"; defaults to left for the first column and right for the rest, which is what numeric tables want.

 None
escape bool

escape LaTeX specials in cells. Turn it off to pass math through.

 True
booktabs bool

use \toprule/\midrule/\bottomrule (needs the booktabs package); otherwise plain \hline.

 True
Source code in xwm/plots/tables.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def latex_table(
    headers: Sequence[str],
    rows: Sequence[Row],
    *,
    caption: str | None = None,
    label: str | None = None,
    align: str | None = None,
    float_format: str = "{:.4f}",
    escape: bool = True,
    booktabs: bool = True,
) -> str:
    """A ``tabular`` (optionally wrapped in ``table``) using booktabs rules.

    Args:
        align: column spec such as ``"lrrr"``; defaults to left for the first
            column and right for the rest, which is what numeric tables want.
        escape: escape LaTeX specials in cells. Turn it off to pass math through.
        booktabs: use ``\\toprule``/``\\midrule``/``\\bottomrule`` (needs the
            ``booktabs`` package); otherwise plain ``\\hline``.
    """
    body = _cells(rows, float_format)
    head = [str(h) for h in headers]
    if escape:
        head = [escape_latex(h) for h in head]
        body = [[escape_latex(c) for c in row] for row in body]
    if align is None:
        align = "l" + "r" * (len(head) - 1)
    if len(align) != len(head):
        raise ValueError(f"align {align!r} has {len(align)} columns, headers have {len(head)}")

    top, mid, bottom = (
        (r"\toprule", r"\midrule", r"\bottomrule") if booktabs else (r"\hline",) * 3
    )
    lines = [f"\\begin{{tabular}}{{{align}}}", top, " & ".join(head) + r" \\", mid]
    lines += [" & ".join(row) + r" \\" for row in body]
    lines += [bottom, r"\end{tabular}"]
    tabular = "\n".join(lines)

    if caption is None and label is None:
        return tabular
    wrapped = [r"\begin{table}[t]", r"\centering", tabular]
    if caption is not None:
        wrapped.append(f"\\caption{{{escape_latex(caption) if escape else caption}}}")
    if label is not None:
        wrapped.append(f"\\label{{{label}}}")
    wrapped.append(r"\end{table}")
    return "\n".join(wrapped)

markdown_table 

markdown_table(headers: Sequence[str], rows: Sequence[Row], *, float_format: str = '{:.4f}') -> str

A GitHub-flavoured Markdown table, column-aligned for readability.

For printing to a terminal or pasting into a README; save_table writes JSON and LaTeX to disk, not this.

Source code in xwm/plots/tables.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def markdown_table(
    headers: Sequence[str],
    rows: Sequence[Row],
    *,
    float_format: str = "{:.4f}",
) -> str:
    """A GitHub-flavoured Markdown table, column-aligned for readability.

    For printing to a terminal or pasting into a README; :func:`save_table`
    writes JSON and LaTeX to disk, not this.
    """
    body = _cells(rows, float_format)
    head = [str(h) for h in headers]
    widths = [
        max(len(head[i]), *(len(r[i]) for r in body)) if body else len(head[i])
        for i in range(len(head))
    ]
    line = "| " + " | ".join(h.ljust(w) for h, w in zip(head, widths, strict=True)) + " |"
    rule = "| " + " | ".join("-" * w for w in widths) + " |"
    out = [line, rule]
    out += [
        "| " + " | ".join(c.ljust(w) for c, w in zip(row, widths, strict=True)) + " |"
        for row in body
    ]
    return "\n".join(out)

save_json 

save_json(path: str | Path, payload: Any) -> Path

Write payload as pretty-printed JSON, creating parent directories.

Source code in xwm/plots/tables.py 
199
200
201
202
203
204
def save_json(path: str | Path, payload: Any) -> Path:
    """Write ``payload`` as pretty-printed JSON, creating parent directories."""
    path = Path(path).with_suffix(".json")
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(jsonable(payload), indent=2, sort_keys=False) + "\n")
    return path

save_metrics 

save_metrics(path_stem: str | Path, metrics: Mapping[str, Any]) -> Path

Write a flat mapping of scalar metrics to <stem>.json.

Source code in xwm/plots/tables.py 
207
208
209
def save_metrics(path_stem: str | Path, metrics: Mapping[str, Any]) -> Path:
    """Write a flat mapping of scalar metrics to ``<stem>.json``."""
    return save_json(Path(path_stem), dict(metrics))

save_table 

save_table(path_stem: str | Path, headers: Sequence[str], rows: Sequence[Row], *, caption: str | None = None, label: str | None = None, align: str | None = None, float_format: str = '{:.4f}') -> dict[str, Path]

Write <stem>.json (full precision) and <stem>.tex (formatted).

Returns a mapping from extension to the path written.

Source code in xwm/plots/tables.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def save_table(
    path_stem: str | Path,
    headers: Sequence[str],
    rows: Sequence[Row],
    *,
    caption: str | None = None,
    label: str | None = None,
    align: str | None = None,
    float_format: str = "{:.4f}",
) -> dict[str, Path]:
    """Write ``<stem>.json`` (full precision) and ``<stem>.tex`` (formatted).

    Returns a mapping from extension to the path written.
    """
    stem = Path(path_stem)
    stem.parent.mkdir(parents=True, exist_ok=True)

    json_path = save_json(stem, table_to_dict(headers, rows, caption=caption, label=label))
    tex = stem.with_suffix(".tex")
    tex.write_text(
        latex_table(
            headers, rows, caption=caption, label=label, align=align, float_format=float_format
        )
        + "\n"
    )
    return {"json": json_path, "tex": tex}

table_to_dict 

table_to_dict(headers: Sequence[str], rows: Sequence[Row], *, caption: str | None = None, label: str | None = None) -> dict[str, Any]

The JSON payload for a table: full-precision values keyed by column.

Source code in xwm/plots/tables.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def table_to_dict(
    headers: Sequence[str],
    rows: Sequence[Row],
    *,
    caption: str | None = None,
    label: str | None = None,
) -> dict[str, Any]:
    """The JSON payload for a table: full-precision values keyed by column."""
    columns = [str(h) for h in headers]
    payload: dict[str, Any] = {
        "columns": columns,
        "rows": [
            {column: jsonable(value) for column, value in zip(columns, row, strict=True)}
            for row in rows
        ],
    }
    if caption is not None:
        payload["caption"] = caption
    if label is not None:
        payload["label"] = label
    return payload

plot_frames 

plot_frames(frames: Array, *, titles=None, max_frames: int = 12, scale: float = 1.3, suptitle: str | None = None)

Draw a strip of (T, C, H, W) frames.

scale is inches per panel. Text scales with it: a strip sized to show a 768 px render natively is ~30 inches wide, where 8 pt titles are unreadable and a fixed-height suptitle lands on top of them.

Source code in xwm/plots/visualize.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def plot_frames(
    frames: Array,
    *,
    titles=None,
    max_frames: int = 12,
    scale: float = 1.3,
    suptitle: str | None = None,
):
    """Draw a strip of ``(T, C, H, W)`` frames.

    ``scale`` is inches per panel. Text scales with it: a strip sized to show a
    768 px render natively is ~30 inches wide, where 8 pt titles are unreadable
    and a fixed-height suptitle lands on top of them.
    """
    plt = pyplot()
    n = min(frames.shape[0], max_frames)
    title_size = max(8.0, 6.0 * scale)
    suptitle_size = max(10.0, 7.0 * scale)
    height = scale + 0.4
    with plot_style(1):
        fig, axes = plt.subplots(1, n, figsize=(scale * n, height))
        axes = np.atleast_1d(axes)
        for i in range(n):
            axes[i].imshow(_to_hwc(frames[i]), cmap=SEQUENTIAL_CMAP)
            axes[i].axis("off")
            axes[i].set_title(
                str(titles[i]) if titles is not None else f"t={i}", fontsize=title_size
            )
        if suptitle:
            fig.suptitle(suptitle, fontsize=suptitle_size)
            # Reserve the suptitle's own height, in figure fractions, so it
            # cannot overlap the panel titles at any aspect ratio.
            reserved = min(0.4, (suptitle_size * 2.0 / 72.0) / height)
            fig.tight_layout(rect=(0.0, 0.0, 1.0, 1.0 - reserved))
        else:
            fig.tight_layout()
    return fig

plot_latent_pca 

plot_latent_pca(z: Array, *, labels: Array | None = None, ax=None, label_name: str = 'value', title: str | None = None)

Scatter embeddings on their first two principal components.

labels is a continuous quantity (an agent coordinate, a joint angle), so it is encoded with the viridis ramp and a colourbar -- magnitude, not identity. If the structure the encoder learned corresponds to that quantity, the scatter shows a smooth gradient rather than a blob.

Source code in xwm/plots/visualize.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def plot_latent_pca(
    z: Array,
    *,
    labels: Array | None = None,
    ax=None,
    label_name: str = "value",
    title: str | None = None,
):
    """Scatter embeddings on their first two principal components.

    ``labels`` is a *continuous* quantity (an agent coordinate, a joint angle),
    so it is encoded with the viridis ramp and a colourbar -- magnitude, not
    identity. If the structure the encoder learned corresponds to that
    quantity, the scatter shows a smooth gradient rather than a blob.
    """
    plt = pyplot()
    flat = z.reshape(-1, z.shape[-1])
    if labels is not None:
        n_labels = np.asarray(labels).reshape(-1).shape[0]
        if n_labels != flat.shape[0]:
            raise ValueError(
                f"got {n_labels} labels for {flat.shape[0]} points (input shape "
                f"{tuple(z.shape)} flattens to {flat.shape[0]} rows). Pool token "
                "sequences to one vector per sample first, e.g. z.mean(axis=-2)."
            )
    centred = flat - jnp.mean(flat, axis=0, keepdims=True)
    _, _, vt = jnp.linalg.svd(centred, full_matrices=False)
    proj = np.asarray(centred @ vt[:2].T)
    with plot_style(1):
        ax = ax or plt.subplots(figsize=(5, 4.2))[1]
        if labels is None:
            from .style import viridis_colors

            ax.scatter(proj[:, 0], proj[:, 1], s=14, color=viridis_colors(1)[0], alpha=0.75)
        else:
            scatter = ax.scatter(
                proj[:, 0],
                proj[:, 1],
                c=np.asarray(labels).reshape(-1),
                s=14,
                cmap=SEQUENTIAL_CMAP,
                alpha=0.9,
            )
            ax.figure.colorbar(scatter, ax=ax, label=label_name)
        ax.set_xlabel("PC 1")
        ax.set_ylabel("PC 2")
        ax.grid(axis="both", visible=True)
        if title:
            ax.set_title(title)
    return ax

plot_mask 

plot_mask(frame: Array, mask_idx: Array, grid: tuple[int, int], *, ax=None)

Overlay a token mask on an image, dimming the masked patches.

Parameters:

Name Type Description Default
frame Array

(C, H, W) image.

 required
mask_idx Array

flat token indices considered visible.

 required
grid tuple[int, int]

(gh, gw) patch grid the indices refer to.

 required
Source code in xwm/plots/visualize.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def plot_mask(frame: Array, mask_idx: Array, grid: tuple[int, int], *, ax=None):
    """Overlay a token mask on an image, dimming the masked patches.

    Args:
        frame: ``(C, H, W)`` image.
        mask_idx: flat token indices considered *visible*.
        grid: ``(gh, gw)`` patch grid the indices refer to.
    """
    plt = pyplot()
    gh, gw = grid
    visible = np.zeros(gh * gw, dtype=bool)
    visible[np.asarray(mask_idx)] = True
    img = _to_hwc(frame)
    h, w = img.shape[:2]
    alpha = np.kron(visible.reshape(gh, gw), np.ones((h // gh, w // gw)))
    shaded = img * (0.25 + 0.75 * alpha[..., None] if img.ndim == 3 else 0.25 + 0.75 * alpha)
    with plot_style(1):
        ax = ax or plt.subplots(figsize=(3, 3))[1]
        ax.imshow(np.clip(shaded, 0, 1), cmap=SEQUENTIAL_CMAP)
        ax.set_title(f"{int(visible.sum())}/{gh * gw} tokens visible", fontsize=9)
        ax.axis("off")
    return ax

plot_rollout 

plot_rollout(true_frames: Array, imagined: Array | None = None, *, max_frames: int = 10)

Compare a real trajectory against an imagined one, frame by frame.

imagined is optional because a latent world model has nothing to render; pass decoded frames only if you have a decoder. Otherwise use this for the ground-truth strip and report latent distances numerically.

Source code in xwm/plots/visualize.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def plot_rollout(true_frames: Array, imagined: Array | None = None, *, max_frames: int = 10):
    """Compare a real trajectory against an imagined one, frame by frame.

    ``imagined`` is optional because a latent world model has nothing to render;
    pass decoded frames only if you have a decoder. Otherwise use this for the
    ground-truth strip and report latent distances numerically.
    """
    plt = pyplot()
    rows = 1 if imagined is None else 2
    n = min(true_frames.shape[0], max_frames)
    with plot_style(1):
        fig, axes = plt.subplots(rows, n, figsize=(1.3 * n, 1.5 * rows), squeeze=False)
        for i in range(n):
            axes[0][i].imshow(_to_hwc(true_frames[i]), cmap=SEQUENTIAL_CMAP)
            axes[0][i].set_title(f"t={i}", fontsize=8)
            if imagined is not None:
                axes[1][i].imshow(_to_hwc(imagined[i]), cmap=SEQUENTIAL_CMAP)
        for row, name in zip(axes, ("observed", "imagined"), strict=False):
            for j, cell in enumerate(row):
                cell.set_xticks([])
                cell.set_yticks([])
                for spine in cell.spines.values():
                    spine.set_visible(False)
                if j == 0:
                    cell.set_ylabel(name, fontsize=9)
        fig.tight_layout()
    return fig