def layout(node):
'''
Set Nodestyle for tree.
'''
#add face with color per genus (circle at outer ring of the tree)
if node.is_leaf():
if 'strepto' in node.name:
genus = ete3.CircleFace(radius=5, color='#d6604d')
ete3.faces.add_face_to_node(genus,
node,
column=1,
position="aligned")
if 'lacto' in node.name:
genus = ete3.CircleFace(radius=5, color='#2166ac')
ete3.faces.add_face_to_node(
genus,
node,
column=1,
position="aligned",
)
if 'flori' in node.name:
genus = ete3.CircleFace(radius=5, color='#636363')
ete3.faces.add_face_to_node(genus,
node,
column=1,
position="aligned")
def _node_layout(node):
global _circle_tested
node.img_style["size"] = 0
if node.is_leaf():
if node.name not in annotation:
print('Got unknown leaf "%s"' % (node.name))
return
size = annotation[node.name]
else:
children = COMMA.join(sorted([leaf.name for leaf in node]))
if children not in annotation:
print('Got unknown node ' + children)
return
size = annotation[children]
dimension = 4 * math.sqrt(size)
labeld = {
'text': "%d" % (size),
'color': 'white',
'font': 'Helvetica',
'size': 8
}
if size % 2:
clabel = dict(labeld)
if size == 35:
if not _circle_tested:
_circle_tested = True
clabel['text'] = 'Ly'
clabel['size'] = 12
thisFace = ete.CircleFace(dimension / 2,
"steelblue",
"circle",
label=clabel)
elif size == 43:
clabel['size'] = 6
del clabel['color']
thisFace = ete.CircleFace(dimension / 2,
"steelblue",
"sphere",
label=clabel)
else:
thisFace = ete.CircleFace(dimension / 2,
"steelblue",
"sphere",
label="%d" % (size))
else:
thisFace = ete.RectFace(dimension,
dimension,
'green',
'blue',
label=labeld)
thisFace.opacity = 0.7
ete.add_face_to_node(thisFace, node, column=0, position="float")
textF = ete.TextFace(str(size), fsize=12, fgcolor="steelblue")
ete.add_face_to_node(textF, node, column=0, position="aligned")
def fam_size_layout(node):
try:
node_color = self.graphics_options[node.event]
except:
node_color = self.graphics_options['=']
relative_size = 8 * (node.fam_size / max(family.fam_sizes))
cf = ete3.CircleFace(
radius=relative_size * self.graphics_options['scale'],
color=node_color,
style='circle'
)
cf.opacity = self.graphics_options['opacity']
node.add_face(cf, column=10, position='float')
# add the family size number and asterisks to the figure
famsize_str = '{}{}\n'.format(
get_pvalue_asterisks(node),
node.fam_size,
)
tf = ete3.TextFace(famsize_str, fsize=4, fgcolor=node_color)
node.add_face(tf, column=1, position='float')
# remove the silly default blue node dot
nstyle = ete3.NodeStyle()
nstyle['size'] = 0
node.set_style(nstyle)
return
def set_node_style(node, status, n_gl_sets, ref_label=None):
if status != 'internal':
if status not in scolors:
raise Exception('status \'%s\' not in scolors' % status)
node.img_style['bgcolor'] = scolors[status]
if status not in used_colors:
used_colors[status] = scolors[status]
if glutils.is_novel(node.name):
node.add_face(ete3.CircleFace(args.novel_dot_size, scolors['novel']), column=1) #, position='float') # if args.leaf_names else 'branch')
# linewidth = 2
# node.img_style['hz_line_width'] = linewidth
# node.img_style['vt_line_width'] = linewidth
names = status.split('-&-')
if node.is_leaf():
if args.pie_chart_faces and len(names) > 1:
pcf = ete3.PieChartFace(percents=[100./len(names) for _ in range(len(names))], width=args.leafheight, height=args.leafheight, colors=[scolors[n] for n in names], line_color=None)
# pcf = ete3.StackedBarFace(percents=[100./len(names) for _ in range(len(names))], width=30, height=50, colors=[scolors[n] for n in names], line_color=None)
node.add_face(pcf, column=0, position='aligned')
elif len(names) == 1 and names[0] in used_faces:
node.add_face(ete3.RectFace(width=5, height=args.leafheight, bgcolor=used_faces[names[0]], fgcolor=None), column=0, position='aligned')
elif n_gl_sets > 2:
rectnames = [n for n in names if n in used_faces]
node.add_face(ete3.StackedBarFace(percents=[100./len(names) for _ in range(len(rectnames))], width=5 * len(rectnames), height=args.leafheight, colors=[used_faces[rn] for rn in rectnames], line_color=None), column=0, position='aligned')
else: # every leaf has to have a face, so that every leaf takes up the same vertical space
node.add_face(ete3.RectFace(width=1, height=args.leafheight, bgcolor=None, fgcolor=None), column=0, position='aligned')
def prior_layout(node):
if 'prior' in node.features:
C = ete3.CircleFace(
radius=20 * np.sqrt(node.prior),
color="RoyalBlue",
style="circle",
#label={'text': "%.5f" % node.prior,
# 'fontsize': 5}
)
C.opacity = 0.5
ete3.faces.add_face_to_node(C, node, 0, position="float")
def style_type_node(color="Black"):
"""Specify the appearance of Type nodes."""
face = ete3.CircleFace(
radius=TYPE_NODE_RADIUS,
color=TYPE_NODE_COLOR_DICT.get(node.name, "White"),
style="circle",
label={
"text":
node.name,
"color":
color,
"fontsize": (TYPE_NODE_FONT_SIZE_FILE if self._render_type
== "file" else TYPE_NODE_FONT_SIZE_BROWSE),
},
)
set_face_margin(face)
node.add_face(face, column=0, position="branch-right")
def process_tree(self, t):
# Add faces
if self.attribute:
values = set([getattr(l, self.attribute) for l in t.get_leaves()])
colours = get_colour_set(len(values))
colour_map = dict(zip(values, colours))
for l in t.iter_leaves():
mycolour = colour_map[getattr(l, self.attribute)]
if not self.dummy:
l.add_face(
ete3.CircleFace(radius=10,
color=mycolour,
style="sphere"), 0)
# Apply labels
if not self.dummy:
for l in t.iter_leaves():
l.add_face(ete3.TextFace(getattr(l, self.label)), 1)
# Plot or save
if self.output:
kw = {}
if self.height or self.width:
kw["h"] = self.height
kw["w"] = self.width
kw["units"] = self.units
kw["dpi"] = self.dpi
if self.multiple:
base, ext = os.path.splitext(self.output)
filename = base + ("_%06d" % (self.n + 1)) + ext
else:
filename = self.output
if not self.dummy:
t.render(filename, ultrametric, tree_style=self.ts, **kw)
else: # pragma: no cover
if not self.dummy:
t.show(ultrametric, tree_style=self.ts)
self.n += 1
if self.multiple:
return None
else:
raise StopIteration
def fam_size_piechart_layout(node):
'''
the PieChart layout function, defined in local scope so it
can access class attributes (graphics options).
'''
if not node.is_root():
if self.count_all_expansions:
n_exp = node.expansion
n_con = node.decrease
else:
n_exp = node.sig_expansions
n_con = node.sig_expansions
if hasattr(self, 'lambda_colors'):
circle_color = self.lambda_colors[node.lambda_group]
else:
circle_color = "blue"
# add a text that shows expansions & contractions, e.g. +10/-20
exp_cnt_txt = ete3.TextFace(
'+{} -{}\n'.format(int(n_exp), int(n_con)), fsize=6,
fgcolor=circle_color
)
pos = 'aligned' if node.is_leaf() else 'float'
# add a circle that shows the average expansion
ete3.faces.add_face_to_node(exp_cnt_txt, node, 1, position=pos)
# add average expansion info:
scale_factor = 1 + node.avg_expansion
avg_exp = '{:+}'.format(round(node.avg_expansion, 2))
circle = ete3.CircleFace(radius=9 * scale_factor,
color=circle_color,
label={'text': avg_exp, 'color': 'white',
'fontsize': 3 + (2.25*scale_factor)})
circle.opacity = self.graphics_options['opacity']
ete3.faces.add_face_to_node(circle, node, 2, position='float')
nstyle = ete3.NodeStyle()
nstyle['size'] = 0
node.set_style(nstyle)
return
def my_layout(node):
seqmeta = annotations.get(node.name)
# handle leaves
if seqmeta and not re.compile("naive").match(node.name):
leaf_style(node, seqmeta, tp_colors, highlight_node)
# Deal with naive and true internal nodes
else:
# Have to set all node sizes to 0 or we end up distorting the x-axis
nstyle = ete3.NodeStyle()
nstyle['size'] = 0
node.set_style(nstyle)
# Highlight just those nodes in the seedlineage which have nonzero branch lengths, or bad things
if node in seed_lineage and node.dist > 0:
position = "float"
# Note; we use circleface instead of node style to avoid borking the x-axis
circle_face = ete3.CircleFace(radius=5,
color='brown',
style="circle")
ete3.add_face_to_node(circle_face,
node,
column=0,
position=position)
def write_legend(plotdir):
def get_leg_name(status):
if args.legends is not None and status in args.glslabels:
lname = args.legends[args.glslabels.index(status)]
elif status == 'both':
if len(args.glsfnames) == 2:
lname = 'both'
elif len(args.glsfnames) == 3:
lname = 'two'
else:
raise Exception('wtf %d' % len(args.glsfnames))
elif status == 'all':
if len(args.glsfnames) == 2:
lname = 'both'
elif len(args.glsfnames) == 3:
lname = 'all three'
else:
raise Exception('wtf %d' % len(args.glsfnames))
else:
lname = status
return lname
def add_stuff(status, leg_name, color):
legfo[leg_name] = color
if status in used_faces:
facefo[leg_name] = used_faces[status]
legfo, facefo = {}, {}
if args.ref_label is not None:
for status, color in simu_colors.items():
add_stuff(status, status, color)
else:
added_two_method_color = False
for status, color in used_colors.items():
if '-&-' in status:
for substatus in status.split(
'-&-'
): # arg, have to handle cases where the single one isn't in there
if get_leg_name(substatus) not in legfo:
add_stuff(substatus, get_leg_name(substatus),
scolors[substatus])
if not added_two_method_color:
leg_name = get_leg_name('both')
added_two_method_color = True
else:
continue
else:
leg_name = get_leg_name(status)
add_stuff(status, leg_name, color)
# figure out the order we want 'em in
lnames = sorted(legfo.keys())
for status in ['both', 'all']:
if get_leg_name(status) in lnames:
lnames.remove(get_leg_name(status))
lnames.append(get_leg_name(status))
etree = ete3.ClusterTree() #'(a);')
tstyle = ete3.TreeStyle()
tstyle.show_scale = False
# tstyle.show_leaf_name = False
# for node in etree.traverse():
# print node.name
# node.add_face(ete3.CircleFace(args.novel_dot_size, scolors['novel']), column=1) #, position='float') # if args.leaf_names else 'branch')
dummy_column = 0
pic_column = 1
text_column = 2
leg_title_height = 1.5 * args.leafheight # if args.legend_title is not None else 0.75 * args.leafheight
for icol in range(text_column + 1): # add a top border
tstyle.title.add_face(ete3.RectFace(0.9 * args.leafheight,
0.9 * args.leafheight,
fgcolor=None,
bgcolor=None),
column=icol)
tstyle.title.add_face(ete3.TextFace(' ', fsize=leg_title_height),
column=dummy_column) # adds a left border
if args.legend_title is not None:
tstyle.title.add_face(
ete3.TextFace('', fsize=leg_title_height), column=pic_column
) # keeps the first legend entry from getting added on this line
tstyle.title.add_face(
ete3.TextFace(args.legend_title,
fsize=leg_title_height,
fgcolor='black',
bold=True),
column=text_column
) # add an empty title so there's some white space at the top, even with no actual title text
for leg_name in lnames:
color = legfo[leg_name]
size_factor = 2.
if leg_name in facefo:
tstyle.title.add_face(
ete3.StackedBarFace([80., 20.],
width=size_factor * args.leafheight,
height=size_factor * args.leafheight,
colors=[color, facefo[leg_name]],
line_color='black'),
column=pic_column
) # looks like maybe they reversed fg/bg kwarg names
else:
tstyle.title.add_face(
ete3.RectFace(size_factor * args.leafheight,
size_factor * args.leafheight,
fgcolor='black',
bgcolor=color),
column=pic_column
) # looks like maybe they reversed fg/bg kwarg names
tstyle.title.add_face(ete3.TextFace(' ' + leg_name,
fsize=args.leafheight,
fgcolor='black'),
column=text_column)
tstyle.title.add_face(ete3.CircleFace(1.5 * args.novel_dot_size,
scolors['novel']),
column=pic_column)
tstyle.title.add_face(
ete3.TextFace('novel allele', fsize=args.leafheight),
column=text_column
) # keeps the first legend entry from getting added on this line
etree.render(plotdir + '/legend.svg', tree_style=tstyle)
scolors = {
'ok': 'DarkSeaGreen',
'missing': 'IndianRed',
'spurious': 'IndianRed',
'data': 'LightSteelBlue',
'both': 'LightGrey',
}
metafos = heads.read_metadata('kate-qrs')
for ds in metafos:
if 'LN1' in ds or 'LN2' in ds:
scolors[ds] = '#85ad98' # green
elif 'LN4' in ds or 'LN3' in ds:
scolors[ds] = '#94a3d1' # blue
faces = {
'missing': ete3.CircleFace(10, 'white'),
'spurious': ete3.CircleFace(10, 'black')
}
def get_cmdfos(cmdstr, workdir, outfname):
return [{'cmd_str': cmdstr, 'workdir': workdir, 'outfname': outfname}]
# ----------------------------------------------------------------------------------------
def make_tree(all_genes, workdir, use_cache=False):
aligned_fname = workdir + '/all-aligned.fa'
raxml_label = 'xxx'
raxml_output_fnames = [
'%s/RAxML_%s.%s' % (workdir, fn, raxml_label)
for fn in ['parsimonyTree', 'log', 'result', 'info', 'bestTree']
