Images, heatmaps and colormaps

Correlation matrices, images and gridded data are all 2D arrays, and a well-chosen colormap is the difference between insight and illusion.

ax.imshow(A) draws a 2D array as a grid of colored pixels. Three behaviors to know:

Row 0 is at the top by default (image convention). Pass origin="lower" for matrices that should read like a graph.
The aspect ratio is fixed so pixels are square; aspect="auto" lets the image fill the Axes.
3D arrays of shape (rows, cols, 3) or (..., 4) are treated as RGB(A) images, not mapped through a colormap.

ax.pcolormesh(X, Y, C) draws a grid of quadrilaterals with arbitrary, possibly uneven coordinates, which suits scientific grids better than pixel positions.

Values are mapped to colors by a norm and a colormap. vmin/vmax fix the range so several panels share a scale. Choose the colormap by data type:

sequential ("viridis", "magma", "cividis"): ordered values from low to high;
diverging ("RdBu_r", "coolwarm"): values around a meaningful center like 0; pair with vmin=-m, vmax=m or colors.CenteredNorm();
qualitative ("tab10"): categories, not magnitudes.

Avoid "jet" and rainbow maps: their uneven brightness creates fake boundaries and they fail in grayscale.

Always add fig.colorbar(im, ax=ax); without it the colors have no meaning. For annotated heatmaps (numbers in each cell), loop over the cells with ax.text, or let seaborn's heatmap do it.

Example

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

rng = np.random.default_rng(0)
data = rng.normal(size=(200, 4))
data[:, 1] += data[:, 0]                  # make columns 0 and 1 correlated
corr = np.corrcoef(data, rowvar=False)    # 4 x 4 matrix

fig, ax = plt.subplots()
im = ax.imshow(corr, cmap="RdBu_r", vmin=-1, vmax=1)
fig.colorbar(im, ax=ax, label="correlation")
ax.set_xticks(range(4), labels=list("ABCD"))
ax.set_yticks(range(4), labels=list("ABCD"))
for i in range(4):
    for j in range(4):
        ax.text(j, i, f"{corr[i, j]:.2f}", ha="center", va="center", fontsize=8)

print(corr.shape, round(corr[0, 1], 2))  # (4, 4) 0.72
print(im.get_clim())                     # (-1.0, 1.0)
fig.savefig("corr.png")

Remember

imshow puts row 0 at the top; use origin='lower' for graph-like matrices
Sequential maps for magnitudes, diverging maps centred on a meaningful value
Fix vmin/vmax to compare panels, and always add a colorbar

Watch out for

Using a rainbow colormap such as jet, which invents visual boundaries.

Check yourself

You plot a correlation matrix (values from -1 to 1). Which setup is most appropriate?

Official docs: matplotlib.org/stable/users/explain/colors/colormaps.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks