genome_diversity: dpmix_plot.py comparison

comparison dpmix_plot.py @ 31:a631c2f6d913

Update to Miller Lab devshed revision 3c4110ffacc3

author	Richard Burhans <burhans@bx.psu.edu>
date	Fri, 20 Sep 2013 13:25:27 -0400
parents	8997f2ca8c7a
children

comparison

equal deleted inserted replaced

-:4188853b940b
+:a631c2f6d913
 #!/usr/bin/env python
 import os
 import sys
 import math
 import matplotlib as mpl
 mpl.use('PDF')
+from matplotlib.backends.backend_pdf import PdfPages
 import matplotlib.pyplot as plt
 from matplotlib.path import Path
 import matplotlib.patches as patches
 ################################################################################
 #   labels.append('{0:.1f}'.format(vals[-1]/math.pow(10, digits)))
 return vals, labels
 ################################################################################
+################################################################################
+################################################################################
+################################################################################
+def space_for_legend(plot_params):
+space = 0.0
+legend_states = plot_params['legend_states']
+if legend_states:
+ind_space = plot_params['ind_space']
+ind_height = plot_params['ind_height']
+space += len(legend_states) * (ind_space + ind_height) - ind_space
+return space
+################################################################################
+def space_for_chroms(plot_params, chroms, individuals, data):
+space_dict = {}
+chrom_height = plot_params['chrom_height']
+ind_space = plot_params['ind_space']
+ind_height = plot_params['ind_height']
+for chrom in chroms:
+space_dict[chrom] = chrom_height
+individual_count = 0
+for individual in individuals:
+if individual in data[chrom]:
+individual_count += 1
+space_dict[chrom] += individual_count * (ind_space + ind_height)
+return space_dict
+################################################################################
 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, used_states = parse_input_file(input_file)
-if populations == 3:
+## populate chrom_len
-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]
 #check_chroms(chroms, chrom_len, input_dbkey)
 check_data(data, chrom_len, input_dbkey)
-## units below are inches
+## plot parameters
-top_space = 0.10
+plot_params = {
-chrom_space = 0.25
+'plot_dpi':        300,
-chrom_height = 0.25
+'page_width':     8.50,
-ind_space = 0.10
+'page_height':   11.00,
-ind_height = 0.25
+'top_margin':     0.10,
+'bottom_margin':  0.10,
-total_height = 0.0
+'chrom_space':    0.25,
+'chrom_height':   0.25,
-## make a legend
+'ind_space':      0.10,
-## only print out states that are
+'ind_height':     0.25,
+'legend_space':   0.10
+}
+## in the 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:
+plot_params['legend_states'] = legend_states
-total_height += len(legend_states) * (ind_space + ind_height)
-total_height += (top_space - ind_space)
+## choose the correct make_state_rectangle method
-at_top = False
+if populations == 3:
-else:
+plot_params['rectangle_method'] = make_state_rectangle_3pop
-at_top = True
+elif populations == 2:
+plot_params['rectangle_method'] = make_state_rectangle_2pop
-for chrom in chroms:
-if at_top:
+pdf_pages = PdfPages(output_file)
-total_height += (top_space + chrom_height)
-at_top = False
+	## generate a list of chroms for each page
+needed_for_legend = space_for_legend(plot_params)
+needed_for_chroms = space_for_chroms(plot_params, chroms, individuals, data)
+chrom_space_per_page = plot_params['page_height']
+chrom_space_per_page -= plot_params['top_margin'] + plot_params['bottom_margin']
+chrom_space_per_page -= needed_for_legend + plot_params['legend_space']
+chrom_space_per_page -= plot_params['chrom_space']
+chroms_left = chroms[:]
+pages = []
+space_left = chrom_space_per_page
+chrom_list = []
+while chroms_left:
+chrom = chroms_left.pop(0)
+space_needed = needed_for_chroms[chrom] + plot_params['chrom_space']
+if (space_needed > chrom_space_per_page):
+print >> sys.stderr, 'Multipage chroms not yet supported'
+sys.exit(1)
+		## sometimes 1.9 - 1.9 < 0 (-4.4408920985e-16)
+		## so, we make sure it's not more than a millimeter over
+if space_left - space_needed > -0.04:
+chrom_list.append(chrom)
+space_left -= space_needed
 else:
-total_height += (top_space + chrom_space + chrom_height)
+pages.append(chrom_list[:])
+chrom_list = []
-individual_count = 0
+chroms_left.insert(0, chrom)
-for individual in individuals:
+space_left = chrom_space_per_page
-if individual in data[chrom]:
-individual_count += 1
+############################################################################
-total_height += individual_count * (ind_space + ind_height)
+plot_dpi = plot_params['plot_dpi']
-width = 7.5
+page_width = plot_params['page_width']
-height = math.ceil(total_height)
+page_height = plot_params['page_height']
+top_margin = plot_params['top_margin']
-bottom = 1.0
+ind_space = plot_params['ind_space']
+ind_height = plot_params['ind_height']
-fig = plt.figure(figsize=(width, height))
+make_state_rectangle = plot_params['rectangle_method']
+legend_space = plot_params['legend_space']
-if legend_states:
+chrom_space = plot_params['chrom_space']
-at_top = True
+chrom_height = plot_params['chrom_height']
-for state in sorted(legend_states):
-if at_top:
+for page in pages:
-bottom -= (top_space + ind_height)/height
+fig = plt.figure(figsize=(page_width, page_height), dpi=plot_dpi)
-at_top = False
+bottom = 1.0 - (top_margin/page_height)
+# print legend
+if legend_states:
+top = True
+for state in sorted(legend_states):
+if top:
+bottom -= ind_height/page_height
+top = False
+else:
+bottom -= (ind_space + ind_height)/page_height
+ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/page_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/page_height], frame_on=False)
+plt.axis('off')
+plt.text(0.0, 0.5, state2name[state], fontsize=10, ha='left', va='center')
+bottom -= legend_space/page_height
+# print chroms
+top = True
+for chrom in page:
+length = chrom_len[chrom]
+vals, labels = tick_foo(0, length)
+if top:
+bottom -= chrom_height/page_height
+top = False
 else:
-bottom -= (ind_space + ind_height)/height
+bottom -= (chrom_space + chrom_height)/page_height
-## add code here to draw legend
-# [left, bottom, width, height]
+ax = fig.add_axes([0.0, bottom, 1.0, chrom_height/page_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)
-if at_top:
-bottom -= (top_space + chrom_height)/height
-at_top = False
-else:
-bottom -= (top_space + chrom_space + chrom_height)/height
-if not for_webb:
-ax = fig.add_axes([0.0, bottom, 1.0, chrom_height/height])
 plt.axis('off')
 plt.text(0.5, 0.5, chrom, fontsize=14, ha='center')
 individual_count = 0
 for individual in individuals:
 if individual in data[chrom]:
 individual_count += 1
 i = 0
 for individual in individuals:
 if individual in data[chrom]:
 i += 1
+bottom -= (ind_space + ind_height)/page_height
-bottom -= (ind_space + ind_height)/height
-if not for_webb:
+ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/page_height])
-# [left, bottom, width, height]
-ax1 = fig.add_axes([0.0, bottom, 0.09, ind_height/height])
 plt.axis('off')
 plt.text(1.0, 0.5, individual, fontsize=10, ha='right', va='center')
-# [left, bottom, width, height]
-ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/height], frame_on=False)
+ax2 = fig.add_axes([0.10, bottom, 0.88, ind_height/page_height], frame_on=False)
 ax2.set_xlim(0, length)
 ax2.set_ylim(0, 1)
-if i != individual_count:
-plt.axis('off')
+if i != individual_count:
-else:
+plt.axis('off')
-if not for_webb:
+else:
 ax2.tick_params(top=False, left=False, right=False, labelleft=False)
 ax2.set_xticks(vals)
 ax2.set_xticklabels(labels)
-else:
-plt.axis('off')
+for p1, p2, state in sorted(data[chrom][individual]):
-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)
-for patch in make_state_rectangle(p1, p2, state, chrom, individual):
-ax2.add_patch(patch)
+# extend last state to end of chrom
+if p2 < length:
-plt.savefig(output_file)
+for patch in make_state_rectangle(p2, length, state, chrom, individual):
+ax2.add_patch(patch)
+pdf_pages.savefig(fig)
+plt.close(fig)
+pdf_pages.close()
 ################################################################################
 if __name__ == '__main__':
-input_dbkey, input_file, output_file, galaxy_data_index_dir = sys.argv[1:5]
+make_dpmix_plot('loxAfr3', 'output.dat', 'output2_files/picture.pdf', '/scratch/galaxy/home/oocyte/galaxy_oocyte/tool-data', state2name={0: 'heterochromatin', 1: 'reference', 2: 'asian'}, populations=2)
-make_dpmix_plot(input_dbkey, input_file, output_file, galaxy_data_index_dir)
+#    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 @ 31:a631c2f6d913