vcf_subset_by_chromosome.py

#!/usr/local/bin/python
# Python 2.7.6
# vcf_subset_by_chromosome.py
# a python script that filters .vcf files for the CHROM field, i.e. extracts RAD-tags by their names
# Copyright (C) 2017, ETH Zurich, Mathias Scharmann
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# 		
# If you use this code please cite:
#
# "Scharmann M, Grafe TU, Metali F, Widmer A. (2017) Sex-determination 
# and sex chromosomes are shared across the radiation of dioecious 
# Nepenthes pitcher plants. XXX"
# AND/OR
# "Scharmann M, Metali F, Grafe TU, Widmer A. (2017) Divergence with 
# gene flow among multiple sympatric carnivorous Nepenthes pitcher 
# plants is linked to trap morphology. XXX"
# 	
# contact: mathias.scharmann[-at-]env.ethz.ch or msph52[-at-]gmail.com


# usage example
# python vcf_subset_by_chromosome.py --vcf raw.01.vcf.tmp6.recode.vcf --exclude_chrom fufuchroms.txt


import argparse
import os

########################## HEAD

# this function checks if file exists and exits script if not
def extant_file(x):
	"""
	'Type' for argparse - checks that file exists but does not open.
	"""
	
	if not os.path.exists(x):
		print "Error: {0} does not exist".format(x)
		exit()
	x = str(x)
	return x

# breaks script if non-UNIX linebreaks in input file popmap
def linebreak_check(x):

	if "\r" in open(x, "rb").readline():
		print "Error: classic mac (CR) or DOS (CRLF) linebreaks in {0}".format(x)
		exit()

######

# parses arguments from command line
def get_commandline_arguments ():
	
	parser = argparse.ArgumentParser()
	
	parser.add_argument("--vcf", required=True, type=extant_file,
		help="name/path of vcf input file", metavar="FILE")
	parser.add_argument("--exclude_chrom", required=True, type=extant_file,
		help="name/path of the file with names of chromosomes to exclude", metavar="FILE")
	
	args = parser.parse_args()
	
	# finish
	return args

################################## CORE


def parse_vcf(vcf_file, excl_chroms):
	
	try:
		os.remove(vcf_file+"subset.vcf")
	except OSError:
		None
	
	# get total lines:
	num_lines = sum(1 for line in open(vcf_file))

	# now start the main business of walking through the vcf:	
	out_lines = []
	
	with open(vcf_file, "r") as INFILE:
		
		with open(vcf_file+"subset.vcf", "a") as OUTFILE:
		
			linecnt = 0
			tlinecnt = 0
			retained_cnt = 0
			for line in INFILE:
				linecnt += 1
				tlinecnt += 1
#				print tlinecnt
				if line.startswith("#"):
					out_lines.append(line.strip("\n"))
				else:
					# throw away unwanted chromosomes:
					fields = line.strip("\n").split("\t")
					if fields[0] not in excl_chroms:
						
						# throw away also all non-biallelic SNPs:
						if "," not in fields[4]:
							out_lines.append(line.strip("\n"))
							retained_cnt += 1
					
				if linecnt == 100000:
					# write to file and flush memory
					OUTFILE.write("\n".join(out_lines)+"\n")
					linecnt = 0
					out_lines = []
				elif tlinecnt == num_lines:
					OUTFILE.write("\n".join(out_lines))
	print "retained\t" + str(retained_cnt)		
				
################################## MAIN


	
args = 	get_commandline_arguments ()
	
#print args

excl_chrom = set()
with open(args.exclude_chrom, "r") as INFILE:
	for line in INFILE:
		excl_chrom.add(line.strip("\n"))

parse_vcf(args.vcf, excl_chrom)

	
print "Done!"