forked from mathii/gdc
-
Notifications
You must be signed in to change notification settings - Fork 0
/
vcf2eigenstrat.py
139 lines (119 loc) · 4.45 KB
/
vcf2eigenstrat.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
# Convert a vcf file to eigenstrat format
# removes multi-alleleic and indel sites.
# usage: python vcf2eigenstrat.py -v vcf_file.vcf(.gz) -o out_root
# will generate out_root.[snp,ind,geno].
# removed multiallelic sites and indels
# Deals with haploid cases including mixed haploid/diplod like X as well.
# -i option is a .ind file to get population names and sex.
from __future__ import division
import sys, getopt, gdc, pdb
################################################################################
def parse_options():
"""
Options are described by the help() function
"""
options ={ "vcf":None, "out":"out", "ref":None, "indAsPop":False, "indmap":None }
try:
opts, args = getopt.getopt(sys.argv[1:], "v:o:r:i:", ["vcf", "out", "ref", "indmap", "indAsPop"])
print opts, args
except Exception as err:
print str(err)
sys.exit()
for o, a in opts:
print o,a
if o in ["-v","--vcf"]: options["vcf"] = a
if o in ["-r","--ref"]: options["ref"] = a
if o in ["-i","--ind"]: options["indmap"] = a
if o in ["--indAsPop"]: options["indAsPop"] = True
elif o in ["-o","--out"]: options["out"] = a
print "found options:"
print options
return options
################################################################################
def main(options):
"""
Convert vcf to eigenstrat format (ind, snp and geno files)
"""
vcf=gdc.open2(options["vcf"])
snp, ind, geno = [open(options["out"]+x, "w") for x in [".snp", ".ind", ".geno"]]
removed={"multiallelic":0, "indel":0}
count=0
if options["indmap"]:
pop_map={}
sex_map={}
ind_map_file=open(options["indmap"], "r")
for line in ind_map_file:
bits=line[:-1].split()
pop_map[bits[0]]=bits[2]
sex_map[bits[0]]=bits[1]
ind_map_file.close()
for line in vcf:
if line[:2]=="##": # Comment line
next
elif line[:6]=="#CHROM": # Header line
inds=line.split()[9:]
if options["ref"]:
ind.write(options["ref"]+"\tU\tREF\n")
if options["indmap"]:
for indi in inds:
ind.write(indi+"\t"+sex_map.get(indi, "U")+"\t"+pop_map.get(indi, "POP")+"\n")
elif options["indAsPop"]:
for indi in inds:
ind.write(indi+"\tU\t"+indi+"\n")
else:
for indi in inds:
ind.write(indi+"\tU\tPOP\n")
else: # data
bits=line.split()
if "," in bits[4]:
removed["indel"]+=1
continue
if len(bits[3])!=1 or len(bits[4])!=1:
removed["multiallelic"]+=1
continue
else:
if bits[2]==".":
bits[2]=bits[0]+":"+bits[1]
snp.write(" ".join([bits[2], bits[0], "0.0", bits[1], bits[3], bits[4]])+"\n")
geno_string=""
if options["ref"]:
geno_string="2"
for gt in bits[9:]:
geno_string+=decode_gt_string(gt)
geno.write(geno_string+"\n")
count+=1
[f.close for f in [ind, snp, geno]]
print "Done. Wrote "+str(count) + " sites"
print "Excluded " + str(sum(removed.values())) + " sites"
for key in removed:
print "Excluded " + str(removed[key]) + " " + key
return
################################################################################
def decode_gt_string(gt_string):
"""
Tries to work out the genotype from a vcf genotype entry. 9 for missing [or not in {0,1,2}]
"""
gt=gt_string.split(":")[0]
if len(gt)==1:
if gt=="0": # haploid
return "2"
elif gt=="1":
return "0"
else:
return "9"
elif len(gt)==3:
if gt[0]=="0" and gt[2]=="0":
return "2"
if gt[0]=="0" and gt[2]=="1":
return "1"
if gt[0]=="1" and gt[2]=="0":
return "1"
if gt[0]=="1" and gt[2]=="1":
return "0"
else:
return "9"
raise Exception("Unknown genotype: "+gt)
################################################################################
if __name__=="__main__":
options=parse_options()
main(options)