New plotter for phylogenetic trees: plotting

Here is a first pass at a plotter, following up on the last post.

It's ugly code yet, but I thought I would show you what I have. What I need to do is really think harder about how this module should be organized. But you can compare the result with what we got here.

[UPDATE: Modified the code to be easily customizable. Next up: more extensive testing.]
UPDATE2: Modified the code yet again.]

output:

alternate name dict None default e node dot color k dot color dict None e node bar color r e node bar color list None e node dot size 75 e node dots visible True e node label default color k e node label size 14 e node labels visible True figure type png horizontal axis visible True i node bar color k i node dot color magenta i node dots visible False i node vertical bar color k label width factor 1.1 line width 2 node label color dict None r node dot color orange using alternate names False using e node specific colors False vertical axis visible False

code listing:

import matplotlib.pyplot as plt import tree_utils as tu def plot(D,attr=None,ofn=None): if not attr: attr = get_default_attr() SZ = attr['e_node_dot_size'] lw = attr['line_width'] maxx, maxy = D['meta']['max_xy'] all_node_names = D['meta']['all_node_names'] e_node_names = D['meta']['e_node_names'] i_node_names = D['meta']['i_node_names'] # actual node dictionaries e_node_dicts = [D[n] for n in e_node_names] i_node_dicts = [D[n] for n in i_node_names] # extract the values we need e_node_positions = [(nD['x'],nD['y']) for nD in e_node_dicts] e_node_lengths = [(nD['dist_to_parent']) for nD in e_node_dicts] i_node_positions = [(nD['x'],nD['y']) for nD in i_node_dicts] i_node_lengths = [(nD['dist_to_parent']) for nD in i_node_dicts] i_node_verticals = [(nD['y_bott'], nD['y_top']) for nD in i_node_dicts] # external nodes for i in range(len(e_node_names)): x,y = e_node_positions[i] d = e_node_lengths[i] xleft = x - d # bars L = attr['e_node_bar_color_list'] if L: c = L[i] else: c = attr['e_node_bar_color'] plt.plot([xleft,x],[y,y],color=c,lw=lw,zorder=1) # dots name = e_node_names[i] if attr['e_node_dots_visible']: if attr['using_e_node_specific_colors']: c = attr['dot_color_dict'][name] else: c = attr['default_e_node_dot_color'] plt.scatter(x,y,color=c,s=SZ,zorder=2) # internal nodes, root is i = 0 for i in range(len(i_node_names))[1:]: x,y = i_node_positions[i] d = i_node_lengths[i] y_bott, y_top = i_node_verticals[i] x0 = x - d # bars c = attr['i_node_bar_color'] plt.plot([x0,x],[y,y],color=c,lw=lw,zorder=1) # verticals c = attr['i_node_vertical_bar_color'] plt.plot([x,x],[y_bott, y_top],color=c,lw=lw,zorder=1) # dots if attr['i_node_dots_visible']: c = attr['i_node_dot_color'] plt.scatter(x,y,color=c,s=SZ,zorder=2) # root i = 0 x,y = i_node_positions[i] y_bott, y_top = i_node_verticals[i] # verticals c = attr['i_node_vertical_bar_color'] plt.plot([x,x],[y_bott, y_top],color=c,lw=lw,zorder=1) # dots if attr['i_node_dots_visible']: c = attr['r_node_dot_color'] plt.scatter(x,y,color=c,s=SZ,zorder=2) L = [len(name) for name in e_node_names] max_label_width = maxx/100.0 * max(L) max_label_width *= attr['label_width_factor'] # external node labels if attr['e_node_labels_visible']: for i in range(len(e_node_names)): x,y = e_node_positions[i] dx = maxx/100.0 * 4 name = e_node_names[i] if attr['using_alternate_names']: s = attr['alternate_name_dict'][name] else: s = name if attr['using_e_node_specific_colors']: c = attr['node_label_color_dict'][name] else: c = attr['e_node_label_default_color'] plt.text(x + dx, y, s, fontname = 'Helvetica', fontsize = attr['e_node_label_size'], color = c, ha = 'left',va = 'center') ax = plt.axes() if not attr['vertical_axis_visible']: ax.yaxis.set_visible(False) if not attr['horizontal_axis_visible']: ax.xaxis.set_visible(False) ax.set_xlim(-maxx/10.0,maxx*1.1 + max_label_width) ax.set_ylim(-maxy/10.0,maxy*1.1) if ofn: plt.savefig(ofn + '.' + attr['figure_type']) else: plt.savefig('example.' + attr['figure_type']) def get_default_attr(): attr = dict() attr['e_node_dot_size'] = 75 attr['line_width'] = 2 attr['figure_type'] = 'png' attr['horizontal_axis_visible'] = True attr['vertical_axis_visible'] = False attr['e_node_bar_color'] = 'r' attr['e_node_bar_color_list'] = None attr['e_node_dots_visible'] = True attr['e_node_labels_visible'] = True attr['default_e_node_dot_color'] = 'k' attr['e_node_label_size'] = 14 attr['e_node_label_default_color'] = 'k' attr['label_width_factor'] = 1.1 attr['i_node_bar_color'] = 'k' attr['i_node_vertical_bar_color'] = 'k' attr['i_node_dots_visible'] = False attr['i_node_dot_color'] = 'magenta' attr['r_node_dot_color'] = 'orange' attr['using_alternate_names'] = False attr['alternate_name_dict'] = None attr['using_e_node_specific_colors'] = False attr['node_label_color_dict'] = None attr['dot_color_dict'] = None return attr def print_defaults(): attr = get_default_attr() N = max([len(k) for k in attr.keys()]) for k in sorted(attr.keys()): v = attr[k] k = k.replace('_',' ') print k.rjust(N), ' ', v if __name__ == '__main__': fn = 'tree.txt' file_data = tu.load_data(fn) D = tu.make_tree_dict(ts=file_data) plot(D) print_defaults()