matplotlib: a plot you can read, saved to a file, from a script or a server

The two-object model

matplotlib has two ways of being used, and most confusion comes from mixing them. The quick way, plt.plot(x, y), draws on an implicit "current figure". The explicit way creates the objects and draws on them by name:

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(steps, loss)
ax.set_xlabel("training step")
ax.set_ylabel("loss")
ax.set_title("Loss per step")
fig.savefig("loss.png", dpi=150, bbox_inches="tight")

fig is the canvas; ax is a set of axes on it. Everything you draw goes through ax; everything about the file goes through fig. Learn this form and the implicit one will still make sense when you read it in other people's code, while the reverse is not true.

bbox_inches="tight" stops the labels being cut off at the edge, which is otherwise the first thing that goes wrong.

The plots you will actually make

A line over time or steps — a loss curve, cost per day:

ax.plot(steps, loss, label="train")
ax.plot(steps, val_loss, label="validation")
ax.legend()

A histogram — token counts per message, latencies:

ax.hist(tokens, bins=50)
ax.set_xlabel("tokens per message")

A scatter — two quantities per item, looking for a relationship:

ax.scatter(prompt_tokens, latency_s, s=8, alpha=0.4)

s is marker size and alpha is transparency; both matter at a few thousand points, where solid dots become a blob.

A bar chart — a number per category, from a groupby:

by_model = df.groupby("model")["cost_usd"].sum().sort_values()
ax.barh(by_model.index, by_model.values)

Horizontal bars fit long category names. Sort before plotting; an unsorted bar chart hides the ranking it exists to show.

pandas wraps all of these: df["tokens"].plot.hist(bins=50, ax=ax) and by_model.plot.barh(ax=ax) draw onto the axes you pass, so you can mix the convenience with the control.

Log scale

Loss falls from 4 to 0.4 to 0.04 over training. On a linear axis the last two-thirds of the run is a flat line at the bottom and you cannot see whether it is still improving. Latencies and token counts have a long tail: most messages are 50 tokens and a few are 5,000, and a linear histogram shows one bar.

ax.set_yscale("log")            # loss curve
ax.set_xscale("log")            # histogram of a long-tailed quantity
ax.hist(tokens, bins=np.logspace(1, 4, 40))   # log-spaced bins to match

Reach for a log axis whenever the data spans more than two orders of magnitude. It is the single change that most often turns an unreadable plot into a readable one.

Several plots in one figure

fig, axes = plt.subplots(1, 3, figsize=(15, 4))
axes[0].plot(steps, loss)
axes[1].hist(tokens, bins=40)
axes[2].scatter(x, y, s=6)
for ax in axes:
    ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig("report.png", dpi=150)

subplots(rows, cols) returns an array of axes. tight_layout spaces them so labels do not overlap. When comparing runs, sharey=True puts them on the same scale, which is what makes the comparison honest.

Saving versus showing

plt.show() opens a window, or renders inline in a notebook. On a server, in a cron job, in a test, there is no window, and matplotlib either errors about a missing display or hangs waiting for one that will never appear. A script that produces plots should save them:

import matplotlib
matplotlib.use("Agg")           # before importing pyplot: file output only, no display
import matplotlib.pyplot as plt

Agg is the backend that renders to pixels with no window; it is the right choice for anything not run by a person at a screen. Save PNG for reports and web pages, SVG (savefig("plot.svg")) for anything that will be scaled or edited in Inkscape, PDF for print.

plt.close(fig) after saving releases memory. A loop that makes a hundred figures without closing them holds a hundred figures and prints a warning about it.

What makes a plot checkable

Axis labels with units. A title that says what the data is and when it was taken. A legend when there is more than one series. Sensible limits: ax.set_ylim(0, None) for a quantity that cannot be negative, so the eye is not tricked by a zoomed baseline. The reader should be able to look at the figure with no surrounding text and know what they are looking at — including you, in a month.

Colour is not the way to carry meaning alone; around one in twelve men cannot separate red from green. Use different markers or line styles as well when the series must be distinguished, and the default palette, which was chosen with this in mind.

Free, and everywhere

matplotlib is free, runs on a phone under Pydroid or Termux, and is what pandas, seaborn and most scientific Python produce their figures with. seaborn adds statistical plots on top of it — a heatmap of a confusion matrix in one call — and plotly produces interactive charts for a browser; both are free, and both give you a figure you can still touch through matplotlib's axes.

Try this now

Plot the cost-per-model bar chart from your groupby, the token histogram from your messages on a log x-axis, and a loss curve from any training log you can find, as three panels in one figure saved to PNG with Agg. Then run the script with no display — DISPLAY= python plot.py on Linux — to confirm it works headless.