55 of 100: Treemap in matplotlib
At the beginning of the year I challenged myself to create all 100 visualizations using python and matplotlib from the 1 dataset,100 visualizations project and I am sharing with you the code for all the visualizations.
Note: Data Viz Project is copyright Ferdio and available under a Creative Commons Attribution – Non Commercial – No Derivatives 4.0 International license. I asked Ferdio and they told me they used a Design tool to create all the plots.
Collaborate
There are a ton of improvements that can be made on the code, so let me know in the comments any improvements you make and I will update the post accordingly!
To be improved: I didnt do this myself, I got help from stackoverflow. I need to go through the code to add the legends and data labels in the correct places.
This is the original viz that we are trying to recreate in matplotlib:

Plot the chart
import numpy as np
import matplotlib.pyplot as plt
import squarify
from matplotlib import colormaps
from matplotlib.colors import to_rgba
DEFAULT_COLORS = list(zip(colormaps["tab20"].colors[::2],
colormaps["tab20"].colors[1::2]))
def color_to_grayscale(color):
# Adapted from: https://stackoverflow.com/a/689547/2912349
r, g, b, a = to_rgba(color)
return (0.299 * r + 0.587 * g + 0.114 * b) * a
class PairedTreeMap:
def __init__(self, values, colors=DEFAULT_COLORS, labels=None, ax=None, bbox=(0, 0, 200, 100)):
"""
Draw a treemap of value pairs.
values : list[tuple[float, float]]
A list of value pairs.
colors : list[tuple[RGBA, RGBA]]
The corresponding color pairs. Defaults to light/dark tab20 matplotlib color pairs.
labels : list[str]
The labels, one for each pair.
ax : matplotlib.axes._axes.Axes
The matplotlib axis instance to draw on.
bbox : tuple[float, float, float, float]
The (x, y) origin and (width, height) extent of the treemap.
"""
self.ax = self.initialize_axis(ax)
self.rects = self.get_layout(values, bbox)
self.artists = list(self.draw(self.rects, values, colors, self.ax))
if labels:
self.labels = list(self.add_labels(self.rects, labels, values, colors, self.ax))
def get_layout(self, values, bbox):
maxima = np.max(values, axis=1)
order = np.argsort(maxima)[::-1]
normalized_maxima = squarify.normalize_sizes(maxima[order], *bbox[2:])
rects = squarify.padded_squarify(normalized_maxima, *bbox)
reorder = np.argsort(order)
return [rects[ii] for ii in reorder]
def initialize_axis(self, ax=None):
if ax is None:
fig, ax = plt.subplots()
ax.set_aspect("equal")
ax.axis("off")
return ax
def _get_artist_pair(self, rect, value_pair, color_pair):
x, y, w, h = rect["x"], rect["y"], rect["dx"], rect["dy"]
(small, large), (color_small, color_large) = zip(*sorted(zip(value_pair, color_pair)))
ratio = np.sqrt(small / large)
return (plt.Rectangle((x, y), w, h, color=color_large, zorder=1),
plt.Rectangle((x, y), w * ratio, h * ratio, color=color_small, zorder=2))
def draw(self, rects, values, colors, ax):
for rect, value_pair, color_pair in zip(rects, values, colors):
large_patch, small_patch = self._get_artist_pair(rect, value_pair, color_pair)
ax.add_patch(large_patch)
ax.add_patch(small_patch)
yield(large_patch, small_patch)
ax.autoscale_view()
def add_labels(self, rects, labels, values, colors, ax):
for rect, label, value_pair, color_pair in zip(rects, labels, values, colors):
x, y, w, h = rect["x"], rect["y"], rect["dx"], rect["dy"]
# decide a fontcolor based on background brightness
(small, large), (color_small, color_large) = zip(*sorted(zip(value_pair, color_pair)))
ratio = small / large
background_brightness = color_to_grayscale(color_large) if ratio < 0.33 else color_to_grayscale(color_small) # i.e. 0.25 + some fudge
fontcolor = "white" if background_brightness < 0.5 else "black"
yield ax.text(x + w/2, y + h/2, label, va="center", ha="center", color=fontcolor)
if __name__ == "__main__":
values = [
(4, 10),
(13, 15),
(5, 8),
]
colors = [
("red", "coral"),
("royalblue", "cornflowerblue"),
("darkslategrey", "gray"),
]
labels = [
"Denmark",
"Sweden",
"Norway"
]
PairedTreeMap(values, colors=colors, labels=labels, bbox=(0, 0, 100, 100))
plt.show()
The result:

Reader Interactions