Coverage for /home/deng/Projects/metatree_drawer/metatreedrawer/treeprofiler/layouts/taxon_layouts.py: 29%
167 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
2from ete4.smartview import RectFace, CircleFace, SeqMotifFace, TextFace, OutlineFace, LegendFace
3from collections import OrderedDict
4from treeprofiler.src import utils
5paried_color = ["red", "darkblue", "darkgreen", "darkyellow", "violet", "mediumturquoise", "sienna", "lightCoral", "lightSkyBlue", "indigo", "tan", "coral", "olivedrab", "teal"]
7#collapse in layout
8#kingdom, phylum, class, order, family, genus, species, subspecies
9def get_level(node, level=0):
10 if node.is_root:
11 return level
12 else:
13 return get_level(node.up, level + 1)
15def summary(nodes):
16 "Return a list of names summarizing the given list of nodes"
17 return list(OrderedDict((first_name(node), None) for node in nodes).keys())
19def first_name(tree):
20 "Return the name of the first node that has a name"
22 sci_names = []
23 for node in tree.traverse('preorder'):
24 if node.is_leaf:
25 sci_name = node.props.get('sci_name')
26 sci_names.append(sci_name)
28 return next(iter(sci_names))
30class TaxaClade(TreeLayout):
31 def __init__(self, name, level, rank, color_dict, legend=True):
32 super().__init__(name, aligned_faces=True)
34 self.activate = False
35 self.name = name
36 self.column = level
37 self.rank = rank
38 self.color_dict = color_dict
39 self.legend = legend
41 def set_tree_style(self, tree, tree_style):
42 super().set_tree_style(tree, tree_style)
43 if self.legend:
44 if self.color_dict:
45 tree_style.add_legend(title=self.rank,
46 variable='discrete',
47 colormap=self.color_dict,
48 )
50 # def set_node_style(self, node):
51 # if not node.is_root and node.props.get('rank') == self.rank:
52 # if node.props.get('sci_name'):
53 # node.sm_style["bgcolor"] = self.color_dict[node.props.get('sci_name')] # highligh clade
55 def set_node_style(self, node):
56 named_lineage = node.props.get('named_lineage', None)
57 if named_lineage:
58 for clade, color in self.color_dict.items():
59 if clade in named_lineage:
60 node.sm_style["hz_line_color"] = color
61 node.sm_style["hz_line_width"] = 2
62 node.sm_style["vt_line_color"] = color
63 node.sm_style["vt_line_width"] = 2
64 #node.sm_style["draw_descendants"] = False
65 node.sm_style["outline_color"] = color
66 break
68 if not node.is_leaf:
69 # Collapsed face
70 if node.props.get('rank') == self.rank:
71 if node.props.get('sci_name') is not None:
72 sci_name = node.props.get('sci_name')
73 color = self.color_dict.get(sci_name, 'gray')
74 node.add_face(TextFace(sci_name, padding_x=2, color = color),
75 position="branch_right", column=1, collapsed_only=True)
78 # if not node.is_root and node.props.get('rank') == self.rank:
79 # if node.props.get('sci_name'):
80 # color = self.color_dict[node.props.get('sci_name')]
81 # node.sm_style["hz_line_color"] = color
82 # node.sm_style["hz_line_width"] = 2
83 # node.sm_style["vt_line_color"] = color
84 # node.sm_style["vt_line_width"] = 2
85 # #node.sm_style["draw_descendants"] = False
86 # node.sm_style["outline_color"] = color
91class LayoutSciName(TreeLayout):
92 def __init__(self, name="Scientific name", color_dict={}):
93 super().__init__(name, aligned_faces=True)
94 self.color_dict = color_dict
96 def set_node_style(self, node):
97 if node.is_leaf:
98 sci_name = node.props.get('sci_name')
99 prot_id = node.name
101 rank_colordict = self.color_dict.get(node.props.get('rank'),'')
102 if rank_colordict:
103 color = rank_colordict.get(sci_name, 'gray')
104 else:
105 color = 'gray'
106 node.add_face(TextFace(sci_name, color = color, padding_x=2, min_fsize=4, max_fsize=25),
107 column=0, position="branch_right")
109 if len(prot_id) > 40:
110 prot_id = prot_id[0:37] + " ..."
112 #node.add_face(TextFace(prot_id, color = 'Gray', padding_x=2), column = 2, position = "aligned")
113 else:
114 # Collapsed face
115 names = summary(node.children)
116 texts = names if len(names) < 6 else (names[:3] + ['...'] + names[-2:])
117 for i, text in enumerate(texts):
118 sci_name = node.props.get('sci_name')
119 rank_colordict = self.color_dict.get(node.props.get('rank'),'')
120 if rank_colordict:
121 color = rank_colordict.get(sci_name, 'gray')
122 else:
123 color = 'gray'
124 node.add_face(TextFace(text, padding_x=2, color = color, min_fsize=4, max_fsize=25),
125 position="branch_right", column=1, collapsed_only=True)
127class TaxaRectangular(TreeLayout):
128 def __init__(self, name="Last common ancestor", rank=None, color_dict={}, rect_width=20, column=0, legend=True ):
129 super().__init__(name, aligned_faces=True)
131 self.active = True
132 self.rank = rank
133 self.color_dict = color_dict
134 self.rect_width = rect_width
135 self.column = column
137 def set_tree_style(self, tree, tree_style):
138 super().set_tree_style(tree, tree_style)
139 if self.legend:
140 if self.color_dict:
141 tree_style.add_legend(title='TaxaRectangular_'+self.rank,
142 variable='discrete',
143 colormap=self.color_dict,
144 )
146 def set_node_style(self, node):
147 node_rank = node.props.get('rank')
148 node_sciname = node.props.get('sci_name')
149 if node_sciname and (node_rank == self.rank):
150 lca = node_sciname
151 color = self.color_dict.get(lca, 'lightgray')
152 level = get_level(node, level=self.column)
153 tooltip = ""
154 if node.name:
155 tooltip += f'<b>{node_sciname}</b><br>'
156 if lca:
157 tooltip += f'rank: {node_rank}<br>'
158 tooltip += f'sci_name: {lca}<br>'
160 lca_face = RectFace(self.rect_width, None, text = lca, color=color, padding_x=1, padding_y=1, tooltip=tooltip)
161 lca_face.rotate_text = True
162 node.add_face(lca_face, position='aligned', column=level)
163 node.add_face(lca_face, position='aligned', column=level,
164 collapsed_only=True)
165 # else:
166 # named_lineage = node.props.get('named_lineage', None)
167 # lca = next((elem for elem in named_lineage if elem in self.taxa_list), None)
168 # if lca:
169 # color = self.color_dict.get(lca, 'lightgray')
170 # #level = get_level(node, level=self.column)
171 # level = self.column
172 # tooltip = ""
173 # if node.name:
174 # tooltip += f'<b>{node.name}</b><br>'
175 # if lca:
176 # tooltip += f'rank: {node_rank}<br>'
177 # tooltip += f'sci_name: {lca}<br>'
179 # lca_face = RectFace(self.rect_width, None, text = lca, color=color, padding_x=1, padding_y=1, tooltip=tooltip)
180 # lca_face.rotate_text = True
181 # node.add_face(lca_face, position='aligned', column=level)
182 # node.add_face(lca_face, position='aligned', column=level,
183 # collapsed_only=True)
185class TaxaCollapse(TreeLayout):
186 def __init__(self, name="Last common ancestor", rank=None, color_dict={}, rect_width=20, column=0, padding_x=1, padding_y=0, legend=True ):
187 super().__init__(name, aligned_faces=True)
189 self.active = True
190 self.rank = rank
191 self.color_dict=color_dict
192 self.rect_width = rect_width
193 self.column = column
194 self.padding_x = padding_x
195 self.padding_y = padding_y
197 self.taxa_list = list(color_dict.keys())
199 def set_tree_style(self, tree, tree_style):
200 super().set_tree_style(tree, tree_style)
201 super().set_tree_style(tree, tree_style)
202 text = TextFace(" ", min_fsize=10, max_fsize=15, padding_x=self.padding_x, width=self.rect_width, rotation=315)
203 tree_style.aligned_panel_header.add_face(text, column=self.column)
204 if self.legend:
205 if self.color_dict:
206 tree_style.add_legend(title='TaxaRectangular_'+self.rank,
207 variable='discrete',
208 colormap=self.color_dict,
209 )
211 def set_node_style(self, node):
212 node_rank = node.props.get('rank')
213 node_sciname = node.props.get('sci_name')
214 named_lineage = node.props.get('named_lineage', None)
215 #lca = next((elem for elem in named_lineage if elem in self.taxa_list), None)
216 lca_value = node.props.get('lca')
217 if lca_value:
218 lca_dict = utils.string_to_dict(lca_value)
219 lca = lca_dict.get(self.rank, None)
220 if lca:
221 color = self.color_dict.get(lca, 'lightgray')
222 #level = get_level(node, level=self.column)
223 level = self.column
224 tooltip = ""
225 if node.name:
226 tooltip += f'<b>{lca}</b><br>'
227 if lca:
228 tooltip += f'rank: {self.rank}<br>'
229 tooltip += f'sci_name: {lca}<br>'
231 lca_face = RectFace(self.rect_width, None, text = lca, color=color, padding_x=1, padding_y=1, tooltip=tooltip)
232 lca_face.rotate_text = True
233 node.sm_style["draw_descendants"] = False
235 node.add_face(lca_face, position='aligned', column=level)
236 node.add_face(lca_face, position='aligned', column=level,
237 collapsed_only=True)
238 node.add_face(TextFace(lca, color = color, padding_x=2),
239 column=1, position="branch_right", collapsed_only=True)
241 # elif node_sciname and (node_rank == self.rank):
242 # lca = node_sciname
243 # color = self.color_dict.get(lca, 'lightgray')
244 # level = self.column
245 # tooltip = ""
246 # if node.name:
247 # tooltip += f'<b>{node.name}</b><br>'
248 # if lca:
249 # tooltip += f'rank: {node_rank}<br>'
250 # tooltip += f'sci_name: {lca}<br>'
252 # lca_face = RectFace(self.rect_width, None, text = lca, color=color, padding_x=1, padding_y=1, tooltip=tooltip)
253 # lca_face.rotate_text = True
254 # node.sm_style["draw_descendants"] = False
255 # node.add_face(lca_face, position='aligned', column=level)
256 # node.add_face(lca_face, position='aligned', column=level,
257 # collapsed_only=True)
259class LayoutEvolEvents(TreeLayout):
260 def __init__(self, name="Evolutionary events",
261 prop="evoltype",
262 speciation_color="blue",
263 duplication_color="red", node_size = 2,
264 legend=True):
265 super().__init__(name)
267 self.prop = prop
268 self.speciation_color = speciation_color
269 self.duplication_color = duplication_color
270 self.node_size = node_size
271 self.legend = legend
273 self.active = True
275 def set_tree_style(self, tree, tree_style):
276 super().set_tree_style(tree, tree_style)
277 if self.legend:
278 colormap = { "Speciation event": self.speciation_color,
279 "Duplication event": self.duplication_color }
280 tree_style.add_legend(title=self.name,
281 variable="discrete",
282 colormap=colormap)
284 def set_node_style(self, node):
285 if not node.is_leaf:
286 if node.props.get(self.prop, "") == "S":
287 node.sm_style["fgcolor"] = self.speciation_color
288 node.sm_style["size"] = self.node_size
290 elif node.props.get(self.prop, "") == "D":
291 node.sm_style["fgcolor"] = self.duplication_color
292 node.sm_style["size"] = self.node_size