Coverage for /home/deng/Projects/metatree_drawer/metatreedrawer/treeprofiler/layouts/text_layouts.py: 26%

299 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-08-07 10:33 +0200

1from collections import OrderedDict, namedtuple 

2from math import pi 

3import random, string 

4 

5from ete4.smartview import TreeStyle, NodeStyle, TreeLayout, PieChartFace 

6from ete4.smartview import Face, RectFace, CircleFace, SeqMotifFace, TextFace, OutlineFace 

7from ete4.smartview.renderer.draw_helpers import draw_text, draw_line, draw_array, draw_rect 

8from treeprofiler.layouts.general_layouts import get_piechartface, get_stackedbarface 

9 

10""" 

11label_layout, colorbranch_layout, rectangular_layout  

12""" 

13Box = namedtuple('Box', 'x y dx dy') # corner and size of a 2D shape 

14 

15class AlignLinkFace(Face): 

16 def __init__(self, width=70, height=None, 

17 stroke_color='gray', stroke_width=0.5, 

18 line_type=1, opacity=0.8): 

19 """Line types: 0 solid, 1 dotted, 2 dashed""" 

20 

21 Face.__init__(self, padding_x=0, padding_y=0) 

22 

23 self.line = None 

24 self.width = width 

25 self.height = height 

26 self.stroke_color = stroke_color 

27 self.stroke_width = stroke_width 

28 self.type = line_type; 

29 self.opacity = opacity 

30 

31 self.always_drawn = True 

32 

33 def __name__(self): 

34 return "AlignLinkFace" 

35 

36 def _compute_bounding_box(self, 

37 drawer, 

38 point, size, 

39 dx_to_closest_child, 

40 bdx, bdy, 

41 bdy0, bdy1, 

42 pos, row, 

43 n_row, n_col, 

44 dx_before, dy_before): 

45 

46 if drawer.NPANELS > 1 and drawer.viewport and pos == 'branch_right': 

47 x, y = point 

48 dx, dy = size 

49 p1 = (x + bdx + dx_before, y + dy/2) 

50 if drawer.TYPE == 'rect': 

51 p2 = (drawer.viewport.x + drawer.viewport.dx, y + dy/2) 

52 else: 

53 aligned = sorted(drawer.tree_style.aligned_grid_dxs.items()) 

54 # p2 = (drawer.node_size(drawer.tree)[0], y + dy/2) 

55 if not len(aligned): 

56 return Box(0, 0, 0, 0) 

57 p2 = (aligned[0][1] - bdx, y + dy/2) 

58 if p1[0] > p2[0]: 

59 return Box(0, 0, 0, 0) 

60 p1, p2 = cartesian(p1), cartesian(p2) 

61 

62 self.line = (p1, p2) 

63 

64 return Box(0, 0, 0, 0) # Should not take space 

65 

66 def get_box(self): 

67 return Box(0, 0, 0, 0) # Should not take space 

68 

69 def fits(self): 

70 return True 

71 

72 def _draw(self, drawer): 

73 

74 if drawer.NPANELS < 2: 

75 return None 

76 

77 style = { 

78 'type': self.type, 

79 'stroke': self.stroke_color, 

80 'stroke-width': self.stroke_width, 

81 'opacity': self.opacity, 

82 } 

83 if drawer.panel == 0 and drawer.viewport and\ 

84 (self.node.is_leaf or self.node.is_collapsed)\ 

85 and self.line: 

86 p1, p2 = self.line 

87 yield draw_line(p1, p2, 'align-link', style=style) 

88 

89 def get_random_string(self, length): 

90 """ Generates random string to nameless trees """ 

91 letters = string.ascii_lowercase 

92 result_str = ''.join(random.choice(letters) for i in range(length)) 

93 return result_str 

94 

95 

96 def compute_bounding_box(self, 

97 drawer, 

98 point, size, 

99 dx_to_closest_child, 

100 bdx, bdy, 

101 bdy0, bdy1, 

102 pos, row, 

103 n_row, n_col, 

104 dx_before, dy_before): 

105 

106 box = super().compute_bounding_box( 

107 drawer, 

108 point, size, 

109 dx_to_closest_child, 

110 bdx, bdy, 

111 bdy0, bdy1, 

112 pos, row, 

113 n_row, n_col, 

114 dx_before, dy_before) 

115 

116 x, y, dx, dy = box 

117 zx, zy = self.zoom 

118 zx = 1 if self.stretch\ 

119 and pos.startswith('aligned')\ 

120 and drawer.TYPE != 'circ'\ 

121 else zx 

122 

123 r = (x or 1e-10) if drawer.TYPE == 'circ' else 1 

124 

125 def get_dimensions(max_width, max_height): 

126 if not (max_width or max_height): 

127 return 0, 0 

128 if (type(max_width) in (int, float) and max_width <= 0) or\ 

129 (type(max_height) in (int, float) and max_height <= 0): 

130 return 0, 0 

131 

132 width = self.width / zx if self.width is not None else None 

133 height = self.height / zy if self.height is not None else None 

134 

135 if width is None: 

136 return max_width or 0, min(height or float('inf'), max_height) 

137 if height is None: 

138 return min(width, max_width or float('inf')), max_height 

139 

140 hw_ratio = height / width 

141 

142 if max_width and width > max_width: 

143 width = max_width 

144 height = width * hw_ratio 

145 if max_height and height > max_height: 

146 height = max_height 

147 if not self.stretch or drawer.TYPE == 'circ': 

148 width = height / hw_ratio 

149 

150 height /= r # in circular drawer 

151 return width, height 

152 

153 max_dy = dy * r # take into account circular mode 

154 

155 if pos == 'branch_right': 

156 width, height = get_dimensions(dx, max_dy) 

157 box = (x, y + (dy - height) / 2, width, height) 

158 

159 elif pos.startswith('aligned'): 

160 width, height = get_dimensions(None, dy) 

161 # height = min(dy, (self.height - 2 * self.padding_y) / zy) 

162 # width = min(self.width - 2 * self.padding_x) / zx 

163 

164 if pos == 'aligned_bottom': 

165 y = y + dy - height 

166 elif pos == 'aligned_top': 

167 y = y 

168 else: 

169 y = y + (dy - height) / 2 

170 

171 box = (x, y, width, height) 

172 

173 self._box = Box(*box) 

174 return self._box 

175 

176 def draw(self, drawer): 

177 self._check_own_variables() 

178 

179 circ_drawer = drawer.TYPE == 'circ' 

180 style = { 

181 'fill': self.stroke_color, 

182 'opacity': self.opacity, 

183 'stroke': self.stroke_color, 

184 'stroke-width': self.stroke_width 

185 } 

186 if circ_drawer: 

187 rect_id = self.get_random_string(10) 

188 style['id'] = rect_id 

189 

190 yield draw_rect(self._box, 

191 self.name, 

192 style=style, 

193 ) 

194 

195 

196class LayoutText(TreeLayout): 

197 def __init__(self, name, column, color_dict, text_prop, width=70, min_fsize=5, max_fsize=15, padding_x=1, padding_y=0, legend=True, aligned_faces=True): 

198 super().__init__(name, aligned_faces=aligned_faces) 

199 self.aligned_faces = True 

200 self.text_prop = text_prop 

201 self.column = column 

202 self.color_dict = color_dict 

203 self.internal_prop = text_prop+'_counter' 

204 self.legend = legend 

205 self.width = width 

206 self.height = None 

207 self.min_fsize = min_fsize 

208 self.max_fsize = max_fsize 

209 self.absence_color = "#EBEBEB" 

210 self.padding_x = padding_x 

211 self.padding_y = padding_y 

212 

213 def set_tree_style(self, tree, tree_style): 

214 super().set_tree_style(tree, tree_style) 

215 text = TextFace(self.text_prop, min_fsize=self.min_fsize, max_fsize=self.max_fsize, padding_x=self.padding_x, width=self.width, rotation=315) 

216 tree_style.aligned_panel_header.add_face(text, column=self.column) 

217 

218 if self.legend: 

219 if self.color_dict: 

220 tree_style.add_legend(title=self.text_prop, 

221 variable='discrete', 

222 colormap=self.color_dict 

223 ) 

224 

225 def set_node_style(self, node): 

226 if node.is_leaf and node.props.get(self.text_prop): 

227 prop_text = node.props.get(self.text_prop) 

228 if prop_text: 

229 if type(prop_text) == list: 

230 prop_text = ",".join(prop_text) 

231 else: 

232 pass 

233 if self.color_dict: 

234 prop_face = TextFace(prop_text, color=self.color_dict.get(prop_text, 'black'),min_fsize=self.min_fsize, max_fsize=self.max_fsize, padding_x=self.padding_x, width=self.width ) 

235 else: 

236 prop_face = TextFace(prop_text, color='black', min_fsize=self.min_fsize, max_fsize=self.max_fsize, padding_x=self.padding_x, width=self.width ) 

237 node.add_face(prop_face, column=self.column, position="aligned") 

238 

239 elif node.is_leaf and node.props.get(self.internal_prop): 

240 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

241 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=False) 

242 

243 elif node.props.get(self.internal_prop): 

244 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

245 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=True) 

246 # piechart_face = get_piechartface(node, self.internal_prop, self.color_dict, radius=25) 

247 # node.add_face(piechart_face, column = self.column, position = "branch_right", collapsed_only=False) 

248 # node.add_face(piechart_face, column = self.column, position = "branch_right", collapsed_only=True) 

249 else: 

250 #prop_face = CircleFace(radius=self.radius, color='grey', padding_x=self.padding_x, padding_y=self.padding_y) 

251 prop_face = RectFace(width=self.width, height=self.height, color=self.absence_color, \ 

252 padding_x=self.padding_x , padding_y=self.padding_y, tooltip=None) 

253 node.add_face(prop_face, column=self.column, position="aligned", collapsed_only=True) 

254 

255class LayoutColorbranch(TreeLayout): 

256 def __init__(self, name, column, color_dict, text_prop, legend=True, width=70, padding_x=1, padding_y=0): 

257 super().__init__(name) 

258 self.aligned_faces = True 

259 self.text_prop = text_prop 

260 self.column = column 

261 self.color_dict = color_dict 

262 self.internal_prop = text_prop+'_counter' 

263 self.legend = legend 

264 self.height = None 

265 self.absence_color = "#EBEBEB" 

266 self.width = width 

267 self.padding_x = padding_x 

268 self.padding_y = padding_y 

269 

270 def set_tree_style(self, tree, tree_style): 

271 super().set_tree_style(tree, tree_style) 

272 text = TextFace(self.text_prop, min_fsize=5, max_fsize=15, padding_x=self.padding_x, width=self.width, rotation=315) 

273 tree_style.aligned_panel_header.add_face(text, column=self.column) 

274 if self.legend: 

275 if self.color_dict: 

276 tree_style.add_legend(title=self.text_prop, 

277 variable='discrete', 

278 colormap=self.color_dict 

279 ) 

280 

281 def set_node_style(self, node): 

282 prop_text = node.props.get(self.text_prop) 

283 if prop_text is not None: 

284 if type(prop_text) == list: 

285 prop_text = ",".join(prop_text) 

286 else: 

287 pass 

288 

289 if self.color_dict: 

290 node.add_face(TextFace(node.name, color = self.color_dict.get(prop_text,""), 

291 padding_x=self.padding_x),column=0, position="branch_right") 

292 node.add_face(TextFace(node.name, color = self.color_dict.get(prop_text,""), 

293 padding_x=self.padding_x),column=self.column, position="branch_right", collapsed_only=True) 

294 

295 node.sm_style["hz_line_color"] = self.color_dict.get(prop_text,"") 

296 node.sm_style["hz_line_width"] = 2 

297 node.sm_style["vt_line_color"] = self.color_dict.get(prop_text,"") 

298 node.sm_style["vt_line_width"] = 2 

299 node.sm_style['outline_color'] = self.color_dict.get(prop_text,"") 

300 node.add_face(RectFace(width=self.width, height=None, color=self.absence_color, \ 

301 padding_x=self.padding_x , padding_y=self.padding_y, tooltip=None),column=self.column, position="aligned") 

302 

303 elif node.is_leaf and node.props.get(self.internal_prop): 

304 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

305 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=False) 

306 

307 

308 if node.props.get(self.internal_prop): 

309 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

310 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=True) 

311 

312 else: 

313 prop_face = RectFace(width=self.width, height=self.height, color=self.absence_color, \ 

314 padding_x=self.padding_x , padding_y=self.padding_y, tooltip=None) 

315 node.add_face(prop_face, column=self.column, position="aligned", collapsed_only=True) 

316 

317class LayoutRect(TreeLayout): 

318 def __init__(self, name, column, color_dict, text_prop, width=70, height=None, padding_x=1, padding_y=0, legend=True): 

319 super().__init__(name) 

320 self.aligned_faces = True 

321 self.text_prop = text_prop 

322 self.column = column 

323 self.color_dict = color_dict 

324 self.absence_color = "#EBEBEB" 

325 self.internal_prop = text_prop+'_counter' 

326 self.legend = legend 

327 self.width = width 

328 self.height = height 

329 self.min_fsize = 5 

330 self.max_fsize = 15 

331 self.padding_x = padding_x 

332 self.padding_y = padding_y 

333 

334 def set_tree_style(self, tree, tree_style): 

335 super().set_tree_style(tree, tree_style) 

336 text = TextFace(self.text_prop, min_fsize=self.min_fsize, max_fsize=self.max_fsize, padding_x=self.padding_x, width=self.width, rotation=315) 

337 tree_style.aligned_panel_header.add_face(text, column=self.column) 

338 

339 if self.legend: 

340 if self.color_dict: 

341 self.color_dict['NA'] = self.absence_color 

342 tree_style.add_legend(title=self.text_prop, 

343 variable='discrete', 

344 colormap=self.color_dict 

345 ) 

346 else: 

347 tree_style.add_legend(title=self.text_prop, 

348 variable='discrete', 

349 colormap={'NA':self.absence_color} 

350 ) 

351 def set_node_style(self, node): 

352 if node.is_leaf: 

353 prop_text = node.props.get(self.text_prop) 

354 if prop_text: 

355 if type(prop_text) == list: 

356 prop_text = ",".join(prop_text) 

357 else: 

358 pass 

359 tooltip = "" 

360 if node.name: 

361 tooltip += f'<b>{node.name}</b><br>' 

362 if self.text_prop: 

363 tooltip += f'<br>{self.text_prop}: {prop_text}<br>' 

364 

365 if self.color_dict: 

366 color = self.color_dict.get(str(prop_text),"") 

367 prop_face = RectFace(width=self.width, height=self.height, color=color, \ 

368 padding_x=self.padding_x , padding_y=self.padding_y, tooltip=tooltip) 

369 node.add_face(prop_face, column=self.column, position="aligned") 

370 else: 

371 #prop_face = CircleFace(radius=self.radius, color='grey', padding_x=self.padding_x, padding_y=self.padding_y) 

372 prop_face = RectFace(width=self.width, height=self.height, text="NA", color=self.absence_color, \ 

373 padding_x=self.padding_x , padding_y=self.padding_y, tooltip=None) 

374 node.add_face(prop_face, column=self.column, position="aligned") 

375 

376 elif node.is_leaf and node.props.get(self.internal_prop): 

377 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

378 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=False) 

379 

380 elif node.props.get(self.internal_prop): 

381 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

382 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=True) 

383 

384 else: 

385 

386 prop_face = RectFace(width=self.width, height=self.height, color=self.absence_color, \ 

387 padding_x=self.padding_x , padding_y=self.padding_y, tooltip=None) 

388 node.add_face(prop_face, column=self.column, position="aligned", collapsed_only=True) 

389 

390class LayoutPiechart(TreeLayout): 

391 def __init__(self, name, color_dict, text_prop, radius=20, padding_x=1, padding_y=0, legend=True, aligned_faces=True): 

392 super().__init__(name, aligned_faces=aligned_faces) 

393 self.aligned_faces = True 

394 

395 self.text_prop = text_prop+"_counter" 

396 self.internal_prop = text_prop+'_counter' 

397 self.color_dict = color_dict 

398 self.radius = radius 

399 self.padding_x = padding_x 

400 self.padding_y = padding_y 

401 self.legend = legend 

402 

403 def set_tree_style(self, tree, tree_style): 

404 super().set_tree_style(tree, tree_style) 

405 if self.legend: 

406 if self.color_dict: 

407 tree_style.add_legend(title=self.text_prop, 

408 variable='discrete', 

409 colormap=self.color_dict 

410 ) 

411 

412 def set_node_style(self, node): 

413 if not node.is_leaf: 

414 if node.props.get(self.internal_prop): 

415 piechart_face = get_piechartface(node, self.internal_prop, self.color_dict, radius=self.radius) 

416 node.add_face(piechart_face, column = 1, position = "branch_right", collapsed_only=False) 

417 node.add_face(piechart_face, column = 1, position = "branch_right", collapsed_only=True) 

418 

419class LayoutBackground(TreeLayout): 

420 def __init__(self, name, color_dict, text_prop, width=70, column=0, 

421 padding_x=1, padding_y=0, legend=True, aligned_faces=True): 

422 super().__init__(name, aligned_faces=aligned_faces) 

423 

424 self.aligned_faces = True 

425 

426 self.text_prop = text_prop 

427 self.internal_prop = text_prop+'_counter' 

428 self.color_dict = color_dict 

429 self.width = width 

430 self.column = column 

431 self.padding_x = padding_x 

432 self.padding_y = padding_y 

433 self.absence_color = "#EBEBEB" 

434 self.legend = legend 

435 

436 def set_tree_style(self, tree, tree_style): 

437 super().set_tree_style(tree, tree_style) 

438 if self.legend: 

439 if self.color_dict: 

440 self.color_dict['NA'] = self.absence_color 

441 tree_style.add_legend(title=self.text_prop, 

442 variable='discrete', 

443 colormap=self.color_dict 

444 ) 

445 

446 def set_node_style(self, node): 

447 prop_text = node.props.get(self.text_prop) 

448 if prop_text: 

449 if type(prop_text) == list: 

450 prop_text = ",".join(prop_text) 

451 else: 

452 pass 

453 tooltip = "" 

454 if node.name: 

455 tooltip += f'<b>{node.name}</b><br>' 

456 if self.text_prop: 

457 tooltip += f'<br>{self.text_prop}: {prop_text}<br>' 

458 

459 if self.color_dict: 

460 color = self.color_dict.get(str(prop_text), self.absence_color) 

461 align_link_face = AlignLinkFace(width=self.width*2, height=None, 

462 stroke_color=color, stroke_width=2, line_type=0, opacity=0.8) 

463 node.sm_style["bgcolor"] = color 

464 node.sm_style["fgopacity"] = 0.5 

465 node.sm_style['outline_color'] = color 

466 node.add_face(align_link_face, 

467 position='branch_right', 

468 #column=self.column, 

469 ) 

470 elif node.is_leaf and node.props.get(self.internal_prop): 

471 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

472 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=False) 

473 

474 if node.props.get(self.internal_prop): 

475 stackedbar_face = get_stackedbarface(node, self.internal_prop, self.color_dict, width=self.width, padding_x=self.padding_x, padding_y=self.padding_y) 

476 node.add_face(stackedbar_face, column = self.column, position = "aligned", collapsed_only=True) 

477 # else: 

478 # prop_face = RectFace(width=self.width, height=None, color=self.absence_color, \ 

479 # padding_x=self.padding_x , padding_y=self.padding_y, tooltip=None) 

480 # node.add_face(prop_face, column=self.column, position="aligned", collapsed_only=True)