Coverage for /home/deng/Projects/metatree_drawer/metatreedrawer/treeprofiler/layouts/conditional_layouts.py: 23%
135 statements
« prev ^ index » next coverage.py v7.2.7, created at 2024-08-07 10:33 +0200
« prev ^ index » next coverage.py v7.2.7, created at 2024-08-07 10:33 +0200
1from ete4.smartview import TreeStyle, NodeStyle, TreeLayout, PieChartFace
2from ete4.smartview import (RectFace, CircleFace, SeqMotifFace, TextFace, OutlineFace, \
3 SelectedFace, SelectedCircleFace, SelectedRectFace, LegendFace)
4from treeprofiler.layouts.general_layouts import get_heatmapface, get_aggregated_heatmapface
5from treeprofiler.src.utils import to_code, call, counter_call, check_nan
6# for boolean layouts
7try:
8 from distutils.util import strtobool
9except ImportError:
10 from treeprofiler.src.utils import strtobool
12# branch thicken, background highlighted to purple
13def highlight_layout(conditions, level, prop2type={}, color='purple'):
14 conditional_output = to_code(conditions)
15 def layout_fn(node):
16 final_call = False
17 for condition in conditional_output:
18 #normal
20 op = condition[1]
21 if op == 'in':
22 value = condition[0]
23 prop = condition[2]
24 datatype = prop2type.get(prop)
25 final_call = call(node, prop, datatype, op, value)
27 elif ':' in condition[0] :
28 internal_prop, leaf_prop = condition[0].split(':')
29 value = condition[2]
30 datatype = prop2type[internal_prop]
31 final_call = counter_call(node, internal_prop, leaf_prop, datatype, op, value)
32 else:
33 prop = condition[0]
34 value = condition[2]
35 datatype = prop2type.get(prop)
36 final_call = call(node, prop, datatype, op, value)
38 if final_call == False:
39 break
40 else:
41 continue
43 if final_call:
45 #prop_face = SelectedRectFace(name='prop')
46 node.sm_style["bgcolor"] = color # highligh clade
47 #node.sm_style["hz_line_width"] = 5
48 #node.add_face(prop_face, column=level, position = "branch_right")
49 while (node):
50 node = node.up
51 if node:
52 node.sm_style["hz_line_width"] = 5
53 #node.sm_style["hz_line_color"] = color
54 #node.add_face(OutlineFace(color=color), column=level, collapsed_only=True)
55 return layout_fn
56 return
58# conditional collapse layouts
59def collapsed_by_layout(conditions, level, prop2type={}, color='red'):
60 conditional_output = to_code(conditions)
61 def layout_fn(node):
62 final_call = False
63 for condition in conditional_output:
64 #normal
66 op = condition[1]
67 if op == 'in':
68 value = condition[0]
69 prop = condition[2]
70 datatype = prop2type.get(prop)
71 final_call = call(node, prop, datatype, op, value)
73 elif ':' in condition[0] :
74 internal_prop, leaf_prop = condition[0].split(':')
75 value = condition[2]
76 datatype = prop2type[internal_prop]
77 final_call = counter_call(node, internal_prop, leaf_prop, datatype, op, value)
78 else:
79 prop = condition[0]
80 value = condition[2]
81 datatype = prop2type.get(prop)
82 final_call = call(node, prop, datatype, op, value)
84 if final_call == False:
85 break
86 else:
87 continue
88 if final_call:
89 if not node.is_root:
90 node.sm_style["draw_descendants"] = False
91 node.sm_style["outline_color"] = color
92 return layout_fn
93 return
95class LayoutBinary(TreeLayout):
96 def __init__(self, name=None, level=1, color='#E60A0A', \
97 bool_prop=None, reverse=False, aggregate=False, \
98 max_count=0, \
99 radius=25, padding_x=1, padding_y=0, width=70, \
100 legend=True):
101 super().__init__(name)
102 self.aligned_faces = True
103 self.bool_prop = bool_prop
104 self.column = level
105 self.color = color
106 self.negative_color = '#EBEBEB'
107 self.internal_prop = bool_prop+'_counter'
108 self.reverse = reverse
109 self.aggregate = aggregate
110 self.max_count = max_count
111 self.radius = radius
112 self.padding_x = padding_x
113 self.padding_y = padding_y
114 self.legend = legend
115 self.width = width
116 self.height = None
117 self.min_fsize = 5
118 self.max_fsize = 10
121 # def set_tree_style(self, tree, tree_style):
122 # super().set_tree_style(tree, tree_style)
123 # text = TextFace(self.name, max_fsize=11, padding_x=1)
124 # tree_style.aligned_panel_header.add_face(text, column=self.column)
125 def update_header_width(self):
126 return
128 def set_tree_style(self, tree, tree_style):
129 super().set_tree_style(tree, tree_style)
130 text = TextFace(self.bool_prop, min_fsize=10, max_fsize=15, padding_x=self.padding_x, width=self.width, rotation=315)
131 tree_style.aligned_panel_header.add_face(text, column=self.column)
132 if self.legend:
134 if self.reverse:
135 title = 'ReverseBinary_' + self.bool_prop
136 colormap = {
137 "False": self.color,
138 "True" : self.negative_color,
139 "NA": 'white'
140 }
141 tree_style.add_legend(title=title,
142 variable='discrete',
143 colormap=colormap,
144 )
145 else:
146 title = 'Binary_' + self.bool_prop
147 colormap = {
148 "True": self.color,
149 "False" : self.negative_color,
150 "NA": 'white'
151 }
152 tree_style.add_legend(title=title,
153 variable='discrete',
154 colormap=colormap,
155 )
158 def set_node_style(self, node):
159 # need to correct
160 if node.is_leaf and node.props.get(self.bool_prop):
161 #if node.props.get(self.bool_prop):
162 prop_bool = node.props.get(self.bool_prop)
164 if not check_nan(prop_bool):
165 str2bool = strtobool(prop_bool)
166 tooltip = ""
167 if node.name:
168 tooltip += f'<b>{node.name}</b><br>'
169 if self.bool_prop:
170 tooltip += f'<br>{self.bool_prop}: {node.props.get(self.bool_prop)}<br>'
172 if self.reverse:
173 if not bool(str2bool):
174 #prop_face = CircleFace(radius=self.radius, color=self.color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
175 prop_face = RectFace(width=self.width, height=self.height, color=self.color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
176 node.add_face(prop_face, column=self.column, position = "aligned")
177 else:
178 #prop_face = CircleFace(radius=self.radius, color=self.negative_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
179 prop_face = RectFace(width=self.width, height=self.height, color=self.negative_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
180 node.add_face(prop_face, column=self.column, position = "aligned")
181 else:
182 if bool(str2bool):
183 #prop_face = CircleFace(radius=self.radius, color=self.color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
184 prop_face = RectFace(width=self.width, height=self.height, color=self.color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
185 node.add_face(prop_face, column=self.column, position = "aligned")
186 else:
187 #prop_face = CircleFace(radius=self.radius, color=self.negative_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
188 prop_face = RectFace(width=self.width, height=self.height, color=self.negative_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
189 node.add_face(prop_face, column=self.column, position = "aligned")
190 else: #mising
191 prop_face = RectFace(width=self.width, height=self.height, text="NA", color=self.negative_color, padding_x=self.padding_x, padding_y=self.padding_y, stroke_color=self.negative_color, tooltip=None)
192 node.add_face(prop_face, column=self.column, position = "aligned")
193 # else:
194 # prop_face = RectFace(width=self.width, height=self.height, text="NA", color=self.negative_color, padding_x=self.padding_x, padding_y=self.padding_y, stroke_color=self.negative_color, tooltip=None)
195 # node.add_face(prop_face, column=self.column, position = "aligned")
197 elif node.is_leaf and node.props.get(self.internal_prop):
198 if self.aggregate:
199 heatmapFace = get_aggregated_heatmapface(node, self.internal_prop, max_color=self.color, width=self.width, height=self.height, padding_x=self.padding_x, max_count=self.max_count)
200 else:
201 heatmapFace = get_heatmapface(node, self.internal_prop, max_color=self.color, width=self.width, height=self.height, padding_x=self.padding_x,)
202 node.add_face(heatmapFace, column = self.column, position = "aligned", collapsed_only=False)
204 elif node.props.get(self.internal_prop):
205 if self.aggregate:
206 heatmapFace = get_aggregated_heatmapface(node, self.internal_prop, max_color=self.color, width=self.width, height=self.height, padding_x=self.padding_x, max_count=self.max_count)
207 else:
208 heatmapFace = get_heatmapface(node, self.internal_prop, max_color=self.color, width=self.width, height=self.height, padding_x=self.padding_x,reverse=self.reverse)
209 node.add_face(heatmapFace, column = self.column, position = "aligned", collapsed_only=True)
210 # else:
211 # prop_face = RectFace(width=self.width, height=self.height, text="NA", color=self.negative_color, padding_x=self.padding_x, padding_y=self.padding_y, stroke_color=self.negative_color, tooltip=None)
212 # node.add_face(prop_face, column=self.column, position = "aligned")