1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798 |
- #!/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)
-
- polycount = 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
- # 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 :
- #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)
|