diff SNP_Mapping.py @ 5:fb735600b4ef

Uploaded
author gregory-minevich
date Tue, 27 Mar 2012 11:30:05 -0400
parents 30fa4d84e84c
children 7d6bccd5f88c
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/SNP_Mapping.py	Tue Mar 27 11:30:05 2012 -0400
@@ -0,0 +1,341 @@
+#!/usr/bin/python
+
+import re
+import sys
+import optparse
+import csv
+import re
+from decimal import *
+from rpy import *
+
+def main():
+	csv.field_size_limit(1000000000)
+
+	parser = optparse.OptionParser()
+	parser.add_option('-p', '--sample_pileup', dest = 'sample_pileup', action = 'store', type = 'string', default = None, help = "Sample pileup from mpileup")
+	parser.add_option('-v', '--haw_vcf', dest = 'haw_vcf', action = 'store', type = 'string', default = None, help = "vcf file of Hawaiian SNPs")
+	parser.add_option('-l', '--loess_span', dest = 'loess_span', action = 'store', type = 'float', default = .01, help = "Loess span")
+	parser.add_option('-d', '--d_yaxis', dest = 'd_yaxis', action = 'store', type = 'float', default = .7, help = "y-axis upper limit for dot plot")  
+	parser.add_option('-y', '--h_yaxis', dest = 'h_yaxis', action = 'store', type = 'int', default = 500, help = "y-axis upper limit for dot plot")   
+	parser.add_option('-c', '--points_color', dest = 'points_color', action = 'store', type = 'string', default = "gray27", help = "Color for data points") 
+	parser.add_option('-k', '--loess_color', dest = 'loess_color', action = 'store', type = 'string', default = "red", help = "Color for loess regression line")        
+	parser.add_option('-z', '--standardize', dest = 'standardize', default= 'false', help = "Standardize X-axis")
+
+	parser.add_option('-o', '--output', dest = 'output', action = 'store', type = 'string', default = None, help = "Output file name")
+	parser.add_option('-s', '--location_plot_output', dest = 'location_plot_output', action = 'store', type = 'string', default = "SNP_Mapping_Plot.pdf", help = "Output file name of SNP plots by chromosomal location")
+
+	#For plotting with map units on the X-axis instead of physical distance
+	#parser.add_option('-u', '--mpu_plot_output', dest = 'mpu_plot_output', action = 'store', type = 'string', default = None, help = "Output file name of SNP plots by map unit location")
+	(options, args) = parser.parse_args()
+
+	haw_snps = build_haw_snp_dictionary(haw_vcf = options.haw_vcf)	
+	pileup_info = parse_pileup(sample_pileup = options.sample_pileup, haw_snps = haw_snps)
+	output_pileup_info(output = options.output, pileup_info = pileup_info)
+
+	#output plot with all ratios
+	output_scatter_plots_by_location(location_plot_output = options.location_plot_output, pileup_info = pileup_info, loess_span=options.loess_span, d_yaxis=options.d_yaxis, h_yaxis=options.h_yaxis, points_color=options.points_color, loess_color=options.loess_color, standardize =options.standardize)
+
+	#For plotting with map units on the X-axis instead of physical distance)
+	#output_scatter_plots_by_mapping_units(mpu_plot_output = options.mpu_plot_output, pileup_info = pileup_info)
+
+def skip_headers(reader = None, i_file = None):
+	# count headers
+	comment = 0
+	while reader.next()[0].startswith('#'):
+		comment = comment + 1
+	
+	# skip headers
+	i_file.seek(0)
+	for i in range(0, comment):
+		reader.next()
+
+def location_comparer(location_1, location_2):
+	chr_loc_1 = location_1.split(':')[0]
+	pos_loc_1 = int(location_1.split(':')[1])
+
+	chr_loc_2 = location_2.split(':')[0]
+	pos_loc_2 = int(location_2.split(':')[1])
+
+	if chr_loc_1 == chr_loc_2:
+		if pos_loc_1 < pos_loc_2:
+			return -1
+		elif pos_loc_1 == pos_loc_1:
+			return 0
+		elif pos_loc_1 > pos_loc_2:
+			return 1
+	elif chr_loc_1 < chr_loc_2:
+		return -1
+	elif chr_loc_1 > chr_loc_2:
+		return 1
+
+def output_pileup_info(output = None, pileup_info = None):
+	o_file = open(output, 'wb')
+	writer = csv.writer(o_file, delimiter = '\t')
+
+	writer.writerow(["#Chr\t", "Pos\t", "ID\t", "Alt Count\t", "Ref Count\t", "Read Depth\t", "Ratio\t", "Mapping Unit"])
+
+	location_sorted_pileup_info_keys = sorted(pileup_info.keys(), cmp=location_comparer)
+
+	for location in location_sorted_pileup_info_keys:
+		alt_allele_count, ref_allele_count, read_depth, ratio, haw_snp_id, mapping_unit = pileup_info[location]
+		
+		location_info = location.split(':')
+		chromosome = location_info[0]
+		position = location_info[1]
+
+		writer.writerow([chromosome, position, haw_snp_id, alt_allele_count, ref_allele_count, read_depth, ratio, mapping_unit])
+
+	o_file.close()
+
+def output_scatter_plots_by_location(location_plot_output = None, pileup_info = None, loess_span="", d_yaxis="", h_yaxis="", points_color="", loess_color="", standardize=None):
+	i = {}
+	ii = {}
+	iii = {}
+	iv = {}
+	v = {}
+	x = {}
+
+	breaks = { 'I' : 16 , 'II' : 16,  'III' : 14, 'IV' : 18, 'V' : 21, 'X' : 18 }
+
+	for location in pileup_info:
+		chromosome = location.split(':')[0]
+		position = location.split(':')[1]
+
+		alt_allele_count, ref_allele_count, read_depth, ratio, haw_snp_id, mapping_unit = pileup_info[location]
+
+		if chromosome == "I": 
+       	        	i[position] = ratio
+               	elif chromosome == "II": 
+               		ii[position] = ratio
+                elif chromosome == "III": 
+                	iii[position] = ratio
+               	elif chromosome == "IV": 
+               		iv[position] = ratio
+                elif chromosome == "V": 
+                	v[position] = ratio
+               	elif chromosome == "X": 
+               		x[position] = ratio
+
+	x_label = "Location (Mb)"
+	filtered_label = ''
+
+
+	try:
+        	r.pdf(location_plot_output, 8, 8)
+		if i:
+		        plot_data(chr_dict = i, chr = "I" + filtered_label, x_label = "Location (Mb)", divide_position = True, draw_secondary_grid_lines = True, loess_span=loess_span, d_yaxis=d_yaxis, h_yaxis=h_yaxis, points_color=points_color, loess_color=loess_color, breaks = breaks["I"], standardize=standardize)
+		if ii:
+        		plot_data(chr_dict = ii, chr = "II" + filtered_label, x_label = "Location (Mb)", divide_position = True, draw_secondary_grid_lines = True, loess_span=loess_span, d_yaxis=d_yaxis, h_yaxis=h_yaxis, points_color=points_color, loess_color=loess_color, breaks = breaks["II"], standardize=standardize)
+		if iii:
+		        plot_data(chr_dict = iii, chr = "III" + filtered_label, x_label = "Location (Mb)", divide_position = True, draw_secondary_grid_lines = True, loess_span=loess_span, d_yaxis=d_yaxis, h_yaxis=h_yaxis, points_color=points_color, loess_color=loess_color, breaks = breaks["III"], standardize=standardize)
+        	if iv:
+			plot_data(chr_dict = iv, chr = "IV" + filtered_label, x_label = "Location (Mb)", divide_position = True, draw_secondary_grid_lines = True, loess_span=loess_span, d_yaxis=d_yaxis, h_yaxis=h_yaxis, points_color=points_color, loess_color=loess_color, breaks = breaks["IV"], standardize=standardize)
+		if v:
+		        plot_data(chr_dict = v, chr = "V" + filtered_label, x_label = "Location (Mb)", divide_position = True, draw_secondary_grid_lines = True, loess_span=loess_span, d_yaxis=d_yaxis, h_yaxis=h_yaxis, points_color=points_color, loess_color=loess_color, breaks = breaks["V"], standardize=standardize)
+		if x:
+        		plot_data(chr_dict = x, chr = "X" + filtered_label, x_label = "Location (Mb)", divide_position = True, draw_secondary_grid_lines = True, loess_span=loess_span, d_yaxis=d_yaxis, h_yaxis=h_yaxis, points_color=points_color, loess_color=loess_color, breaks = breaks["X"], standardize=standardize)
+
+	        r.dev_off()
+    	except Exception as inst:
+        	print inst
+        	print "There was an error creating the location plot pdf... Please try again"
+
+def output_scatter_plots_by_mapping_units(mpu_plot_output = None, pileup_info = None):
+	i = {}
+	ii = {}
+	iii = {}
+	iv = {}
+	v = {}
+	x = {}
+	
+	for location in pileup_info:
+		chromosome = location.split(':')[0]
+		position = location.split(':')[1]
+
+		alt_allele_count, ref_allele_count, read_depth, ratio, haw_snp_id, mapping_unit = pileup_info[location]
+
+		if chromosome == "I": 
+       	        	i[mapping_unit] = ratio
+               	elif chromosome == "II": 
+               		ii[mapping_unit] = ratio
+                elif chromosome == "III": 
+        	        iii[mapping_unit] = ratio
+                elif chromosome == "IV": 
+               		iv[mapping_unit] = ratio
+	        elif chromosome == "V": 
+                	v[mapping_unit] = ratio
+               	elif chromosome == "X": 
+               		x[mapping_unit] = ratio
+
+	x_label = "Map Units"
+
+	try:
+        	r.pdf(mpu_plot_output, 8, 8)
+	        plot_data(chr_dict = i, chr = "I", x_label = "Map Units")
+        	plot_data(chr_dict = ii, chr = "II", x_label = "Map Units")
+	        plot_data(chr_dict = iii, chr = "III", x_label = "Map Units")
+        	plot_data(chr_dict = iv, chr = "IV", x_label = "Map Units")
+	        plot_data(chr_dict = v, chr = "V", x_label = "Map Units")
+        	plot_data(chr_dict = x, chr = "X", x_label = "Map Units")
+	        r.dev_off()
+    	except Exception as inst:
+        	print inst
+        	print "There was an error creating the map unit plot pdf... Please try again"
+
+
+def plot_data(chr_dict =  None, chr = "", x_label = "", divide_position = False, draw_secondary_grid_lines = False, loess_span=None, d_yaxis=None, h_yaxis=None, points_color="", loess_color="", breaks = None, standardize= None):
+	ratios = "c("
+	positions = "c("
+	z_ratios = "c("
+	z_positions = "c("
+
+	for position in chr_dict:
+		ratio = chr_dict[position]
+		if divide_position:
+		       	position = float(position) / 1000000.0
+	        positions = positions + str(position) + ", "
+		ratios = ratios + str(ratio) + ", "
+		if ratio == 0:
+			if divide_position:
+			       	z_position = float(position) / 1000000.0
+		        z_positions = z_positions + str(position) + ", "
+			z_ratios = z_ratios + str(ratio) + ", "
+
+    
+	if len(ratios) == 2:
+		ratios = ratios + ")"
+	else:
+		ratios = ratios[0:len(ratios) - 2] + ")"
+
+	if len(z_ratios) == 2:
+		z_ratios = z_ratios + ")"
+	else:
+		z_ratios = z_ratios[0:len(z_ratios) - 2] + ")"
+	
+
+	if len(positions) == 2:
+		positions = positions + ")"
+	else:
+		positions = positions[0:len(positions) - 2] + ")"
+
+	if len(z_positions) == 2:
+		z_positions = z_positions + ")"
+	else:
+		z_positions = z_positions[0:len(z_positions) - 2] + ")"
+
+	r("x <- " + positions)
+	r("y <- " + ratios)
+
+	r("xz <- " + z_positions)
+	r("yz <- " + z_ratios)
+
+	if (standardize=='true'):    
+		r("plot(x, y, cex=0.60, xlim=c(0,21), main='LG " + chr + "', xlab= '" + x_label + "', ylim = c(0, %f " %d_yaxis + "), ylab='HA Ratio [Hawaiian Reads / Total Read Depth]', pch=18, col='"+ points_color +"')")
+		r("lines(loess.smooth(x, y, span = %f "%loess_span + "), lwd=5, col='"+ loess_color +"')")
+		r("axis(1, at=seq(0, 21, by=1), labels=FALSE, tcl=-0.5)")
+		r("axis(1, at=seq(0, 21, by=0.5), labels=FALSE, tcl=-0.25)")
+		r("axis(2, at=seq(floor(min(y)), 1, by=0.1), labels=FALSE, tcl=-0.2)")
+	elif (standardize=='false'):
+		r("plot(x, y, cex=0.60, main='LG " + chr + "', xlab= '" + x_label + "', ylim = c(0, %f " %d_yaxis + "), ylab='HA Ratio [Hawaiian Reads / Total Read Depth]', pch=18, col='"+ points_color +"')")
+		r("lines(loess.smooth(x, y, span = %f "%loess_span + "), lwd=5, col='"+ loess_color +"')")    
+		r("axis(1, at=seq(0, as.integer( ' " + str(breaks) + " '), by=1), labels=FALSE, tcl=-0.5)")
+		r("axis(1, at=seq(0, as.integer( ' " + str(breaks) + " '), by=0.5), labels=FALSE, tcl=-0.25)")	
+		r("axis(2, at=seq(floor(min(y)), 1, by=0.1), labels=FALSE, tcl=-0.2)")
+
+		
+	if draw_secondary_grid_lines:
+		r("abline(h = seq(floor(min(y)), 1, by=0.1), v = seq(floor(min(x)), length(x), by= 1), col='gray')")
+	else:
+		r("grid(lty = 1, col = 'gray')")
+
+	if (standardize=='true'):
+		r("hist(xz, col='darkgray', xlim=c(0,21), xlab='Location (Mb)', ylab='Frequency SNP Positions Where HA Ratio=0 ', ylim=c(0, %f " %h_yaxis + "), breaks = seq(0, as.integer( ' " + str(breaks) + " '), by=1), main='LG " + chr + "')")
+		r("hist(xz, add=TRUE, col=rgb(1, 0, 0, 1), xlim=c(0,21), xlab='Location (Mb)', ylab='Frequency SNP Positions Where HA Ratio=0 ', ylim=c(0, %f " %h_yaxis + "), breaks = seq(0, as.integer( ' " + str(breaks) + " '), by=.5), main='Chr " + chr + "')")
+		r("axis(1, at=seq(0, 21, by=1), labels=FALSE, tcl=-0.5)")
+		r("axis(1, at=seq(0, 21, by=0.5), labels=FALSE, tcl=-0.25)")
+	elif (standardize=='false'):
+		r("hist(xz, col='darkgray', xlab='Location (Mb)', ylab='Frequency SNP Positions Where HA Ratio=0 ', ylim=c(0, %f " %h_yaxis + "), breaks = seq(0, as.integer( ' " + str(breaks) + " '), by=1), main='LG " + chr + "')")
+		r("hist(xz, add=TRUE, col=rgb(1, 0, 0, 1), xlab='Location (Mb)', ylab='Frequency SNP Positions Where HA Ratio=0 ', ylim=c(0, %f " %h_yaxis + "), breaks = seq(0, as.integer( ' " + str(breaks) + " '), by=.5), main='Chr " + chr + "')")	
+		r("axis(1, at=seq(0, as.integer( ' " + str(breaks) + " '), by=1), labels=FALSE, tcl=-0.5)")
+		r("axis(1, at=seq(0, as.integer( ' " + str(breaks) + " '), by=0.5), labels=FALSE, tcl=-0.25)")
+
+
+def build_haw_snp_dictionary(haw_vcf = None):
+	haw_snps = {}
+	
+	i_file = open(haw_vcf, 'rU')
+	reader = csv.reader(i_file, delimiter = '\t')
+
+	skip_headers(reader = reader, i_file = i_file)
+	
+	for row in reader:
+		#print row
+		chromosome = row[0].upper()
+		chromosome = re.sub("chr", "", chromosome, flags = re.IGNORECASE)
+		chromosome = re.sub("CHROMOSOME_", "", chromosome, flags = re.IGNORECASE)
+
+		position = row[1]
+		haw_snp_id = row[2]
+		ref_allele = row[3]
+		alt_allele = row[4]
+        	info = row[7]
+            
+        	mapping_unit = info.replace("MPU=", "")
+
+		location = chromosome + ":" + position
+		haw_snps[location] = (alt_allele, haw_snp_id, mapping_unit)
+
+	i_file.close()
+
+	return haw_snps
+
+def parse_pileup(sample_pileup = None, haw_snps = None):
+	i_file = open(sample_pileup, 'rU')
+	reader = csv.reader(i_file, delimiter = '\t', quoting = csv.QUOTE_NONE)	
+
+	pileup_info = {}
+
+	for row in reader:
+		chromosome = row[0].upper()
+		chromosome = re.sub("chr", "", chromosome, flags = re.IGNORECASE)
+		chromosome = re.sub("CHROMOSOME_", "", chromosome, flags = re.IGNORECASE)
+
+		position = row[1]
+		ref_allele = row[2]
+		read_depth = row[3]
+		read_bases = row[4]
+
+		location = chromosome + ":" + position
+		if location in haw_snps:
+			alt_allele, haw_snp_id, mapping_unit = haw_snps[location]
+			ref_allele_count, alt_allele_count = parse_read_bases(read_bases = read_bases, alt_allele = alt_allele)
+		
+			getcontext().prec = 6	
+			ratio = Decimal(alt_allele_count) / Decimal(read_depth)
+		
+			pileup_info[location] = (alt_allele_count, ref_allele_count, read_depth, ratio, haw_snp_id, mapping_unit)
+		
+			#debug line
+			#print chromosome, position, read_depth, ref_allele_count, alt_allele_count, ratio, id
+
+	i_file.close()
+
+	return pileup_info
+
+def parse_read_bases(read_bases = None, alt_allele = None):
+	read_bases = re.sub('\$', '', read_bases)
+	read_bases = re.sub('\^[^\s]', '', read_bases)
+
+	ref_allele_matches = re.findall("\.|\,", read_bases)
+	ref_allele_count = len(ref_allele_matches)
+
+	alt_allele_matches = re.findall(alt_allele, read_bases, flags = re.IGNORECASE)
+	alt_allele_count = len(alt_allele_matches)
+
+	#debug line
+	#print read_bases, alt_allele, alt_allele_count, ref_allele_count
+
+	return ref_allele_count, alt_allele_count
+
+if __name__ == "__main__":
+	main()