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

Concept Flow - Heatmap with plt.pcolormesh

Prepare 2D data array

↓

Create X, Y grid coordinates

↓

Call plt.pcolormesh(X, Y, data)

↓

Color cells based on data values

↓

Add colorbar for scale

↓

Show heatmap plot

This flow shows how to create a heatmap by preparing data, mapping it on a grid, coloring cells, and displaying the plot.

Execution Sample

Matplotlib

import numpy as np
import matplotlib.pyplot as plt

Z = np.array([[1, 2], [3, 4]])
plt.pcolormesh(Z, cmap='viridis')
plt.colorbar()
plt.show()

This code creates a simple 2x2 heatmap using pcolormesh with a color scale.

Execution Table

Step	Action	Input/State	Output/Result
1	Create 2D data array Z	None	[[1, 2], [3, 4]]
2	Call plt.pcolormesh(Z, cmap='viridis')	Z=[[1,2],[3,4]]	Grid cells colored by values
3	Add colorbar	Plot with colored cells	Color scale shown on side
4	Show plot	Plot ready	Heatmap displayed on screen
5	End	Plot displayed	Execution stops

💡 Plot is shown and program ends

Variable Tracker

Variable	Start	After Step 1	After Step 2	After Step 3	Final
Z	None	[[1, 2], [3, 4]]	[[1, 2], [3, 4]]	[[1, 2], [3, 4]]	[[1, 2], [3, 4]]
Plot	None	None	Heatmap grid created	Colorbar added	Displayed

Key Moments - 3 Insights

Why does plt.pcolormesh require a 2D array for data?

Because each cell in the heatmap corresponds to one value in the 2D array, as shown in execution_table step 2 where the array Z defines cell colors.

What does the colorbar represent in the heatmap?

The colorbar shows the mapping from data values to colors, helping interpret the heatmap colors, as added in step 3 of the execution_table.

Why do we call plt.show() at the end?

plt.show() displays the plot window so we can see the heatmap, as indicated in step 4 where the heatmap is displayed on screen.

Visual Quiz - 3 Questions

Look at the execution_table at step 2, what is the input data used for coloring the heatmap cells?

A [[4, 3], [2, 1]]

B [[1, 2], [3, 4]]

C [1, 2, 3, 4]

D None

💡 Hint

Check the 'Input/State' column at step 2 in execution_table.

At which step in the execution_table is the colorbar added to the plot?

A Step 1

B Step 2

C Step 3

D Step 4

💡 Hint

Look for the action mentioning 'Add colorbar' in execution_table.

If we change the data array Z to [[5, 6], [7, 8]], how would the heatmap colors change?

A Colors would represent higher values, shifting the color scale

B Colors would stay the same

C Heatmap would become empty

D Colorbar would disappear

💡 Hint

Refer to variable_tracker for Z values and how they map to colors in execution_table step 2.

Concept Snapshot

plt.pcolormesh(data, cmap) creates a heatmap by coloring grid cells based on 2D data values.
Data must be a 2D array where each value colors one cell.
Add plt.colorbar() to show the color scale.
Call plt.show() to display the heatmap plot.
Useful for visualizing matrices or grids of values.

Full Transcript

This visual execution trace shows how to create a heatmap using matplotlib's plt.pcolormesh. First, a 2D data array Z is created with values. Then plt.pcolormesh is called with Z and a color map to color each grid cell according to its value. Next, a colorbar is added to show the mapping from values to colors. Finally, plt.show() displays the heatmap plot on screen. Variables like Z and the plot state are tracked through each step. Key moments clarify why data must be 2D, the role of the colorbar, and the need to call plt.show(). The quiz tests understanding of data input, colorbar addition, and effect of changing data values on colors. This step-by-step trace helps beginners see how heatmaps are built and displayed visually.

Practice

1. What does plt.pcolormesh do in matplotlib?

easy

A. It creates a scatter plot with colored dots.

B. It plots a line graph connecting points.

C. It creates a heatmap by coloring a grid based on data values.

D. It draws bar charts with colored bars.

Solution

1. Step 1: Understand the function purpose

   plt.pcolormesh colors a grid of cells based on the values in a 2D array, making a heatmap.

2. Step 2: Compare with other plot types

   Line graphs connect points, scatter plots show dots, and bar charts show bars, none of which match pcolormesh.

3. Final Answer:

   It creates a heatmap by coloring a grid based on data values. -> Option C

4. Quick Check:

   Heatmap = pcolormesh = A [OK]

Hint: pcolormesh colors grids for heatmaps, not lines or bars [OK]

Common Mistakes:

• Confusing pcolormesh with line or scatter plots
• Thinking pcolormesh plots points instead of colored cells
• Assuming pcolormesh creates bar charts

2. Which of these is the correct way to add a colorbar to a heatmap created with plt.pcolormesh?

easy

A. plt.colorbar()

B. plt.add_colorbar()

C. plt.show_colorbar()

D. plt.colorbar_add()

Solution

1. Step 1: Recall matplotlib colorbar syntax

   The correct function to add a colorbar is plt.colorbar().

2. Step 2: Check other options

   Functions like plt.add_colorbar(), plt.show_colorbar(), or plt.colorbar_add() do not exist in matplotlib.

3. Final Answer:

   plt.colorbar() -> Option A

4. Quick Check:

   Colorbar function = plt.colorbar() = B [OK]

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

Common Mistakes:

• Using non-existent functions like plt.add_colorbar()
• Forgetting to add a colorbar after pcolormesh
• Trying to add colorbar before plotting heatmap

3. What will the following code display?

import matplotlib.pyplot as plt
import numpy as np

Z = np.array([[1, 2], [3, 4]])
plt.pcolormesh(Z)
plt.colorbar()
plt.show()

medium

A. An error because Z is not a 1D array.

B. A line plot connecting points 1, 2, 3, 4.

C. A scatter plot with 4 points colored by values.

D. A 2x2 heatmap with colors representing values 1 to 4 and a colorbar.

Solution

1. Step 1: Understand the input array Z

   Z is a 2x2 numpy array with values from 1 to 4.

2. Step 2: What does plt.pcolormesh(Z) do?

   It creates a heatmap grid colored by the values in Z, so 4 colored squares.

3. Step 3: Adding plt.colorbar()

   This adds a color scale legend showing how colors map to values 1 to 4.

4. Final Answer:

   A 2x2 heatmap with colors representing values 1 to 4 and a colorbar. -> Option D

5. Quick Check:

   2x2 heatmap + colorbar = C [OK]

Hint: pcolormesh with 2D array shows colored grid matching array shape [OK]

Common Mistakes:

• Thinking pcolormesh needs 1D array only
• Confusing heatmap with line or scatter plots
• Forgetting colorbar shows value-color mapping

4. Identify the error in this code snippet:

import matplotlib.pyplot as plt
import numpy as np

Z = np.array([[1, 2], [3, 4]])
plt.pcolormesh(Z)
plt.colorbar
plt.show()

medium

A. Missing parentheses after plt.colorbar causing no colorbar to show.

B. Z should be a 1D array, not 2D.

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

D. plt.pcolormesh() requires x and y coordinates explicitly.

Solution

1. Step 1: Check plt.colorbar usage

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

2. Step 2: Effect of missing parentheses

   Without calling the function, no colorbar appears on the plot.

3. Final Answer:

   Missing parentheses after plt.colorbar causing no colorbar to show. -> Option A

4. Quick Check:

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

Hint: Always add () after plt.colorbar to display it [OK]

Common Mistakes:

• Forgetting () after plt.colorbar
• Thinking 2D array is invalid input
• Calling plt.show() too early
• Assuming x,y needed for simple heatmap

5. You want to create a heatmap of a 3x3 matrix with values from 0 to 8 using plt.pcolormesh. Which code correctly sets the x and y axis labels to 'X Axis' and 'Y Axis' and adds a colorbar?

hard

A. import matplotlib.pyplot as plt
import numpy as np
Z = np.arange(9).reshape(3,3)
plt.pcolormesh(Z)
plt.set_xlabel('X Axis')
plt.set_ylabel('Y Axis')
plt.colorbar()
plt.show()

B. import matplotlib.pyplot as plt
import numpy as np
Z = np.arange(9).reshape(3,3)
plt.pcolormesh(Z)
plt.xlabel('X Axis')
plt.ylabel('Y Axis')
plt.colorbar()
plt.show()

C. import matplotlib.pyplot as plt
import numpy as np
Z = np.arange(9).reshape(3,3)
plt.pcolormesh(Z)
plt.xlabel = 'X Axis'
plt.ylabel = 'Y Axis'
plt.colorbar()
plt.show()

D. import matplotlib.pyplot as plt
import numpy as np
Z = np.arange(9).reshape(3,3)
plt.pcolormesh(Z)
plt.axis_labels('X Axis', 'Y Axis')
plt.colorbar()
plt.show()

Solution

1. Step 1: Check axis label functions

   The correct functions to label axes are plt.xlabel() and plt.ylabel().

2. Step 2: Verify colorbar and show usage

   plt.colorbar() must be called with parentheses, and plt.show() displays the plot.

3. Step 3: Identify incorrect options

   import matplotlib.pyplot as plt import numpy as np Z = np.arange(9).reshape(3,3) plt.pcolormesh(Z) plt.set_xlabel('X Axis') plt.set_ylabel('Y Axis') plt.colorbar() plt.show() uses non-existent plt.set_xlabel(), C assigns strings instead of calling functions, D uses a non-existent plt.axis_labels().

4. Final Answer:

   import matplotlib.pyplot as plt import numpy as np Z = np.arange(9).reshape(3,3) plt.pcolormesh(Z) plt.xlabel('X Axis') plt.ylabel('Y Axis') plt.colorbar() plt.show() -> Option B

5. Quick Check:

   Use plt.xlabel(), plt.ylabel(), plt.colorbar() [OK]

Hint: Use plt.xlabel() and plt.ylabel() to label axes [OK]

Common Mistakes:

• Using plt.set_xlabel() instead of plt.xlabel()
• Assigning labels instead of calling functions
• Forgetting parentheses on plt.colorbar()
• Using non-existent axis_labels() function