Heatmap with plt.imshow in Matplotlib - Step-by-Step Execution

Concept Flow - Heatmap with plt.imshow

Prepare 2D data array

↓

Call plt.imshow(data)

↓

Render heatmap colors

↓

Add colorbar (optional)

↓

Show plot with plt.show()

The flow shows how a 2D data array is passed to plt.imshow, which renders it as a colored heatmap, optionally adding a colorbar, then displays the plot.

Execution Sample

Matplotlib

import matplotlib.pyplot as plt
import numpy as np
data = np.array([[1, 2], [3, 4]])
heatmap_image = plt.imshow(data, cmap='viridis')
colorbar = plt.colorbar()
plt.show()

This code creates a 2x2 heatmap from a small 2D array using the 'viridis' color map and shows a colorbar.

Execution Table

Step	Action	Input/Parameter	Result/Output
1	Create 2D numpy array	[[1, 2], [3, 4]]	data array ready
2	Call plt.imshow	data, cmap='viridis'	Heatmap image object created
3	Call plt.colorbar	None	Colorbar linked to heatmap
4	Call plt.show	None	Heatmap with colorbar displayed
5	End	Plot window closed or script ends	Execution stops

💡 Plot displayed and script ends, no further steps

Variable Tracker

Variable	Start	After Step 1	After Step 2	After Step 3	Final
data	None	[[1 2] [3 4]]	[[1 2] [3 4]]	[[1 2] [3 4]]	[[1 2] [3 4]]
heatmap_image	None	None	AxesImage object	AxesImage object	AxesImage object
colorbar	None	None	None	Colorbar object	Colorbar object

Key Moments - 3 Insights

Why does plt.imshow need a 2D array as input?

plt.imshow uses the 2D array values to assign colors for each cell in the heatmap, as shown in execution_table step 2 where the data array is passed to create the heatmap image.

What does the cmap parameter do in plt.imshow?

The cmap parameter sets the color scheme for the heatmap. In step 2, 'viridis' is used to map data values to colors, making the heatmap visually meaningful.

Why do we call plt.colorbar after plt.imshow?

plt.colorbar links a color scale to the heatmap image, helping interpret colors. This is shown in step 3 where the colorbar is created after the heatmap image exists.

Visual Quiz - 3 Questions

Look at the execution table, what is the type of the object created by plt.imshow at step 2?

A Figure object

B Colorbar object

C AxesImage object

D Numpy array

💡 Hint

Check the 'Result/Output' column in step 2 of the execution_table.

At which step is the colorbar linked to the heatmap?

A Step 3

B Step 1

C Step 2

D Step 4

💡 Hint

Look for when plt.colorbar is called in the execution_table.

If the data array was 3x3 instead of 2x2, how would the heatmap change?

A The heatmap would stay 2x2 but colors would change

B The heatmap would have 3 rows and 3 columns of colored cells

C The heatmap would be a single color

D The code would error out

💡 Hint

Refer to variable_tracker for how data shape affects the heatmap size.

Concept Snapshot

plt.imshow(data, cmap='color_map')
- Displays 2D array as colored heatmap
- cmap sets color scheme (e.g., 'viridis')
- Use plt.colorbar() to add legend
- plt.show() renders the plot
- Data shape controls heatmap grid size

Full Transcript

This visual execution traces how to create a heatmap using matplotlib's plt.imshow. First, a 2D numpy array is prepared as data. Then plt.imshow is called with this data and a color map to create a heatmap image object. Next, plt.colorbar adds a color scale legend linked to the heatmap. Finally, plt.show displays the heatmap plot. Variables like the data array, heatmap image, and colorbar objects are tracked through each step. Key beginner questions about input data, color maps, and colorbars are answered referencing the execution steps. A quiz tests understanding of object types, step order, and data shape effects. The snapshot summarizes the main commands and their roles in making a heatmap with plt.imshow.

Practice

1. What does the plt.imshow function do in matplotlib when given a 2D array?

easy

A. It prints the array values as text on the plot.

B. It plots a line graph of the array values.

C. It creates a heatmap by showing the array values as colors.

D. It creates a pie chart from the array values.

Solution

1. Step 1: Understand plt.imshow input

   The function takes a 2D array and maps its values to colors in a grid.

2. Step 2: Recognize heatmap behavior

   This color mapping creates a heatmap, showing patterns visually.

3. Final Answer:

   It creates a heatmap by showing the array values as colors. -> Option C

4. Quick Check:

   plt.imshow = heatmap colors [OK]

Hint: plt.imshow shows 2D data as colored grids [OK]

Common Mistakes:

• Thinking plt.imshow plots lines or text
• Confusing heatmap with pie chart
• Assuming plt.imshow prints numbers

2. Which of these is the correct way to add a color bar to a heatmap created with plt.imshow?

easy

A. plt.colorbar()

B. plt.addcolor()

C. plt.showcolor()

D. plt.colorheat()

Solution

1. Step 1: Recall color bar function

   Matplotlib uses plt.colorbar() to add a color scale next to images.

2. Step 2: Check other options

   Other options are not valid matplotlib functions.

3. Final Answer:

   plt.colorbar() -> Option A

4. Quick Check:

   Color bar function = plt.colorbar() [OK]

Hint: Use plt.colorbar() to add color scale [OK]

Common Mistakes:

• Using non-existent functions like plt.addcolor()
• Forgetting to call plt.colorbar() after plt.imshow()
• Trying to add color bar before plotting heatmap

3. What will the following code display?

import matplotlib.pyplot as plt
import numpy as np

data = np.array([[1, 2], [3, 4]])
plt.imshow(data, cmap='hot')
plt.colorbar()
plt.show()

medium

A. A scatter plot of points (1,2) and (3,4) with hot colors.

B. A line plot of the numbers 1, 2, 3, 4 with a color bar.

C. An error because cmap='hot' is invalid.

D. A 2x2 heatmap with colors from black to red to yellow, and a color bar showing values 1 to 4.

Solution

1. Step 1: Understand plt.imshow with cmap='hot'

   The data array is 2x2, so plt.imshow creates a 2x2 grid colored by values using the 'hot' color map (black to red to yellow).

2. Step 2: Effect of plt.colorbar()

   The color bar shows the scale from the minimum value 1 to maximum 4, helping interpret colors.

3. Final Answer:

   A 2x2 heatmap with colors from black to red to yellow, and a color bar showing values 1 to 4. -> Option D

4. Quick Check:

   plt.imshow + cmap='hot' + colorbar = heatmap with hot colors [OK]

Hint: cmap='hot' colors low to high values black-red-yellow [OK]

Common Mistakes:

• Thinking plt.imshow plots lines or scatter
• Believing cmap='hot' causes error
• Ignoring color bar meaning

4. The code below is intended to show a heatmap but does not display the color bar. What is the error?

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(3,3)
plt.imshow(data)
plt.colorbar
plt.show()

medium

A. plt.colorbar is missing parentheses, so color bar is not shown.

B. plt.imshow needs a cmap argument to show color bar.

C. plt.show() must come before plt.colorbar().

D. np.random.rand(3,3) creates invalid data for heatmap.

Solution

1. Step 1: Identify plt.colorbar usage

   The code uses plt.colorbar without parentheses, so the function is not called.

2. Step 2: Correct function call

   Adding parentheses plt.colorbar() calls the function and displays the color bar.

3. Final Answer:

   plt.colorbar is missing parentheses, so color bar is not shown. -> Option A

4. Quick Check:

   Call plt.colorbar() with () to show color bar [OK]

Hint: Always call plt.colorbar() with parentheses [OK]

Common Mistakes:

• Forgetting parentheses on plt.colorbar
• Thinking cmap is required for color bar
• Calling plt.show() before plt.colorbar

5. You have a 4x4 matrix with values from 0 to 15. You want to create a heatmap using plt.imshow that shows the values clearly with a color bar, but you want the color scale to always be from 0 to 20 regardless of data. Which code snippet correctly does this?

hard

A. plt.imshow(matrix, vmin=0, vmax=15)
plt.colorbar()
plt.show()

B. plt.imshow(matrix, vmin=0, vmax=20)
plt.colorbar()
plt.show()

C. plt.imshow(matrix, cmap='gray')
plt.colorbar()
plt.show()

D. plt.imshow(matrix)
plt.colorbar(vmin=0, vmax=20)
plt.show()

Solution

1. Step 1: Set color scale limits in plt.imshow

   Use vmin=0 and vmax=20 in plt.imshow to fix color scale from 0 to 20.

2. Step 2: Add color bar normally

   Call plt.colorbar() to show the color scale matching the fixed limits.

3. Final Answer:

   plt.imshow(matrix, vmin=0, vmax=20) plt.colorbar() plt.show() -> Option B

4. Quick Check:

   Set vmin/vmax in imshow, then add colorbar [OK]

Hint: Use vmin and vmax in plt.imshow to fix color scale [OK]

Common Mistakes:

• Setting vmin/vmax in plt.colorbar instead of plt.imshow
• Not setting vmax to 20, so scale auto adjusts
• Using wrong vmin/vmax values