From 0e25049a58223f217e969fd91151e9d1a92be7e9 Mon Sep 17 00:00:00 2001
From: tforest <thomas.forest@college-de-france.fr>
Date: Mon, 18 Oct 2021 17:11:51 +0200
Subject: [PATCH] folded version

---
 vcf_to_sfs.py | 37 +++++++++++++++++++++++--------------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/vcf_to_sfs.py b/vcf_to_sfs.py
index 686c8f6..5f27794 100755
--- a/vcf_to_sfs.py
+++ b/vcf_to_sfs.py
@@ -11,31 +11,40 @@ import sys
 
 # default folded SFS
 folded = True
+diploid = True
+
+# PARAM : Nb of indiv
+n = int(sys.argv[2])
+
+if diploid and not folded:
+    n *= 2
+
+SFS_values = dict.fromkeys(range(n),0)
+
 
 with gzip.open(sys.argv[1], "rb") as inputgz:
     line = inputgz.readline()
     genotypes = []
-    SFS_values = {}
+    #SFS_values = {}
     while line:
         line = line.decode('utf-8').strip()
         if not line.startswith("##") and not line.startswith("#"):
             FORMAT = line.split("\t")[8:9]
             SAMPLES = line.split("\t")[9:]
             snp_genotypes = []
+            allele_counts = {}
             for sample in SAMPLES:
                 # for UNPHASED data
-                smpl_genotype = [int(a) for a in sample.split(':')[0].split('/') if a != '.']
-                #if not folded:
-                    
-                print(smpl_genotype)
-                nb_alleles = len(set(smpl_genotype))
-                snp_genotypes.append(nb_alleles)
-            print(snp_genotypes)
-            nb_derived_allele = len([val for val in snp_genotypes if val != 0])
-            print("nb derived allele", nb_derived_allele)
-            if nb_derived_allele not in SFS_values.keys():
-                SFS_values[nb_derived_allele] = 1
-            else:
-                SFS_values[nb_derived_allele] += 1
+                smpl_genotype = [int(a) for a in sample.split(':')[0].split('/') if a != '.']                    
+                nb_alleles = set(smpl_genotype)
+                snp_genotypes += smpl_genotype
+            for k in set(snp_genotypes):
+                allele_counts[snp_genotypes.count(k)] = k
+            if folded :
+                for count in allele_counts.keys():
+                    if count <= len(snp_genotypes)/2 :
+                        SFS_values[count-1] += 1                        
+                    else:
+                        SFS_values[len(snp_genotypes)-count] += 1
         line = inputgz.readline()
         print(SFS_values)