Coverage for /home/deng/Projects/metatree_drawer/metatreedrawer/treeprofiler/layouts/general_layouts.py: 18%

141 statements  

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

1from __future__ import annotations 

2import Bio 

3from Bio import AlignIO 

4from Bio.Align import MultipleSeqAlignment 

5from Bio.Align.AlignInfo import SummaryInfo 

6import numpy as np 

7from distutils.util import strtobool 

8import matplotlib as mpl 

9import matplotlib.colors as mcolors 

10from itertools import chain 

11import math 

12 

13from ete4.smartview import TreeStyle, NodeStyle, TreeLayout, PieChartFace, LegendFace, RectFace 

14from ete4.smartview.renderer.draw_helpers import * 

15 

16from treeprofiler.src.utils import to_code, call, counter_call, check_nan 

17from treeprofiler.src import utils 

18 

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

20 

21def get_piechartface(node, prop, color_dict=None, radius=20, tooltip=None): 

22 pair_delimiter = "--" 

23 item_seperator = "||" 

24 piechart_data = [] 

25 counter_props = node.props.get(prop).split(item_seperator) 

26 for counter_prop in counter_props: 

27 k, v = counter_prop.split(pair_delimiter) 

28 piechart_data.append([k,float(v),color_dict.get(k,None),None]) 

29 

30 if piechart_data: 

31 piechart_face = PieChartFace(radius=radius, data=piechart_data, padding_x=5, tooltip=tooltip) 

32 

33 return piechart_face 

34 else: 

35 return None 

36 

37def get_aggregated_heatmapface(node, prop, min_color="#EBEBEB", max_color="#971919", tooltip=None, 

38 width=70, height=None, padding_x=1, padding_y=0, count_missing=True, max_count=0): 

39 counter_props = node.props.get(prop).split('||') 

40 total = 0 

41 positive = 0 

42 for counter_prop in counter_props: 

43 k, v = counter_prop.split('--') 

44 if count_missing: 

45 if not check_nan(k): 

46 if strtobool(k): 

47 positive += float(v) 

48 total += float(v) # Consider missing data in total 

49 else: 

50 if not check_nan(k): 

51 total += float(v) # Doesn't consider missing data in total 

52 if strtobool(k): 

53 positive += float(v) 

54 

55 total = int(total) 

56 # ratio = positive / total if total != 0 else 0 

57 # if ratio < 0.05 and ratio != 0: # Show minimum color for too low 

58 # ratio = 0.05 

59 

60 # Adjust the maximum color based on 'total' to simulate darkening 

61 adjusted_max_color = utils.make_color_darker_scaled(max_color, positive, max_count, base=10, scale_factor=10) 

62 #adjusted_max_color = make_color_darker(max_color, darkening_factor=0.01) # Example factor 

63 #gradient_color = color_gradient(min_color, adjusted_max_color, mix=ratio) 

64 

65 if not tooltip: 

66 tooltip = f'<b>{node.name}</b><br>' if node.name else '' 

67 if prop: 

68 tooltip += f'<br>{prop}: {positive} / {total} <br>' 

69 if positive == 0: 

70 aggregateFace = RectFace(width=width, text=int(positive), height=height, color=min_color, padding_x=padding_x, padding_y=padding_y, tooltip=tooltip) 

71 else: 

72 aggregateFace = RectFace(width=width, text=int(positive), height=height, color=adjusted_max_color, padding_x=padding_x, padding_y=padding_y, tooltip=tooltip) 

73 return aggregateFace 

74 

75def get_heatmapface(node, prop, min_color="#EBEBEB", max_color="#971919", tooltip=None, width=70, height=None, padding_x=1, padding_y=0, count_missing=True, reverse=False): 

76 counter_props = node.props.get(prop).split('||') 

77 total = 0 

78 positive = 0 

79 for counter_prop in counter_props: 

80 k, v = counter_prop.split('--') 

81 if count_missing: 

82 if not check_nan(k): 

83 if strtobool(k): 

84 positive = float(v) 

85 total += float(v) # here consider missing data in total 

86 else: 

87 if not check_nan(k): 

88 total += float(v) # here doesn't consider missing data in total 

89 if strtobool(k): 

90 positive = float(v) 

91 

92 total = int(total) 

93 if total != 0: 

94 ratio = positive / total 

95 else: 

96 ratio = 0 

97 

98 if reverse: 

99 ratio = 1 - ratio 

100 

101 if ratio < 0.05 and ratio != 0: # show minimum color for too low 

102 ratio = 0.05 

103 

104 c1 = min_color 

105 c2 = max_color 

106 gradient_color = utils.color_gradient(c1, c2, mix=ratio) 

107 text = f"{positive} / {total}" 

108 # gradientFace = RectFace(width=100,height=50,text="%.1f" % (ratio*100), color=gradient_color,  

109 # padding_x=1, padding_y=1) 

110 

111 if not tooltip: 

112 if node.name: 

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

114 else: 

115 tooltip = '' 

116 if prop: 

117 tooltip += '<br>{}: {} / {} <br>'.format(prop, positive, total) 

118 

119 gradientFace = RectFace(width=width, height=height, 

120 #text=text,  

121 color=gradient_color, 

122 padding_x=padding_x, padding_y=padding_y, tooltip=tooltip) 

123 return gradientFace 

124 

125 

126SeqRecord = Bio.SeqRecord.SeqRecord 

127def get_consensus_seq(filename: Path | str, threshold=0.7) -> SeqRecord: 

128 #https://stackoverflow.com/questions/73702044/how-to-get-a-consensus-of-multiple-sequence-alignments-using-biopython 

129 common_alignment = MultipleSeqAlignment( 

130 chain(*AlignIO.parse(filename, "fasta")) 

131 ) 

132 summary = SummaryInfo(common_alignment) 

133 consensus = summary.dumb_consensus(threshold, "-") 

134 return consensus 

135 

136def get_stackedbarface(node, prop, color_dict=None, width=70, height=None, padding_x=1, padding_y=0, tooltip=None): 

137 pair_delimiter = "--" 

138 item_seperator = "||" 

139 stackedbar_data = [] 

140 absence_color = "#EBEBEB" 

141 counter_props = node.props.get(prop).split(item_seperator) 

142 tooltip = "" 

143 total = 0 

144 

145 for counter_prop in counter_props: 

146 k, v = counter_prop.split(pair_delimiter) 

147 if v: 

148 total += float(v) 

149 stackedbar_data.append([k,float(v),color_dict.get(k,absence_color),None]) 

150 

151 if stackedbar_data: 

152 tooltip = "" 

153 if node.name: 

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

155 

156 if counter_props: 

157 for counter_prop in counter_props: 

158 k, v = counter_prop.split(pair_delimiter) 

159 tooltip += f'<b>{k}</b>: {v}/{int(total)}<br>' 

160 

161 stackedbar_face = StackedBarFace(width=width, height=None, data=stackedbar_data, padding_x=padding_x, padding_y=padding_y, tooltip=tooltip) 

162 

163 return stackedbar_face 

164 else: 

165 return None 

166 

167 

168 

169class StackedBarFace(RectFace): 

170 def __init__(self, width, height, data=None, name="", opacity=0.7, 

171 min_fsize=6, max_fsize=15, ftype='sans-serif', 

172 padding_x=1, padding_y=0, tooltip=None): 

173 

174 RectFace.__init__(self, width=width, height=height, name=name, color=None, 

175 min_fsize=min_fsize, max_fsize=max_fsize, 

176 padding_x=padding_x, padding_y=padding_y, tooltip=tooltip) 

177 

178 self.width = width 

179 self.height = height 

180 

181 # data = [ [name, value, color, tooltip], ... ] 

182 # self.data = [ 

183 # ['first', 10, 'red', None], 

184 # ['second', 40, 'blue', None], 

185 # ['green', 50, 'green', None] 

186 # ] 

187 self.data = data 

188 

189 

190 def __name__(self): 

191 return "StackedBarFace" 

192 

193 def draw(self, drawer): 

194 

195 # Draw RectFace if only one datum 

196 

197 if len(self.data) == 1: 

198 self.color = self.data[0][2] 

199 yield from RectFace.draw(self, drawer) 

200 

201 else: 

202 total_value = sum(d[1] for d in self.data) 

203 start_x, start_y, dx, dy = self._box 

204 

205 for i in range(len(self.data)): 

206 i_value = self.data[i][1] 

207 color = self.data[i][2] 

208 

209 if i > 0: 

210 start_x += new_dx # start with where last segment ends 

211 

212 new_dx = i_value/total_value * dx # width of segment 

213 

214 self._box = Box(start_x,start_y,new_dx,dy) 

215 style = { 'fill': color } 

216 yield draw_rect(self._box, 

217 self.name, 

218 style=style, 

219 tooltip=self.tooltip) 

220 

221# def color_gradient(c1, c2, mix=0): 

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

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

224# c1 = np.array(mpl.colors.to_rgb(c1)) 

225# c2 = np.array(mpl.colors.to_rgb(c2)) 

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

227 

228# def make_color_darker_log(hex_color, total, base=10): 

229# """Darkens the hex color based on a logarithmic scale of the total.""" 

230# # Calculate darkening factor using a logarithmic scale 

231# darkening_factor = math.log(1 + total, base) / 50 # Adjust base and divisor as needed 

232# return make_color_darker(hex_color, darkening_factor) 

233 

234# def make_color_darker(hex_color, darkening_factor): 

235# """Darkens the hex color by a factor. Simplified version for illustration.""" 

236# # Simple darkening logic for demonstration 

237# c = mcolors.hex2color(hex_color) # Convert hex to RGB 

238# darker_c = [max(0, x - darkening_factor) for x in c] # Darken color 

239# return mcolors.to_hex(darker_c) 

240 

241# def make_color_darker_scaled(hex_color, positive, maximum, base=10, scale_factor=10, min_darkness=0.6): 

242# """ 

243# Darkens the hex color based on the positive count, maximum count, and a scaling factor. 

244 

245# :param hex_color: The original color in hex format. 

246# :param positive: The current count. 

247# :param maximum: The maximum count achievable, corresponding to the darkest color. 

248# :param base: The base for the logarithmic calculation, affecting darkening speed. 

249# :param scale_factor: Factor indicating how much darker the color can get at the maximum count. 

250# :param min_darkness: The minimum darkness level allowed. 

251# :return: The darkened hex color. 

252# """ 

253# if positive > maximum: 

254# raise ValueError("Positive count cannot exceed the maximum specified.") 

255 

256# # Calculate the normalized position of 'positive' between 0 and 'maximum' 

257# normalized_position = positive / maximum if maximum != 0 else 0 

258 

259# # Calculate the logarithmic scale position 

260# log_position = math.log(1 + normalized_position * (scale_factor - 1), base) / math.log(scale_factor, base) 

261 

262# # Ensure the log_position respects the min_darkness threshold 

263# if log_position >= min_darkness: 

264# log_position = min_darkness 

265 

266# # Convert hex to RGB 

267# rgb = mcolors.hex2color(hex_color) 

268 

269# # Apply the darkening based on log_position 

270# darkened_rgb = [(1 - log_position) * channel for channel in rgb] 

271 

272# return mcolors.to_hex(darkened_rgb)