Images and 2D data

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 heatmap or contour plot 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.0

Thresholds 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 units

Reductions 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.
• extent sets the axis values, so ticks show metres rather than pixel indices.
• A colorbar 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 with vmin=-m, vmax=m.
• Avoid jet and rainbow: 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 argmax and 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_area for an area.
• imshow(..., origin="lower", extent=...) plus a labelled colorbar; contour for 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.