Lesson 12

Heatmaps

Turn a grid of numbers into colour with sns.heatmap: pivot long data into rows and columns, write values with annot and fmt, and choose a fitting colour map.

Intermediate25 min

A heatmap colours every cell of a table, so patterns across rows and columns jump out before you read a single number. It needs a grid, not long data: the row labels, the column labels, and one value per cell. A correlation matrix is the classic example.

A three by three correlation heatmap of week, hours, and score with values in each cell

Build the grid first

pivot_table turns the long study table into a grid of mean scores: tracks down the side, weeks across the top. This is exactly the shape heatmap wants.

python

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

rng = np.random.default_rng(26)
study = pd.DataFrame({
    "track": rng.choice(["Python", "Data", "Web"], 120),
    "mode": rng.choice(["online", "classroom"], 120),
    "week": rng.integers(1, 7, 120),
})
study["hours"] = (rng.uniform(1, 6, 120) + study["week"] * 0.8).round(1)
study["score"] = (35 + study["hours"] * 4.5 + rng.normal(0, 7, 120)).clip(0, 100).round()

grid = study.pivot_table(index="track", columns="week", values="score", aggfunc="mean")
print(grid.round(0))

Output:

text

week       1     2     3     4     5     6
track                                     
Data    53.0  63.0  62.0  68.0  70.0  73.0
Python  50.0  57.0  59.0  65.0  62.0  78.0
Web     56.0  58.0  65.0  66.0  66.0  67.0

Colour and numbers

annot=True writes each value in its cell and fmt=".0f" rounds it to a whole number. cmap picks the colours; a sequential map such as "mako" suits values that only go one way, from low to high.

python

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

rng = np.random.default_rng(26)
study = pd.DataFrame({
    "track": rng.choice(["Python", "Data", "Web"], 120),
    "mode": rng.choice(["online", "classroom"], 120),
    "week": rng.integers(1, 7, 120),
})
study["hours"] = (rng.uniform(1, 6, 120) + study["week"] * 0.8).round(1)
study["score"] = (35 + study["hours"] * 4.5 + rng.normal(0, 7, 120)).clip(0, 100).round()

grid = study.pivot_table(index="track", columns="week", values="score", aggfunc="mean")
ax = sns.heatmap(grid, annot=True, fmt=".0f", cmap="mako", linewidths=1)
ax.set_title("Mean score by track and week")
plt.show()

print(grid.shape)

Output:

text

(3, 6)

A correlation matrix

corr() gives every pair of numeric columns a value from -1 to 1. For values that can go either way, use a diverging map with center=0, and fix vmin=-1 and vmax=1 so the colours mean the same thing in every chart.

python

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

rng = np.random.default_rng(26)
study = pd.DataFrame({
    "track": rng.choice(["Python", "Data", "Web"], 120),
    "mode": rng.choice(["online", "classroom"], 120),
    "week": rng.integers(1, 7, 120),
})
study["hours"] = (rng.uniform(1, 6, 120) + study["week"] * 0.8).round(1)
study["score"] = (35 + study["hours"] * 4.5 + rng.normal(0, 7, 120)).clip(0, 100).round()

corr = study[["week", "hours", "score"]].corr()
sns.heatmap(corr, annot=True, fmt=".2f", cmap="vlag", center=0, vmin=-1, vmax=1, square=True)
plt.show()

print(corr.round(2))

Output:

text

       week  hours  score
week   1.00   0.72   0.58
hours  0.72   1.00   0.82
score  0.58   0.82   1.00

Hide the repeated half

A correlation matrix is a mirror image, and the diagonal is always 1. A mask of True values hides cells; np.triu builds one for the upper triangle so each pair appears once.

python

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

rng = np.random.default_rng(26)
study = pd.DataFrame({
    "track": rng.choice(["Python", "Data", "Web"], 120),
    "mode": rng.choice(["online", "classroom"], 120),
    "week": rng.integers(1, 7, 120),
})
study["hours"] = (rng.uniform(1, 6, 120) + study["week"] * 0.8).round(1)
study["score"] = (35 + study["hours"] * 4.5 + rng.normal(0, 7, 120)).clip(0, 100).round()

corr = study[["week", "hours", "score"]].corr()
mask = np.triu(np.ones_like(corr, dtype=bool))
sns.heatmap(corr, mask=mask, annot=True, fmt=".2f", cmap="vlag", center=0,
            vmin=-1, vmax=1, cbar_kws={"label": "correlation"})
plt.show()

print(mask)

Output:

text

[[ True  True  True]
 [False  True  True]
 [False False  True]]

Counts as a heatmap

pd.crosstab counts rows for every pair of categories, and a heatmap of it shows which combinations are common. fmt="d" prints whole numbers because counts are integers.

python

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

rng = np.random.default_rng(26)
study = pd.DataFrame({
    "track": rng.choice(["Python", "Data", "Web"], 120),
    "mode": rng.choice(["online", "classroom"], 120),
    "week": rng.integers(1, 7, 120),
})
study["hours"] = (rng.uniform(1, 6, 120) + study["week"] * 0.8).round(1)
study["score"] = (35 + study["hours"] * 4.5 + rng.normal(0, 7, 120)).clip(0, 100).round()

counts = pd.crosstab(study["track"], study["mode"])
sns.heatmap(counts, annot=True, fmt="d", cmap="rocket_r")
plt.show()

print(counts)

Output:

text

mode    classroom  online
track                    
Data           17      17
Python         15      20
Web            22      29

Key takeaways

heatmap needs a grid: pivot or crosstab long data first
annot=True writes values; fmt=".0f" or "d" formats them
Use a sequential cmap for one-way values, diverging with center=0 otherwise
Mask the upper triangle of a correlation matrix with np.triu
Next: Regression Plots
Previous: Strip, Swarm, and catplot

Loading playground…

Check your understanding

Heatmaps Quiz

What shape of data does heatmap need?

1 / 10

← Strip, Swarm, and catplot

Regression Plots →