What you'll learn
- Plot a heatmap and contours of 2D data
- Build 2D coordinate grids with meshgrid and keep rows along y
- Find peaks and areas above a threshold in a 2D array
- Choose sequential or diverging colormaps for the data
Build 2D grids with meshgrid, find peaks and thresholds in 2D arrays, and show them as heatmaps and contour plots.
Why this matters
A lot of scientific data is two-dimensional: a temperature map across a heated plate, a stress field in a beam from finite-element software, a detector image from a microscope or X-ray camera, a potential-energy surface, a gel scan. In NumPy these are just 2D arrays. The skills are the same each time: build the coordinate grid, compute or load the values, find the interesting features (peaks, hot spots, regions above a threshold), and show the result as a heatmapA plot that shows a 2D array as colours, read with a labelled colorbar. or contour plotA plot whose lines join points of equal value, like a topographic map. that a reader can interpret.
Core idea: a grid of coordinates
np.meshgrid turns two 1D coordinate arrays into two 2D arrays, one holding the x-coordinate of every grid point and one holding the y-coordinate:
import numpy as np
x = np.linspace(-2, 2, 201) # 201 columns
y = np.linspace(-1, 1, 101) # 101 rows
X, Y = np.meshgrid(x, y) # both have shape (101, 201)Note the shape: rows follow y, columns follow x, so Z[i, j] is the value at (x[j], y[i]). This "row = y" convention matches how images are stored and is the single biggest source of confusion with 2D data.
With the grid in hand, any formula becomes a 2D field by broadcasting:
# A Gaussian hot spot centred at (0.5, -0.2) on a 20 °C plate
Z = 20 + 60 * np.exp(-((X - 0.5)**2 + (Y + 0.2)**2) / (2 * 0.3**2))Finding features
np.argmax returns the position of the largest value in the flattened array. Convert it back to a row and column with np.unravel_index:
i, j = np.unravel_index(np.argmax(Z), Z.shape)
print(x[j], y[i], Z[i, j]) # 0.5 -0.2 80.0Thresholds work with Boolean masks, exactly as in 1D:
hot = Z > 50 # 2D Boolean array
print(hot.mean()) # fraction of the plate above 50 °C
cell_area = (x[1] - x[0]) * (y[1] - y[0])
print(hot.sum() * cell_area) # hot area in square unitsReductions take an axis. Z.mean(axis=0) averages over rows (giving a profile along x); Z.max(axis=1) gives the hottest value in each row. A line profile through the peak is just a slice: Z[i, :].
Heatmaps with imshow and pcolormesh
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(5, 3))
im = ax.imshow(Z, origin="lower", extent=[x[0], x[-1], y[0], y[-1]],
cmap="viridis", aspect="auto")
fig.colorbar(im, ax=ax, label="Temperature (°C)")
ax.set_xlabel("x (m)")
ax.set_ylabel("y (m)")Three arguments matter:
origin="lower"puts row 0 at the bottom, as on a graph. The default ("upper") is image convention, with row 0 at the top. Getting this wrong flips your plot vertically.extentsets the axis values, so ticks show metres rather than pixel indices.- A colorbarThe scale beside a heatmap that maps colours to values; it needs a label and unit. with a label and unit is required; a heatmap without one is unreadable.
ax.pcolormesh(X, Y, Z, shading="auto") does the same job and also handles grids that are not evenly spaced.
Contours
Contour lines join points of equal value, like a topographic map:
cs = ax.contour(X, Y, Z, levels=[30, 40, 50, 60, 70], colors="white", linewidths=0.8)
ax.clabel(cs, fmt="%d °C", fontsize=7)contourf fills between levels. Contours are better than colour when readers need to read actual values, and they print well in greyscale. Overlaying a few contour lines on a heatmap often gives the best of both.
Colormaps
- Sequential maps (
viridis,magma,cividis) for values that go from low to high, like temperature or intensity. - Diverging maps (
RdBu_r,coolwarm) for values with a meaningful centre, like residuals or tension versus compression. Centre them on zero withvmin=-m, vmax=m. - Avoid
jetandrainbow: they create bands that look like features and are hard for colour-blind readers.
Fix vmin and vmax when comparing several images side by side, or each panel gets its own scale and the comparison is meaningless.
Common mistakes
- Swapping x and y.
Z[i, j]is(y[i], x[j]). Check with an asymmetric test pattern. - Forgetting
origin="lower", so the plot is upside down compared with your coordinates. - No
extent, so the axes show pixel numbers instead of physical units. - Comparing panels with different colour scales.
- Using
argmaxand treating the result as a column index: it is a flat index until you unravel it.
Quick recap
X, Y = np.meshgrid(x, y)gives 2D coordinate arrays of shape(len(y), len(x)).- Peaks:
np.unravel_index(np.argmax(Z), Z.shape). Thresholds: Boolean masks,.mean()for a fraction,.sum() * cell_areafor an area. imshow(..., origin="lower", extent=...)plus a labelled colorbar;contourfor lines of equal value.- Sequential colormaps for magnitudes, diverging ones for signed data.
Try it
In the Do tab you will build a temperature field on a grid, find its hottest point, and measure the area above a threshold.
Key terms (5)
- meshgrid
- A NumPy function that turns 1D x and y coordinates into 2D arrays holding the coordinates of every grid point.
- heatmap
- A plot that shows a 2D array as colours, read with a labelled colorbar.
- contour plot
- A plot whose lines join points of equal value, like a topographic map.
- unravel_index
- A NumPy function that converts a flat index, such as argmax's result, into a row and column.
- colorbar
- The scale beside a heatmap that maps colours to values; it needs a label and unit.