Last active
June 10, 2021 16:33
-
-
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.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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