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

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 

11 

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 

19 

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) 

26 

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) 

37 

38 if final_call == False: 

39 break 

40 else: 

41 continue 

42 

43 if final_call: 

44 

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 

57 

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 

65 

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) 

72 

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) 

83 

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 

94 

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 

119 

120 

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 

127 

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: 

133 

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 ) 

156 

157 

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) 

163 

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

171 

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

196 

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) 

203 

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