def parseTree(root):
tree = Tree()
tree.name = root['Name']
tree.add_face(TextFace(root['Split'], fgcolor="red"), column=0, position="branch-bottom")
if root['Children']:
for child in root['Children']:
tree.children.append(parseTree(child))
return tree
def to_tree_node(self):
t = Tree(f"{self.function_type};", format=1)
for child in self.children:
t.add_child(child.to_tree_node())
tf = TextFace(f"{self.function_type}")
tf.rotation = -90
t.add_face(tf, column=1, position="branch-top")
return t
def to_tree_node(self):
t = Tree("IF;", format=1)
t.add_child(self.children[1].to_tree_node())
t.add_child(self.children[0].to_tree_node())
t.add_child(self.children[2].to_tree_node())
tf = TextFace("IF")
tf.rotation = -90
t.add_face(tf, column=1, position="branch-top")
return t
def plot_event_tree(self):
t = Tree("((1:0.001, 0.1:0.23, 0.001, >0.001), NIE);")
t = Tree(
"((D: 0.723274, F: 0.567784)1.000000: 0.067192, (B: 0.279326, H: 0.756049)1.000000: 0.807788);"
)
t.support = 0.1
ts = TreeStyle()
ts.show_leaf_name = True
ts.show_branch_length = True
ts.show_branch_support = True
t.add_face(TextFace(" hola "), column=0, position='branch-right')
t.show(tree_style=ts)
def frogTree(startstate):
global frogCount
frogCount = 0
t = Tree()
root = t.add_child(name=startstate)
t.add_face(TextFace(startstate), column=0)
frogalgo(root, root)
messagebox.showinfo("Tree Generation Done",
"number of states:" + str(frogCount))
states.set("Tree Generated! number of states:" + str(frogCount))
ts = TreeStyle()
ts.show_leaf_name = False
ts.show_scale = False
t.show(tree_style=ts)
def get_example_tree():
t = Tree("((a,b),c);")
right_c0_r0 = TextFace("right_col0_row0")
right_c0_r1 = TextFace("right_col0_row1")
right_c1_r0 = TextFace("right_col1_row0")
right_c1_r1 = TextFace("right_col1_row1")
right_c1_r2 = TextFace("right_col1_row2")
top_c0_r0 = TextFace("top_col0_row0")
top_c0_r1 = TextFace("top_col0_row1")
bottom_c0_r0 = TextFace("bottom_col0_row0")
bottom_c1_r0 = TextFace("bottom_col1_row0")
aligned_c0_r0 = TextFace("aligned_col0_row0")
aligned_c0_r1 = TextFace("aligned_col0_row1")
aligned_c1_r0 = TextFace("aligned_col1_row0")
aligned_c1_r1 = TextFace("aligned_col1_row1")
t.add_face(right_c0_r1, column=1, position="branch-right")
t.add_face(right_c0_r0, column=0, position="branch-right")
t.add_face(right_c1_r2, column=2, position="branch-right")
t.add_face(right_c1_r1, column=1, position="branch-right")
t.add_face(right_c1_r0, column=0, position="branch-right")
t.add_face(top_c0_r1, column=1, position="branch-top")
t.add_face(top_c0_r0, column=0, position="branch-top")
t.add_face(bottom_c0_r0, column=0, position="branch-bottom")
t.add_face(bottom_c1_r0, column=1, position="branch-bottom")
for leaf in t.iter_leaves():
leaf.add_face(aligned_c0_r1, 0, "aligned")
leaf.add_face(aligned_c0_r0, 0, "aligned")
leaf.add_face(aligned_c1_r1, 0, "aligned")
leaf.add_face(aligned_c1_r0, 0, "aligned")
return t, TreeStyle()
def get_example_tree():
t = Tree("((a,b),c);")
right_c0_r0 = TextFace("right_col0_row0")
right_c0_r1 = TextFace("right_col0_row1")
right_c1_r0 = TextFace("right_col1_row0")
right_c1_r1 = TextFace("right_col1_row1")
right_c1_r2 = TextFace("right_col1_row2")
top_c0_r0 = TextFace("top_col0_row0")
top_c0_r1 = TextFace("top_col0_row1")
bottom_c0_r0 = TextFace("bottom_col0_row0")
bottom_c1_r0 = TextFace("bottom_col1_row0")
aligned_c0_r0 = TextFace("aligned_col0_row0")
aligned_c0_r1 = TextFace("aligned_col0_row1")
aligned_c1_r0 = TextFace("aligned_col1_row0")
aligned_c1_r1 = TextFace("aligned_col1_row1")
t.add_face(right_c0_r1, column=1, position="branch-right")
t.add_face(right_c0_r0, column=0, position="branch-right")
t.add_face(right_c1_r2, column=2, position="branch-right")
t.add_face(right_c1_r1, column=1, position="branch-right")
t.add_face(right_c1_r0, column=0, position="branch-right")
t.add_face(top_c0_r1, column=1, position="branch-top")
t.add_face(top_c0_r0, column=0, position="branch-top")
t.add_face(bottom_c0_r0, column=0, position="branch-bottom")
t.add_face(bottom_c1_r0, column=1, position="branch-bottom")
for leaf in t.iter_leaves():
leaf.add_face(aligned_c0_r1, 0, "aligned")
leaf.add_face(aligned_c0_r0, 0, "aligned")
leaf.add_face(aligned_c1_r1, 0, "aligned")
leaf.add_face(aligned_c1_r0, 0, "aligned")
return t, TreeStyle()
def tree_to_png(self, filepath: str):
nodes = list()
for i in range(self.leafs.shape[0]):
thisNode = Tree()
thisNode.add_face(
faces.BarChartFace(
self.leafs[i].detach().cpu().numpy(),
labels=self.action_labels,
min_value=0.0,
max_value=self.leafs[i].max().detach().cpu().numpy() + 1e-7,
), 0)
nodes.append(thisNode)
for d in range(self.depth-1):
for node_i in range(self.nodes_beta[d].shape[0]):
thisNode = Tree()
thisNode.add_child(nodes.pop(1))
thisNode.add_child(nodes.pop(0))
beta = F.softmax(self.nodes_beta[d][node_i].squeeze(), 0
).detach().cpu().numpy()
phi = self.nodes_phi[d][node_i].squeeze().detach().cpu().item()
thisNode.add_face(
faces.BarChartFace(
beta,
labels=self.labels,
min_value=0.0,
max_value=1.0
), 0)
thisNode.add_face(
faces.TextFace('phi: {0:.3f}'.format(phi)), 0)
nodes.append(thisNode)
if filepath is not None:
nodes[0].render(filepath)
return nodes[0]
def partitionTreeSet(N):
if N == 1:
x = Tree(";",format=100)
x.add_features(value=N, name=str(N))
xFace = styleFace(x.name)
x.add_face(xFace,column=0,position="branch-top")
return (x,)
else:
y = ()
base = Tree(";",format=100)
base.dist = 1
for k in range(lam(N)):
left = partitionTreeSet(N-(k+1))
right = partitionTreeSet(k+1)
for l in left:
for r in right:
l.dist = 1
r.dist = 1
z = base.copy()
z.dist = 1
z.add_features(value=N, name=str(N))
z.add_child(l.copy())
z.add_child(r.copy())
zFace = styleFace(z.name)
z.add_face(zFace,column=0,position="branch-top")
y = y + (z,)
return y
def add_Edges(node, tree, rootNode):
def calculate_thickness(child, style):
thickness = (child.total_children / rootNode.total_children) * 60
if child.rank == "Superkingdom" or child.rank == "Kingdom": # just to offset MASSIVE thick lines since those two
thickness = 3 # will typically contain >90% of all results
if thickness < 3: # just to offset later nodes having invisible borders
thickness = 3
style["vt_line_width"] = thickness
style["hz_line_width"] = thickness
return style #unecessary since it passes by reference, but it might be useful in future
def highlight_tax(child, style):
ranks = {'Name':'FireBrick',
'Species':'Crimson',
'Genus':'Chocolate',
'Family':'Gold',
'Order':'LawnGreen',
'Class':'LimeGreen',
'Phylum':'Turquoise',
'Kingdom':'SteelBlue',
'Superkingdom':'Plum'}
style["bgcolor"] = ranks[child.rank]
return style #unecessary since it passes by reference, but it might be useful in future
for child in node.child_nodes:
if child.rank is not ("Name"): # Restricting it to above Species means it won't crash rendering a tree with thousands of results
newNode = Tree(name=child.node_name)
# Set borders
name = TextFace(child.node_name)
name.border.width = 1
name.border.color = '#ffffff'
# Add name, total children, and rank.
newNode.add_face(name,column=0,position="branch-right")
children = "[" + str(child.total_children) + "]"
newNode.add_face(TextFace(children),column=0, position="branch-bottom")
newNode.add_face(TextFace(child.rank),column=0, position="branch-top")
# Set colour based on rank and thickness based on total number of node's children.
style = NodeStyle()
highlight_tax(child, style)
calculate_thickness(child, style)
newNode.set_style(style)
# Keep iterating
tree.add_child(newNode)
add_Edges(child, newNode, rootNode)
def hard_tree_to_png(self, filepath: str):
nodes = list()
leaf_tmpl = '{}: y_{}'
for i in range(self.leafs.shape[0]):
thisNode = Tree()
if self.continuous:
thisNode.add_face(
faces.BarChartFace(
self.leafs[i].detach().cpu().numpy(),
min_value=0.0,
max_value=self.leafs[i].max().detach().cpu().numpy() + 1e-7,
labels=self.action_labels
), 0)
else:
max_leaf_idx = np.argmax(self.leafs[i].detach().cpu().numpy())
thisNode.add_face(faces.TextFace(
leaf_tmpl.format(
self.action_labels[max_leaf_idx],
max_leaf_idx)), 0)
nodes.append(thisNode)
node_tmpl = '{}: x_{} >= {}'
for d in range(self.depth-1):
for node_i in range(self.nodes_beta[d].shape[0]):
thisNode = Tree()
thisNode.add_child(nodes.pop(1))
thisNode.add_child(nodes.pop(0))
beta = F.softmax(self.nodes_beta[d][node_i].squeeze(), 0
).detach().cpu().numpy()
phi = self.nodes_phi[d][node_i].squeeze().detach().cpu().item()
max_beta_idx = np.argmax(beta)
thisNode.add_face(faces.TextFace(node_tmpl.format(
self.labels[max_beta_idx],
max_beta_idx,
phi)), 0)
nodes.append(thisNode)
if filepath is not None:
nodes[0].render(filepath,)
return nodes[0]
def node_viz(self):
thisNode = Tree()
thisNode.set_style(NodeStyle(shape='square'))
if self.child2 is not None:
thisNode.add_child(self.child2.node_viz())
if self.child1 is not None:
thisNode.add_child(self.child1.node_viz())
if self.child1 is not None and self.child2 is not None:
thisNode.add_face(
faces.BarChartFace(
self.beta.squeeze().detach().cpu().numpy(),
labels=self.labels,
min_value=0.0,
max_value=self.beta.max().detach().cpu().item()), 0)
thisNode.add_face(
faces.TextFace("phi: {0:.3f}".format(self.phi.item())), 0)
else:
thisNode.add_face(
faces.BarChartFace(
self.beta.squeeze().detach().cpu().numpy(),
min_value=0.0,
max_value=self.beta.max().detach().cpu().item()), 0)
return thisNode
from ete3 import Tree, TreeStyle, TextFace # @UndefinedVariable
rand_newick = "((water, estuarine_open_water, inland_water)Water_bodies, artificial_surface, cloud, shadow, forest, grassland)Root;"
rand_tree = "rand_tree"
with open(rand_tree, 'w') as TREE:
TREE.write(rand_newick)
# Reading tree t1 and t2
t1 = Tree(rand_newick, format=8) # @UndefinedVariable
#t2 = Tree(rand_tree) # @UndefinedVariable
ts = TreeStyle()
ts.show_leaf_name = True
ts.branch_vertical_margin = 20
ts.title.add_face(TextFace("Auckland LCDB", fsize=20), column=0)
ts.scale = 50
t1.add_face(TextFace("Root "), column=0, position = "branch-top")
for node in t1.traverse():
if node.name == "Water_bodies":
node.add_face(TextFace("Water_bodies"), column=0, position = "branch-top")
print(t1)
t1.render("mytree.png", tree_style=ts, dpi=300)
aligned_c1_r0 = TextFace("aligned_col1_row0")
aligned_c1_r1 = TextFace("aligned_col1_row1")
all_faces = [right_c0_r0, right_c0_r1, right_c1_r0, right_c1_r1, right_c1_r2, top_c0_r0, \
top_c0_r1, bottom_c0_r0, bottom_c0_r1, aligned_c0_r0, aligned_c0_r1,\
aligned_c1_r0, aligned_c1_r1]
# set a border in all faces
for f in all_faces:
# f.margin_border.width = 1
f.margin_bottom = 5
f.margin_top = 5
f.margin_right = 10
t.add_face(right_c0_r0, column=0, position="branch-right")
t.add_face(right_c0_r1, column=0, position="branch-right")
t.add_face(right_c1_r0, column=1, position="branch-right")
t.add_face(right_c1_r1, column=1, position="branch-right")
t.add_face(right_c1_r2, column=1, position="branch-right")
t.add_face(top_c0_r0, column=0, position="branch-top")
t.add_face(top_c0_r1, column=0, position="branch-top")
t.add_face(bottom_c0_r0, column=0, position="branch-bottom")
t.add_face(bottom_c0_r1, column=0, position="branch-bottom")
a = t&"a"
a.set_style(NodeStyle())
a.img_style["bgcolor"] = "lightgreen"
aligned_c1_r0 = TextFace("aligned_col1_row0")
aligned_c1_r1 = TextFace("aligned_col1_row1")
all_faces = [right_c0_r0, right_c0_r1, right_c1_r0, right_c1_r1, right_c1_r2, top_c0_r0, \
top_c0_r1, bottom_c0_r0, bottom_c0_r1, aligned_c0_r0, aligned_c0_r1,\
aligned_c1_r0, aligned_c1_r1]
# set a border in all faces
for f in all_faces:
f.margin_border.width = 1
f.margin_bottom = 5
f.margin_top = 5
f.margin_right = 10
t.add_face(right_c0_r0, column=0, position="branch-right")
t.add_face(right_c0_r1, column=0, position="branch-right")
t.add_face(right_c1_r0, column=1, position="branch-right")
t.add_face(right_c1_r1, column=1, position="branch-right")
t.add_face(right_c1_r2, column=1, position="branch-right")
t.add_face(top_c0_r0, column=0, position="branch-top")
t.add_face(top_c0_r1, column=0, position="branch-top")
t.add_face(bottom_c0_r0, column=0, position="branch-bottom")
t.add_face(bottom_c0_r1, column=0, position="branch-bottom")
a = t&"a"
a.set_style(NodeStyle())
a.img_style["bgcolor"] = "lightgreen"
Each area represents a table in which faces can be added through the TreeNode.add_face() method. For instance,
if two text labels want to be drawn bellow the branch line of a given node,
a pair of TextFace faces can be created and added to the columns 0 and 1 of the branch-bottom area:
refer ot annotate_node_ete.py for graphical display
'''
from ete3 import Tree, TreeStyle, TextFace, faces
t = Tree("((a,b),c);")
# Basic tree style
ts = TreeStyle()
# ts.show_leaf_name = True
# Add two text faces to different columns
t.add_face(TextFace("hola "), column=0, position="branch-right")
t.add_face(TextFace("mundo!"), column=1, position="branch-bottom")
t.add_face(TextFace("here"), column=0, position="branch-bottom")
# t.show(tree_style=ts)
# t.render("/groups/itay_mayrose/udiland/crispys_tomato/families/filtering_scripts/face.pdf")
t = Tree("((a,b),c);")
print(t)
def layout(node):
if len(node) > 1:
faces.add_face_to_node(TextFace("hola"),
node,
column=0,
position="branch-right")
def get_example_tree():
t = Tree()
t.populate(10)
# Margins, alignment, border, background and opacity can now be set for any face
rs1 = faces.TextFace("branch-right\nmargins&borders",
fsize=12,
fgcolor="#009000")
rs1.margin_top = 10
rs1.margin_bottom = 50
rs1.margin_left = 40
rs1.margin_right = 40
rs1.border.width = 1
rs1.background.color = "lightgreen"
rs1.inner_border.width = 0
rs1.inner_border.line_style = 1
rs1.inner_border.color = "red"
rs1.opacity = 0.6
rs1.hz_align = 2 # 0 left, 1 center, 2 right
rs1.vt_align = 1 # 0 left, 1 center, 2 right
br1 = faces.TextFace("branch-right1", fsize=12, fgcolor="#009000")
br2 = faces.TextFace("branch-right3", fsize=12, fgcolor="#009000")
# New face positions (branch-top and branch-bottom)
bb = faces.TextFace("branch-bottom 1", fsize=8, fgcolor="#909000")
bb2 = faces.TextFace("branch-bottom 2", fsize=8, fgcolor="#909000")
bt = faces.TextFace("branch-top 1", fsize=6, fgcolor="#099000")
# And faces can also be used as headers or foot notes of aligned
# columns
t1 = faces.TextFace("Header Face", fsize=12, fgcolor="#aa0000")
t2 = faces.TextFace("Footer Face", fsize=12, fgcolor="#0000aa")
# Attribute faces can now contain prefix and suffix fixed text
aligned = faces.AttrFace("name",
fsize=12,
fgcolor="RoyalBlue",
text_prefix="Aligned (",
text_suffix=")")
# horizontal and vertical alignment per face
aligned.hz_align = 1 # 0 left, 1 center, 2 right
aligned.vt_align = 1
# Node style handling is no longer limited to layout functions. You
# can now create fixed node styles and use them many times, save them
# or even add them to nodes before drawing (this allows to save and
# reproduce an tree image design)
style = NodeStyle()
style["fgcolor"] = "Gold"
style["shape"] = "square"
style["size"] = 15
style["vt_line_color"] = "#ff0000"
t.set_style(style)
# add a face to the style. This face will be render in any node
# associated to the style.
fixed = faces.TextFace("FIXED branch-right", fsize=11, fgcolor="blue")
t.add_face(fixed, column=1, position="branch-right")
# Bind the precomputed style to the root node
# ETE 2.1 has now support for general image properties
ts = TreeStyle()
# You can add faces to the tree image (without any node
# associated). They will be used as headers and foot notes of the
# aligned columns (aligned faces)
ts.aligned_header.add_face(t1, column=0)
ts.aligned_header.add_face(t1, 1)
ts.aligned_header.add_face(t1, 2)
ts.aligned_header.add_face(t1, 3)
t1.hz_align = 1 # 0 left, 1 center, 2 right
t1.border.width = 1
ts.aligned_foot.add_face(t2, column=0)
ts.aligned_foot.add_face(t2, 1)
ts.aligned_foot.add_face(t2, 2)
ts.aligned_foot.add_face(t2, 3)
t2.hz_align = 1
# Set tree image style. Note that aligned header and foot is only
# visible in "rect" mode.
ts.mode = "r"
ts.scale = 10
for node in t.traverse():
# If node is a leaf, add the nodes name and a its scientific
# name
if node.is_leaf():
node.add_face(aligned, column=0, position="aligned")
node.add_face(aligned, column=1, position="aligned")
node.add_face(aligned, column=3, position="aligned")
else:
node.add_face(bt, column=0, position="branch-top")
node.add_face(bb, column=0, position="branch-bottom")
node.add_face(bb2, column=0, position="branch-bottom")
node.add_face(br1, column=0, position="branch-right")
node.add_face(rs1, column=0, position="branch-right")
node.add_face(br2, column=0, position="branch-right")
return t, ts
def get_example_tree():
t = Tree()
t.populate(10)
# Margins, alignment, border, background and opacity can now be set for any face
rs1 = faces.TextFace("branch-right\nmargins&borders",
fsize=12, fgcolor="#009000")
rs1.margin_top = 10
rs1.margin_bottom = 50
rs1.margin_left = 40
rs1.margin_right = 40
rs1.border.width = 1
rs1.background.color = "lightgreen"
rs1.inner_border.width = 0
rs1.inner_border.line_style = 1
rs1.inner_border.color= "red"
rs1.opacity = 0.6
rs1.hz_align = 2 # 0 left, 1 center, 2 right
rs1.vt_align = 1 # 0 left, 1 center, 2 right
br1 = faces.TextFace("branch-right1", fsize=12, fgcolor="#009000")
br2 = faces.TextFace("branch-right3", fsize=12, fgcolor="#009000")
# New face positions (branch-top and branch-bottom)
bb = faces.TextFace("branch-bottom 1", fsize=8, fgcolor="#909000")
bb2 = faces.TextFace("branch-bottom 2", fsize=8, fgcolor="#909000")
bt = faces.TextFace("branch-top 1", fsize=6, fgcolor="#099000")
# And faces can also be used as headers or foot notes of aligned
# columns
t1 = faces.TextFace("Header Face", fsize=12, fgcolor="#aa0000")
t2 = faces.TextFace("Footer Face", fsize=12, fgcolor="#0000aa")
# Attribute faces can now contain prefix and suffix fixed text
aligned = faces.AttrFace("name", fsize=12, fgcolor="RoyalBlue",
text_prefix="Aligned (", text_suffix=")")
# horizontal and vertical alignment per face
aligned.hz_align = 1 # 0 left, 1 center, 2 right
aligned.vt_align = 1
# Node style handling is no longer limited to layout functions. You
# can now create fixed node styles and use them many times, save them
# or even add them to nodes before drawing (this allows to save and
# reproduce an tree image design)
style = NodeStyle()
style["fgcolor"] = "Gold"
style["shape"] = "square"
style["size"] = 15
style["vt_line_color"] = "#ff0000"
t.set_style(style)
# add a face to the style. This face will be render in any node
# associated to the style.
fixed = faces.TextFace("FIXED branch-right", fsize=11, fgcolor="blue")
t.add_face(fixed, column=1, position="branch-right")
# Bind the precomputed style to the root node
# ETE 2.1 has now support for general image properties
ts = TreeStyle()
# You can add faces to the tree image (without any node
# associated). They will be used as headers and foot notes of the
# aligned columns (aligned faces)
ts.aligned_header.add_face(t1, column = 0)
ts.aligned_header.add_face(t1, 1)
ts.aligned_header.add_face(t1, 2)
ts.aligned_header.add_face(t1, 3)
t1.hz_align = 1 # 0 left, 1 center, 2 right
t1.border.width = 1
ts.aligned_foot.add_face(t2, column = 0)
ts.aligned_foot.add_face(t2, 1)
ts.aligned_foot.add_face(t2, 2)
ts.aligned_foot.add_face(t2, 3)
t2.hz_align = 1
# Set tree image style. Note that aligned header and foot is only
# visible in "rect" mode.
ts.mode = "r"
ts.scale = 10
for node in t.traverse():
# If node is a leaf, add the nodes name and a its scientific
# name
if node.is_leaf():
node.add_face(aligned, column=0, position="aligned")
node.add_face(aligned, column=1, position="aligned")
node.add_face(aligned, column=3, position="aligned")
else:
node.add_face(bt, column=0, position="branch-top")
node.add_face(bb, column=0, position="branch-bottom")
node.add_face(bb2, column=0, position="branch-bottom")
node.add_face(br1, column=0, position="branch-right")
node.add_face(rs1, column=0, position="branch-right")
node.add_face(br2, column=0, position="branch-right")
return t, ts
def build_tree(output,
input_seq,
StOS=-3.0,
EOS=-2.,
PAD=-4.,
threshold=0.0,
optimal=None):
from ete3 import TextFace, NodeStyle
root = Tree(name=str(input_seq), dist=1.0)
_label = TextFace(str(input_seq))
root.add_face(_label, column=0)
nodes = {'StOS': root}
EOS_style = NodeStyle()
StOS_style = NodeStyle()
irrelevant_style = NodeStyle()
style = NodeStyle()
decoded_style = NodeStyle()
opt_style = NodeStyle()
# style
style['vt_line_color'] = '#000000'
style['hz_line_color'] = '#000000'
style['vt_line_type'] = 0
style['hz_line_type'] = 0
#decoded
decoded_style['vt_line_color'] = '#000000'
decoded_style['hz_line_color'] = '#000000'
decoded_style['vt_line_type'] = 0
decoded_style['hz_line_type'] = 0
decoded_style['bgcolor'] = 'DarkSeaGreen'
# EOS
EOS_style['vt_line_color'] = '#0000ff'
EOS_style['hz_line_color'] = '#0000ff'
EOS_style['vt_line_type'] = 2
EOS_style['hz_line_type'] = 2
# StOS
StOS_style['vt_line_color'] = '#000000'
StOS_style['hz_line_color'] = '#00ff00'
StOS_style['vt_line_type'] = 0
StOS_style['hz_line_type'] = 0
StOS_style['hz_line_width'] = 5
# cut
irrelevant_style['vt_line_color'] = '#ffdddd'
irrelevant_style['hz_line_color'] = '#ffdddd'
irrelevant_style['vt_line_type'] = 1
irrelevant_style['hz_line_type'] = 1
for tree in output:
# format(1-tree[3],'.2f')
try:
_parent = str(tree[1]) + str(tree[2] - 1)
node_name = str(tree[1] + [tree[0]]) + str(tree[2])
if tree[0] == EOS:
_name = '[EOS]'
elif tree[0] == PAD:
_name = '[PAD]'
else:
_name = str([int(tree[0])])
_label = TextFace(_name)
if tree[1:3] == [[StOS], 1] or tree[1:3] == [[str(StOS)], 1]:
# support = tree[3]
nodes[node_name] = nodes['StOS'].add_child(name=_name,
dist=1,
support=1)
nodes[node_name].add_feature('prob', tree[3])
nodes[node_name].add_face(_label, column=0)
nodes['StOS'].set_style(StOS_style)
else:
_prob = tree[3] * nodes[_parent].support
_prob = 1
if tree[0] == EOS or tree[0] == PAD:
nodes[node_name] = nodes[_parent].add_child(name=_name,
dist=1,
support=_prob)
nodes[node_name].add_feature('prob', tree[3])
if _prob < threshold:
nodes[node_name].set_style(irrelevant_style)
else:
nodes[node_name].set_style(EOS_style)
nodes[node_name].add_face(_label, column=0)
else:
nodes[node_name] = nodes[_parent].add_child(name=_name,
dist=1,
support=_prob)
nodes[node_name].add_feature('prob', tree[3])
if _prob < threshold:
nodes[node_name].set_style(irrelevant_style)
else:
nodes[node_name].add_face(_label, column=0)
nodes[node_name].set_style(style)
except:
pass
return root, nodes
def to_tree_node(self):
t = Tree(f"{self.children[0]};", format=1)
tf = TextFace(f"{self.children[0]}")
tf.rotation = -90
t.add_face(tf, column=1, position="branch-bottom")
return t
