swp2 plotting improve

sfs_tools.py

@@ -208,6 +208,7 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
         #     else:
         #         sfs_val.append(ksi)
         if transformed:
+            ylab = r'$ \phi_i $'
             if folded:
                 val = ((k*(2*n - k)) / (2*n))*(ksi)
             else:
@@ -216,6 +217,9 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
             val = ksi
         sfs_val.append(val)
 
+        if not transformed and not normalized:
+            ylab = r'$ \eta_i $'
+            
     #terminal case, same for folded or unfolded
     if transformed:
         last_bin = list(sfs.values())[n-1] * n/2
@@ -223,19 +227,21 @@ def barplot_sfs(sfs,  xlab, ylab, folded=True, title = "Barplot", transformed =
         last_bin = list(sfs.values())[n-1]
     sfs_val[-1] = last_bin
     if normalized:
-        ylab = "Fraction of SNPs"
+        #ylab = "Fraction of SNPs "
+        ylab = r'$ \phi_i $'
         sum_val = sum(sfs_val)
         for k, sfs_bin in enumerate(sfs_val):
             sfs_val[k] = sfs_bin / sum_val
+        
         #print(sum(sfs_val))
     #build the plot
-    title = title+" (n="+str(len(sfs_val)+1)+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
+    title = title+" (n="+str(len(sfs_val))+") [folded="+str(folded)+"]"+" [transformed="+str(transformed)+"]"
     print("SFS =", sfs)
     if folded:
         xlab = "Minor allele frequency"
     if transformed:
-        print("Transformed SFS ( n =",len(sfs_val)+1, ") :", sfs_val)
-        plt.axhline(y=1/n, color='r', linestyle='-')
+        print("Transformed SFS ( n =",len(sfs_val), ") :", sfs_val)
+        #plt.axhline(y=1/n, color='r', linestyle='-')
     else:
         if normalized:
             # then plot a theoritical distribution as 1/i

swp2.py

@@ -29,15 +29,17 @@ def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen):
     if dim < breaks+1:
         return 0,0,0
     # get n, the last bin of the last group
-    n = int(groups[-1].split(",")[-1])
+    #n = int(groups[-1].split(",")[-1])
+    # revert the list of groups as the most recent times correspond
+    # to the closest and last leafs of the coal. tree.
+    groups = groups[::-1]
     # initiate the dict of times
     t = {}
     # list of thetas 
     theta_L = []
     sum_t = 0
-    #groups.sort(reverse=True)
     for group_nb, group in enumerate(groups):
-        print(group_nb, group)
+        print(group_nb, group, theta_site[group_nb], len(theta_site))
         # store all the thetas one by one, with one theta per group
         theta_L.append(float(theta_site[group_nb]))
         # if the group is of size 1
@@ -69,11 +71,14 @@ def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen):
     
     # build the time x axis
     x = [0]
+    # need to fix: t.values are sorted? and in reverse order because the groups and the times are sorted in opposite order
     for time in range(0, len(t.values())-1):
         x.append(list(t.values())[time])
         x.append(list(t.values())[time])
     x.append(list(t.values())[len(t.values())-1])
 
+    #x = x[::-1]
+    #y = y[::-1]
     #x.sort(reverse=True)
     #y.sort(reverse=True)