correction for sfs transform plot

customgraphics.py

 
 def barplot(x=None, y=None, ylab=None, xlab=None, title=None):
     if x:
+        x = list(x)
+        plt.xticks(x)
         plt.bar(x, y)
     else:
         x = list(range(len(y)))
         plt.bar(x, y)
+        plt.xticks(x)
     if ylab:
         plt.ylabel(ylab)
     if xlab:

sfs_tools.py

 import gzip
 import sys
 import matplotlib.pyplot as plt
+from frst import customgraphics
 
 def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False,
                  strip = False, count_ext = False):
 def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed = False):
     sfs_val = []
     n = len(sfs.values())
-    for k in range(1, n):
-        ksi = list(sfs.values())[k-1]
+    print("n =", n)
+    for k, ksi in sfs.items():
+        #ksi = list(sfs.values())[k-1]
         # k+1 because k starts from 0
         # if folded:
         #     # ?check if 2*n or not?
         #         sfs_val.append(ksi)
         if transformed:
             if folded:
-                sfs_val.append(ksi * k * (2*n - k))
+                #sfs_val.append(ksi * k * (2*n - k))
+                sfs_val.append(((k*(2*n - k)) / (2*n))*ksi) 
             else:
                 sfs_val.append(ksi * k)
         else:
             
     #terminal case, same for folded or unfolded
     if transformed:
-        sfs_val.append(list(sfs.values())[n-1] * n)
+        sfs_val[-1] = list(sfs.values())[n-1] * n
     else:
-         sfs_val.append(list(sfs.values())[n-1])
+        sfs_val[-1] = list(sfs.values())[n-1]
     #build the plot
     title = title+" [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
-    if ylab:
-        plt.ylabel(ylab)
-    if xlab:
-        plt.xlabel(xlab)
-    plt.title(title)
-    plt.bar([i+1 for i in sfs.keys()], sfs_val)
+    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)
     plt.show()
 
 if __name__ == "__main__":