def make_color_legend(filename, colormap, start, end, step, width=100, height=10, display=False): from rasmus import util if filename is None: filename = util.tempfile(".", "colormap", ".svg") temp = True else: temp = False s = svg.Svg(util.open_stream(filename, "w")) s.beginSvg(width, height) xscale = float(width) / (end + step - start) for i in util.frange(start, end + step, step): color = colormap.get(i) s.rect((i - start) * xscale, 0, step * xscale, height, color, color) s.endSvg() s.close() # display if display: os.system("display %s" % filename) # clean up temp files if temp: os.remove(filename)
def drawDistRuler(names, dists, scale=500, padding=10, textsize=12, notchsize=2, labelpadding=5, distsize=9, filename=sys.stdout): """Produce a ruler of pairwise distances""" nameswidth = textsize * max(map(len, names)) out = svg.Svg(util.open_stream(filename, "w")) out.beginSvg(scale * max(dists) + 2 * padding, 2 * padding + nameswidth + 5 * distsize) out.beginTransform(("translate", padding, nameswidth + padding)) # draw ruler out.line(0, 0, scale * max(dists), 0) for name, dist in zip(names, dists): x = scale * dist out.text(name, x + textsize / 2.0, -labelpadding, textsize, angle=-90) out.line(x, notchsize, x, -notchsize) out.text("%.3f" % dist, x + textsize / 2.0, labelpadding + distsize * 3.5, distsize, angle=-90) out.endTransform() out.endSvg()
def make_color_legend(filename, colormap, start, end, step, width=100, height=10): from rasmus import util s = svg.Svg(util.open_stream(filename, "w")) s.beginSvg(width, height) xscale = float(width) / (end + step - start) for i in util.frange(start, end + step, step): color = colormap.get(i) s.rect((i-start) * xscale, 0, step*xscale, height, color, color) s.endSvg()
def draw_tree(tree, brecon, stree, xscale=100, yscale=100, leaf_padding=10, label_size=None, label_offset=None, font_size=12, stree_font_size=20, canvas=None, autoclose=True, rmargin=10, lmargin=100, tmargin=100, bmargin=100, tree_color=(0, 0, 0), tree_trans_color=(0, 0, 0), stree_color=(.3, .7, .3), snode_color=(.2, .2, .7), loss_color=(1, 1, 1), loss_color_border=(.5, .5, .5), dup_color=(0, 0, 1), dup_color_border=(0, 0, 1), trans_color=(1, 1, 0), trans_color_border=(.5, .5, 0), gtrans_color=(1, 0, 0), gtrans_color_border=(.5, 0, 0), event_size=10, snames=None, rootlen=None, stree_width=.8, filename="tree.svg"): '''Takes as input a parasite tree, tree, a reconciliation file, brecon, a host tree, stree, as well as sizes and colors of the trees components and returns a drawing of the reconciliation of the parasite tree on the host tree with event nodes of specified colors''' # set defaults font_ratio = 8. / 11. if label_size is None: label_size = .7 * font_size if sum(x.dist for x in tree.nodes.values()) == 0: legend_scale = False minlen = xscale if snames is None: snames = dict((x, x) for x in stree.leaf_names()) # layout stree slayout = treelib1.layout_tree(stree, xscale, yscale) if rootlen is None: rootlen = .1 * max(l[0] for l in slayout.values()) # setup slayout x, y = slayout[stree.root] slayout[None] = (x - rootlen, y) for node, (x, y) in slayout.items(): slayout[node] = (x + rootlen, y - .5 * yscale) # layout tree ylists = defaultdict(lambda: []) yorders = {} # layout speciations and genes (y) for node in tree.preorder(): if node == list(tree.preorder())[0]: rootNode = node.name yorders[node] = [] for ev in brecon[node]: snode, event, frequency = ev if event == "spec" or event == "gene" or event == "loss": yorders[node].append(len(ylists[snode])) ylists[snode].append(node) # layout dups and transfers (y) for node in tree.postorder(): for ev in brecon[node]: snode, event, frequency = ev if event != "spec" and event != "gene" and event != "loss": # Find number of nodes on a single branch for y-coord v = [ yorders[child] for child in node.children if brecon[child][-1][0] == snode ] if len(v) == 0: yorders[node].append(0) else: yorders[node].append(stats.mean(flatten(v))) # layout node (x) xorders = { } #Dictionary to record number of nodes on a single branch for x-coord branchFrac = {} #Dictionary to record the placement of a node on a branch for node in tree.postorder(): for n in range(len(brecon[node])): snode, event, frequency = brecon[node][n] if event == "spec" or event == "gene" or event == "loss": # Speciation, gene, and loss events happen at host vertices if not node in branchFrac: branchFrac[node] = 0 else: # Transfers and duplications occur on branches v = [branchFrac[child] for child in node.children] if len(v) == 0: branchFrac[node] = 1 else: branchFrac[node] = max(v) + 1 for node in tree.preorder(): xorders[node] = [] for n in range(len(brecon[node])): snode, event, frequency = brecon[node][n] if event == "spec" or event == "gene" or event == "loss": # Speciation, gene, and loss events happen on vertices, not branches xorders[node].append(0) else: if node.parent and containsTransOrDup(node.parent, brecon): # set branchFrac to the branch Frac of the parent, they are # on the same branch branchFrac[node] = branchFrac[node.parent] if containsLoss(node, brecon): # if following a loss, first transfer/duplication event on branch xorders[node].append(1) elif not node.parent: # Root of tree xorders[node].append(0) else: xorders[node].append(maxList(xorders[node.parent]) + 1) # setup layout layout = {None: [slayout[brecon[tree.root][-1][0].parent]]} for node in tree.preorder(): for n in range(len(brecon[node])): snode, event, frequency = brecon[node][n] nx, ny = slayout[snode] px, py = slayout[snode.parent] (npx, npy) = layout[node.parent][-1] # set spacing between nodes on the same branch frac = 50 while branchFrac[node] * frac >= nx - px: frac = frac - 5 # calc x if event == "trans" or event == "gtrans": if npx > px: # transfer parent is farther forward in time than host parent x = npx + frac else: x = px + frac elif event == "dup": x = px + frac else: x = nx # calc y deltay = ny - py slope = deltay / float(nx - px) deltax2 = x - px deltay2 = slope * deltax2 offset = py + deltay2 frac = (yorders[node][n] + 1) / float(max(len(ylists[snode]), 1) + 1) y = offset + (frac - .5) * stree_width * yscale if node in layout: layout[node].append((x, y)) else: layout[node] = [(x, y)] # order brecon nodes temporally brecon[node] = orderLoss(node, brecon, layout) # order layout nodes temporally layout[node] = orderLayout(node, layout) if y > max(l[1] for l in slayout.values()) + 50: print nx, ny print px, py print offset, frac print ylists[snode], yorders[node] print brecon[node] print node, snode, layout[node] # layout label sizes max_label_size = max(len(x.name) for x in tree.leaves()) * font_ratio * font_size max_slabel_size = max( len(x.name) for x in stree.leaves()) * font_ratio * stree_font_size ''' if colormap == None: for node in tree: node.color = (0, 0, 0) else: colormap(tree) if stree and gene2species: recon = phylo.reconcile(tree, stree, gene2species) events = phylo.label_events(tree, recon) losses = phylo.find_loss(tree, stree, recon) else: events = None losses = None # layout tree if layout is None: coords = treelib.layout_tree(tree, xscale, yscale, minlen, maxlen) else: coords = layout ''' xcoords, ycoords = zip(*slayout.values()) maxwidth = max(xcoords) + max_label_size + max_slabel_size maxheight = max(ycoords) + yscale # initialize canvas if canvas is None: canvas = svg.Svg(util.open_stream(filename, "w")) width = int(rmargin + maxwidth + lmargin) height = int(tmargin + maxheight + bmargin) canvas.beginSvg(width, height) canvas.beginStyle("font-family: \"Sans\";") if autoclose == None: autoclose = True else: if autoclose == None: autoclose = False canvas.beginTransform(("translate", lmargin, tmargin)) draw_stree(canvas, stree, slayout, yscale=yscale, stree_width=stree_width, stree_color=stree_color, snode_color=snode_color) # draw stree leaves for node in stree: x, y = slayout[node] if node.is_leaf(): canvas.text(snames[node.name], x + leaf_padding + max_label_size, y + stree_font_size / 2., stree_font_size, fillColor=snode_color) # draw tree for node in tree: containsL = containsLoss(node, brecon) for n in range(len(brecon[node])): x, y = layout[node][n] if containsL == False: # no loss event px, py = layout[node.parent][-1] else: # loss event present if n == 0: # event is loss px, py = layout[node.parent][-1] else: # event stems from loss px, py = layout[node][n - 1] trans = False if node.parent: snode, event, frequency = brecon[node][n] if n == 0: psnode, pevent, pfrequency = brecon[node.parent][-1] # Event stemming from a loss event else: psnode, pevent, pfrequency = brecon[node][n - 1] if pevent == "trans" or pevent == "gtrans": if psnode != snode: trans = True else: trans = False if not trans: canvas.line(x, y, px, py, color=tree_color) # draw the transfer dashed line else: arch = 20 x2 = (x * .5 + px * .5) - arch y2 = (y * .5 + py * .5) x3 = (x * .5 + px * .5) - arch y3 = (y * .5 + py * .5) # draw regular transfer dashed line if pevent == "trans": canvas.write( "<path d='M%f %f C%f %f %f %f %f %f' %s />\n " % (x, y, x2, y2, x3, y3, px, py, " style='stroke-dasharray: 4, 2' " + svg.colorFields(tree_trans_color, (0, 0, 0, 0)))) # draw guilty transfer dashed line else: canvas.write( "<path d='M%f %f C%f %f %f %f %f %f' %s />\n " % (x, y, x2, y2, x3, y3, px, py, " style='stroke-dasharray: 4, 2' " + svg.colorFields(gtrans_color, (0, 0, 0, 0)))) # draw events for node in tree: if node.name == rootNode: x, y = layout[node][0] canvas.polygon((x-20, y, x-50, y+30,x-50, y+15, x-90, y+15, x-90,\ y-15, x-50, y-15, x-50, y-30), strokeColor = (1,.7,.3), \ fillColor = (1,.7,.3)) canvas.text("Root Node", x-88, y+5, font_size+2,\ fillColor = (0,0,0)) for n in range(len(brecon[node])): snode, event, frequency = brecon[node][n] frequency = float(frequency) x, y = layout[node][n] o = event_size / 2.0 if event == "loss": # draw boxes, frequencies of loss events canvas.rect(x - o, y - o, event_size, event_size, fillColor=loss_color, strokeColor=loss_color_border) canvas.text("{:.3f}".format(frequency) + node.name, x - o, y - o, font_size + 2, fillColor=loss_color) if event == "spec": # draw boxes, frequencies of speciation events canvas.rect(x - o, y - o, event_size, event_size, fillColor=(0, 0, 0), strokeColor=(0, 0, 0)) canvas.text("{:.3f}".format(frequency) + node.name, x - o, y - o, font_size + 2, fillColor=(0, 0, 0)) if event == "dup": # draw boxes, frequencies of duplication events canvas.rect(x - o, y - o, event_size, event_size, fillColor=dup_color, strokeColor=dup_color_border) canvas.text("{:.3f}".format(frequency) + node.name, x - o, y - o, font_size + 2, fillColor=dup_color) elif event == "trans": # draw boxes, frequencies of transfer events canvas.rect(x - o, y - o, event_size, event_size, fillColor=trans_color, strokeColor=trans_color_border) canvas.text("{:.3f}".format(frequency) + node.name, x - o, y - o, font_size + 2, fillColor=trans_color) elif event == "gtrans": # draw boxes, frequencies of guilty transfer events canvas.rect(x - o, y - o, event_size, event_size, fillColor=gtrans_color, strokeColor=gtrans_color_border) canvas.text("{:.3f}".format(frequency) + node.name, x - o, y - o, font_size + 2, fillColor=gtrans_color) # draw tree leaves for node in tree: for n in range(len(brecon[node])): x, y = layout[node][n] if node.is_leaf() and brecon[node][n][1] == "gene": canvas.text(node.name, x + leaf_padding, y + font_size / 2., font_size + 2, fillColor=(0, 0, 0)) canvas.endTransform() if autoclose: canvas.endStyle() canvas.endSvg() return canvas
def draw_tree(tree, labels={}, xscale=100, yscale=20, canvas=None, leafPadding=10, leafFunc=lambda x: str(x.name), labelOffset=None, fontSize=10, labelSize=None, minlen=1, maxlen=util.INF, filename=sys.stdout, rmargin=150, lmargin=10, tmargin=0, bmargin=None, colormap=None, stree=None, layout=None, gene2species=None, lossColor=(0, 0, 1), dupColor=(1, 0, 0), eventSize=4, legendScale=False, autoclose=None, extendRoot=True, labelLeaves=True, drawHoriz=True, nodeSize=0): # set defaults fontRatio = 8. / 11. if labelSize == None: labelSize = .7 * fontSize if labelOffset == None: labelOffset = -1 if bmargin == None: bmargin = yscale if sum(x.dist for x in tree.nodes.values()) == 0: legendScale = False minlen = xscale if colormap == None: for node in tree: node.color = (0, 0, 0) else: colormap(tree) if stree and gene2species: recon = phylo.reconcile(tree, stree, gene2species) events = phylo.label_events(tree, recon) losses = phylo.find_loss(tree, stree, recon) else: events = None losses = None if len(labels) > 0 or (stree and gene2species): drawHoriz = True # layout tree if layout is None: coords = treelib.layout_tree(tree, xscale, yscale, minlen, maxlen) else: coords = layout xcoords, ycoords = zip(*coords.values()) maxwidth = max(xcoords) maxheight = max(ycoords) + labelOffset # initialize canvas if canvas == None: canvas = svg.Svg(util.open_stream(filename, "w")) width = int(rmargin + maxwidth + lmargin) height = int(tmargin + maxheight + bmargin) canvas.beginSvg(width, height) if autoclose == None: autoclose = True else: if autoclose == None: autoclose = False # draw tree def walk(node): x, y = coords[node] if node.parent: parentx, parenty = coords[node.parent] else: if extendRoot: parentx, parenty = 0, y else: parentx, parenty = x, y # e.g. no branch # draw branch if drawHoriz: canvas.line(parentx, y, x, y, color=node.color) else: canvas.line(parentx, parenty, x, y, color=node.color) # draw branch labels if node.name in labels: branchlen = x - parentx lines = str(labels[node.name]).split("\n") labelwidth = max(map(len, lines)) labellen = min(labelwidth * fontRatio * fontSize, max(int(branchlen - 1), 0)) for i, line in enumerate(lines): canvas.text( line, parentx + (branchlen - labellen) / 2., y + labelOffset + (-len(lines) + 1 + i) * (labelSize + 1), labelSize) # draw nodes if nodeSize > 0: canvas.circle(x, y, nodeSize, strokeColor=svg.null, fillColor=node.color) # draw leaf labels or recur if node.is_leaf(): if labelLeaves: canvas.text(leafFunc(node), x + leafPadding, y + fontSize / 2., fontSize, fillColor=node.color) else: if drawHoriz: # draw vertical part of branch top = coords[node.children[0]][1] bot = coords[node.children[-1]][1] canvas.line(x, top, x, bot, color=node.color) # draw children for child in node.children: walk(child) canvas.beginTransform(("translate", lmargin, tmargin)) walk(tree.root) if stree and gene2species: draw_events(canvas, tree, coords, events, losses, lossColor=lossColor, dupColor=dupColor, size=eventSize) canvas.endTransform() # draw legend if legendScale: if legendScale == True: # automatically choose a scale length = maxwidth / float(xscale) order = math.floor(math.log10(length)) length = 10**order drawScale(lmargin, tmargin + maxheight + bmargin - fontSize, length, xscale, fontSize, canvas=canvas) if autoclose: canvas.endSvg() return canvas
def heatmap(matrix, width=20, height=20, colormap=None, filename=None, rlabels=None, clabels=None, display=True, xdir=1, ydir=1, xmargin=0, ymargin=0, labelPadding=2, labelSpacing=4, mincutoff=None, maxcutoff=None, showVals=False, formatVals=str, valColor=black, clabelsAngle=270, clabelsPadding=None, rlabelsAngle=0, rlabelsPadding=None, colors=None, strokeColors=None, valAnchor="start", close=True): from rasmus import util if display and (not close): raise Exception("must close file if display is used") # determine filename if filename is None: filename = util.tempfile(".", "heatmap", ".svg") temp = True else: temp = False # determine colormap if colors is None: if colormap is None: colormap = rainbowColorMap(util.flatten(matrix)) # determine matrix size and orientation nrows = len(matrix) ncols = len(matrix[0]) if xdir == 1: xstart = xmargin ranchor = "end" coffset = width elif xdir == -1: xstart = xmargin + ncols * width ranchor = "start" coffset = 0 else: raise Exception("xdir must be 1 or -1") if ydir == 1: ystart = ymargin roffset = height canchor = "start" elif ydir == -1: ystart = ymargin + nrows * width roffset = 0 canchor = "end" else: raise Exception("ydir must be 1 or -1") # begin svg infile = util.open_stream(filename, "w") s = svg.Svg(infile) s.beginSvg(ncols * width + 2 * xmargin, nrows * height + 2 * ymargin) # draw matrix for i in xrange(nrows): for j in xrange(ncols): if mincutoff and matrix[i][j] < mincutoff: continue if maxcutoff and matrix[i][j] > maxcutoff: continue if colors: color = colors[i][j] else: color = colormap.get(matrix[i][j]) if strokeColors: strokeColor = strokeColors[i][j] else: strokeColor = color s.rect(xstart + xdir * j * width, ystart + ydir * i * height, xdir * width, ydir * height, strokeColor, color) # draw values if showVals: # find text size fontwidth = 7 / 11.0 textsize = [] for i in xrange(nrows): for j in xrange(ncols): if mincutoff and matrix[i][j] < mincutoff: continue if maxcutoff and matrix[i][j] > maxcutoff: continue strval = formatVals(matrix[i][j]) if len(strval) > 0: textsize.append( min(height, width / (float(len(strval)) * fontwidth))) textsize = min(textsize) if valAnchor == "start": xoffset = 0 elif valAnchor == "middle": xoffset = 0.5 elif valAnchor == "end": xoffset = 1 else: raise Exception("anchor not supported: %s" % valAnchor) yoffset = int(ydir == -1) for i in xrange(nrows): for j in xrange(ncols): if mincutoff and matrix[i][j] < mincutoff: continue if maxcutoff and matrix[i][j] > maxcutoff: continue strval = formatVals(matrix[i][j]) s.text(strval, xstart + xdir * (j + xoffset) * width, ystart + ydir * (i + yoffset) * height + height / 2.0 + textsize / 2.0, textsize, fillColor=valColor, anchor=valAnchor) # draw labels if rlabels is not None: assert len(rlabels) == nrows, \ "number of row labels does not equal number of rows" if rlabelsPadding is None: rlabelsPadding = labelPadding for i in xrange(nrows): x = xstart - xdir * rlabelsPadding y = ystart + roffset + ydir * i * height - labelSpacing / 2. s.text(rlabels[i], x, y, height - labelSpacing, anchor=ranchor, angle=rlabelsAngle) if clabels is not None: assert len(clabels) == ncols, \ "number of col labels does not equal number of cols" if clabelsPadding is None: clabelsPadding = labelPadding for j in xrange(ncols): x = xstart + coffset + xdir * j * width - labelSpacing / 2. y = ystart - ydir * clabelsPadding s.text(clabels[j], x, y, width - labelSpacing, anchor=canchor, angle=clabelsAngle) # end svg if close: s.endSvg() s.close() # display matrix if display: #if temp: os.system("display %s" % filename) #else: # os.spawnl(os.P_NOWAIT, "display", "display", filename) # clean up temp files if temp: os.remove(filename) return s
def draw_tree(tree, brecon, stree, xscale=100, yscale=100, leaf_padding=10, label_size=None, label_offset=None, font_size=12, stree_font_size=20, canvas=None, autoclose=True, rmargin=10, lmargin=10, tmargin=0, bmargin=0, tree_color=(0, 0, 0), tree_trans_color=(0, 0, 0), stree_color=(.6, .3, .8), snode_color=(.2, .2, .7), loss_color=(1, 1, 1), loss_color_border=(.5, .5, .5), dup_color=(1, 0, 0), dup_color_border=(.5, 0, 0), trans_color=(0, 1, 0), trans_color_border=(0, .5, 0), event_size=10, snames=None, rootlen=None, stree_width=.8, filename="tree.svg"): # set defaults font_ratio = 8. / 11. if label_size is None: label_size = .7 * font_size #if label_offset is None: # label_offset = -1 if sum(x.dist for x in tree.nodes.values()) == 0: legend_scale = False minlen = xscale if snames is None: snames = dict((x, x) for x in stree.leaf_names()) # layout stree slayout = treelib1.layout_tree(stree, xscale, yscale) if rootlen is None: rootlen = .1 * max(l[0] for l in slayout.values()) # setup slayout x, y = slayout[stree.root] slayout[None] = (x - rootlen, y) for node, (x, y) in slayout.items(): slayout[node] = (x + rootlen, y - .5 * yscale) # layout tree ylists = defaultdict(lambda: []) yorders = {} # layout speciations and genes (y) for node in tree.preorder(): for ev in brecon[node]: snode, event, frequency = ev if event == "spec" or event == "gene" or event == "loss": yorders[node] = len(ylists[snode]) ylists[snode].append(node) # layout dups and transfers (y) for node in tree.postorder(): for ev in brecon[node]: snode, event, frequency = ev if event != "spec" and event != "gene" and event != "loss": v = [ yorders[child] for child in node.children if brecon[child][-1][0] == snode ] if len(v) == 0: yorders[node] = 0 else: yorders[node] = stats.mean(v) # layout node (x) xorders = {} xmax = defaultdict(lambda: 0) for node in tree.postorder(): for ev in brecon[node]: snode, event, frequency = ev if event == "spec" or event == "gene" or event == "loss": xorders[node] = 0 else: v = [ xorders[child] for child in node.children if brecon[child][-1][0] == snode ] if len(v) == 0: xorders[node] = 1 else: xorders[node] = max(v) + 1 xmax[snode] = max(xmax[snode], xorders[node]) # setup layout layout = {None: [slayout[brecon[tree.root][-1][0].parent]]} for node in tree: for ev in brecon[node]: snode, event, frequency = ev nx, ny = slayout[snode] px, py = slayout[snode.parent] # calc x frac = (xorders[node]) / float(xmax[snode] + 1) deltax = nx - px x = nx - frac * deltax # calc y deltay = ny - py slope = deltay / float(deltax) deltax2 = x - px deltay2 = slope * deltax2 offset = py + deltay2 frac = (yorders[node] + 1) / float(max(len(ylists[snode]), 1) + 1) y = offset + (frac - .5) * stree_width * yscale if node in layout: layout[node].append((x, y)) else: layout[node] = [(x, y)] brecon[node] = orderLoss(node, brecon, layout) print "Brecon = ", brecon[node] layout[node] = orderLayout(node, layout) print "Layout = ", layout[node] if y > max(l[1] for l in slayout.values()) + 50: print nx, ny print px, py print offset, frac print ylists[snode], yorders[node] print brecon[node] print node, snode, layout[node] # layout label sizes max_label_size = max(len(x.name) for x in tree.leaves()) * font_ratio * font_size max_slabel_size = max( len(x.name) for x in stree.leaves()) * font_ratio * stree_font_size ''' if colormap == None: for node in tree: node.color = (0, 0, 0) else: colormap(tree) if stree and gene2species: recon = phylo.reconcile(tree, stree, gene2species) events = phylo.label_events(tree, recon) losses = phylo.find_loss(tree, stree, recon) else: events = None losses = None # layout tree if layout is None: coords = treelib.layout_tree(tree, xscale, yscale, minlen, maxlen) else: coords = layout ''' xcoords, ycoords = zip(*slayout.values()) maxwidth = max(xcoords) + max_label_size + max_slabel_size maxheight = max(ycoords) + .5 * yscale # initialize canvas if canvas is None: canvas = svg.Svg(util.open_stream(filename, "w")) width = int(rmargin + maxwidth + lmargin) height = int(tmargin + maxheight + bmargin) canvas.beginSvg(width, height) canvas.beginStyle("font-family: \"Sans\";") if autoclose == None: autoclose = True else: if autoclose == None: autoclose = False canvas.beginTransform(("translate", lmargin, tmargin)) draw_stree(canvas, stree, slayout, yscale=yscale, stree_width=stree_width, stree_color=stree_color, snode_color=snode_color) # draw stree leaves for node in stree: x, y = slayout[node] if node.is_leaf(): canvas.text(snames[node.name], x + leaf_padding + max_label_size, y + stree_font_size / 2., stree_font_size, fillColor=snode_color) # draw tree for node in tree: containsL = containsLoss(node, brecon) for n in range(len(brecon[node])): # print brecon[node] x, y = layout[node][n] # print layout[node] if containsL == False: px, py = layout[node.parent][-1] else: if brecon[node][n][1] == "loss": px, py = layout[node.parent][-1] else: px, py = layout[node][n - 1] trans = False if node.parent: for ev in brecon[node]: snode, event, frequency = ev psnode = brecon[node.parent][-1][0] while snode: if psnode == snode: break snode = snode.parent else: trans = True if not trans: canvas.line(x, y, px, py, color=tree_color) else: arch = 20 x2 = (x * .5 + px * .5) - arch y2 = (y * .5 + py * .5) x3 = (x * .5 + px * .5) - arch y3 = (y * .5 + py * .5) canvas.write("<path d='M%f %f C%f %f %f %f %f %f' %s />\n " % (x, y, x2, y2, x3, y3, px, py, " style='stroke-dasharray: 4, 2' " + svg.colorFields(tree_trans_color, (0, 0, 0, 0)))) # draw events for node in tree: for n in range(len(brecon[node])): snode, event, frequency = brecon[node][n] x, y = layout[node][n] o = event_size / 2.0 if event == "loss": canvas.rect(x - o, y - o, event_size, event_size, fillColor=loss_color, strokeColor=loss_color_border) canvas.text(frequency, x - o, y - o, font_size, fillColor=(1, 1, 1)) if event == "spec": canvas.text(frequency, slayout[snode][0] - leaf_padding / 2, slayout[snode][1] - font_size, font_size, fillColor=(0, 0, 0)) if event == "dup": canvas.rect(x - o, y - o, event_size, event_size, fillColor=dup_color, strokeColor=dup_color_border) canvas.text(frequency, x - o, y - o, font_size, fillColor=dup_color) elif event == "trans": canvas.rect(x - o, y - o, event_size, event_size, fillColor=trans_color, strokeColor=trans_color_border) canvas.text(frequency, x - o, y - o, font_size, fillColor=trans_color) # draw tree leaves for node in tree: for n in range(len(brecon[node])): x, y = layout[node][n] if node.is_leaf() and containsLoss(node, brecon) == False: canvas.text(node.name, x + leaf_padding, y + font_size / 2., font_size, fillColor=(0, 0, 0)) canvas.endTransform() if autoclose: canvas.endStyle() canvas.endSvg() return canvas
def draw_tree(tree, stree, extra, xscale=100, yscale=100, leaf_padding=10, label_size=None, label_offset=None, font_size=12, stree_font_size=20, canvas=None, autoclose=True, rmargin=10, lmargin=10, tmargin=0, bmargin=0, stree_color=(.4, .4, 1), snode_color=(.2, .2, .7), event_size=10, rootlen=None, stree_width=.8, filename=sys.stdout, labels=None, slabels=None): recon = extra["species_map"] loci = extra["locus_map"] order = extra["order"] # setup color map all_loci = sorted(set(loci.values())) num_loci = len(all_loci) colormap = util.rainbow_color_map(low=0, high=num_loci - 1) locus_color = {} for ndx, locus in enumerate(all_loci): locus_color[locus] = colormap.get(ndx) # set defaults font_ratio = 8. / 11. if label_size is None: label_size = .7 * font_size #if label_offset is None: # label_offset = -1 if sum(x.dist for x in tree.nodes.values()) == 0: legend_scale = False minlen = xscale snames = dict((x, x) for x in stree.leaf_names()) if labels is None: labels = {} if slabels is None: slabels = {} # layout stree slayout = treelib.layout_tree(stree, xscale, yscale) if rootlen is None: rootlen = .1 * max(l[0] for l in slayout.values()) # setup slayout x, y = slayout[stree.root] slayout[None] = (x - rootlen, y) for node, (x, y) in slayout.items(): slayout[node] = (x + rootlen, y - .5 * yscale) # layout tree ylists = defaultdict(lambda: []) yorders = {} # layout speciations and genes (y) events = phylo.label_events(tree, recon) for node in tree.preorder(): snode = recon[node] event = events[node] if event == "spec" or event == "gene": yorders[node] = len(ylists[snode]) ylists[snode].append(node) # layout internal nodes (y) for node in tree.postorder(): snode = recon[node] event = events[node] if event != "spec" and event != "gene": v = [yorders[child] for child in node.children] yorders[node] = stats.mean(v) # layout node (x) xorders = {} xmax = defaultdict(lambda: 0) for node in tree.postorder(): snode = recon[node] event = events[node] if event == "spec" or event == "gene": xorders[node] = 0 else: v = [xorders[child] for child in node.children] xorders[node] = max(v) + 1 xmax[snode] = max(xmax[snode], xorders[node]) ## # initial order ## xpreorders = {} ## for node in tree.postorder(): ## snode = recon[node] ## event = events[node] ## if event == "spec" or event == "gene": ## xpreorders[node] = 0 ## else: ## v = [xpreorders[child] for child in node.children] ## xpreorders[node] = max(v) + 1 #### print node.name, xpreorders[node] ## # hack-ish approach : shift x until order is satisfied ## def shift(node, x): ## xpreorders[node] += x ## for child in node.children: ## if events[child] != "spec": ## shift(child, x) ## satisfied = False ## while not satisfied: ## satisfied = True ## for snode, d in order.iteritems(): ## for plocus, lst in d.iteritems(): ## # test each pair ## for m, node1 in enumerate(lst): ## x1 = xpreorders[node1] ## for node2 in lst[m+1:]: ## x2 = xpreorders[node2] #### print node1, node2, x1, x2 ## if x2 < x1: ## # violation - shift all descendants in the sbranch ## satisfied = False #### print 'violation', node1, node2, x1, x2, x1-x2+1 ## shift(node2, x1-x2+1) ## break ## # finally, "normalize" xorders ## xorders = {} ## xmax = defaultdict(lambda: 0) ## for node in tree.postorder(): ## snode = recon[node] ## xorders[node] = xpreorders[node] ## xmax[snode] = max(xmax[snode], xorders[node]) #### print node.name, xpreorders[node] # setup layout layout = {None: slayout[None]} for node in tree: snode = recon[node] nx, ny = slayout[snode] px, py = slayout[snode.parent] # calc x frac = (xorders[node]) / float(xmax[snode] + 1) deltax = nx - px x = nx - frac * deltax # calc y deltay = ny - py slope = deltay / float(deltax) deltax2 = x - px deltay2 = slope * deltax2 offset = py + deltay2 frac = (yorders[node] + 1) / float(max(len(ylists[snode]), 1) + 1) y = offset + (frac - .5) * stree_width * yscale layout[node] = (x, y) ## if y > max(l[1] for l in slayout.values()) + 50: ## print nx, ny ## print px, py ## print offset, frac ## print ylists[snode], yorders[node] ## print node, snode, layout[node] # layout label sizes max_label_size = max(len(x.name) for x in tree.leaves()) * font_ratio * font_size max_slabel_size = max( len(x.name) for x in stree.leaves()) * font_ratio * stree_font_size xcoords, ycoords = zip(*slayout.values()) maxwidth = max(xcoords) + max_label_size + max_slabel_size maxheight = max(ycoords) + .5 * yscale # initialize canvas if canvas is None: canvas = svg.Svg(util.open_stream(filename, "w")) width = int(rmargin + maxwidth + lmargin) height = int(tmargin + maxheight + bmargin) canvas.beginSvg(width, height) canvas.beginStyle("font-family: \"Sans\";") if autoclose == None: autoclose = True else: if autoclose == None: autoclose = False canvas.beginTransform(("translate", lmargin, tmargin)) draw_stree(canvas, stree, slayout, yscale=yscale, stree_width=stree_width, stree_color=stree_color, snode_color=snode_color, slabels=slabels) # draw stree leaves for node in stree: x, y = slayout[node] if node.is_leaf(): canvas.text(snames[node.name], x + leaf_padding + max_label_size, y + stree_font_size / 2., stree_font_size, fillColor=snode_color) # draw tree for node in tree: x, y = layout[node] px, py = layout[node.parent] if node.parent: color = locus_color[loci[node.parent]] else: color = locus_color[loci[tree.root]] canvas.line(x, y, px, py, color=color) # draw tree names for node in tree: x, y = layout[node] px, py = layout[node.parent] if node.is_leaf(): canvas.text(node.name, x + leaf_padding, y + font_size / 2., font_size, fillColor=(0, 0, 0)) if node.name in labels: canvas.text(labels[node.name], x, y, label_size, fillColor=(0, 0, 0)) # draw events for node in tree: if node.parent: locus = loci[node] plocus = loci[node.parent] if locus != plocus: color = locus_color[locus] x, y = layout[node] o = event_size / 2.0 canvas.rect(x - o, y - o, event_size, event_size, fillColor=color, strokeColor=color) canvas.endTransform() if autoclose: canvas.endStyle() canvas.endSvg() return canvas