richard
/
geinf
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

add polyallelic snps

master
tforest 2021-10-22 11:31:06 +02:00
parent 110ba26f7e
commit c1e2cc9192
1 changed files with 45 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
vcf_to_sfs.py
View File

@ -1,6 +1,8 @@
#!/usr/bin/env python3
"""
FOREST Thomas (thomas.forest@college-de-france.fr)
Caution : At the moment for gzipped files only.
ARGS
@ -26,6 +28,8 @@ if diploid and not folded:
# initiate SFS_values with a zeros dict
SFS_values = dict.fromkeys(range(n),0)
polycount = 0
with gzip.open(sys.argv[1], "rb") as inputgz:
line = inputgz.readline()
genotypes = []
@ -34,23 +38,60 @@ with gzip.open(sys.argv[1], "rb") as inputgz:
line = line.decode('utf-8').strip()
# every snp line, not comment or header
if not line.startswith("##") and not line.startswith("#"):
FIELDS = line.split("\t")
# REF is col 4 of VCF
REF = FIELDS[3].split(",")
# ALT is col 5 of VCF
ALT = FIELDS[4].split(",")
FORMAT = line.split("\t")[8:9]
SAMPLES = line.split("\t")[9:]
snp_genotypes = []
allele_counts = {}
allele_counts_list = []
# SKIP the SNP if :
# 1 : missing
# 2 : deletion among REF
# 3 : deletion among ALT
if "./.:." in line \
or len(ALT[0]) > 1 \
or len(REF[0]) > 1:
line = inputgz.readline()
continue
for sample in SAMPLES:
if not phased:
# for UNPHASED data
smpl_genotype = [int(a) for a in sample.split(':')[0].split('/') if a != '.']
else:
# for PHASED
smpl_genotype = [int(a) for a in sample.split(':')[0].split('|') if a != '.']
nb_alleles = set(smpl_genotype)
snp_genotypes += smpl_genotype
# if set(snp_genotypes) > 2:
# polyallelic = set(snp_genotypes)
# else:
# polyallelic = False
polyallelic = len(ALT)
##print(REF, ALT, snp_genotypes)
# skip if all individuals have the same genotype
if len(set(snp_genotypes)) == 1:
line = inputgz.readline()
continue
for k in set(snp_genotypes):
allele_counts[snp_genotypes.count(k)] = k
allele_counts_list.append(snp_genotypes.count(k))
if folded :
SFS_values[min(allele_counts.keys())-1] += 1
#allele_counts_list = list(allele_counts.keys())
##print("ALC", allele_counts_list, "POLY", polyallelic, ALT)
# for al in range(polyallelic-1):
# SFS_values[min(allele_counts_list)-1] += 1/len(ALT)
# allele_counts_list.remove(min(allele_counts_list))
if len(ALT) == 1:
SFS_values[min(allele_counts_list)-1] += 1
else:
for al in range(polyallelic-1):
SFS_values[min(allele_counts_list)-1] += 1/len(ALT)
allele_counts_list.remove(min(allele_counts_list))
polycount += 1
line = inputgz.readline()
print(SFS_values)
print(polycount)