correction of transformed sfs plot

sfs_tools.py

@@ -192,10 +192,10 @@ def sfs_from_parsed_vcf(n, vcf_dict, folded = True, diploid = True, phased = Fal
     return SFS_values, count_pluriall
 
 
-def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed = False):
+def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed = False, normalized = False):
     sfs_val = []
     n = len(sfs.values())
-    print("n =", n)
+    sum_sites = sum(list(sfs.values()))
     for k, ksi in sfs.items():
         #ksi = list(sfs.values())[k-1]
         # k+1 because k starts from 0
@@ -209,24 +209,30 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
         #         sfs_val.append(ksi)
         if transformed:
             if folded:
-                #sfs_val.append(ksi * k * (2*n - k))
-                sfs_val.append(((k*(2*n - k)) / (2*n))*ksi) 
+                val = ((k*(2*n - k)) / (2*n))*(ksi)
             else:
-                sfs_val.append(ksi * k)
+                val = ksi * k
         else:
-             sfs_val.append(ksi)
+            val = ksi
+        sfs_val.append(val)
             
     #terminal case, same for folded or unfolded
     if transformed:
-        sfs_val[-1] = list(sfs.values())[n-1] * n
+        last_bin = list(sfs.values())[n-1] * n/2
     else:
-        sfs_val[-1] = list(sfs.values())[n-1]
+        last_bin = list(sfs.values())[n-1]
+    sfs_val[-1] = last_bin
+    if normalized:
+        ylab = "Fraction of SNPs"
+        sum_val = sum(sfs_val)
+        for k, sfs_bin in enumerate(sfs_val):
+            sfs_val[k] = sfs_bin / sum_val
     #build the plot
-    title = title+" [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
+    title = title+" (n="+str(len(sfs_val)+1)+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
     print("SFS =", sfs)
     if transformed:
-        print("Transformed SFS ( n =",len(sfs_val), ") :", sfs_val)
-    customgraphics.barplot(x = sfs.keys(), y= sfs_val, xlab = xlab, ylab = ylab, title = title)
+        print("Transformed SFS ( n =",len(sfs_val)+1, ") :", sfs_val)
+    customgraphics.barplot(x = [x for x in list(sfs.keys())], y= sfs_val, xlab = xlab, ylab = ylab, title = title)
     plt.show()
 
 if __name__ == "__main__":