Coverage for /home/deng/Projects/metatree_drawer/metatreedrawer/treeprofiler/layouts/staple_layouts.py: 19%

412 statements  

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

1import matplotlib as mpl 

2import numpy as np 

3 

4from ete4.smartview import TreeStyle, NodeStyle, TreeLayout 

5from ete4.smartview import TextFace, Face, ScaleFace, LegendFace, RectFace 

6from ete4.smartview.renderer.draw_helpers import * 

7from treeprofiler.src.utils import random_color, add_suffix 

8 

9import colorsys 

10 

11 

12 

13__all__ = [ "LayoutBarplot" ] 

14 

15def heatmap_gradient(hue, intensity, granularity): 

16 min_lightness = 0.35 

17 max_lightness = 0.9 

18 base_value = intensity 

19 

20 # each gradient must contain 100 lightly descendant colors 

21 colors = [] 

22 rgb2hex = lambda rgb: '#%02x%02x%02x' % rgb 

23 l_factor = (max_lightness-min_lightness) / float(granularity) 

24 l = min_lightness 

25 while l <= max_lightness: 

26 l += l_factor 

27 rgb = rgb2hex(tuple(map(lambda x: int(x*255), colorsys.hls_to_rgb(hue, l, base_value)))) 

28 colors.append(rgb) 

29 

30 colors.append("#ffffff") 

31 return list(reversed(colors)) 

32 

33def color_gradient(c1, c2, mix=0): 

34 """ Fade (linear interpolate) from color c1 (at mix=0) to c2 (mix=1) """ 

35 # https://stackoverflow.com/questions/25668828/how-to-create-colour-gradient-in-python 

36 c1 = np.array(mpl.colors.to_rgb(c1)) 

37 c2 = np.array(mpl.colors.to_rgb(c2)) 

38 return mpl.colors.to_hex((1-mix)*c1 + mix*c2) 

39 

40def swap_pos(pos, angle): 

41 if abs(angle) >= pi / 2: 

42 if pos == 'branch_top': 

43 pos = 'branch_bottom' 

44 elif pos == 'branch_bottom': 

45 pos = 'branch_top' 

46 return pos 

47 

48class LayoutPlot(TreeLayout): 

49 def __init__(self, name=None, prop=None, width=200, size_prop=None, 

50 color_prop=None, color_gradient=None, color="red", colors=None, 

51 position="aligned", column=0, padding_x=10, padding_y=0, size_range=[], 

52 internal_rep='avg', scale=True, legend=True, active=True): 

53 super().__init__(name, 

54 aligned_faces=True if position == "aligned" else False, 

55 legend=legend, active=active) 

56 

57 self.width = width 

58 self.position = position 

59 self.column = column 

60 

61 self.scale = scale 

62 self.padding_x = padding_x 

63 self.padding_y = padding_y 

64 

65 self.internal_rep = internal_rep 

66 self.prop = prop 

67 # if not (size_prop or color_prop): 

68 # raise InvalidUsage("Either size_prop or color_prop required") 

69 

70 self.size_prop = size_prop 

71 self.color_prop = color_prop 

72 self.size_range = size_range 

73 self.size_range = size_range 

74 

75 self.color = color 

76 self.colors = colors 

77 self.color_gradient = color_gradient 

78 if self.color_prop and not self.color_gradient: 

79 self.color_gradient = ("#FFF", self.color) 

80 

81 def set_tree_style(self, tree, tree_style): 

82 super().set_tree_style(tree, tree_style) 

83 def update_vals(metric, node): 

84 p, minval, maxval, uniqvals = vals[metric] 

85 prop = node.props.get(p) 

86 try: 

87 prop = float(prop) # prop by default is string  

88 if type(prop) in [int, float]: 

89 vals[metric][1] = min(minval, prop) 

90 vals[metric][2] = max(maxval, prop) 

91 elif prop is None or prop == "": 

92 return 

93 else: 

94 uniqvals.add(prop) 

95 except: 

96 pass 

97 

98 # only when size_range is not provided, calculate min and max values 

99 if self.size_range == []: 

100 vals = { 

101 "size": [ self.size_prop, 0, 0, set() ], # min, max, unique 

102 "color": [ self.color_prop, 0, 0, set() ] # min, max, unique 

103 } 

104 

105 for node in tree.traverse(): 

106 if self.size_prop: 

107 update_vals("size", node) 

108 

109 if self.color_prop: 

110 update_vals("color", node) 

111 

112 if self.size_prop: 

113 self.size_range = vals["size"][1:3] 

114 # if self.color_prop: 

115 # unique = vals["color"][3] 

116 # if len(unique): 

117 # colors = self.colors or random_color(num=len(unique)) 

118 # if type(colors) == dict: 

119 # self.colors = colors.copy() 

120 # else: 

121 # colors = list(colors) 

122 # self.colors = {} 

123 # for idx, value in enumerate(unique): 

124 # self.colors[value] = colors[idx % len(colors)] 

125 # if self.legend: 

126 # tree_style.add_legend(title=self.name, 

127 # variable="discrete", 

128 # colormap=self.colors) 

129 # else: 

130 # self.color_range = vals["color"][1:3] 

131 # if self.legend: 

132 # tree_style.add_legend(title=self.name,  

133 # variable="continuous", 

134 # value_range=self.color_range, 

135 # color_range=self.color_gradient) 

136 

137 # def get_size(self, node, prop): 

138 # if self.size_range != [0, 0]: 

139 # minval, maxval = self.size_range 

140 # else: 

141 # minval, maxval = 0,1 

142 # return float(node.props.get(prop, 0)) / float(maxval) * self.width 

143 

144 def get_size(self, node, prop): 

145 if not self.size_prop: 

146 return self.width 

147 minval, maxval = self.size_range 

148 return float(node.props.get(prop, 0)) / maxval * self.width 

149 

150 def get_color(self, node): 

151 if not self.color_prop: 

152 return self.color 

153 prop = node.props.get(self.color_prop) 

154 if prop is None: 

155 return None 

156 if self.color_range: 

157 minval, maxval = self.color_range 

158 return color_gradient(*self.color_gradient, (prop - minval) / maxval) 

159 else: 

160 return self.colors.get(prop) 

161 

162 def get_legend(self): 

163 return self.legend 

164 

165class LayoutBarplot(LayoutPlot): 

166 def __init__(self, name=None, prop=None, width=200, size_prop=None, 

167 color_prop=None, position="aligned", column=0, 

168 color_gradient=None, color=None, colors=None, 

169 padding_x=10, padding_y=0, scale=True, legend=True, active=True, 

170 internal_rep='avg', scale_size=None, size_range=None, scale_range=None): 

171 

172 name = name or f'Barplot_{size_prop}_{color_prop}' 

173 super().__init__(name=name, prop=prop, width=width, size_prop=size_prop, 

174 color_prop=color_prop, position=position, column=column, 

175 color_gradient=color_gradient, colors=colors, color=color, 

176 padding_x=padding_x, padding_y=padding_y, scale=scale, size_range=size_range, 

177 legend=legend, active=active, 

178 internal_rep=internal_rep) 

179 

180 def set_tree_style(self, tree, tree_style): 

181 super().set_tree_style(tree, tree_style) 

182 if self.scale and self.size_range: 

183 self.scale_width = self.width 

184 self.scale_range = self.size_range 

185 scale = ScaleFace(width=self.width, scale_range=self.size_range, 

186 formatter='%.2f', 

187 padding_x=self.padding_x, padding_y=2) 

188 text = TextFace(self.prop, max_fsize=15, padding_x=self.padding_x, rotation=315) 

189 tree_style.aligned_panel_header.add_face(scale, column=self.column) 

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

191 

192 if self.legend: 

193 if self.color: 

194 colormap = {self.prop: self.color 

195 } 

196 else: 

197 colormap = self.colors 

198 

199 tree_style.add_legend(title=self.prop, 

200 variable='discrete', 

201 colormap=colormap 

202 ) 

203 

204 def get_color(self, node, color_prop, color_dict): 

205 if color_dict and color_prop: 

206 prop = node.props.get(color_prop) 

207 if prop is None: 

208 return None 

209 else: 

210 return color_dict.get(prop, None) 

211 else: 

212 return self.color 

213 

214 def set_node_style(self, node): 

215 internal_prop = self.prop + '_' + self.internal_rep 

216 if node.props.get(self.prop) is not None: 

217 if node.is_leaf: 

218 width = self.get_size(node, self.prop) 

219 color = self.get_color(node, self.color_prop, self.colors) 

220 tooltip = "" 

221 if node.name: 

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

223 if self.size_prop: 

224 tooltip += f'<br>{self.prop}: {node.props.get(self.prop)}<br>' 

225 if self.color_prop: 

226 tooltip += f'<br>{self.color_prop}: {color}<br>' 

227 face = RectFace(width, None, color=color, 

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

229 node.add_face(face, position=self.position, column=self.column, 

230 collapsed_only=False) 

231 

232 elif node.is_leaf and node.props.get(internal_prop): 

233 width = self.get_size(node, internal_prop) 

234 color = self.get_color(node, self.color_prop, self.colors) 

235 tooltip = "" 

236 if node.name: 

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

238 if self.size_prop: 

239 tooltip += f'<br>{self.prop}: {node.props.get(internal_prop)}<br>' 

240 if self.color_prop: 

241 tooltip += f'<br>{self.color_prop}: {color}<br>' 

242 face = RectFace(width, None, color=color, 

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

244 node.add_face(face, position=self.position, column=self.column, 

245 collapsed_only=False) 

246 

247 elif node.props.get(internal_prop): 

248 width = self.get_size(node, internal_prop) 

249 color = self.get_color(node, self.color_prop, self.colors) 

250 tooltip = "" 

251 if node.name: 

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

253 if self.size_prop: 

254 tooltip += f'<br>{self.size_prop}: {node.props.get(internal_prop)}<br>' 

255 if self.color_prop: 

256 tooltip += f'<br>{self.color_prop}: {color}<br>' 

257 face = RectFace(width, None, color=color, 

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

259 node.add_face(face, position=self.position, column=self.column, 

260 collapsed_only=True) 

261 

262class LayoutHeatmap(TreeLayout): 

263 def __init__(self, name=None, column=0, width=70, height=None, 

264 padding_x=1, padding_y=0, heatmap_prop=None, internal_rep=None, 

265 value_color=None, value_range=[], color_range=None, minval=0, maxval=None, 

266 absence_color="#EBEBEB", 

267 legend=True): 

268 

269 super().__init__(name) 

270 self.aligned_faces = True 

271 

272 self.heatmap_prop = heatmap_prop 

273 self.internal_prop = add_suffix(heatmap_prop, internal_rep) 

274 self.column = column 

275 self.value_color = value_color 

276 self.value_range = value_range 

277 self.color_range = color_range 

278 self.absence_color = absence_color 

279 self.maxval = maxval 

280 self.minval = minval 

281 

282 self.width = width 

283 self.height = height 

284 self.padding_x = padding_x 

285 self.padding_y = padding_y 

286 

287 def set_tree_style(self, tree, tree_style): 

288 super().set_tree_style(tree, tree_style) 

289 

290 text = TextFace(self.heatmap_prop, padding_x=self.padding_x, width=self.width, rotation=315) 

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

292 

293 if self.legend: 

294 tree_style.add_legend(title=self.heatmap_prop, 

295 variable='continuous', 

296 value_range=self.value_range , 

297 color_range=[ 

298 self.color_range.get(20), 

299 self.color_range.get(10), 

300 self.color_range.get(1), 

301 ] 

302 ) 

303 def set_node_style(self, node): 

304 heatmap_num = node.props.get(self.heatmap_prop) 

305 if heatmap_num is not None and heatmap_num != 'NaN': 

306 heatmap_num = float(heatmap_num) 

307 if node.is_leaf: 

308 # heatmap 

309 tooltip = "" 

310 if node.name: 

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

312 if self.heatmap_prop: 

313 tooltip += f'<br>{self.heatmap_prop}: {heatmap_num}<br>' 

314 

315 gradient_color = self.value_color.get(heatmap_num) 

316 

317 if gradient_color: 

318 identF = RectFace(width=self.width, height=self.height, text="%.2f" % (float(heatmap_num)), \ 

319 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip) 

320 node.add_face(identF, column = self.column, position = 'aligned') 

321 

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

323 heatmap_num = node.props.get(self.internal_prop) 

324 heatmap_num = float(heatmap_num) 

325 # heatmap 

326 tooltip = "" 

327 if node.name: 

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

329 if self.heatmap_prop: 

330 tooltip += f'<br>{self.heatmap_prop}: {heatmap_num}<br>' 

331 

332 gradient_color = self.value_color.get(heatmap_num) 

333 

334 if gradient_color: 

335 identF = RectFace(width=self.width, height=self.height, text="%.2f" % (float(heatmap_num)), \ 

336 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip) 

337 node.add_face(identF, column = self.column, position = 'aligned', collapsed_only=True) 

338 

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

340 heatmap_num = node.props.get(self.internal_prop) 

341 heatmap_num = float(heatmap_num) 

342 # heatmap 

343 tooltip = "" 

344 if node.name: 

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

346 if self.heatmap_prop: 

347 tooltip += f'<br>{self.heatmap_prop}: {heatmap_num}<br>' 

348 

349 gradient_color = self.value_color.get(heatmap_num) 

350 

351 if gradient_color: 

352 identF = RectFace(width=self.width, height=self.height, text="%.2f" % (float(heatmap_num)), \ 

353 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip) 

354 node.add_face(identF, column = self.column, position = 'aligned', collapsed_only=True) 

355 

356 else: 

357 # heatmap 

358 tooltip = "" 

359 if node.name: 

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

361 if self.heatmap_prop: 

362 tooltip += f'<br>{self.heatmap_prop}: {heatmap_num}<br>' 

363 

364 identF = RectFace(width=self.width, height=self.height, text=heatmap_num, color=self.absence_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=None) 

365 node.add_face(identF, column = self.column, position = 'aligned', collapsed_only=False) 

366 

367class LayoutHeatmapOld(TreeLayout): 

368 def __init__(self, name=None, column=0, width=70, height=None, padding_x=1, padding_y=0, \ 

369 internal_rep=None, prop=None, maxval=100, minval=0, mean_val=0, std_val=0, \ 

370 norm_method='min-max', color_dict=None, legend=True): 

371 super().__init__(name) 

372 self.aligned_faces = True 

373 self.num_prop = prop 

374 self.column = column 

375 self.color_dict = color_dict 

376 self.maxval = maxval 

377 self.minval = minval 

378 self.mean_val = mean_val 

379 self.std_val = std_val 

380 self.norm_method = norm_method 

381 

382 self.internal_prop = add_suffix(prop, internal_rep) 

383 self.width = width 

384 self.height = height 

385 self.padding_x = padding_x 

386 self.padding_y = padding_y 

387 self.min_fsize = 5 

388 self.max_fsize = 15 

389 

390 def set_tree_style(self, tree, tree_style): 

391 super().set_tree_style(tree, tree_style) 

392 text = TextFace(self.num_prop, min_fsize=self.min_fsize, max_fsize=self.max_fsize, padding_x=self.padding_x, width=self.width, rotation=315) 

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

394 

395 if self.legend: 

396 tree_style.add_legend(title=self.num_prop, 

397 variable='continuous', 

398 value_range=[self.minval, self.maxval], 

399 color_range=[ 

400 self.color_dict[20], 

401 self.color_dict[10], 

402 self.color_dict[1] 

403 ] 

404 ) 

405 

406 def min_max_normalize(self, value): 

407 return (value - self.minval) / (self.maxval - self.minval) 

408 

409 def mean_normalize(self, value): 

410 return (value - self.mean_val) / (self.maxval - self.minval) 

411 

412 def z_score_normalize(self, value): 

413 return (value - self.mean_val) / self.std_val 

414 

415 def _get_color(self, search_value, norm_method='min-max'): 

416 num = len(self.color_dict) 

417 search_value = float(search_value) 

418 if norm_method == "min-max": 

419 normalized_value = self.min_max_normalize(search_value) 

420 index_values = np.linspace(0, 1, num) 

421 elif norm_method == "mean": 

422 normalized_value = self.mean_normalize(search_value) 

423 index_values = np.linspace(-1, 1, num) 

424 elif norm_method == "zscore": 

425 normalized_value = self.z_score_normalize(search_value) 

426 index_values = np.linspace(-3, 3, num) 

427 else: 

428 raise ValueError("Unsupported normalization method.") 

429 index = np.abs(index_values - normalized_value).argmin() + 1 

430 #index = np.abs(index_values - search_value).argmin() + 1 

431 return self.color_dict.get(index, "") 

432 

433 def set_node_style(self, node): 

434 if node.props.get(self.num_prop) is not None: 

435 if node.is_leaf: 

436 # heatmap 

437 tooltip = "" 

438 if node.name: 

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

440 if self.num_prop: 

441 tooltip += f'<br>{self.num_prop}: {node.props.get(self.num_prop)}<br>' 

442 

443 gradient_color = self._get_color(node.props.get(self.num_prop), norm_method=self.norm_method) 

444 if gradient_color: 

445 identF = RectFace(width=self.width, height=self.height, text="%.2f" % (float(node.props.get(self.num_prop))), \ 

446 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip) 

447 else: # for miss data 

448 identF = RectFace(width=self.width, height=self.height, text="NA", 

449 color="", 

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

451 

452 node.add_face(identF, column = self.column, position = 'aligned') 

453 

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

455 # heatmap 

456 tooltip = "" 

457 if node.name: 

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

459 if self.num_prop: 

460 tooltip += f'<br>{self.internal_prop}: {node.props.get(self.internal_prop)}<br>' 

461 

462 gradient_color = self._get_color(node.props.get(self.internal_prop), norm_method=self.norm_method) 

463 if gradient_color: 

464 identF = RectFace(width=self.width, height=self.height, text="%.2f" % (float(node.props.get(self.internal_prop))), \ 

465 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip) 

466 else: # for miss data 

467 identF = RectFace(width=self.width, height=self.height, text="NA", 

468 color="", 

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

470 node.add_face(identF, column = self.column, position = 'aligned', collapsed_only=True) 

471 

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

473 # heatmap 

474 tooltip = "" 

475 if node.name: 

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

477 if self.num_prop: 

478 tooltip += f'<br>{self.internal_prop}: {node.props.get(self.internal_prop)}<br>' 

479 

480 gradient_color = self._get_color(node.props.get(self.internal_prop), norm_method=self.norm_method) 

481 if gradient_color: 

482 identF = RectFace(width=self.width, height=self.height, text="%.2f" % (float(node.props.get(self.internal_prop))), \ 

483 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip) 

484 else: # for miss data 

485 identF = RectFace(width=self.width, height=self.height, text="NA", 

486 color="", 

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

488 node.add_face(identF, column = self.column, position = 'aligned', collapsed_only=True) 

489 

490class LayoutBranchScore(TreeLayout): 

491 def __init__(self, name, color_dict, score_prop, internal_rep=None, \ 

492 value_range=None, color_range=None, show_score=False, legend=True, active=True): 

493 super().__init__(name) 

494 self.aligned_faces = True 

495 self.score_prop = score_prop 

496 if internal_rep: 

497 self.internal_prop = add_suffix(score_prop, internal_rep) 

498 else: 

499 self.internal_prop = None 

500 self.color_dict = color_dict 

501 self.legend = legend 

502 self.absence_color = "black" 

503 self.value_range = value_range 

504 self.color_range = color_range 

505 self.show_score = show_score 

506 self.line_width = 3 

507 self.active = active 

508 

509 def set_tree_style(self, tree, tree_style): 

510 if self.legend: 

511 if self.color_dict: 

512 tree_style.add_legend(title=self.name, 

513 variable='continuous', 

514 value_range=self.value_range, 

515 color_range=self.color_range, 

516 ) 

517 

518 # def _get_color_for_score(self, search_value): 

519 # num = len(self.color_dict) 

520 # index_values = np.linspace(self.value_range[0], self.value_range[1], num) 

521 # index = np.abs(index_values - search_value).argmin() + 1 

522 # return self.color_dict.get(index, self.absence_color) 

523 

524 def set_node_style(self, node): 

525 prop_score = node.props.get(self.score_prop) 

526 if prop_score is not None: 

527 prop_score = float(prop_score) 

528 node.sm_style["hz_line_color"] = self.color_dict.get(prop_score) 

529 node.sm_style["hz_line_width"] = self.line_width 

530 node.sm_style["vt_line_color"] = self.color_dict.get(prop_score) 

531 node.sm_style["vt_line_width"] = self.line_width 

532 node.sm_style["outline_color"] = self.color_dict.get(prop_score) 

533 

534 if self.show_score: 

535 node.add_face( 

536 TextFace("%.2f" % (float(prop_score)), color=self.color_dict.get(prop_score)), 

537 position="branch_bottom") 

538 

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

540 prop_score = node.props.get(self.internal_prop) 

541 prop_score = float(prop_score) 

542 node.sm_style["hz_line_color"] = self.color_dict.get(prop_score) 

543 node.sm_style["hz_line_width"] = self.line_width 

544 node.sm_style["vt_line_color"] = self.color_dict.get(prop_score) 

545 node.sm_style["vt_line_width"] = self.line_width 

546 node.sm_style["outline_color"] = self.color_dict.get(prop_score) 

547 

548 if self.show_score: 

549 node.add_face( 

550 TextFace("%.2f" % (float(prop_score)), color=self.color_dict.get(prop_score)), 

551 position="branch_bottom") 

552 

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

554 prop_score = node.props.get(self.internal_prop) 

555 prop_score = float(prop_score) 

556 node.sm_style["hz_line_color"] = self.color_dict.get(prop_score) 

557 node.sm_style["hz_line_width"] = self.line_width 

558 node.sm_style["vt_line_color"] = self.color_dict.get(prop_score) 

559 node.sm_style["vt_line_width"] = self.line_width 

560 node.sm_style["outline_color"] = self.color_dict.get(prop_score) 

561 

562 if self.show_score: 

563 node.add_face( 

564 TextFace("%.2f" % (float(prop_score)), color=self.color_dict.get(prop_score)), 

565 position="branch_bottom") 

566 

567class LayoutBubble(TreeLayout): 

568 def __init__(self, name=None, prop=None, position="aligned", 

569 column=0, color=None, max_radius=10, abs_maxval=None, 

570 padding_x=2, padding_y=0, 

571 scale=True, legend=True, active=True, 

572 internal_rep='avg'): 

573 

574 name = name or f'Barplot_{size_prop}_{color_prop}' 

575 super().__init__(name) 

576 

577 self.aligned_faces = True 

578 self.num_prop = prop 

579 self.internal_prop = add_suffix(prop, internal_rep) 

580 

581 self.column = column 

582 self.position = position 

583 self.color = color 

584 self.positive_color = "#ff0000" 

585 self.negative_color = "#0000ff" 

586 self.internal_rep = internal_rep 

587 self.max_radius = float(max_radius) 

588 self.abs_maxval = float(abs_maxval) 

589 

590 self.padding_x = padding_x 

591 self.padding_y = padding_y 

592 

593 self.legend = legend 

594 self.active = active 

595 

596 def set_tree_style(self, tree, tree_style): 

597 super().set_tree_style(tree, tree_style) 

598 text = TextFace(self.num_prop, min_fsize=5, max_fsize=15, padding_x=self.padding_x, rotation=315) 

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

600 colormap = { 

601 "positive": self.positive_color, 

602 "negative": self.negative_color 

603 } 

604 if self.legend: 

605 tree_style.add_legend(title=self.num_prop, 

606 variable='discrete', 

607 colormap=colormap 

608 ) 

609 

610 def _get_bubble_size(self, search_value): 

611 search_value = abs(float(search_value)) 

612 bubble_size = search_value / self.abs_maxval * self.max_radius 

613 return bubble_size 

614 

615 

616 def set_node_style(self, node): 

617 number = node.props.get(self.num_prop) 

618 if number is not None: 

619 bubble_size = self._get_bubble_size(number) 

620 #bubble_size = self.max_radius 

621 if number > 0: 

622 bubble_color = self.positive_color 

623 else: 

624 bubble_color = self.negative_color 

625 node.sm_style["size"] = bubble_size 

626 node.sm_style["fgcolor"] = bubble_color 

627 

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

629 if number > 0: 

630 bubble_color = self.positive_color 

631 else: 

632 bubble_color = self.negative_color 

633 bubble_size = self._get_bubble_size(number) 

634 node.sm_style["size"] = bubble_size 

635 node.sm_style["fgcolor"] = bubble_color 

636 

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

638 number = node.props.get(self.internal_prop) 

639 if number > 0: 

640 bubble_color = self.positive_color 

641 else: 

642 bubble_color = self.negative_color 

643 bubble_size = self._get_bubble_size(number) 

644 node.sm_style["size"] = bubble_size 

645 node.sm_style["fgcolor"] = bubble_color