corr. vcf_to_sfs

vcf_to_sfs.py

@@ -17,9 +17,7 @@ import sys
 
 def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False):
 
-    """
-    Multiplication de deux nombres entiers.
-    Cette fonction ne sert pas à grand chose.
+    """ Returns an SFS from a VCF file.
 
     Parameters
     ----------
@@ -42,7 +40,8 @@ def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, ver
         n *= 2
     # initiate SFS_values with a zeros dict
     SFS_values = dict.fromkeys(range(n),0)
-
+    # store nb polyallellic sites
+    polyall = 0
     with gzip.open(vcf_file, "rb") as inputgz:
         line = inputgz.readline()
         genotypes = []
@@ -88,17 +87,13 @@ def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, ver
                     allele_counts[snp_genotypes.count(k)] = k
                     allele_counts_list.append(snp_genotypes.count(k))
                 if folded and len(ALT) >= 2:
-                    pass
-                    # TODO - work in progress
-                    # 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))
+                    polyall += 1
                 else:
                     SFS_values[min(allele_counts_list)-1] += 1
             line = inputgz.readline()
             if verbose:
                 print(SFS_values)
-    return SFS_values
+    return SFS_values, polyall
 
 if __name__ == "__main__":
             
@@ -106,8 +101,10 @@ if __name__ == "__main__":
         print("Need 2 args")
         exit(0)
 
+    # PARAM : vcf_file
+    vcf_file = sys.argv[1]
     # PARAM : Nb of indiv
     n = int(sys.argv[2])
 
-    sfs = sfs_from_vcf(n, sys.argv[1], folded = True, diploid = True, phased = False)
+    sfs, nb_polyall = sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False)
     print(sfs)