82 lines
2.5 KiB
Python
Executable File
82 lines
2.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
"""
|
|
FOREST Thomas (thomas.forest@college-de-france.fr)
|
|
|
|
Caution : At the moment for gzipped files only.
|
|
|
|
ARGS
|
|
--------
|
|
|
|
usage : vcf_to_sfs.py VCF.gz nb_indiv
|
|
|
|
"""
|
|
|
|
import gzip
|
|
import sys
|
|
|
|
# default folded SFS
|
|
folded = True
|
|
diploid = True
|
|
phased = False
|
|
|
|
# PARAM : Nb of indiv
|
|
n = int(sys.argv[2])
|
|
|
|
if diploid and not folded:
|
|
n *= 2
|
|
# initiate SFS_values with a zeros dict
|
|
SFS_values = dict.fromkeys(range(n),0)
|
|
|
|
with gzip.open(sys.argv[1], "rb") as inputgz:
|
|
line = inputgz.readline()
|
|
genotypes = []
|
|
while line:
|
|
# decode gzipped binary lines
|
|
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
|
|
# 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 and len(ALT) >= 2:
|
|
for al in range(len(ALT)-1):
|
|
SFS_values[min(allele_counts_list)-1] += 1/len(ALT)
|
|
allele_counts_list.remove(min(allele_counts_list))
|
|
else:
|
|
SFS_values[min(allele_counts_list)-1] += 1
|
|
line = inputgz.readline()
|
|
print(SFS_values)
|