Subplots
16 min
plt.subplots(rows, cols) creates a grid of Axes objects for
multiple plots in one figure — genuinely useful for comparing several
charts side by side. A real gotcha worth knowing up front: the shape of
the returned axes array depends on the grid shape itself. A full
grid (both rows and cols greater than 1) returns a proper 2D array;
a single row OR a single column gets automatically "squeezed" down to a
flat 1D array instead.
def create_subplot_grid(rows, cols):
fig, axes = plt.subplots(rows, cols)
if rows == 1 or cols == 1:
return len(axes)
return (len(axes), len(axes[0]))
Code that blindly does axes[0][0] will crash with a TypeError the moment someone calls it with rows=1 -- this is a real, commonly-hit bug in matplotlib code, and exactly why libraries built on top of it (like seaborn) often normalize this shape difference away internally before you ever see it.
Write `create_subplot_grid(rows, cols)`: call `plt.subplots(rows, cols)` to get a grid of axes. If `rows == 1` or `cols == 1`, `plt.subplots` returns a flat 1D array — return `len(axes)`. Otherwise it returns a 2D array — return `(len(axes), len(axes[0]))`.
Why does `plt.subplots(rows, cols)` return a genuinely DIFFERENT shaped object (1D array vs. 2D array) depending on whether rows or cols equals 1?
You can create a subplot grid and understand matplotlib's dimension-squeezing behavior for single-row or single-column grids, a real gotcha that trips up naive indexing code.