Mercurial > repos > miller-lab > genome_diversity

comparison dpmix_plot.py @ 27:8997f2ca8c7a

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Update to Miller Lab devshed revision bae0d3306d3b
author Richard Burhans <burhans@bx.psu.edu>
date Mon, 15 Jul 2013 10:47:35 -0400
parents 4b6590dd7250
children a631c2f6d913
comparison
equal deleted inserted replaced
26:91e835060ad2 27:8997f2ca8c7a
35 def parse_input_file(input_file): 35 def parse_input_file(input_file):
36 chroms = [] 36 chroms = []
37 individuals = [] 37 individuals = []
38 data = {} 38 data = {}
39 chrom_len = {} 39 chrom_len = {}
40 used_states = []
40 41
41 with open(input_file) as fh: 42 with open(input_file) as fh:
42 for line in fh: 43 for line in fh:
43 line = line.strip() 44 line = line.strip()
44 if line: 45 if line:
45 elems = line.split() 46 elems = line.split()
46 chrom = elems[0] 47 chrom = elems[0]
47 p1, p2, state = map(int, elems[1:4]) 48 p1, p2, state = map(int, elems[1:4])
48 id = elems[4] 49 id = elems[4]
49 50
51 if state not in used_states:
52 used_states.append(state)
53
50 if chrom not in chroms: 54 if chrom not in chroms:
51 chroms.append(chrom) 55 chroms.append(chrom)
52 56
53 if id not in individuals: 57 if id not in individuals:
54 individuals.append(id) 58 individuals.append(id)
58 data[chrom][id].append((p1, p2, state)) 62 data[chrom][id].append((p1, p2, state))
59 63
60 if p2 > chrom_len.setdefault(chrom, 0): 64 if p2 > chrom_len.setdefault(chrom, 0):
61 chrom_len[chrom] = p2 65 chrom_len[chrom] = p2
62 66
63 return chroms, individuals, data, chrom_len 67 return chroms, individuals, data, chrom_len, used_states
64 68
65 def check_chroms(chroms, chrom_len, dbkey): 69 def check_chroms(chroms, chrom_len, dbkey):
66 error = 0 70 error = 0
67 for chrom in chroms: 71 for chrom in chroms:
68 if chrom not in chrom_len: 72 if chrom not in chrom_len:
110 def make_split_rectangle(p1, p2, top_color, bottom_color): 114 def make_split_rectangle(p1, p2, top_color, bottom_color):
111 patch1 = make_rectangle(p1, p2, bottom_color, top=0.5) 115 patch1 = make_rectangle(p1, p2, bottom_color, top=0.5)
112 patch2 = make_rectangle(p1, p2, top_color, bottom=0.5) 116 patch2 = make_rectangle(p1, p2, top_color, bottom=0.5)
113 return [patch1, patch2] 117 return [patch1, patch2]
114 118
115 def make_state_rectangle(p1, p2, state, chrom, individual): 119 def make_state_rectangle_2pop(p1, p2, state, chrom, individual):
120 p1_color = 'r'
121 p2_color = 'g'
122 heterochromatin_color = '#c7c7c7'
123
116 if state == 0: 124 if state == 0:
117 return [ make_rectangle(p1, p2, 'r') ] 125 return [ make_rectangle(p1, p2, heterochromatin_color) ]
118 elif state == 1: 126 elif state == 1:
119 return make_split_rectangle(p1, p2, 'r', 'g') 127 return [ make_rectangle(p1, p2, p1_color) ]
120 elif state == 2: 128 elif state == 2:
121 return [ make_rectangle(p1, p2, 'g') ] 129 return [ make_rectangle(p1, p2, p2_color) ]
122 elif state == 3: 130 elif state == 3:
123 return [ make_rectangle(p1, p2, '#c7c7c7') ] 131 return make_split_rectangle(p1, p2, p1_color, p2_color)
132 else:
133 print >> sys.stderr, "Unknown state: {0}: {1} {2} {3} {4}".format(state, chrom, p1, p2, state, individual)
134 sys.exit(1)
135
136 def make_state_rectangle_3pop(p1, p2, state, chrom, individual):
137 p1_color = 'r'
138 p2_color = 'g'
139 p3_color = 'b'
140 heterochromatin_color = '#c7c7c7'
141
142 if state == 0:
143 return [ make_rectangle(p1, p2, heterochromatin_color) ]
144 if state == 1:
145 return [ make_rectangle(p1, p2, p1_color) ]
146 if state == 2:
147 return [ make_rectangle(p1, p2, p2_color) ]
148 if state == 3:
149 return [ make_rectangle(p1, p2, p3_color) ]
150 if state == 4:
151 return make_split_rectangle(p1, p2, p1_color, p2_color)
152 if state == 5:
153 return make_split_rectangle(p1, p2, p1_color, p3_color)
154 if state == 6:
155 return make_split_rectangle(p1, p2, p2_color, p3_color)
124 else: 156 else:
125 print >> sys.stderr, "Unknown state: {0}: {1} {2} {3} {4}".format(state, chrom, p1, p2, state, individual) 157 print >> sys.stderr, "Unknown state: {0}: {1} {2} {3} {4}".format(state, chrom, p1, p2, state, individual)
126 sys.exit(1) 158 sys.exit(1)
127 159
128 def nicenum(num, round=False): 160 def nicenum(num, round=False):
193 225
194 return vals, labels 226 return vals, labels
195 227
196 ################################################################################ 228 ################################################################################
197 229
198 def make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir): 230 def make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir, state2name=None, populations=3):
199 fs_chrom_len = build_chrom_len_dict(input_dbkey, galaxy_data_index_dir) 231 fs_chrom_len = build_chrom_len_dict(input_dbkey, galaxy_data_index_dir)
200 chroms, individuals, data, chrom_len = parse_input_file(input_file) 232 chroms, individuals, data, chrom_len, used_states = parse_input_file(input_file)
233
234 if populations == 3:
235 make_state_rectangle = make_state_rectangle_3pop
236 elif populations == 2:
237 make_state_rectangle = make_state_rectangle_2pop
238 else:
239 pass
201 240
202 for chrom in chrom_len.keys(): 241 for chrom in chrom_len.keys():
203 if chrom in fs_chrom_len: 242 if chrom in fs_chrom_len:
204 chrom_len[chrom] = fs_chrom_len[chrom] 243 chrom_len[chrom] = fs_chrom_len[chrom]
205 244
212 chrom_height = 0.25 251 chrom_height = 0.25
213 ind_space = 0.10 252 ind_space = 0.10
214 ind_height = 0.25 253 ind_height = 0.25
215 254
216 total_height = 0.0 255 total_height = 0.0
217 at_top = True 256
257 ## make a legend
258 ## only print out states that are
259 ## 1) in the data
260 ## - AND -
261 ## 2) in the state2name map
262 ## here, we only calculate the space needed
263 legend_states = []
264 if state2name is not None:
265 for state in used_states:
266 if state in state2name:
267 legend_states.append(state)
268
269 if legend_states:
270 total_height += len(legend_states) * (ind_space + ind_height)
271 total_height += (top_space - ind_space)
272 at_top = False
273 else:
274 at_top = True
275
218 for chrom in chroms: 276 for chrom in chroms:
219 if at_top: 277 if at_top:
220 total_height += (top_space + chrom_height) 278 total_height += (top_space + chrom_height)
221 at_top = False 279 at_top = False
222 else: 280 else:
233 291
234 bottom = 1.0 292 bottom = 1.0
235 293
236 fig = plt.figure(figsize=(width, height)) 294 fig = plt.figure(figsize=(width, height))
237 295
238 at_top = True 296 if legend_states:
297 at_top = True
298 for state in sorted(legend_states):
299 if at_top:
300 bottom -= (top_space + ind_height)/height
301 at_top = False
302 else:
303 bottom -= (ind_space + ind_height)/height
304 ## add code here to draw legend
305 # [left, bottom, width, height]
306 ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/height])
307 plt.axis('off')
308 ax1.set_xlim(0, 1)
309 ax1.set_ylim(0, 1)
310 for patch in make_state_rectangle(0, 1, state, 'legend', state2name[state]):
311 ax1.add_patch(patch)
312
313 ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/height], frame_on=False)
314 plt.axis('off')
315 plt.text(0.0, 0.5, state2name[state], fontsize=10, ha='left', va='center')
316 else:
317 at_top = True
318
239 for_webb = False 319 for_webb = False
240 320
241 for chrom in chroms: 321 for chrom in chroms:
242 length = chrom_len[chrom] 322 length = chrom_len[chrom]
243 vals, labels = tick_foo(0, length) 323 vals, labels = tick_foo(0, length)
281 ax2.set_xticks(vals) 361 ax2.set_xticks(vals)
282 ax2.set_xticklabels(labels) 362 ax2.set_xticklabels(labels)
283 else: 363 else:
284 plt.axis('off') 364 plt.axis('off')
285 for p1, p2, state in sorted(data[chrom][individual]): 365 for p1, p2, state in sorted(data[chrom][individual]):
366 #for patch in make_state_rectangle(p1, p2, state, chrom, individual):
286 for patch in make_state_rectangle(p1, p2, state, chrom, individual): 367 for patch in make_state_rectangle(p1, p2, state, chrom, individual):
287 ax2.add_patch(patch) 368 ax2.add_patch(patch)
288 369
289 plt.savefig(output_file) 370 plt.savefig(output_file)
290 371
293 if __name__ == '__main__': 374 if __name__ == '__main__':
294 input_dbkey, input_file, output_file, galaxy_data_index_dir = sys.argv[1:5] 375 input_dbkey, input_file, output_file, galaxy_data_index_dir = sys.argv[1:5]
295 make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir) 376 make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir)
296 sys.exit(0) 377 sys.exit(0)
297 378
379 ## notes
380 # 1) pass in a state to name mapping
381 # 2) only print out names for states which exist in the data, and are in the state to name mapping