forked from biopython/biopython
-
Notifications
You must be signed in to change notification settings - Fork 0
/
_utils.py
501 lines (432 loc) · 20.2 KB
/
_utils.py
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
# Copyright (C) 2009 by Eric Talevich (eric.talevich@gmail.com)
# This code is part of the Biopython distribution and governed by its
# license. Please see the LICENSE file that should have been included
# as part of this package.
"""Utilities for handling, displaying and exporting Phylo trees.
Third-party libraries are loaded when the corresponding function is called.
"""
__docformat__ = "restructuredtext en"
import math
import sys
def to_networkx(tree):
"""Convert a Tree object to a networkx graph.
The result is useful for graph-oriented analysis, and also interactive
plotting with pylab, matplotlib or pygraphviz, though the resulting diagram
is usually not ideal for displaying a phylogeny.
Requires NetworkX version 0.99 or later.
"""
try:
import networkx
except ImportError:
from Bio import MissingPythonDependencyError
raise MissingPythonDependencyError(
"Install NetworkX if you want to use to_networkx.")
# NB (1/2010): the networkx API stabilized at v.1.0
# 1.0+: edges accept arbitrary data as kwargs, weights are floats
# 0.99: edges accept weight as a string, nothing else
# pre-0.99: edges accept no additional data
# Ubuntu Lucid LTS uses v0.99, let's support everything
if networkx.__version__ >= '1.0':
def add_edge(graph, n1, n2):
graph.add_edge(n1, n2, weight=n2.branch_length or 1.0)
# Copy branch color value as hex, if available
if hasattr(n2, 'color') and n2.color is not None:
graph[n1][n2]['color'] = n2.color.to_hex()
elif hasattr(n1, 'color') and n1.color is not None:
# Cascading color attributes
graph[n1][n2]['color'] = n1.color.to_hex()
n2.color = n1.color
# Copy branch weight value (float) if available
if hasattr(n2, 'width') and n2.width is not None:
graph[n1][n2]['width'] = n2.width
elif hasattr(n1, 'width') and n1.width is not None:
# Cascading width attributes
graph[n1][n2]['width'] = n1.width
n2.width = n1.width
elif networkx.__version__ >= '0.99':
def add_edge(graph, n1, n2):
graph.add_edge(n1, n2, (n2.branch_length or 1.0))
else:
def add_edge(graph, n1, n2):
graph.add_edge(n1, n2)
def build_subgraph(graph, top):
"""Walk down the Tree, building graphs, edges and nodes."""
for clade in top:
graph.add_node(clade.root)
add_edge(graph, top.root, clade.root)
build_subgraph(graph, clade)
if tree.rooted:
G = networkx.DiGraph()
else:
G = networkx.Graph()
G.add_node(tree.root)
build_subgraph(G, tree.root)
return G
def draw_graphviz(tree, label_func=str, prog='twopi', args='',
node_color='#c0deff', **kwargs):
"""Display a tree or clade as a graph, using the graphviz engine.
Requires NetworkX, matplotlib, Graphviz and either PyGraphviz or pydot.
The third and fourth parameters apply to Graphviz, and the remaining
arbitrary keyword arguments are passed directly to networkx.draw(), which
in turn mostly wraps matplotlib/pylab. See the documentation for Graphviz
and networkx for detailed explanations.
The NetworkX/matplotlib parameters are described in the docstrings for
networkx.draw() and pylab.scatter(), but the most reasonable options to try
are: *alpha, node_color, node_size, node_shape, edge_color, style,
font_size, font_color, font_weight, font_family*
:Parameters:
label_func : callable
A function to extract a label from a node. By default this is str(),
but you can use a different function to select another string
associated with each node. If this function returns None for a node,
no label will be shown for that node.
The label will also be silently skipped if the throws an exception
related to ordinary attribute access (LookupError, AttributeError,
ValueError); all other exception types will still be raised. This
means you can use a lambda expression that simply attempts to look
up the desired value without checking if the intermediate attributes
are available:
>>> Phylo.draw_graphviz(tree, lambda n: n.taxonomies[0].code)
prog : string
The Graphviz program to use when rendering the graph. 'twopi'
behaves the best for large graphs, reliably avoiding crossing edges,
but for moderate graphs 'neato' looks a bit nicer. For small
directed graphs, 'dot' may produce a normal-looking cladogram, but
will cross and distort edges in larger graphs. (The programs 'circo'
and 'fdp' are not recommended.)
args : string
Options passed to the external graphviz program. Normally not
needed, but offered here for completeness.
Example
-------
>>> import pylab
>>> from Bio import Phylo
>>> tree = Phylo.read('ex/apaf.xml', 'phyloxml')
>>> Phylo.draw_graphviz(tree)
>>> pylab.show()
>>> pylab.savefig('apaf.png')
"""
try:
import networkx
except ImportError:
from Bio import MissingPythonDependencyError
raise MissingPythonDependencyError(
"Install NetworkX if you want to use to_networkx.")
G = to_networkx(tree)
try:
# NetworkX version 1.8 or later (2013-01-20)
Gi = networkx.convert_node_labels_to_integers(G,
label_attribute='label')
int_labels = {}
for integer, nodeattrs in Gi.node.items():
int_labels[nodeattrs['label']] = integer
except TypeError:
# Older NetworkX versions (before 1.8)
Gi = networkx.convert_node_labels_to_integers(G,
discard_old_labels=False)
int_labels = Gi.node_labels
try:
posi = networkx.graphviz_layout(Gi, prog, args=args)
except ImportError:
raise MissingPythonDependencyError(
"Install PyGraphviz or pydot if you want to use draw_graphviz.")
def get_label_mapping(G, selection):
"""Apply the user-specified node relabeling."""
for node in G.nodes():
if (selection is None) or (node in selection):
try:
label = label_func(node)
if label not in (None, node.__class__.__name__):
yield (node, label)
except (LookupError, AttributeError, ValueError):
pass
if 'nodelist' in kwargs:
labels = dict(get_label_mapping(G, set(kwargs['nodelist'])))
else:
labels = dict(get_label_mapping(G, None))
kwargs['nodelist'] = list(labels.keys())
if 'edge_color' not in kwargs:
kwargs['edge_color'] = [isinstance(e[2], dict) and
e[2].get('color', 'k') or 'k'
for e in G.edges(data=True)]
if 'width' not in kwargs:
kwargs['width'] = [isinstance(e[2], dict) and
e[2].get('width', 1.0) or 1.0
for e in G.edges(data=True)]
posn = dict((n, posi[int_labels[n]]) for n in G)
networkx.draw(G, posn, labels=labels, node_color=node_color, **kwargs)
def draw_ascii(tree, file=sys.stdout, column_width=80):
"""Draw an ascii-art phylogram of the given tree.
The printed result looks like::
_________ Orange
______________|
| |______________ Tangerine
______________|
| | _________________________ Grapefruit
_| |_________|
| |______________ Pummelo
|
|__________________________________ Apple
:Parameters:
file : file-like object
File handle opened for writing the output drawing.
column_width : int
Total number of text columns used by the drawing.
"""
taxa = tree.get_terminals()
# Some constants for the drawing calculations
max_label_width = max(len(str(taxon)) for taxon in taxa)
drawing_width = column_width - max_label_width - 1
drawing_height = 2 * len(taxa) - 1
def get_col_positions(tree):
"""Create a mapping of each clade to its column position."""
depths = tree.depths()
# If there are no branch lengths, assume unit branch lengths
if not max(depths.values()):
depths = tree.depths(unit_branch_lengths=True)
# Potential drawing overflow due to rounding -- 1 char per tree layer
fudge_margin = int(math.ceil(math.log(len(taxa), 2)))
cols_per_branch_unit = ((drawing_width - fudge_margin)
/ float(max(depths.values())))
return dict((clade, int(round(blen*cols_per_branch_unit + 0.5)))
for clade, blen in depths.items())
def get_row_positions(tree):
positions = dict((taxon, 2*idx) for idx, taxon in enumerate(taxa))
def calc_row(clade):
for subclade in clade:
if subclade not in positions:
calc_row(subclade)
positions[clade] = ((positions[clade.clades[0]] +
positions[clade.clades[-1]]) // 2)
calc_row(tree.root)
return positions
col_positions = get_col_positions(tree)
row_positions = get_row_positions(tree)
char_matrix = [[' ' for x in range(drawing_width)]
for y in range(drawing_height)]
def draw_clade(clade, startcol):
thiscol = col_positions[clade]
thisrow = row_positions[clade]
# Draw a horizontal line
for col in range(startcol, thiscol):
char_matrix[thisrow][col] = '_'
if clade.clades:
# Draw a vertical line
toprow = row_positions[clade.clades[0]]
botrow = row_positions[clade.clades[-1]]
for row in range(toprow+1, botrow+1):
char_matrix[row][thiscol] = '|'
# NB: Short terminal branches need something to stop rstrip()
if (col_positions[clade.clades[0]] - thiscol) < 2:
char_matrix[toprow][thiscol] = ','
# Draw descendents
for child in clade:
draw_clade(child, thiscol+1)
draw_clade(tree.root, 0)
# Print the complete drawing
for idx, row in enumerate(char_matrix):
line = ''.join(row).rstrip()
# Add labels for terminal taxa in the right margin
if idx % 2 == 0:
line += ' ' + str(taxa[idx//2])
file.write(line + '\n')
file.write('\n')
def draw(tree, label_func=str, do_show=True, show_confidence=True,
# For power users
axes=None, branch_labels=None, *args, **kwargs):
"""Plot the given tree using matplotlib (or pylab).
The graphic is a rooted tree, drawn with roughly the same algorithm as
draw_ascii.
Additional keyword arguments passed into this function are used as pyplot
options. The input format should be in the form of:
pyplot_option_name=(tuple), pyplot_option_name=(tuple, dict), or
pyplot_option_name=(dict).
Example using the pyplot options 'axhspan' and 'axvline':
>>> Phylo.draw(tree, axhspan=((0.25, 7.75), {'facecolor':'0.5'}),
... axvline={'x':'0', 'ymin':'0', 'ymax':'1'})
Visual aspects of the plot can also be modified using pyplot's own functions
and objects (via pylab or matplotlib). In particular, the pyplot.rcParams
object can be used to scale the font size (rcParams["font.size"]) and line
width (rcParams["lines.linewidth"]).
:Parameters:
label_func : callable
A function to extract a label from a node. By default this is str(),
but you can use a different function to select another string
associated with each node. If this function returns None for a node,
no label will be shown for that node.
do_show : bool
Whether to show() the plot automatically.
show_confidence : bool
Whether to display confidence values, if present on the tree.
axes : matplotlib/pylab axes
If a valid matplotlib.axes.Axes instance, the phylogram is plotted
in that Axes. By default (None), a new figure is created.
branch_labels : dict or callable
A mapping of each clade to the label that will be shown along the
branch leading to it. By default this is the confidence value(s) of
the clade, taken from the ``confidence`` attribute, and can be
easily toggled off with this function's ``show_confidence`` option.
But if you would like to alter the formatting of confidence values,
or label the branches with something other than confidence, then use
this option.
"""
try:
import matplotlib.pyplot as plt
except ImportError:
try:
import pylab as plt
except ImportError:
from Bio import MissingPythonDependencyError
raise MissingPythonDependencyError(
"Install matplotlib or pylab if you want to use draw.")
import matplotlib.collections as mpcollections
# Arrays that store lines for the plot of clades
horizontal_linecollections = []
vertical_linecollections = []
# Options for displaying branch labels / confidence
def conf2str(conf):
if int(conf) == conf:
return str(int(conf))
return str(conf)
if not branch_labels:
if show_confidence:
def format_branch_label(clade):
if hasattr(clade, 'confidences'):
# phyloXML supports multiple confidences
return '/'.join(conf2str(cnf.value)
for cnf in clade.confidences)
if clade.confidence:
return conf2str(clade.confidence)
return None
else:
def format_branch_label(clade):
return None
elif isinstance(branch_labels, dict):
def format_branch_label(clade):
return branch_labels.get(clade)
else:
assert callable(branch_labels), \
"branch_labels must be either a dict or a callable (function)"
format_branch_label = branch_labels
# Layout
def get_x_positions(tree):
"""Create a mapping of each clade to its horizontal position.
Dict of {clade: x-coord}
"""
depths = tree.depths()
# If there are no branch lengths, assume unit branch lengths
if not max(depths.values()):
depths = tree.depths(unit_branch_lengths=True)
return depths
def get_y_positions(tree):
"""Create a mapping of each clade to its vertical position.
Dict of {clade: y-coord}.
Coordinates are negative, and integers for tips.
"""
maxheight = tree.count_terminals()
# Rows are defined by the tips
heights = dict((tip, maxheight - i)
for i, tip in enumerate(reversed(tree.get_terminals())))
# Internal nodes: place at midpoint of children
def calc_row(clade):
for subclade in clade:
if subclade not in heights:
calc_row(subclade)
# Closure over heights
heights[clade] = (heights[clade.clades[0]] +
heights[clade.clades[-1]]) / 2.0
if tree.root.clades:
calc_row(tree.root)
return heights
x_posns = get_x_positions(tree)
y_posns = get_y_positions(tree)
# The function draw_clade closes over the axes object
if axes is None:
fig = plt.figure()
axes = fig.add_subplot(1, 1, 1)
elif not isinstance(axes, plt.matplotlib.axes.Axes):
raise ValueError("Invalid argument for axes: %s" % axes)
def draw_clade_lines(use_linecollection=False, orientation='horizontal',
y_here=0, x_start=0, x_here=0, y_bot=0, y_top=0,
color='black', lw='.1'):
"""Create a line with or without a line collection object.
Graphical formatting of the lines representing clades in the plot can be
customized by altering this function.
"""
if (use_linecollection==False and orientation=='horizontal'):
axes.hlines(y_here, x_start, x_here, color=color, lw=lw)
elif (use_linecollection==True and orientation=='horizontal'):
horizontal_linecollections.append(mpcollections.LineCollection(
[[(x_start, y_here), (x_here, y_here)]], color=color, lw=lw),)
elif (use_linecollection==False and orientation=='vertical'):
axes.vlines(x_here, y_bot, y_top, color=color)
elif (use_linecollection==True and orientation=='vertical'):
vertical_linecollections.append(mpcollections.LineCollection(
[[(x_here, y_bot), (x_here, y_top)]], color=color, lw=lw),)
def draw_clade(clade, x_start, color, lw):
"""Recursively draw a tree, down from the given clade."""
x_here = x_posns[clade]
y_here = y_posns[clade]
# phyloXML-only graphics annotations
if hasattr(clade, 'color') and clade.color is not None:
color = clade.color.to_hex()
if hasattr(clade, 'width') and clade.width is not None:
lw = clade.width * plt.rcParams['lines.linewidth']
# Draw a horizontal line from start to here
draw_clade_lines(use_linecollection=True, orientation='horizontal',
y_here=y_here, x_start=x_start, x_here=x_here, color=color, lw=lw)
# Add node/taxon labels
label = label_func(clade)
if label not in (None, clade.__class__.__name__):
axes.text(x_here, y_here, ' %s' % label, verticalalignment='center')
# Add label above the branch (optional)
conf_label = format_branch_label(clade)
if conf_label:
axes.text(0.5*(x_start + x_here), y_here, conf_label,
fontsize='small', horizontalalignment='center')
if clade.clades:
# Draw a vertical line connecting all children
y_top = y_posns[clade.clades[0]]
y_bot = y_posns[clade.clades[-1]]
# Only apply widths to horizontal lines, like Archaeopteryx
draw_clade_lines(use_linecollection=True, orientation='vertical',
x_here=x_here, y_bot=y_bot, y_top=y_top, color=color, lw=lw)
# Draw descendents
for child in clade:
draw_clade(child, x_here, color, lw)
draw_clade(tree.root, 0, 'k', plt.rcParams['lines.linewidth'])
# If line collections were used to create clade lines, here they are added
# to the pyplot plot.
for i in horizontal_linecollections:
axes.add_collection(i)
for i in vertical_linecollections:
axes.add_collection(i)
# Aesthetics
if hasattr(tree, 'name') and tree.name:
axes.set_title(tree.name)
axes.set_xlabel('branch length')
axes.set_ylabel('taxa')
# Add margins around the tree to prevent overlapping the axes
xmax = max(x_posns.values())
axes.set_xlim(-0.05 * xmax, 1.25 * xmax)
# Also invert the y-axis (origin at the top)
# Add a small vertical margin, but avoid including 0 and N+1 on the y axis
axes.set_ylim(max(y_posns.values()) + 0.8, 0.2)
# Parse and process key word arguments as pyplot options
for key, value in kwargs.items():
try:
# Check that the pyplot option input is iterable, as required
[i for i in value]
except TypeError:
raise ValueError('Keyword argument "%s=%s" is not in the format '
'pyplot_option_name=(tuple), pyplot_option_name=(tuple, dict),'
' or pyplot_option_name=(dict) '
% (key, value))
if isinstance(value, dict):
getattr(plt, str(key))(**dict(value))
elif not (isinstance(value[0], tuple)):
getattr(plt, str(key))(*value)
elif (isinstance(value[0], tuple)):
getattr(plt, str(key))(*value[0], **dict(value[1]))
if do_show:
plt.show()