put_vcf_genotypes_into_fasta_reference_no_discard_unmapped_sites.py

#!/usr/bin/env python
"""
Unwrap fasta file so that each sequence takes up only one line.

Usage:
	%program  input_file  output_file
"""

import os, sys, argparse

# parses command line arguments
def get_commandline_arguments ():
		
	parser = argparse.ArgumentParser()
	
	parser.add_argument("--vcf", required=True, dest="vcf", type=extant_file, help="vcf genotypes", metavar="FILE")
	parser.add_argument("--ref_fasta", required=True, dest="ref_fasta", type=extant_file, help="fasta reference sequences", metavar="FILE")
	parser.add_argument("--out", required=True, help="name and path of output file", metavar="FILENAME")

	args = parser.parse_args()
	
	return args

# checks if file exists and breaks 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

########

def read_wrapped_or_unwrapped_fasta (infile):
	
	outlines = []
	with open(infile, "r") as INFILE:
	
		first_id = INFILE.readline().strip("\n").strip(">")
		outlines.append(first_id)
		seq = ""
	
		for line in INFILE:
			line_clean = line.strip("\n")
			if line_clean.startswith(">"):
				outlines.append(seq)
				outlines.append(line_clean.strip(">"))
				seq = ""
			else:
				if len(line_clean) > 0:
					seq += line_clean
	
		# append last seq
		outlines.append(seq)
	
	i=0
	j=1
	out_dict = {}
	for x in range(len(outlines)/2):
		out_dict[outlines[i]] = outlines[j]
		i += 2
		j += 2
	
	return out_dict


def write_phylip (concat_dict, outfile):
	
	with open(outfile, "w") as OUTFILE:
		ntaxa = len(concat_dict.keys())
		len_align = len(concat_dict[concat_dict.keys()[0]]) # gets value of first element in dictionary -> this OK since all seqs have same length
		header = str(ntaxa) + " " + str(len_align) + "\n"
		OUTFILE.write(header)
		for sample, seq in concat_dict.items():
			out_line = sample + "    " + seq + "\n"
			OUTFILE.write(out_line)
	OUTFILE.close()

def parse_vcf(vcf_file):
	
	
	with open(vcf_file, "r") as INFILE:
		for line in INFILE:
			if line.startswith("##"):
				continue
			if line.startswith("#"):
				header_line = line.lstrip("#").strip("\n").split("\t")	
				print header_line
				break
	
	samples = sorted(header_line[9:])

### freebayes vcf line:	
#	E3379_L96	43	.	ATCG	ATTG,ATCA,GTCA,GTCG	4179.3	.	AB=BLABLA	0/0:1:
	
	
### STACKS vcf line:
#	#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	sample_203-ampullaria-Brunei-3
#	un	2312	25	G	A	.	PASS	NS=17;AF=0.529,0.471	GT:DP:GL	1/1:32:.,.,.	

	samples_vcf_idx = {}
	for sample in samples:
		print sample
		vcf_idx = header_line.index(sample)
		samples_vcf_idx[sample] = vcf_idx
#	print popdict_vcf_idx

# now start the main business of walking through the vcf:	
	out_dict = {x:{} for x in samples}
	with open(vcf_file, "r") as INFILE:
		
		snpcnt = 0
		for line in INFILE:
			if (snpcnt / 1000.0).is_integer():
				print snpcnt
			if line.startswith("#"):
				continue
			if len(line) < 2: # empty lines or so
				continue
			fields = line.strip("\n").split("\t")
			if len(fields[3]) > 1:
				continue # exclude variants that are not SNPs; theres no way to code those things in a phylip format!
			
			snpcnt += 1
			variants = [fields[3]] + fields[4].split(",") 
			
			contig = fields[0]
			pos = fields[1]
			
			for sample in samples:
				idxes = [int(x) for x in fields[samples_vcf_idx[sample]].split(":")[0].split("/") if not x == "." ]					
				if len(idxes) > 0:
					alleles = [variants[x] for x in idxes ]
					# get IUPAC code if necessary (heterozygotes)
					outnuc = get_IUPAC_amb_code ( "".join(sorted(alleles)) ) 
				else:
					outnuc = "-"
				try:
					out_dict[sample][contig][pos] = outnuc
				except KeyError:
					out_dict[sample][contig] = {pos : outnuc}
	
		return out_dict
	

def get_IUPAC_amb_code (nucs):
	
	# nucs must be a "".joined string of the sorted list of nucleotides
	
	# the IUPAC ambiguity codes:
	# the keys in this IUPAC ambiguity code dictionary are produced by:
	# "".join(sorted(list("NUCLEOTIDES"))) -> consistent lookup possible without additional loops!
	combs = {'AC':'M', 'GT':'K', 'CG':'S', 'AT':'W', 'AG':'R', 'CT':'Y',
	'ACG':'V', 'CGT':'B', 'AGT':'D', 'ACT':'H', 'ACGT':'N'
	}

	if len(set(nucs)) > 1:
		# sometimes theres already an N, so need to except that and return N instead!
		try:		
			out_code = combs[nucs]
		except KeyError:
			out_code = "N"
	else:
		out_code = list(nucs)[0]
	return out_code


def replace_str_index(text,index=0,replacement=''):
    return '%s%s%s'%(text[:index],replacement,text[index+1:])

def genotypes_into_fasta ( refgenome_dict, samples_genotype_dict ):
	
	out_dict = {}
	for contig, positions in samples_genotype_dict.items():
		refseq = refgenome_dict[contig]
		
# 		for pos
# 		
# 		if len(refseq) < max( [ int(x) for x in positions.keys() ] ):
# 			print "FUCK!", len(refseq), [ int(x) for x in positions.keys() ]
# 			print contig, refseq
#			exit()
		
#		print positions
		newseq = refseq
#		if not "-" in positions.values():
		for idx, nuc in positions.items():
			newseq = replace_str_index(newseq, int(idx)-1, nuc)  ### we substract 1 because vcf format has 1-based offset whereas python has 0-based offset
#			print idx, nuc
		out_dict[contig] = newseq
		
	return out_dict

				
################################## MAIN


######################

args = get_commandline_arguments ()

refgenome_file = args.ref_fasta
vcf_file = args.vcf
outfilename = args.out


ref_fasta = read_wrapped_or_unwrapped_fasta (refgenome_file)
vcf_dict = parse_vcf(vcf_file)

output_dict = {}
for sample, gts in vcf_dict.items():
	print sample
	individualised_fasta_seqs = genotypes_into_fasta ( ref_fasta, gts )
	output_dict[sample] = individualised_fasta_seqs


all_contigs = vcf_dict[ vcf_dict.keys()[0] ].keys()
print "loci / contigs / alignment partitions:	", len(all_contigs)


concat_dict = {}
for i in output_dict.keys():
	models_record = []
	start_idx = 1
	end_idx = 0

	concat_seq = ""
	for l in all_contigs:
		concat_seq += output_dict[i][l]
		end_idx += len( output_dict[i][l] )
		models_record.append( "DNA," + l + "=" + str(start_idx) + "-" + str(end_idx) )
		start_idx += len( output_dict[i][l] )
	print i, len(concat_seq)
	concat_dict[i] = concat_seq

with open(args.out + ".model.txt", "w") as OUTF:
	OUTF.write( "\n".join( models_record ) + "\n")

write_phylip (concat_dict, outfilename)
print "Done!"