Jul 31, 2013

Playing with TreeMaps

Just started trying to implement a TreeMap visualization in Python, using Matplotlib. I found some sample code that got me started and led to some of the original academic papers on the topic.

I've got the basics working, with a simple Tree class that I wrote to track nodes and leaves, along with weights and a place to store metadata. The Tree class just provides a wrapper around an array but helps in tracking parent/ child relationships and working out the relative weights of nodes.

class Tree(object):

    def __init__(self, parent=None, weight=None, name=None):
        self.parent = parent
        self.children = []
        self.name = name
        self.weight = weight
        self.changed = False
        if self.parent:
            self.parent.add_child(self)


    def add_child(self, child):
        self.changed = True
        self.children.append(child)


    def __iter__(self):
        for child in self.children:
            yield child


    def is_leaf(self):
        return len(self.children) == 0


    def get_weight(self, recalculate = False):

        if (recalculate and self.children) or not self.weight or self.changed:
            self.weight = 0
            self.changed = False
            for child in self.children:
                self.weight += child.get_weight(recalculate)

        return self.weight

    def get_normalized_weight(self):
        return self.weight/ float(self.parent.weight)

The actual drawing algorithm is straightforward too - for each level in the hierarchy, divide up the available space between all the nodes, weighted by their size and then flip between horizontal or vertical packing at each level of the hierarchy. Currently I'm putting the leaf node weight into the center of each rectangle as an annotation. This will probably need to change with additional nodes, as they will get too small to see. Mouseover tooltips or a datacursor that updates when a node is selected will probably be more useful.

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

The full source is available on github. Trees can be implemented by creating all the nodes, one instance at a time.

from Tree import HueTree as t

a = t()
b = t(a)
e = t(b, 1, 'e')
f = t(b, 2, 'f')

c = t(a, 3, 'c')

d = t(a)
g = t(d)
j = t(g, 1, 'j')
k = t(g, 1, 'k')
h = t(d, 4, 'h')
l = t(d)
lprime = t(l, 1, 'l')
m = t(l, 1, 'm')
n = t(l, 1, 'n')
o = t(l, 1, 'o')

TreeMap(a).show()

Simple Treemap

Here I'm using the relative weight of the node within a given point in the hierarchy to specify the Hue value from an Hue/Saturation/Value triplet that gets converted to an RGB colour.

The long form way of instantiating a tree gets quite cumbersome, so it is also possible to define a tree just as a nested set of tuples, and then use a helper function, make_tree to construct the actual Tree objects.

import Tree
from TreeMap import TreeMap

short_map = (((1, 2), 3, ((1, 1), 4, ((32,34,1,2), 1, 1, 
                    (1, 2, 4, 5,(2,2,(2,(2,(1,1,1,1,(3,2,12),1)))))))))

x=Tree.make_tree(short_map, TreeType = Tree.Tree)
TreeMap(x).show()

Things look good with a small number of nodes, but when you start adding more leaves, you find that skinny rectangles start to dominate the plot. In this case, I'm using the relative weight as a gray value, rather than varying the hue. Other algorithms are available for how the space is subdivided that can tend to give a more square structure to the resulting map - I'm starting to investigate implementing the map using those algorithms to see how that looks.

Skinny Treemap

There are comments.

Comments !