Skip to content

Instantly share code, notes, and snippets.

@epilys
Last active June 10, 2021 16:33
Show Gist options
  • Save epilys/4df67027dc5db661969b086bac56385f to your computer and use it in GitHub Desktop.
Save epilys/4df67027dc5db661969b086bac56385f to your computer and use it in GitHub Desktop.
Draw treemap with hatch patterns in vanilla matplotlib + python3. It can be easily modified to use colors instead of hatches.
Display the source blob
Display the rendered blob
Raw
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Display the source blob
Display the rendered blob
Raw
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
import matplotlib
matplotlib.rcParams["hatch.linewidth"] = 0.9
from matplotlib import pylab
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from functools import reduce
import io
# Hatch reference: https://matplotlib.org/stable/gallery/shapes_and_collections/hatch_style_reference.html#sphx-glr-gallery-shapes-and-collections-hatch-style-reference-py
# Doubling the size of a string (i.e. '+'*2) makes the hatches thicker
# Use matplotlib.rcParams['hatch.linewidth'] to set the linewidth
def hatches():
list_ = ["/", "\\", "|", "-", "+", "x", "o", "O", ".", "*"]
i = 0
wraps = 0
while True:
index = i % len(list_)
yield list_[i] * (wraps + 1)
i = i + 1
if i >= len(list_):
i -= len(list_)
wraps += 1
"""
Code adapted from https://scipy-cookbook.readthedocs.io/items/Matplotlib_TreeMap.html
"""
class Treemap:
"""
Treemap builder using pylab.
Uses algorithm straight from http://hcil.cs.umd.edu/trs/91-03/91-03.html
James Casbon 29/7/2006
"""
def __init__(self, title, tree, labels):
"""example trees:
tree= ((5,(3,5)), 4, (5,2,(2,3,(3,2,2)),(3,3)), (3,2) )
tree = ((6, 5))
"""
self.done = False
self.hatches_gen = hatches()
self.tree = tree
self.labels = labels
self.title = title
def compute(self):
def new_size_method():
size_cache = {}
def _size(thing):
if isinstance(thing, int):
return thing
if thing in size_cache:
return size_cache[thing]
else:
size_cache[thing] = reduce(int.__add__, [_size(x) for x in thing])
return size_cache[thing]
return _size
self.size_method = new_size_method()
self.ax = pylab.subplot(111, aspect="equal")
pylab.subplots_adjust(left=0, right=1, top=1, bottom=0)
self.ax.set_xticks([])
self.ax.set_yticks([])
self.iter_method = iter
self.rectangles = []
self.addnode(self.tree, lower=[0, 0], upper=[1, 1], axis=0)
i = 0
for (n, r) in self.rectangles:
if isinstance(n, int):
label = str(self.labels[i])
i += 1
rx, ry = r.get_xy()
cx = rx + r.get_width() / 2.0
cy = ry + r.get_height() / 2.0
cx = rx + r.get_width() / 2.0
cy = ry + r.get_height() / 2.0
self.ax.annotate(
f"{label}",
(cx, cy),
color="k",
backgroundcolor="w",
weight="bold",
fontsize=9,
ha="center",
va="center",
)
self.ax.set_xlabel(self.title)
self.done = True
def as_svg(self):
matplotlib.use("Agg") # Use pixel canvas backend instead of GUI backend
if not self.done:
self.compute()
svg = io.BytesIO()
plt.savefig(svg, format="svg")
return svg.getvalue().decode(encoding="UTF-8").strip()
def addnode(self, node, lower=[0, 0], upper=[1, 1], axis=0):
axis = axis % 2
hatch = self.draw_rectangle(lower, upper, node)
width = upper[axis] - lower[axis]
try:
for child in self.iter_method(node):
upper[axis] = lower[axis] + (
width * float(self.size_method(child))
) / self.size_method(node)
self.addnode(child, list(lower), list(upper), axis + 1)
lower[axis] = upper[axis]
except TypeError:
pass
def draw_rectangle(self, lower, upper, node):
h = None
if isinstance(node, int):
h = next(self.hatches_gen)
r = Rectangle(
lower,
upper[0] - lower[0],
upper[1] - lower[1],
edgecolor="k",
fill=False,
hatch=h,
)
self.ax.add_patch(r)
self.rectangles.append((node, r))
return h
from itertools import islice
import math
import random
from datamap import Treemap
from matplotlib import pylab
def random_label():
word_file = "/usr/share/dict/words"
words = [
w
for w in open(word_file).read().splitlines()
if not w.endswith("'s") and len(w) < 4 and len(w) > 1 and w.islower()
]
random.shuffle(words)
for w in words:
yield w
if __name__ == "__main__":
tree = ((5, (3, 5)), 4, (5, 2, (2, 3, (3, 2, 2)), (3, 3)), (3, 2))
labels = list(islice(random_label(), 15))
t = Treemap(f"tree {tree}", tree, labels)
t.compute()
pylab.show()
tree = (5, 4, 3, 2)
labels = [str(n + 1) for n in range(0, 5)]
t = Treemap(f"tree {tree}", tree, labels)
t.compute()
pylab.show()
"""
Show that the area ratios are correct:
"""
a = 500
b = 156
c = 301
total = a + b + c
tree = (a, b, c)
labels = [str(n) for n in tree]
t = Treemap(f"tree {tree}", tree, labels)
t.compute()
pylab.show()
get_area = lambda r: r.get_height() * r.get_width()
for (i, node) in enumerate(tree):
ratio = (node * 1.0) / total
rectangle_ratio = get_area(next(x for x in t.rectangles if x[0] == node)[1])
if not math.isclose(ratio, rectangle_ratio):
print(
"For node",
node,
"the ratio",
ratio,
"differs from rectangle ratio",
rectangle_ratio,
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment