← Back to branch summary

~ubuntu-branches/ubuntu/natty/python-cogent/natty

« back to all changes in this revision

Viewing changes to cogent/draw/dendrogram.py

  • Committer: Bazaar Package Importer
  • Author(s): Steffen Moeller
  • Date: 2010-12-04 22:30:35 UTC
  • mfrom: (1.1.1 upstream)
  • Revision ID: james.westby@ubuntu.com-20101204223035-j11kinhcrrdgg2p2
Tags: 1.5-1
* Bumped standard to 3.9.1, no changes required.
* New upstream version.
  - major additions to Cookbook
  - added AlleleFreqs attribute to ensembl Variation objects.
  - added getGeneByStableId method to genome objects.
  - added Introns attribute to Transcript objects and an Intron class.
  - added Mann-Whitney test and a Monte-Carlo version
  - exploratory and confirmatory period estimation techniques (suitable for
    symbolic and continuous data)
  - Information theoretic measures (AIC and BIC) added
  - drawing of trees with collapsed nodes
  - progress display indicator support for terminal and GUI apps
  - added parser for illumina HiSeq2000 and GAiix sequence files as 
    cogent.parse.illumina_sequence.MinimalIlluminaSequenceParser.
  - added parser to FASTQ files, one of the output options for illumina's
    workflow, also added cookbook demo.
  - added functionality for parsing of SFF files without the Roche tools in
    cogent.parse.binary_sff
  - thousand fold performance improvement to nmds
  - >10-fold performance improvements to some Table operations

Show diffs side-by-side

added added

removed removed

Lines of Context:
cogent/draw/dendrogram.py
22
22
#  - orientation switch
23
23
# Layout gets more complicated for rooted tree styles if dy is allowed to vary,
24
24
# and constant-y is suitable for placing alongside a sequence alignment anyway.
 
25
from cogent.core.tree import TreeNode
25
26
import rlg2mpl
26
27
import matplotlib.colors
27
 
from matplotlib.patches import PathPatch
 
28
from matplotlib.patches import PathPatch, Polygon
28
29
from matplotlib.path import Path
29
30
from matplotlib.text import Text
30
31
import numpy
34
35
__credits__ = ["Peter Maxwell", "Gavin Huttley", "Rob Knight",
35
36
                    "Zongzhi Liu", "Daniel McDonald"]
36
37
__license__ = "GPL"
37
 
__version__ = "1.4.1"
 
38
__version__ = "1.5.0"
38
39
__maintainer__ = "Peter Maxwell"
39
40
__email__ = "pm67nz@gmail.com"
40
41
__status__ = "Production"
41
42
 
 
43
to_rgb = matplotlib.colors.colorConverter.to_rgb
 
44
 
42
45
def _sign(x):
43
46
    """Returns True if x is positive, False otherwise."""
44
47
    return x and x/abs(x)
45
48
 
46
 
def _treemaxof(param):
47
 
    """Returns maximum value of a parameter in a tree."""
48
 
    def _max(node, child_results):
49
 
        params = node.edge.params
50
 
        if param in params:
51
 
            return max([params[param]] + child_results)
52
 
        elif child_results:
53
 
            return max(child_results)
54
 
        else:
55
 
            return None
56
 
    return _max
57
 
 
 
49
def _first_non_none(values):
 
50
    for item in values:
 
51
        if item is not None:
 
52
            return item
 
53
            
58
54
def SimpleColormap(color0, color1, name=None):
59
55
    """Linear interpolation between any two colours"""
60
 
    to_rgb = matplotlib.colors.colorConverter.to_rgb
61
56
    c0 = to_rgb(color0)
62
57
    c1 = to_rgb(color1)
63
58
    cn = ['red', 'green', 'blue']
123
118
        set to 0 (False) or 1 (True) to highlight the mammals.
124
119
    
125
120
    4.  You have f(node) -> color. Pass in f as edge_color_callback.
 
121
    
 
122
    Alternatively, set the Color attribute of the dendrogram edges.
126
123
    """
127
124
    if callback_returns_name is not None:
128
125
        pass # give deprecation warning?, no longer needed
129
 
        
 
126
    
 
127
    edge_color = to_rgb(edge_color)
 
128
    highlight_color = to_rgb(highlight_color)
 
129
    
130
130
    if edge_color_callback is not None:
131
 
        return edge_color_callback
 
131
        return lambda edge:edge_color_callback(edge)
132
132
    elif shade_param:
133
133
        if cmap is None:
134
134
            cmap = SimpleColormap(edge_color, highlight_color)
135
 
        return ScalarColormapShading(shade_param, min_value, max_value, cmap)
 
135
        return ScalarColormapShading(
 
136
                shade_param, min_value, max_value, cmap)
136
137
    else:
137
138
        return lambda edge:edge_color
138
 
    
 
139
        
139
140
    
140
141
class MatplotlibRenderer(object):
141
142
    """Returns a matplitlib render including font size, stroke width, etc.
145
146
    stroke width is not yet implemented by should be.
146
147
    """
147
148
    def __init__(self, font_size=None, stroke_width=3, label_pad=None, **kw):
148
 
        self.edge_color = makeColorCallback(**kw)
 
149
        self.calculated_edge_color = makeColorCallback(**kw)
149
150
        self.text_opts = {}
150
151
        if font_size is not None:
151
152
            self.text_opts['fontsize'] = font_size
153
154
        if stroke_width is not None:
154
155
            self.line_opts['linewidth'] = stroke_width
155
156
        if label_pad is None:
156
 
            label_pad = self.stringWidth(' ')
 
157
            label_pad = 8
157
158
        self.labelPadDistance = label_pad
158
159
    
 
160
    def edge_color(self, edge):
 
161
        if edge.Color is None:
 
162
            return self.calculated_edge_color(edge.original)
 
163
        else:
 
164
            return edge.Color
 
165
        
159
166
    def line(self, x1, y1, x2, y2, edge=None):
160
 
        path = Path([(x1, y1), (x2, y2)], [Path.MOVETO, Path.LINETO])
161
167
        opts = self.line_opts.copy()
162
168
        if edge is not None:
163
169
            opts['edgecolor'] = self.edge_color(edge)
 
170
        path = Path([(x1, y1), (x2, y2)], [Path.MOVETO, Path.LINETO])
164
171
        return PathPatch(path,  **opts)
165
172
    
166
 
    def string(self, x, y, string, ha=None, va=None, rotation=None):
 
173
    def polygon(self, vertices, color):
 
174
        opts = self.line_opts.copy()
 
175
        opts['color'] = color
 
176
        return Polygon(vertices, **opts)
 
177
        
 
178
    def string(self, x, y, string, ha=None, va=None, rotation=None, color=None):
167
179
        opts = self.text_opts.copy()
168
180
        if ha is not None:
169
181
            opts['ha'] = ha
171
183
            opts['va'] = va
172
184
        if rotation is not None:
173
185
            opts['rotation'] = rotation
 
186
        if color is not None:
 
187
            opts['color'] = color
174
188
        return Text(x, y, string, **opts)
175
 
    
176
 
    def stringWidth(self, text):
177
 
        #rlg2mpl.String(0, 0, text, **self.text_opts).getEast()
178
 
        return self.text_opts.get('fontSize',10) * .8 * len(text) # not very accurate!
179
 
    
180
 
    def stringHeight(self, text):
181
 
        return self.text_opts.get('fontSize', 10) * len(text.split('\n'))
182
 
    
 
189
        
183
190
class DendrogramLabelStyle(object):
184
191
    """Label options"""
185
192
    
215
222
        return self.edgeLabelCallback(edge)
216
223
    
217
224
    def getNodeLabel(self, edge):
218
 
        if self.showInternalLabels or not edge.Children:
 
225
        if edge.Name is not None:
219
226
            return edge.Name
 
227
        elif self.showInternalLabels or not edge.Children:
 
228
            return edge.original.Name
220
229
        else:
221
230
            return ""
222
 
        
223
 
class _Dendrogram(rlg2mpl.Drawable):
224
 
    # One of these for each tree edge.  Attributes:
225
 
    #    edge - the real tree node
226
 
    #    children - Dendrograms of self.edge.Children
 
231
 
 
232
def ValidColorProperty(real_name, doc='A color name or other spec'):
 
233
    """Can only be set to Null or a valid color"""
 
234
    def getter(obj):
 
235
        return getattr(obj, real_name, None)
 
236
    def setter(obj, value):
 
237
        if value is not None: to_rgb(value)
 
238
        setattr(obj, real_name, value)
 
239
    def deleter(obj):
 
240
        setattr(obj, real_name, None)
 
241
    return property(getter, setter, deleter, doc)
 
242
 
 
243
class _Dendrogram(rlg2mpl.Drawable, TreeNode):
 
244
    # One of these for each tree edge.  Extra attributes:
227
245
    #    depth - distance from root to bottom of edge
228
246
    #    height - max distance from a decendant leaf to top of edge
229
247
    #    width - number of decendant leaves
237
255
    # code reuse - vertical drawing, new tree styles, new graphics
238
256
    # libraries etc.
239
257
    
240
 
    def __init__(self, edge):
241
 
        if hasattr(edge, 'TreeRoot'):
242
 
            edge = edge.TreeRoot
243
 
        self.edge = edge
244
 
        self.children = [self.__class__(c) for c in edge.Children]
245
 
        self._up_to_date = False
 
258
    aspect_distorts_lengths = True
 
259
 
 
260
    def __init__(self, edge, use_lengths=True):
 
261
        children = [type(self)(child) for child in edge.Children]
 
262
        TreeNode.__init__(self, Params=edge.params.copy(), Children=children, 
 
263
            Name=("" if children else edge.Name))
 
264
        self.Length = edge.Length
 
265
        self.original = edge  # for edge_color_callback
 
266
        self.Collapsed = False
 
267
        self.use_lengths_default = use_lengths
 
268
    
 
269
    # Colors are properties so that invalid color names are caught immediately
 
270
    Color = ValidColorProperty('_Color', 'Color of line segment')
 
271
    NameColor = ValidColorProperty('_NameColor', 'Color of node name')
 
272
    CladeColor = ValidColorProperty('_CladeColor', 'Color of collapsed descendants')
246
273
    
247
274
    def __repr__(self):
248
275
        return '%s %s %s %s' % (
249
 
                self.depth, self.length, self.height, self.children)
250
 
    
251
 
    def postorder(self, f):
252
 
        rs = [child.postorder(f) for child in self.children]
253
 
        return f(self, rs)
254
 
    
255
 
    def updateGeometry(self, renderer, use_lengths=True,
256
 
            depth=None, track_coordinates=None):
 
276
                self.depth, self.length, self.height, self.Children)
 
277
    
 
278
    def updateGeometry(self, use_lengths, depth=None, track_coordinates=None):
257
279
        """Calculate tree node attributes such as height and depth.
258
280
        Despite the name this first pass is ignorant of issues like
259
281
        scale and orientation"""
260
282
        
261
 
        if self.edge.Length is None or not use_lengths:
 
283
        if self.Length is None or not use_lengths:
262
284
            if depth is None:
263
285
                self.length = 0
264
286
            else:
265
287
                self.length = 1
266
288
        else:
267
 
            self.length = self.edge.Length
 
289
            self.length = self.Length
268
290
        
269
291
        self.depth = (depth or 0) + self.length
270
292
        
271
 
        if self.children:
272
 
            for c in self.children:
273
 
                c.updateGeometry(renderer, use_lengths, self.depth, track_coordinates)
274
 
            self.height = max([c.height for c in self.children]) + self.length
275
 
            self.leafcount  = sum([c.leafcount for c in self.children])
276
 
            self.edgecount  = sum([c.edgecount for c in self.children]) + 1
277
 
            self.labelwidth = max([c.labelwidth for c in self.children])
 
293
        children = self.Children
 
294
        if children:
 
295
            for c in children:
 
296
                c.updateGeometry(use_lengths, self.depth, track_coordinates)
 
297
            self.height = max([c.height for c in children]) + self.length
 
298
            self.leafcount  = sum([c.leafcount for c in children])
 
299
            self.edgecount  = sum([c.edgecount for c in children]) + 1
 
300
            self.longest_label = max([c.longest_label for c in children],
 
301
                    key=len)
278
302
        else:
279
303
            self.height = self.length
280
304
            self.leafcount = self.edgecount = 1
281
 
            node_label = self.edge.Name or ''
282
 
            self.labelwidth = renderer.stringWidth(node_label)
 
305
            self.longest_label = self.Name or ''
283
306
        
284
 
        if track_coordinates is not None and self.edge.Name != "root":
285
 
            self.track_y = track_coordinates[self.edge.Name]
 
307
        if track_coordinates is not None and self.Name != "root":
 
308
            self.track_y = track_coordinates[self.Name]
286
309
        else:
287
310
            self.track_y = 0
288
311
 
290
313
        """Return list of [node_name, node_id, x, y, child_ids]"""
291
314
        self.asArtist(height, width)
292
315
        result = []
293
 
        def _f(node, child_results):
294
 
            result.append([node.edge.Name, id(node), node.x2, node.y2] + [map(id, node.children)])
295
 
        self.postorder(_f)
 
316
        for node in self.postorder(include_self=True):
 
317
            result.append([node.Name, id(node), node.x2, node.y2] + [map(id, node.Children)])
296
318
        return result
297
319
    
298
320
    def makeFigure(self, width=None, height=None, margin=.25, use_lengths=None, **kw):
308
330
        ax.set_yticks([])
309
331
        if use_lengths is None:
310
332
            use_lengths = self.use_lengths_default
 
333
        else:
 
334
            pass # deprecate setting use_lengths here?
311
335
        if use_lengths and self.aspect_distorts_lengths:
312
336
            ax.set_aspect('equal')
313
337
        g = self.asArtist(width, height, use_lengths=use_lengths, 
317
341
    
318
342
    def asArtist(self, width, height, margin=20, use_lengths=None,
319
343
            scale_bar="left", show_params=None, show_internal_labels=False,
320
 
            label_template=None, edge_label_callback=None, **kw):
321
 
        """A reportlab drawing"""
322
 
        
323
 
        label_style = DendrogramLabelStyle(
324
 
                show_params = show_params,
325
 
                show_internal_labels = show_internal_labels,
326
 
                label_template = label_template,
327
 
                edge_label_callback = edge_label_callback,
328
 
                )
329
 
        if kw.get('shade_param', None) is not None and \
330
 
                kw.get('max_value', None) is None:
331
 
            kw['max_value'] = self.postorder(_treemaxof(kw['shade_param']))
332
 
        renderer = MatplotlibRenderer(**kw)
333
 
        self.updateGeometry(renderer, use_lengths=use_lengths)
 
344
            label_template=None, edge_label_callback=None, shade_param=None, 
 
345
            max_value=None, font_size=None, **kw):
 
346
        
 
347
        if use_lengths is None:
 
348
            use_lengths = self.use_lengths_default
 
349
        self.updateGeometry(use_lengths=use_lengths)
 
350
        
334
351
        if width <= 2 * margin:
335
352
            raise ValueError('%spt not wide enough for %spt margins' %
336
353
                    (width, margin))
339
356
                    (height, margin))
340
357
        width -= 2 * margin
341
358
        height -= 2 * margin
342
 
        scale = self.updateCoordinates(width, height)
 
359
 
 
360
        label_length = len(self.longest_label)
 
361
        label_width = label_length * 0.8 * (font_size or 10) # not very accurate
 
362
        (left_labels, right_labels) = self.labelMargins(label_width)
 
363
        total_label_width = left_labels + right_labels
 
364
        if width < total_label_width:
 
365
            raise ValueError('%spt not wide enough for ""%s"' %
 
366
                    (width, self.longest_label))
 
367
    
 
368
        scale = self.updateCoordinates(width-total_label_width, height)
 
369
 
 
370
        if shade_param is not None and max_value is None:
 
371
            for edge in self.postorder(include_self=True):
 
372
                sp = edge.params.get(shade_param, None)
 
373
                if max_value is None or sp > max_value:
 
374
                    max_value = sp
 
375
        renderer = MatplotlibRenderer(shade_param=shade_param, 
 
376
                max_value=max_value, font_size=font_size, **kw)
 
377
 
 
378
        labelopts = {}
 
379
        for labelopt in ['show_params', 'show_internal_labels', 
 
380
                'label_template', 'edge_label_callback']:
 
381
            labelopts[labelopt] = locals()[labelopt]
 
382
        label_style = DendrogramLabelStyle(**labelopts)
 
383
 
343
384
        ss = self._draw(renderer, label_style)
344
385
        if use_lengths:
345
386
            # Placing the scale properly might take some work,
348
389
            if scale_bar == "right":
349
390
                x1, x2 = (width-scale*unit, width)
350
391
            elif scale_bar == "left":
351
 
                x1, x2 = (0, scale*unit)
 
392
                x1, x2 = (-left_labels, scale*unit-left_labels)
352
393
            else:
353
394
                assert not scale_bar, scale_bar
354
395
            if scale_bar:
356
397
                ss.append(renderer.string((x1+x2)/2, 5, str(unit), va='bottom', ha='center'))
357
398
        
358
399
        g = rlg2mpl.Group(*ss)
359
 
        g.translate(margin, margin)
 
400
        g.translate(margin+left_labels, margin)
360
401
        return g
361
402
    
362
403
    def _draw(self, renderer, label_style):
363
404
        g = []
364
405
        g += self._draw_edge(renderer, label_style)
365
 
        for child in self.children:
366
 
            g += child._draw(renderer, label_style)
367
 
        g += self._draw_node_label(renderer, label_style)
 
406
        if self.Collapsed:
 
407
            g += self._draw_collapsed_clade(renderer, label_style)
 
408
        else:
 
409
            g += self._draw_node(renderer, label_style)
 
410
            for child in self.Children:
 
411
                g += child._draw(renderer, label_style)
 
412
            g += self._draw_node_label(renderer, label_style)
368
413
        return g
369
414
    
370
 
    def _draw_edge(self, renderer, label_style):
 
415
    def _draw_node(self, renderer, label_style):
371
416
        g = []
372
417
        # Joining line for square form
373
 
        if self.children:
374
 
            cys = [c.y1 for c in self.children] + [self.y2]
 
418
        if self.Children:
 
419
            cys = [c.y1 for c in self.Children] + [self.y2]
375
420
            if max(cys) > min(cys):
376
 
                g.append(renderer.line(self.x2, min(cys), self.x2, max(cys), self.edge))
377
 
        
 
421
                g.append(renderer.line(self.x2, min(cys), self.x2, max(cys), self))
 
422
        return g
 
423
    
 
424
    def _draw_edge(self, renderer, label_style):
 
425
        g = []
378
426
        if ((self.x1, self.y1) == (self.x2, self.y2)):
379
427
            # avoid labeling zero length line, eg: root
380
428
            return g
381
 
            
 
429
 
382
430
        # Main line
383
 
        g.append(renderer.line(self.x1, self.y1, self.x2, self.y2, self.edge))
 
431
        g.append(renderer.line(self.x1, self.y1, self.x2, self.y2, self))
384
432
        
385
433
        # Edge Label
386
 
        text = label_style.getEdgeLabel(self.edge)
 
434
        text = label_style.getEdgeLabel(self)
387
435
        if text:
388
436
            midx, midy = (self.x1+self.x2)/2, (self.y1+self.y2)/2
389
437
            if self.x1 == self.x2:
397
445
        return g
398
446
    
399
447
    def _draw_node_label(self, renderer, label_style):
400
 
        text = label_style.getNodeLabel(self.edge)
401
 
        (x, ha, y, va) = self.getLabelCoordinates(text, renderer)
402
 
        return [renderer.string(x, y, text, ha=ha, va=va)]
403
 
    
404
 
 
405
 
class _RootedDendrogramForm(object):
406
 
    """A rooted dendrogram form defines how lengths get mapped to X and Y coodinates.
407
 
    
408
 
    _RootedDendrogramStyle subclasses provide yCoords and xCoords, which examine
 
448
        text = label_style.getNodeLabel(self)
 
449
        color = self.NameColor
 
450
        (x, ha, y, va) = self.getLabelCoordinates(text, renderer)
 
451
        return [renderer.string(x, y, text, ha=ha, va=va, color=color)]
 
452
        
 
453
    def _draw_collapsed_clade(self, renderer, label_style):
 
454
        text = label_style.getNodeLabel(self)
 
455
        color = _first_non_none([self.CladeColor, self.Color, 'black'])
 
456
        icolor = 'white' if sum(to_rgb(color))/3 < 0.5 else 'black'
 
457
        g = []
 
458
        if not self.Children:
 
459
            return g
 
460
        (l,r,t,b), vertices = self.wedgeVertices()
 
461
        g.append(renderer.polygon(vertices, color))
 
462
        if not b <= self.y2 <= t:
 
463
            # ShelvedDendrogram needs this extra line segment
 
464
            g.append(renderer.line(self.x2, self.y2, self.x2, b, self))
 
465
        (x, ha, y, va) = self.getLabelCoordinates(text, renderer)
 
466
        g.append(renderer.string(
 
467
                (self.x2+r)/2, (t+b)/2, str(self.leafcount), ha=ha, va=va,
 
468
                color=icolor))
 
469
        g.append(renderer.string(
 
470
                x-self.x2+r, y, text, ha=ha, va=va, color=self.NameColor))
 
471
        return g
 
472
    
 
473
    def setCollapsed(self, collapsed=True, label=None, color=None):
 
474
        if color is not None:
 
475
            self.CladeColor = color
 
476
        if label is not None:
 
477
            self.Name = label
 
478
        self.Collapsed = collapsed
 
479
 
 
480
 
 
481
class Dimensions(object):
 
482
    def __init__(self, xscale, yscale, total_tree_height):
 
483
        self.x = xscale
 
484
        self.y = yscale
 
485
        self.height = total_tree_height
 
486
    
 
487
 
 
488
class _RootedDendrogram(_Dendrogram):
 
489
    """_RootedDendrogram subclasses provide yCoords and xCoords, which examine
409
490
    attributes of a node (its length, coodinates of its children) and return
410
491
    a tuple for start/end of the line representing the edge."""
411
 
    aspect_distorts_lengths = True
412
 
    
413
 
    def __init__(self, tree, width, height):
414
 
        self.yscale = 1.0
415
 
        self.yscale = height / self.widthRequiredFor(tree)
416
 
        self.xscale = width / tree.height
417
 
        self.total_tree_height = tree.height
418
 
    
419
 
    def widthRequiredFor(self, node):
420
 
        return node.leafcount * self.yscale
421
 
    
422
 
    def xCoords(self, node, x1):
423
 
        raise NotImplementedError
424
 
    
425
 
    def yCoords(self, node, x1):
426
 
        raise NotImplementedError
427
 
    
 
492
    def labelMargins(self, label_width):
 
493
        return (0, label_width)
 
494
            
 
495
    def widthRequired(self):
 
496
        return self.leafcount
 
497
    
 
498
    def xCoords(self, scale, x1):
 
499
        raise NotImplementedError
 
500
    
 
501
    def yCoords(self, scale, y1):
 
502
        raise NotImplementedError
 
503
        
 
504
    def updateCoordinates(self, width, height):
 
505
        xscale = width / self.height
 
506
        yscale = height / self.widthRequired()
 
507
        scale = Dimensions(xscale, yscale, self.height)
 
508
        
 
509
        # y coords done postorder, x preorder, y first.
 
510
        # so it has to be done in 2 passes.
 
511
        self.update_y_coordinates(scale)
 
512
        self.update_x_coordinates(scale)
 
513
        return xscale
 
514
    
 
515
    def update_y_coordinates(self, scale, y1=None):
 
516
        """The second pass through the tree.  Y coordinates only
 
517
        depend on the shape of the tree and yscale"""
 
518
        if y1 is None:
 
519
            y1 = self.widthRequired() * scale.y
 
520
        child_y = y1
 
521
        for child in self.Children:
 
522
            child.update_y_coordinates(scale, child_y)
 
523
            child_y -= child.widthRequired() * scale.y
 
524
        (self.y1, self.y2) = self.yCoords(scale, y1)
 
525
    
 
526
    def update_x_coordinates(self, scale, x1=0):
 
527
        """For non 'square' styles the x coordinates will depend
 
528
        (a bit) on the y coodinates, so they should be done first"""
 
529
        (self.x1, self.x2) = self.xCoords(scale, x1)
 
530
        for child in self.Children:
 
531
            child.update_x_coordinates(scale, self.x2)
 
532
    
 
533
    def getLabelCoordinates(self, text, renderer):
 
534
        return (self.x2+renderer.labelPadDistance, 'left', self.y2, 'center')
428
535
 
429
 
class SquareDendrogramForm(_RootedDendrogramForm):
 
536
class SquareDendrogram(_RootedDendrogram):
430
537
    aspect_distorts_lengths = False
431
538
    
432
 
    def yCoords(self, node, y1):
433
 
        cys = [c.y1 for c in node.children]
 
539
    def yCoords(self, scale, y1):
 
540
        cys = [c.y1 for c in self.Children]
434
541
        if cys:
435
542
            y2 = (cys[0]+cys[-1]) / 2.0
436
543
        else:
437
 
            y2 = y1 - self.yscale / 2.0
 
544
            y2 = y1 - 0.5 * scale.y
438
545
        return (y2, y2)
439
546
    
440
 
    def xCoords(self, node, x1):
441
 
        dx = self.xscale * node.length
 
547
    def xCoords(self, scale, x1):
 
548
        dx = scale.x * self.length
442
549
        x2 = x1 + dx
443
550
        return (x1, x2)
444
 
    
445
 
 
446
 
class ContemporaneousDendrogramForm(SquareDendrogramForm):
447
 
    def xCoords(self, node, x1):
448
 
        return (x1, (self.total_tree_height-(node.height-node.length))*self.xscale)
449
 
    
450
 
 
451
 
class ShelvedDendrogramForm(ContemporaneousDendrogramForm):
452
 
    def widthRequiredFor(self, node):
453
 
        return node.edgecount * self.yscale
454
 
    
455
 
    def yCoords(self, node, y1):
456
 
        cys = [c.y1 for c in node.children]
457
 
        if cys:
458
 
            y2 = cys[-1] - 1.0 * self.yscale
459
 
        else:
460
 
            y2 = y1    - 0.5 * self.yscale
461
 
        return (y2, y2)
462
 
    
463
 
 
464
 
class AlignedShelvedDendrogramForm(ShelvedDendrogramForm):
465
 
    def __init__(self, tree, width, height):
466
 
        self.yscale = 1.0
467
 
        self.xscale = width / tree.height
468
 
        self.total_tree_height = tree.height
469
 
    
470
 
    def yCoords(self, node, y1):
471
 
        if hasattr(node, 'track_y'):
472
 
            return (node.track_y, node.track_y)
473
 
        else:
474
 
            raise RuntimeError, node.edge.Name
475
 
            return ShelvedDendrogramForm.yCoords(self, node, y1)
476
 
    
477
 
    def widthRequiredFor(self, node):
478
 
        raise RuntimeError, node.edge.Name
479
 
    
480
 
 
481
 
class StraightDendrogramForm(_RootedDendrogramForm):
482
 
    def yCoords(self, node, y1):
 
551
 
 
552
    def wedgeVertices(self):
 
553
        tip_ys = [(c.y2 + self.y2)/2 for c in self.iterTips()]
 
554
        t,b = max(tip_ys), min(tip_ys)
 
555
        cxs = [c.x2 for c in self.iterTips()]
 
556
        l,r = min(cxs), max(cxs)
 
557
        return (l,r,t,b), [(self.x2, b), (self.x2, t), (l, t), (r, b)]
 
558
 
 
559
 
 
560
class StraightDendrogram(_RootedDendrogram):
 
561
    def yCoords(self, scale, y1):
483
562
        # has a side effect of adjusting the child y1's to meet nodes' y2's
484
 
        cys = [c.y1 for c in node.children]
 
563
        cys = [c.y1 for c in self.Children]
485
564
        if cys:
486
565
            y2 = (cys[0]+cys[-1]) / 2.0
487
 
            distances = [child.length for child in node.children]
488
 
            closest_child = node.children[distances.index(min(distances))]
 
566
            distances = [child.length for child in self.Children]
 
567
            closest_child = self.Children[distances.index(min(distances))]
489
568
            dy = closest_child.y1 - y2
490
 
            max_dy = 0.8*max(5, closest_child.length*self.xscale)
 
569
            max_dy = 0.8*max(5, closest_child.length*scale.x)
491
570
            if abs(dy) > max_dy:
492
 
                # 'moved', node.edge.Name, y2, 'to within', max_dy,
493
 
                # 'of', closest_child.edge.Name, closest_child.y1
 
571
                # 'moved', node.Name, y2, 'to within', max_dy,
 
572
                # 'of', closest_child.Name, closest_child.y1
494
573
                y2 = closest_child.y1 - _sign(dy) * max_dy
495
574
        else:
496
 
            y2 = y1 - self.yscale / 2.0
 
575
            y2 = y1 - scale.y / 2.0
497
576
        y1 = y2
498
 
        for child in node.children:
 
577
        for child in self.Children:
499
578
            child.y1 = y2
500
579
        return (y1, y2)
501
580
    
502
 
    def xCoords(self, node, x1):
503
 
        dx = node.length * self.xscale
504
 
        dy = node.y2 - node.y1
 
581
    def xCoords(self, scale, x1):
 
582
        dx = self.length * scale.x
 
583
        dy = self.y2 - self.y1
505
584
        dx = numpy.sqrt(max(dx**2 - dy**2, 1))
506
585
        return (x1, x1 + dx)
507
 
    
508
 
 
509
 
class ContemporaneousStraightDendrogramForm(StraightDendrogramForm):
510
 
    def xCoords(self, node, x1):
511
 
        return (x1, (self.total_tree_height-(node.height-node.length))*self.xscale)
512
 
    
513
 
 
514
 
class _RootedDendrogram(_Dendrogram):
515
 
    def updateCoordinates(self, width, height):
516
 
        if width < self.labelwidth:
517
 
            raise ValueError('%spt not wide enough for %spt wide labels' %
518
 
                    (width, self.labelwidth))
519
 
        width -= self.labelwidth
520
 
        
521
 
        form = self.FormClass(self, width, height)
522
 
        
523
 
        # y coords done postorder, x preorder, y first.
524
 
        # so it has to be done in 2 passes.
525
 
        self.update_y_coordinates(form)
526
 
        self.update_x_coordinates(form)
527
 
        return form.xscale
528
 
    
529
 
    def update_y_coordinates(self, style, y1=None):
530
 
        """The second pass through the tree.  Y coordinates only
531
 
        depend on the shape of the tree and yscale"""
532
 
        if y1 is None:
533
 
            y1 = style.widthRequiredFor(self)
534
 
        child_y = y1
535
 
        for child in self.children:
536
 
            child.update_y_coordinates(style, child_y)
537
 
            child_y -= style.widthRequiredFor(child)
538
 
        (self.y1, self.y2) = style.yCoords(self, y1)
539
 
    
540
 
    def update_x_coordinates(self, style, x1=0):
541
 
        """For non 'square' styles the x coordinates will depend
542
 
        (a bit) on the y coodinates, so they should be done first"""
543
 
        (self.x1, self.x2) = style.xCoords(self, x1)
544
 
        for child in self.children:
545
 
            child.update_x_coordinates(style, self.x2)
546
 
    
547
 
    def getLabelCoordinates(self, text, renderer):
548
 
        return (self.x2+renderer.labelPadDistance, 'left', self.y2, 'center')
549
 
 
550
 
class SquareDendrogram(_RootedDendrogram):
551
 
    FormClass = SquareDendrogramForm
552
 
    aspect_distorts_lengths = FormClass.aspect_distorts_lengths
553
 
    use_lengths_default = True
554
 
 
555
 
class StraightDendrogram(_RootedDendrogram):
556
 
    FormClass = StraightDendrogramForm
557
 
    aspect_distorts_lengths = FormClass.aspect_distorts_lengths
558
 
    use_lengths_default = True
559
 
 
560
 
class ContemporaneousDendrogram(_RootedDendrogram):
561
 
    FormClass = ContemporaneousDendrogramForm
562
 
    aspect_distorts_lengths = FormClass.aspect_distorts_lengths
563
 
    use_lengths_default = False
564
 
 
565
 
class ShelvedDendrogram(_RootedDendrogram):
566
 
    FormClass = ShelvedDendrogramForm
567
 
    aspect_distorts_lengths = FormClass.aspect_distorts_lengths
568
 
    use_lengths_default = False
569
 
 
570
 
class AlignedShelvedDendrogram(_RootedDendrogram):
571
 
    FormClass = AlignedShelvedDendrogramForm
572
 
    aspect_distorts_lengths = FormClass.aspect_distorts_lengths
573
 
    use_lengths_default = False
574
 
    
575
 
    def update_y_coordinates(self, style, y1=None):
576
 
        """The second pass through the tree.  Y coordinates only
577
 
        depend on the shape of the tree and yscale"""
578
 
        for child in self.children:
579
 
            child.update_y_coordinates(style, None)
580
 
        (self.y1, self.y2) = style.yCoords(self, None)
581
 
 
582
 
class ContemporaneousStraightDendrogram(_RootedDendrogram):
583
 
    FormClass = ContemporaneousStraightDendrogramForm
584
 
    aspect_distorts_lengths = FormClass.aspect_distorts_lengths
585
 
    use_lengths_default = False
 
586
 
 
587
    def wedgeVertices(self):
 
588
        tip_ys = [(c.y2 + self.y2)/2 for c in self.iterTips()]
 
589
        t,b = max(tip_ys), min(tip_ys)
 
590
        cxs = [c.x2 for c in self.iterTips()]
 
591
        l,r = min(cxs), max(cxs)
 
592
        vertices = [(self.x2, self.y2), (l, t), (r, b)]
 
593
        return (l,r,t,b), vertices
 
594
 
 
595
class _ContemporaneousMixin(object):
 
596
    """A dendrogram with all of the tips lined up.  
 
597
    Tidy but not suitable for displaying evolutionary distances accurately"""
 
598
 
 
599
    # Overrides init to change default for use_lengths
 
600
    def __init__(self, edge, use_lengths=False):
 
601
        super(_ContemporaneousMixin, self).__init__(edge, use_lengths)
 
602
        
 
603
    def xCoords(self, scale, x1):
 
604
        return (x1, (scale.height-(self.height-self.length))*scale.x)
 
605
 
 
606
class ContemporaneousDendrogram(_ContemporaneousMixin, SquareDendrogram):
 
607
    pass
 
608
    
 
609
class ContemporaneousStraightDendrogram(_ContemporaneousMixin, StraightDendrogram):
 
610
    pass
 
611
 
 
612
 
 
613
class ShelvedDendrogram(ContemporaneousDendrogram):
 
614
    """A dendrogram in which internal nodes also get a row to themselves"""
 
615
    def widthRequired(self):
 
616
        return self.edgecount  # as opposed to tipcount
 
617
    
 
618
    def yCoords(self, scale, y1):
 
619
        cys = [c.y1 for c in self.Children]
 
620
        if cys:
 
621
            y2 = cys[-1] - 1.0 * scale.y
 
622
        else:
 
623
            y2 = y1 - 0.5 * scale.y
 
624
        return (y2, y2)
 
625
 
 
626
class AlignedShelvedDendrogram(ShelvedDendrogram):
 
627
    
 
628
    def update_y_coordinates(self, scale, y1=None):
 
629
        """The second pass through the tree.  Y coordinates only
 
630
        depend on the shape of the tree and yscale"""
 
631
        for child in self.Children:
 
632
            child.update_y_coordinates(scale, None)
 
633
        (self.y1, self.y2) = self.yCoords(scale, None)
 
634
    
 
635
    def yCoords(self, scale, y1):
 
636
        if hasattr(self, 'track_y'):
 
637
            return (self.track_y, self.track_y)
 
638
        else:
 
639
            raise RuntimeError, self.Name
 
640
    
586
641
 
587
642
class UnrootedDendrogram(_Dendrogram):
588
 
    use_lengths_default = True
589
643
    aspect_distorts_lengths = True
 
644
 
 
645
    def labelMargins(self, label_width):
 
646
        return (label_width, label_width)
590
647
    
 
648
    def wedgeVertices(self):
 
649
        tip_dists = [(c.depth-self.depth)*self.scale for c in self.iterTips()]
 
650
        (near, far) = (min(tip_dists), max(tip_dists))
 
651
        a = self.angle - 0.25 * self.wedge
 
652
        (x1, y1) = (self.x2+near*numpy.sin(a), self.y2+near*numpy.cos(a))
 
653
        a = self.angle + 0.25 * self.wedge
 
654
        (x2, y2) = (self.x2+far*numpy.sin(a), self.y2+far*numpy.cos(a))
 
655
        vertices = [(self.x2, self.y2), (x1, y1), (x2, y2)]
 
656
        return (self.x2, (x1+x2)/2, self.y2, (y1+y2)/2), vertices
 
657
 
591
658
    def updateCoordinates(self, width, height):
592
 
        if width < 2*self.labelwidth:
593
 
            raise ValueError('%spt not wide enough for %spt wide labels' %
594
 
                    (width, self.labelwidth))
595
 
        width -= 2*self.labelwidth
596
 
        
597
659
        angle = 2*numpy.pi / self.leafcount
598
660
        # this loop is a horrible brute force hack
599
661
        # there are better (but complex) ways to find
610
672
                best_scale = scale
611
673
                mid_x = width/2-((max(xs)+min(xs))/2)*scale
612
674
                mid_y = height/2-((max(ys)+min(ys))/2)*scale
613
 
                best_args = (scale, mid_x+self.labelwidth, mid_y, direction, angle)
 
675
                best_args = (scale, mid_x, mid_y, direction, angle)
614
676
        self._update_coordinates(*best_args)
615
677
        return best_scale
616
678
    
618
680
        # Constant angle algorithm.  Should add maximim daylight step.
619
681
        (x2, y2) = (x1+self.length*s*numpy.sin(a), y1+self.length*s*numpy.cos(a))
620
682
        (self.x1, self.y1, self.x2, self.y2, self.angle) = (x1, y1, x2, y2, a)
 
683
        if self.Collapsed:
 
684
            self.wedge = self.leafcount * da
 
685
            self.scale = s
 
686
            (l,r,t,b), vertices = self.wedgeVertices()
 
687
            return vertices
 
688
            
621
689
        a -= self.leafcount * da / 2
622
 
        if not self.children:
 
690
        if not self.Children:
623
691
            points = [(x2, y2)]
624
692
        else:
625
693
            points = []
626
 
            for (i,child) in enumerate(self.children):
 
694
            for (i,child) in enumerate(self.Children):
627
695
                ca = child.leafcount * da
628
696
                points += child._update_coordinates(s, x2, y2, a+ca/2, da)
629
697
                a += ca