12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182 |
- #!/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)
|