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

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"] 

6 

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) 

14 

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()) 

18 

19def first_name(tree): 

20 "Return the name of the first node that has a name" 

21 

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) 

27 

28 return next(iter(sci_names)) 

29 

30class TaxaClade(TreeLayout): 

31 def __init__(self, name, level, rank, color_dict, legend=True): 

32 super().__init__(name, aligned_faces=True) 

33 

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 

40 

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 ) 

49 

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 

54 

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 

67 

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) 

76 

77 

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 

87 

88 

89 

90 

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 

95 

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 

100 

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") 

108 

109 if len(prot_id) > 40: 

110 prot_id = prot_id[0:37] + " ..." 

111 

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) 

126 

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) 

130 

131 self.active = True 

132 self.rank = rank 

133 self.color_dict = color_dict 

134 self.rect_width = rect_width 

135 self.column = column 

136 

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 ) 

145 

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>' 

159 

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>' 

178 

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) 

184 

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) 

188 

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 

196 

197 self.taxa_list = list(color_dict.keys()) 

198 

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 ) 

210 

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>' 

230 

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 

234 

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) 

240 

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>' 

251 

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) 

258 

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) 

266 

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 

272 

273 self.active = True 

274 

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) 

283 

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 

289 

290 elif node.props.get(self.prop, "") == "D": 

291 node.sm_style["fgcolor"] = self.duplication_color 

292 node.sm_style["size"] = self.node_size