treemap/TreeMap.py at master · GordonMcGregor/treemap · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
from Tree import Tree
import pylab
from matplotlib.patches import Rectangle


'''
A simple TreeMap visualization in pylab

Gordon McGregor  31/7/2013


Based upon http://wiki.scipy.org/Cookbook/Matplotlib/TreeMap:

Comment from that cookbook example:
		Treemap builder using pylab.
		Uses algorithm straight from http://hcil.cs.umd.edu/trs/91-03/91-03.html
		James Casbon 29/7/2006
'''

class TreeMap(object):

	def __init__(self, root):
		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.add_node(root)


	def add_node(self, node, lower=[0.005,0.005], upper=[0.995,0.995], axis = 0):
		axis = axis % 2
		self.draw_rectangle(lower, upper, node)

		width = upper[axis] - lower[axis]

		for branch in node:
			upper[axis] = lower[axis] + (width * float(branch.get_weight())) / node.get_weight()
			self.add_node(branch, list(lower), list(upper), axis + 1)
			lower[axis] = upper[axis]


	def draw_rectangle(self, lower, upper, node):
		r = Rectangle(lower, upper[0] - lower[0], upper[1]-lower[1],
			edgecolor='k',
			facecolor = (0,0,0))
		self.ax.add_patch(r)
		if node.is_leaf():
			rx, ry = r.get_xy()
			cx = rx + r.get_width()/2.0
			cy = ry + r.get_height()/2.0
			r.set_facecolor( node.get_colour())
			self.ax.annotate(node.get_weight(), (cx, cy), color=(0,0,0), fontsize = 10, ha='center', va='center')
			print node.name, rx, ry, cx, cy


	def show(self):
		pylab.show()


if __name__ == '__main__':

	root = Tree(None, None, 'root')
	c1 = Tree(root, 10, 'c1')
	c2 = Tree(root, 2, 'c2')
	n1 = Tree(root, None, 'n1')
	c3 = Tree(n1, 3, 'c3')
	c4 = Tree(n1, 4, 'c4')
	n2 = Tree(n1, None, 'n2')
	c5 = Tree(n2, 5, 'c5')
	c6 = Tree(n2, 6, 'c6')
	TreeMap(root).show()