From 44449033db789a4db57fada071c7a2cc2f04ffe5 Mon Sep 17 00:00:00 2001
From: tforest <thomas.forest@college-de-france.fr>
Date: Wed, 21 Feb 2024 22:11:24 +0100
Subject: [PATCH] Update SFS plotting function

---
 customgraphics.py |  9 ++++++---
 sfs_tools.py      | 32 ++++++++++++++++++++++----------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/customgraphics.py b/customgraphics.py
index af22486..82503a1 100644
--- a/customgraphics.py
+++ b/customgraphics.py
@@ -200,14 +200,14 @@ def scatter(x, y, ylab=None, xlab=None, title=None):
         plt.title(title)
     plt.show()
 
-def barplot(x=None, y=None, ylab=None, xlab=None, title=None):
+def barplot(x=None, y=None, ylab=None, xlab=None, title=None, label=None, xticks = None, width=1):
     if x:
         x = list(x)
         plt.xticks(x)
-        plt.bar(x, y)
+        plt.bar(x, y, width=width, label=label)
     else:
         x = list(range(len(y)))
-        plt.bar(x, y)
+        plt.bar(x, y, width=width, label=label)
         plt.xticks(x)
     if ylab:
         plt.ylabel(ylab)
@@ -215,6 +215,9 @@ def barplot(x=None, y=None, ylab=None, xlab=None, title=None):
         plt.xlabel(xlab)
     if title:
         plt.title(title)
+    if xticks:
+        plt.xticks(xticks)
+    plt.legend()
     plt.show()
 
 def plot_chrom_continuity(vcf_entries, chr_id, x=None, y=None, outfile = None,
diff --git a/sfs_tools.py b/sfs_tools.py
index 6633d5c..3cfe5bf 100755
--- a/sfs_tools.py
+++ b/sfs_tools.py
@@ -21,6 +21,7 @@ import gzip
 import sys
 import matplotlib.pyplot as plt
 from frst import customgraphics
+import numpy as np
 
 def sfs_from_vcf(n, vcf_file, folded = True, diploid = True, phased = False, verbose = False,
                  strip = False, count_ext = False):
@@ -192,7 +193,7 @@ 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, normalized = False):
+def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed = False, normalized = False, ploidy = 2):
     sfs_val = []
     n = len(sfs.values())
     sum_sites = sum(list(sfs.values()))
@@ -222,7 +223,7 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
             
     #terminal case, same for folded or unfolded
     if transformed:
-        last_bin = list(sfs.values())[n-1] * n/2
+        last_bin = list(sfs.values())[n-1] * n/ploidy
     else:
         last_bin = list(sfs.values())[n-1]
     sfs_val[-1] = last_bin
@@ -235,22 +236,33 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
         
         #print(sum(sfs_val))
     #build the plot
-    title = title+" (n="+str(len(sfs_val))+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
-    print("SFS =", sfs)
     if folded:
         xlab = "Minor allele frequency"
+        n_title = n
+    else:
+        # the spectrum is n-1 long when unfolded
+        n_title = n+1
+    
+    title = title+" (n="+str(n_title)+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
+    print("SFS =", sfs)
+
+    X_axis = list(sfs.keys()) 
+    
+
     if transformed:
-        print("Transformed SFS ( n =",len(sfs_val), ") :", sfs_val)
+        print("Transformed SFS ( n =",n_title, ") :", sfs_val)
         #plt.axhline(y=1/n, color='r', linestyle='-')
+        plt.bar([x+0.2 for x in list(sfs.keys())], [1/n]*n, color='r', linestyle='-', width = 0.4, label= "H0 Theoric constant")
+
     else:
         if normalized:
             # then plot a theoritical distribution as 1/i
-            expected_y = [1/(2*x+1) for x in list(sfs.keys())]
+            sum_expected = sum([(1/(i+1)) for i,x in enumerate(list(sfs.keys()))])
+            expected_y = [(1/(i+1))/sum_expected for i,x in enumerate(list(sfs.keys()))]
+            print(expected_y)
+            plt.bar([x+0.2 for x in list(sfs.keys())], expected_y, color='r', linestyle='-', width = 0.4, label= "H0 Theoric constant")
             print(sum(expected_y))
-            #plt.plot([x for x in list(sfs.keys())], expected_y, color='r', linestyle='-')
-            #print(expected_y)
-            
-    customgraphics.barplot(x = [x for x in list(sfs.keys())], y= sfs_val, xlab = xlab, ylab = ylab, title = title)
+    customgraphics.barplot(x = [x-0.2 for x in X_axis], width=0.4, y= sfs_val, xlab = xlab, ylab = ylab, title = title, label = "H1 Observed spectrum", xticks =list(sfs.keys()) )
     plt.show()
 
 if __name__ == "__main__":