genome_diversity: dpmix_plot.py comparison

comparison dpmix_plot.py @ 27:8997f2ca8c7a

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

-:91e835060ad2
+:8997f2ca8c7a
 def parse_input_file(input_file):
 chroms = []
 individuals = []
 data = {}
 chrom_len = {}
+used_states = []
 with open(input_file) as fh:
 for line in fh:
 line = line.strip()
 if line:
 elems = line.split()
 chrom = elems[0]
 p1, p2, state = map(int, elems[1:4])
 id = elems[4]
+if state not in used_states:
+used_states.append(state)
 if chrom not in chroms:
 chroms.append(chrom)
 if id not in individuals:
 individuals.append(id)
 data[chrom][id].append((p1, p2, state))
 if p2 > chrom_len.setdefault(chrom, 0):
 chrom_len[chrom] = p2
-return chroms, individuals, data, chrom_len
+return chroms, individuals, data, chrom_len, used_states
 def check_chroms(chroms, chrom_len, dbkey):
 error = 0
 for chrom in chroms:
 if chrom not in chrom_len:
 def make_split_rectangle(p1, p2, top_color, bottom_color):
 patch1 = make_rectangle(p1, p2, bottom_color, top=0.5)
 patch2 = make_rectangle(p1, p2, top_color, bottom=0.5)
 return [patch1, patch2]
-def make_state_rectangle(p1, p2, state, chrom, individual):
+def make_state_rectangle_2pop(p1, p2, state, chrom, individual):
+p1_color = 'r'
+p2_color = 'g'
+heterochromatin_color = '#c7c7c7'
 if state == 0:
-return [ make_rectangle(p1, p2, 'r') ]
+return [ make_rectangle(p1, p2, heterochromatin_color) ]
 elif state == 1:
-return make_split_rectangle(p1, p2, 'r', 'g')
+return [ make_rectangle(p1, p2, p1_color) ]
 elif state == 2:
-return [ make_rectangle(p1, p2, 'g') ]
+return [ make_rectangle(p1, p2, p2_color) ]
 elif state == 3:
-return [ make_rectangle(p1, p2, '#c7c7c7') ]
+return make_split_rectangle(p1, p2, p1_color, p2_color)
+else:
+print >> sys.stderr, "Unknown state: {0}: {1} {2} {3} {4}".format(state, chrom, p1, p2, state, individual)
+sys.exit(1)
+def make_state_rectangle_3pop(p1, p2, state, chrom, individual):
+p1_color = 'r'
+p2_color = 'g'
+p3_color = 'b'
+heterochromatin_color = '#c7c7c7'
+if state == 0:
+return [ make_rectangle(p1, p2, heterochromatin_color) ]
+if state == 1:
+return [ make_rectangle(p1, p2, p1_color) ]
+if state == 2:
+return [ make_rectangle(p1, p2, p2_color) ]
+if state == 3:
+return [ make_rectangle(p1, p2, p3_color) ]
+if state == 4:
+return make_split_rectangle(p1, p2, p1_color, p2_color)
+if state == 5:
+return make_split_rectangle(p1, p2, p1_color, p3_color)
+if state == 6:
+return make_split_rectangle(p1, p2, p2_color, p3_color)
 else:
 print >> sys.stderr, "Unknown state: {0}: {1} {2} {3} {4}".format(state, chrom, p1, p2, state, individual)
 sys.exit(1)
 def nicenum(num, round=False):
 return vals, labels
 ################################################################################
-def make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir):
+def make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir, state2name=None, populations=3):
 fs_chrom_len = build_chrom_len_dict(input_dbkey, galaxy_data_index_dir)
-chroms, individuals, data, chrom_len = parse_input_file(input_file)
+chroms, individuals, data, chrom_len, used_states = parse_input_file(input_file)
+if populations == 3:
+make_state_rectangle = make_state_rectangle_3pop
+elif populations == 2:
+make_state_rectangle = make_state_rectangle_2pop
+else:
+pass
 for chrom in chrom_len.keys():
 if chrom in fs_chrom_len:
 chrom_len[chrom] = fs_chrom_len[chrom]
 chrom_height = 0.25
 ind_space = 0.10
 ind_height = 0.25
 total_height = 0.0
-at_top = True
+## make a legend
+## only print out states that are
+##   1) in the data
+##    - AND -
+##   2) in the state2name map
+## here, we only calculate the space needed
+legend_states = []
+if state2name is not None:
+for state in used_states:
+if state in state2name:
+legend_states.append(state)
+if legend_states:
+total_height += len(legend_states) * (ind_space + ind_height)
+total_height += (top_space - ind_space)
+at_top = False
+else:
+at_top = True
 for chrom in chroms:
 if at_top:
 total_height += (top_space + chrom_height)
 at_top = False
 else:
 bottom = 1.0
 fig = plt.figure(figsize=(width, height))
-at_top = True
+if legend_states:
+at_top = True
+for state in sorted(legend_states):
+if at_top:
+bottom -= (top_space + ind_height)/height
+at_top = False
+else:
+bottom -= (ind_space + ind_height)/height
+## add code here to draw legend
+# [left, bottom, width, height]
+ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/height])
+plt.axis('off')
+ax1.set_xlim(0, 1)
+ax1.set_ylim(0, 1)
+for patch in make_state_rectangle(0, 1, state, 'legend', state2name[state]):
+ax1.add_patch(patch)
+ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/height], frame_on=False)
+plt.axis('off')
+plt.text(0.0, 0.5, state2name[state], fontsize=10, ha='left', va='center')
+else:
+at_top = True
 for_webb = False
 for chrom in chroms:
 length = chrom_len[chrom]
 vals, labels = tick_foo(0, length)
 ax2.set_xticks(vals)
 ax2.set_xticklabels(labels)
 else:
 plt.axis('off')
 for p1, p2, state in sorted(data[chrom][individual]):
+#for patch in make_state_rectangle(p1, p2, state, chrom, individual):
 for patch in make_state_rectangle(p1, p2, state, chrom, individual):
 ax2.add_patch(patch)
 plt.savefig(output_file)
 if __name__ == '__main__':
 input_dbkey, input_file, output_file, galaxy_data_index_dir = sys.argv[1:5]
 make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir)
 sys.exit(0)
+## notes
+# 1) pass in a state to name mapping
+# 2) only print out names for states which exist in the data, and are in the state to name mapping

Mercurial > repos > miller-lab > genome_diversity

comparison dpmix_plot.py @ 27:8997f2ca8c7a