diff --git a/vcf_to_sfs.py b/vcf_to_sfs.py index 56f8cc5..4b21b39 100755 --- a/vcf_to_sfs.py +++ b/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)