add polyallelic snps

vcf_to_sfs.py

@@ -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:
-                # for UNPHASED data
-                smpl_genotype = [int(a) for a in sample.split(':')[0].split('/') if a != '.']
-                
+                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)