From f79d48bf5b21c89391f15caf0ff47ad9654f1bca Mon Sep 17 00:00:00 2001
From: tforest <thomas.forest@college-de-france.fr>
Date: Fri, 17 Nov 2023 17:12:27 +0100
Subject: [PATCH] Rectify swp2 plot

---
 swp2.py | 38 ++++++++++++++++----------------------
 1 file changed, 16 insertions(+), 22 deletions(-)

diff --git a/swp2.py b/swp2.py
index 2ef7acd..c38ff07 100644
--- a/swp2.py
+++ b/swp2.py
@@ -1,4 +1,3 @@
-
 import matplotlib.pyplot as plt
 import os
 
@@ -29,17 +28,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])
     # revert the list of groups as the most recent times correspond
     # to the closest and last leafs of the coal. tree.
     groups = groups[::-1]
+    theta_site = theta_site[::-1]
     # initiate the dict of times
     t = {}
-    # list of thetas 
+    # list of thetas
     theta_L = []
     sum_t = 0
     for group_nb, group in enumerate(groups):
-        print(group_nb, group, theta_site[group_nb], len(theta_site))
+        ###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
@@ -51,48 +50,40 @@ def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen):
             i = int(group.split(",")[0])
             j = int(group.split(",")[-1])
         t[i] = 0
+        #t =
         if len(group.split(',')) == 1:
             k = i
             t[i] += ((theta_L[group_nb] ) / (k*(k-1)) * tgen) / mu
         else:
-            for k in range(i, j+1):
+            for k in range(j, i-1, -1 ):
                 t[i] += ((theta_L[group_nb] ) / (k*(k-1)) * tgen) / mu
         # we add the cumulative times at the end
         t[i] += sum_t
         sum_t = t[i]
     # build the y axis (sizes)
     y = []
-    #theta_L.sort(reverse=True)
     for theta in theta_L:
         # with size N = theta/4mu
         size = theta / (4*mu)
         y.append(size)
         y.append(size)
-    
     # 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)
-    
     return x,y,likelihood
 
 def plot_3epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title"):
-    
+
     scenari = {}
     cpt = 0
-    
+
     for file_name in os.listdir(folder_path):
         if os.path.isfile(os.path.join(folder_path, file_name)):
             # Perform actions on each file
-            x,y,likelihood = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks, 
+            x,y,likelihood = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                              tgen = tgen,
                                                              mu = mu)
             if x == 0 or y == 0:
@@ -100,7 +91,7 @@ def plot_3epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title")
             cpt +=1
             scenari[likelihood] = x,y
     print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(breaks)+"\n*******\n")
-    print(cpt, "theta files have been scanned.")
+    print(cpt, "theta file(s) have been scanned.")
     # sort starting by the smallest -log(Likelihood)
     best10_scenari = (sorted(list(scenari.keys())))[:10]
     print("10 greatest Likelihoods", best10_scenari)
@@ -109,9 +100,13 @@ def plot_3epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title")
     my_dpi = 300
     plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
     plt.plot(x, y, 'r-', lw=2, label = 'Lik='+greatest_likelihood)
+    plt.yscale('log')
+    #plt.xscale('log')
+    plt.grid(True,which="both", linestyle='--')
+
     for scenario in best10_scenari[1:]:
         x,y = scenari[scenario]
-        print("\n----  Lik:",scenario,"\n\nt=", x,"\n\nN=",y, "\n\n")
+        #print("\n----  Lik:",scenario,"\n\nt=", x,"\n\nN=",y, "\n\n")
         plt.plot(x, y, '--', lw=1, label = 'Lik='+scenario)
     plt.ylabel("Individuals (N)")
     plt.xlabel("Time (years)")
@@ -119,14 +114,13 @@ def plot_3epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title")
     plt.title(title)
     #plt.gcf().set_size(1000, 500)
     plt.savefig(title+'_b'+str(breaks)+'.pdf')
-    #plt.show()
 
 if __name__ == "__main__":
-            
+
     if len(sys.argv) != 4:
         print("Need 3 args: ThetaFolder MutationRate GenerationTime")
         exit(0)
-        
+
     folder_path = sys.argv[1]
     mu = sys.argv[2]
     tgen = sys.argv[3]