From a92dba2d2519ca012da186f00c691086323ba4d4 Mon Sep 17 00:00:00 2001
From: tforest <thomas.forest@college-de-france.fr>
Date: Fri, 8 Jul 2022 14:15:41 +0200
Subject: [PATCH] correction for sfs transform plot

---
 customgraphics.py |  3 +++
 sfs_tools.py      | 23 ++++++++++++-----------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/customgraphics.py b/customgraphics.py
index 60c833b..af22486 100644
--- a/customgraphics.py
+++ b/customgraphics.py
@@ -202,10 +202,13 @@ def scatter(x, y, ylab=None, xlab=None, title=None):
 
 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:
diff --git a/sfs_tools.py b/sfs_tools.py
index a77642a..225273a 100755
--- a/sfs_tools.py
+++ b/sfs_tools.py
@@ -20,6 +20,7 @@ Rectify SFS comp in parsed funct.
 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):
@@ -194,8 +195,9 @@ def sfs_from_parsed_vcf(n, vcf_dict, folded = True, diploid = True, phased = Fal
 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?
@@ -207,7 +209,8 @@ 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(ksi * k * (2*n - k))
+                sfs_val.append(((k*(2*n - k)) / (2*n))*ksi) 
             else:
                 sfs_val.append(ksi * k)
         else:
@@ -215,17 +218,15 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
             
     #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__":