Mercurial > repos > george-weingart > lefse
comparison home/ubuntu/lefse_to_export/plot_cladogram.py @ 1:db64b6287cd6 draft
Modified datatypes
author | george-weingart |
---|---|
date | Wed, 20 Aug 2014 16:56:51 -0400 |
parents | |
children |
comparison
equal
deleted
inserted
replaced
0:e7cd19afda2e | 1:db64b6287cd6 |
---|---|
1 #!/usr/bin/env python | |
2 | |
3 import os,sys,matplotlib,argparse,string | |
4 matplotlib.use('Agg') | |
5 from pylab import * | |
6 from lefse import * | |
7 import numpy as np | |
8 | |
9 colors = ['r','g','b','m','c',[1.0,0.5,0.0],[0.0,1.0,0.0],[0.33,0.125,0.0],[0.75,0.75,0.75],'k'] | |
10 dark_colors = [[0.4,0.0,0.0],[0.0,0.2,0.0],[0.0,0.0,0.4],'m','c',[1.0,0.5,0.0],[0.0,1.0,0.0],[0.33,0.125,0.0],[0.75,0.75,0.75],'k'] | |
11 | |
12 class CladeNode: | |
13 def __init__(self, name, abundance, viz = True): | |
14 self.id = name | |
15 self.name = name.split(".") | |
16 self.last_name = self.name[-1] | |
17 self.abundance = abundance | |
18 self.pos = (-1.0,-1.0) | |
19 self.children = {} | |
20 self.isleaf = True | |
21 self.color = 'y' | |
22 self.next_leaf = -1 | |
23 self.prev_leaf = -1 | |
24 self.viz = viz | |
25 def __repr__(self): | |
26 return self.last_name | |
27 def add_child(self,node): | |
28 self.isleaf = False | |
29 self.children[node.__repr__()] = node | |
30 def get_children(self): | |
31 ck = sorted(self.children.keys()) | |
32 return [self.children[k] for k in ck] | |
33 def get_color(self): | |
34 return self.color | |
35 def set_color(self,c): | |
36 self.color = c | |
37 def set_pos(self,pos): | |
38 self.pos = pos | |
39 | |
40 def read_params(args): | |
41 parser = argparse.ArgumentParser(description='Cladoplot') | |
42 parser.add_argument('input_file', metavar='INPUT_FILE', type=str, help="tab delimited input file") | |
43 parser.add_argument('output_file', metavar='OUTPUT_FILE', type=str, help="the file for the output image") | |
44 parser.add_argument('--clade_sep',dest="clade_sep", type=float, default=1.5) | |
45 parser.add_argument('--max_lev',dest="max_lev", type=int, default=-1) | |
46 parser.add_argument('--max_point_size',dest="max_point_size", type=float, default=6.0) | |
47 parser.add_argument('--min_point_size',dest="min_point_size", type=float, default=1) | |
48 parser.add_argument('--point_edge_width',dest="markeredgewidth", type=float, default=.25) | |
49 parser.add_argument('--siblings_connector_width',dest="siblings_connector_width", type=float, default=2) | |
50 parser.add_argument('--parents_connector_width',dest="parents_connector_width", type=float, default=0.75) | |
51 parser.add_argument('--radial_start_lev',dest="radial_start_lev", type=int, default=1) | |
52 parser.add_argument('--labeled_start_lev',dest="labeled_start_lev", type=int, default=2) | |
53 parser.add_argument('--labeled_stop_lev',dest="labeled_stop_lev", type=int, default=5) | |
54 parser.add_argument('--abrv_start_lev',dest="abrv_start_lev", type=int, default=3) | |
55 parser.add_argument('--abrv_stop_lev',dest="abrv_stop_lev", type=int, default=5) | |
56 parser.add_argument('--expand_void_lev',dest="expand_void_lev", type=int, default=1) | |
57 parser.add_argument('--class_legend_vis',dest="class_legend_vis", type=int, default=1) | |
58 parser.add_argument('--colored_connector',dest="colored_connectors", type=int, default=1) | |
59 parser.add_argument('--alpha',dest="alpha", type=float, default=0.2) | |
60 parser.add_argument('--title',dest="title", type=str, default="Cladogram") | |
61 parser.add_argument('--sub_clade',dest="sub_clade", type=str, default="") | |
62 parser.add_argument('--title_font_size',dest="title_font_size", type=str, default="14") | |
63 parser.add_argument('--right_space_prop',dest="r_prop", type=float, default=0.1) | |
64 parser.add_argument('--left_space_prop',dest="l_prop", type=float, default=0.1) | |
65 parser.add_argument('--label_font_size',dest="label_font_size", type=str, default="6") | |
66 parser.add_argument('--background_color',dest="back_color", type=str, choices=["k","w"], default="w", help="set the color of the background") | |
67 parser.add_argument('--colored_labels',dest="col_lab", type=int, choices=[0,1], default=1, help="draw the label with class color (1) or in black (0)") | |
68 parser.add_argument('--class_legend_font_size',dest="class_legend_font_size", type=str, default="10") | |
69 parser.add_argument('--dpi',dest="dpi", type=int, default=72) | |
70 parser.add_argument('--format', dest="format", choices=["png","svg","pdf"], default="svg", type=str, help="the format for the output file") | |
71 parser.add_argument('--all_feats', dest="all_feats", type=str, default="") | |
72 args = parser.parse_args() | |
73 return vars(args) | |
74 | |
75 def cmp_names(la,lb): | |
76 if len(la) != len(lb): return False | |
77 for p in [(a,b) for i,a in enumerate(la) for j,b in enumerate(lb) if i == j]: | |
78 if p[0] != p[1]: return False | |
79 return True | |
80 | |
81 def build_tree(father,all_nodes,l,depth,viz): | |
82 cc = [n for n in all_nodes if len(n.name) > len(father.name) and cmp_names(father.name,n.name[:len(father.name)])] | |
83 children = [n for n in cc if len(n.name) == len(father.name)+1] | |
84 if len(children) == 0 and l < depth -1: # !!! | |
85 nc = CladeNode(father.id+"."+father.id.split(".")[-1],1.0,viz) | |
86 father.add_child(nc) | |
87 children.append(nc) | |
88 for child in children: | |
89 build_tree(child,cc,l+1,depth,viz) | |
90 father.add_child(child) | |
91 | |
92 def get_all_nodes(father): | |
93 ret = [father] | |
94 children = father.get_children() | |
95 for c in children: | |
96 ret += get_all_nodes(c) | |
97 return ret | |
98 | |
99 def read_data(input_file,params): | |
100 with open(input_file, 'r') as inp: | |
101 if params['sub_clade'] == "": rows = [line.strip().split()[:-1] for line in inp.readlines() if params['max_lev'] < 1 or line.split()[0].count(".") < params['max_lev']] | |
102 else: rows = [line.split(params['sub_clade']+".")[1].strip().split()[:-1] for line in inp.readlines() if ( params['max_lev'] < 1 or line.split()[0].count(".") < params['max_lev'] ) and line.startswith(params['sub_clade']+".")] | |
103 all_names = [lin[0] for lin in rows] | |
104 to_add = [] | |
105 | |
106 abundances = [float(v) for v in zip(*rows)[1] if v >= 0.0] | |
107 tree = {} | |
108 tree['classes'] = list(set([v[2] for v in rows if len(v)>2])) | |
109 tree['classes'].sort() | |
110 all_nodes = [CladeNode("root."+row[0],float(row[1])) for row in rows] | |
111 | |
112 depth = max([len(n.name) for n in all_nodes]) | |
113 | |
114 n2 = ["_".join(nn.name) for nn in all_nodes] | |
115 for i,nn in enumerate(all_nodes): | |
116 n = nn | |
117 while "_".join(n.name[:-1]) not in n2 and len(n.name) > 1: | |
118 n = CladeNode(".".join(n.name[:-1]),n.abundance) | |
119 all_nodes.append(n) | |
120 n2.append("_".join(n.name)) | |
121 | |
122 cls2 = [] | |
123 if params['all_feats'] != "": | |
124 cls2 = sorted(params['all_feats'].split(":")) | |
125 for i,v in enumerate(rows): | |
126 if len(v)>2: | |
127 if len(cls2) > 0: all_nodes[i].set_color(colors[cls2.index(v[2])%len(colors)]) | |
128 else: | |
129 if v[2].count('rgbcol') > 0: | |
130 ccc = [float(tt) for tt in v[2].split('_')[1:]] | |
131 all_nodes[i].set_color(ccc) | |
132 else: all_nodes[i].set_color(colors[sorted(tree['classes']).index(v[2])%len(colors)]) | |
133 root = CladeNode("root",-1.0) | |
134 root.set_pos((0.0,0.0)) | |
135 | |
136 build_tree(root,all_nodes,0,depth,params['expand_void_lev']==1) | |
137 | |
138 all_nodes = get_all_nodes(root) | |
139 | |
140 tree['root'] = root | |
141 tree['max_abs'] = max(abundances) | |
142 tree['min_abs'] = min(abundances) | |
143 levs = [] | |
144 for i in range(depth): | |
145 depthi = [n for n in all_nodes if len(n.name) == i+1] | |
146 levs.append(len(depthi)) | |
147 tree['nlev'] = levs | |
148 return tree | |
149 | |
150 def add_all_pos(father,n,distn,seps,tsep,mlev,last_leaf=-1,nc=1): | |
151 children = father.get_children() | |
152 leaves = True if children[0].isleaf else False | |
153 for i,child in enumerate(children): | |
154 if leaves: | |
155 n += 1.0 | |
156 men = 0.5 if len(children) == 1 else 0.0 | |
157 child.set_pos((n*distn-men*float(distn)+tsep,(len(father.name))/float(mlev-1))) | |
158 if last_leaf != -1: | |
159 child.prev_leaf = last_leaf | |
160 last_leaf.next_leaf = child | |
161 last_leaf = child | |
162 else: | |
163 ln = n | |
164 ltsep = tsep | |
165 n,tsep,last_leaf = add_all_pos(child,n,distn,seps,tsep,mlev,last_leaf,len(children)) | |
166 nn = (ln + n)*0.5*distn | |
167 ssep = (ltsep + tsep)*0.5 | |
168 if n-ln == 1: | |
169 ssep = ltsep | |
170 child.set_pos((nn+ssep,(len(father.name))/float(mlev-1))) | |
171 tsep += seps[len(father.name)-1] | |
172 return n,tsep,last_leaf | |
173 | |
174 def plot_points(father,params,pt_scale,ax): | |
175 children = father.get_children() | |
176 children.sort(key = lambda a: -int(a.get_color() == 'y')*a.abundance) | |
177 x,r = father.pos[0], father.pos[1] | |
178 for i,child in enumerate(children): | |
179 xc,rc = plot_points(child,params,pt_scale,ax) | |
180 if not father.viz: return x,r | |
181 ps = pt_scale[0]+father.abundance/pt_scale[1]+pt_scale[0] | |
182 col = father.get_color() | |
183 pw = params['markeredgewidth'] if col == 'y' else params['markeredgewidth']*3.0 | |
184 if x==0 and r==0: ax.plot(x,r, 'o',markersize=ps,color=col,markeredgewidth=0.01,markeredgecolor=params['fore_color']) | |
185 else: ax.plot(x,r, 'o',markersize=ps,color=col,markeredgewidth=pw,markeredgecolor=params['fore_color']) | |
186 return x,r | |
187 | |
188 def plot_lines(father,params,depth,ax,xf): | |
189 children = father.get_children() | |
190 x,r = father.pos[0], father.pos[1] | |
191 for i,child in enumerate(children): | |
192 xc,rc = plot_lines(child,params,depth,ax,x) | |
193 if i == 0: x_first, r_first = xc, rc | |
194 if len(father.name) >= depth-params['radial_start_lev']: | |
195 col = params['fore_color'] | |
196 lw=params['parents_connector_width'] | |
197 if not child.viz: continue | |
198 if father.get_color() != 'y' and father.get_color() == child.get_color() and params['colored_connectors']: | |
199 col = child.get_color() | |
200 lw *=2.5 | |
201 if col != params['fore_color']: | |
202 ax.plot([x,xc],[r,rc],"-",color=params['fore_color'],lw=lw*1.5) | |
203 ax.plot([x,xc],[r,rc],"-",color=col,lw=lw) | |
204 | |
205 if not father.viz or (len(children) == 1 and not children[0].viz): return x,r | |
206 if len(father.name) < depth-params['radial_start_lev']: | |
207 col = params['fore_color'] | |
208 lw=params['parents_connector_width'] | |
209 if father.get_color() != 'y': | |
210 f =True | |
211 for child in children: | |
212 if child.get_color() != father.get_color() or not params['colored_connectors']: | |
213 f = False | |
214 break | |
215 if f: | |
216 col = father.get_color() | |
217 lw *= 2.5 | |
218 if not (x==0 and r==0): | |
219 xx = xc if len(children) > 0 else x | |
220 if len(children) == 0: rc = r | |
221 xt = x if len(children)>1 else xx | |
222 if col != params['fore_color']: | |
223 ax.plot([x,xt],[r,rc],"-",color=params['fore_color'],lw=lw*1.5) | |
224 ax.plot([x,xt],[r,rc],"-",color=col,lw=lw) | |
225 if len(children) > 0 and 1 < len(father.name) < depth-params['radial_start_lev']: | |
226 xs = arange(x_first,xc,0.01) | |
227 ys = [rc for t in xs] | |
228 ax.plot(xs,ys,"-",color=col,lw=params['siblings_connector_width'],markeredgecolor=params['fore_color']) | |
229 return x,r | |
230 | |
231 def uniqueid(): | |
232 for l in string.lowercase: yield l | |
233 for l in string.lowercase: | |
234 for i in range(10): | |
235 yield l+str(i) | |
236 i = 0 | |
237 while True: | |
238 yield str(i) | |
239 i += 1 | |
240 | |
241 def plot_names(father,params,depth,ax,u_i,seps): | |
242 children = father.get_children() | |
243 l = len(father.name) | |
244 if len(children)==0: | |
245 if father.prev_leaf == -1 or father.next_leaf == -1: | |
246 fr_0, fr_1 = father.pos[0], father.pos[0] | |
247 else: fr_0, fr_1 = (father.pos[0]+father.prev_leaf.pos[0])*0.5, (father.pos[0]+father.next_leaf.pos[0])*0.5 | |
248 for i,child in enumerate(children): | |
249 fr,to = plot_names(child,params,depth,ax,u_i,seps) | |
250 if i == 0: fr_0 = fr | |
251 fr_1 = to | |
252 if father.get_color() != 'y' and params['labeled_start_lev'] < l <= params['labeled_stop_lev']+1: | |
253 col = father.get_color() | |
254 dd = params['labeled_stop_lev'] - params['labeled_start_lev'] + 1 | |
255 de = depth - 1 | |
256 dim = 1.0/float(de) | |
257 perc_ext = 0.65 if dim > 0.1 else 1.0 | |
258 clto = (de-l+1)*dim+dim*(dd+1-(l-dd-1))*perc_ext | |
259 clto = (de-l+1)*dim+dim*(dd-(l-params['labeled_start_lev'])+1)*perc_ext | |
260 des = float(180.0*(fr_0+fr_1)/np.pi)*0.5-90 | |
261 lab = "" | |
262 txt = father.last_name | |
263 if params['abrv_start_lev'] < l <= params['abrv_stop_lev'] + 1: | |
264 ide = u_i.next() | |
265 lab = str(ide)+": "+father.last_name | |
266 txt = str(ide) | |
267 # ax.bar(fr_0, clto, width = fr_1-fr_0, bottom = float(l-1)/float(depth-1), alpha = params['alpha'], color=col, edgecolor=col) | |
268 ax.bar(fr_0, clto, width = fr_1-fr_0, bottom = float(l-1)/float(de), alpha = params['alpha'], color=col, edgecolor=col) | |
269 ax.bar(0.0, 0.0, width = 0.0, bottom = 0.0, alpha = 1.0, color=col, edgecolor=params['fore_color'], label=lab) | |
270 if l <= params['abrv_stop_lev'] + 1: | |
271 if not params['col_lab']: col = params['fore_color'] | |
272 else: | |
273 if col not in colors: col = params['fore_color'] | |
274 else: col = dark_colors[colors.index(col)%len(dark_colors)] | |
275 ax.text((fr_0+fr_1)*0.5, clto+float(l-1)/float(de)-dim*perc_ext/2.0, txt, size = params['label_font_size'], rotation=des, ha ="center", va="center", color=col) | |
276 return fr_0, fr_1 | |
277 | |
278 def draw_tree(out_file,tree,params): | |
279 plt_size = 7 | |
280 nlev = tree['nlev'] | |
281 pt_scale = (params['min_point_size'],max(1.0,((tree['max_abs']-tree['min_abs']))/(params['max_point_size']-params['min_point_size']))) | |
282 depth = len(nlev) | |
283 sep = (2.0*np.pi)/float(nlev[-1]) | |
284 seps = [params['clade_sep']*sep/float(depth-i+1) for i in range(1,len(tree['nlev'])+1)] | |
285 totseps = sum([s*nlev[i] for i,s in enumerate(seps[:-1])]) | |
286 clade_sep_err = True if totseps > np.pi else False | |
287 while totseps > np.pi: | |
288 params['clade_sep'] *= 0.75 | |
289 seps = [params['clade_sep']*sep/(float(depth-i+1)*0.25) for i in range(1,len(tree['nlev'])+1)] | |
290 totseps = sum([s*nlev[i] for i,s in enumerate(seps[:-1])]) | |
291 if clade_sep_err: print 'clade_sep parameter too large, lowered to',params['clade_sep'] | |
292 | |
293 fig = plt.figure(edgecolor=params['back_color'],facecolor=params['back_color']) | |
294 ax = fig.add_subplot(111, polar=True, frame_on=False, axis_bgcolor=params['back_color'] ) | |
295 plt.subplots_adjust(right=1.0-params['r_prop'],left=params['l_prop']) | |
296 ax.grid(False) | |
297 xticks([]) | |
298 yticks([]) | |
299 | |
300 ds = (2.0*np.pi-totseps)/float(nlev[-1]) | |
301 | |
302 add_all_pos(tree['root'],0.0,ds,seps,0.0,depth) | |
303 | |
304 plot_lines(tree['root'],params,depth,ax,0) | |
305 plot_points(tree['root'],params,pt_scale,ax) | |
306 plot_names(tree['root'],params,depth,ax,uniqueid(),seps) | |
307 | |
308 r = np.arange(0, 3.0, 0.01) | |
309 theta = 2*np.pi*r | |
310 | |
311 def get_col_attr(x): | |
312 return hasattr(x, 'set_color') and not hasattr(x, 'set_facecolor') | |
313 | |
314 h, l = ax.get_legend_handles_labels() | |
315 if len(l) > 0: | |
316 leg = ax.legend(bbox_to_anchor=(1.05, 1), frameon=False, loc=2, borderaxespad=0.,prop={'size':params['label_font_size']}) | |
317 if leg != None: | |
318 gca().add_artist(leg) | |
319 for o in leg.findobj(get_col_attr): | |
320 o.set_color(params['fore_color']) | |
321 | |
322 cll = sorted(tree['classes']) if params['all_feats'] == "" else sorted(params['all_feats'].split(":")) | |
323 nll = [ax.bar(0.0, 0.0, width = 0.0, bottom = 0.0, color=colors[i%len(colors)], label=c) for i,c in enumerate(cll) if c in tree['classes']] | |
324 cl = [c for c in cll if c in tree['classes']] | |
325 | |
326 ax.set_title(params['title'],size=params['title_font_size'],color=params['fore_color']) | |
327 | |
328 if params['class_legend_vis']: | |
329 l2 = legend(nll, cl, loc=2, prop={'size':params['class_legend_font_size']}, frameon=False) | |
330 if l2 != None: | |
331 for o in l2.findobj(get_col_attr): | |
332 o.set_color(params['fore_color']) | |
333 | |
334 plt.savefig(out_file,format=params['format'],facecolor=params['back_color'],edgecolor=params['fore_color'],dpi=params['dpi']) | |
335 plt.close() | |
336 | |
337 if __name__ == '__main__': | |
338 params = read_params(sys.argv) | |
339 params['fore_color'] = 'w' if params['back_color'] == 'k' else 'k' | |
340 clad_tree = read_data(params['input_file'],params) | |
341 draw_tree(params['output_file'],clad_tree,params) | |
342 |