Heatmaps

imshow, colorbar, and a city × month rainfall grid.

imshow draws a 2-D array as a colored grid. A colorbar is the legend for those colors.

Goal

Show a rainfall matrix, label rows and columns, and add a colorbar.

A small grid

# rows: Nairobi, Mombasa, Kisumu  ·  cols: Jan–Apr
rain = np.array(
    [
        [50, 40, 80, 150],
        [20, 15, 30, 90],
        [70, 80, 120, 180],
    ],
    dtype=float,
)
fig, ax = plt.subplots()
image = ax.imshow(rain, cmap="Blues")
fig.colorbar(image, ax=ax, label="mm")
ax.set_xticks(range(4), ["Jan", "Feb", "Mar", "Apr"])
ax.set_yticks(range(3), ["Nairobi", "Mombasa", "Kisumu"])
ax.set_title("Rainfall")
plt.show()

Cell values on top

rain = np.array(
    [
        [50, 40, 80, 150],
        [20, 15, 30, 90],
        [70, 80, 120, 180],
    ],
    dtype=float,
)
fig, ax = plt.subplots()
image = ax.imshow(rain, cmap="Blues")
fig.colorbar(image, ax=ax, label="mm")
ax.set_xticks(range(4), ["Jan", "Feb", "Mar", "Apr"])
ax.set_yticks(range(3), ["Nairobi", "Mombasa", "Kisumu"])
for i in range(rain.shape[0]):
    for j in range(rain.shape[1]):
        ax.text(j, i, f"{rain[i, j]:.0f}", ha="center", va="center", color="0.15")
ax.set_title("Values in cells")
plt.show()

imshow uses row, column like a NumPy array: rain[0, 3] is Nairobi in April (top row, last column).

Pitfall

origin="lower" flips the y-axis to match a math graph. Leave the default (upper) for tables so row 0 stays at the top.