Keeping only thetas with the lowest logLik for breakpoints thetas

swp2.py

@@ -56,12 +56,18 @@ def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_
     #### END of parsing
     # quit this file if the number of dimensions is incorrect
     if dim < breaks+1:
-        return 0,0,0,0,0
+        return 0,0,0,0,0,0
     # get n, the last bin of the last group
     # 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]
+    # store thetas for later use
+    grps = groups.copy()
+    thetas = {}
+    for i in range(len(groups)):
+        grps[i] = grps[i].split(',')
+        thetas[i] = [float(theta_site[i]), grps[i], likelihood]
     # initiate the dict of times
     t = {}
     # list of thetas
@@ -119,7 +125,7 @@ def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_
     #     #     # divide by N0
     #     #     y[i] = y[i]/N0
     #     #     x[i] = x[i]/N0
-    return x,y,likelihood,sfs,L
+    return x,y,likelihood,thetas,sfs,L
 
 def return_x_y_from_stwp_theta_file_as_is(stwp_theta_file, breaks, mu, tgen, relative_theta_scale = False):
     with open(stwp_theta_file, "r") as swp_file:
@@ -163,7 +169,7 @@ def return_x_y_from_stwp_theta_file_as_is(stwp_theta_file, breaks, mu, tgen, rel
 
     for i in range(len(groups)):
         groups[i] = groups[i].split(',')
-        #print(groups[i], len(groups[i]))
+        # print(groups[i], len(groups[i]))
         thetas[i] = [float(theta_site[i]), groups[i], likelihood]
     return thetas, sfs
 
@@ -236,7 +242,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
         breaks = 0
         cpt +=1
         if os.path.isfile(os.path.join(folder_path, file_name)):
-            x, y, likelihood, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
+            x, y, likelihood, theta, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                              tgen = tgen,
                                                              mu = mu, relative_theta_scale = theta_scale)
             SFS_stored = sfs
@@ -246,7 +252,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
                     epochs[breaks] = {}
                 epochs[breaks][likelihood] = x,y
                 breaks += 1
-                x,y,likelihood,sfs,L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
+                x,y,likelihood,theta,sfs,L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks,
                                                                  tgen = tgen,
                                                                   mu = mu, relative_theta_scale = theta_scale)
             if x == 0:
@@ -373,6 +379,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
     if ax is None:
         plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
     print("S", S)
+    # return plots
     return ax
 
 def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, breaks_max = 10, ax = None, n_ticks = 10):
@@ -381,6 +388,7 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
     """
     cpt = 0
     epochs = {}
+    len_sfs = 0
     for file_name in os.listdir(folder_path):
         cpt +=1
         if os.path.isfile(os.path.join(folder_path, file_name)):
@@ -390,7 +398,13 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
                                                                  mu = mu, relative_theta_scale = theta_scale)
                 if thetas == 0:
                     continue
-                epochs[k] = thetas
+                if len(thetas)-1 != k:
+                    continue
+                if k not in epochs.keys():
+                    epochs[k] = {}
+                likelihood = thetas[k][2]
+                epochs[k][likelihood] = thetas
+                #epochs[k] = thetas
     print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(k)+"\n*******\n")
     print(cpt, "theta file(s) have been scanned.")
     # multiple fig
@@ -405,9 +419,16 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         # plt.rcParams['font.size'] = fnt_size
         ax1 = ax[0, 1]
         plt.subplots_adjust(wspace=0.3, hspace=0.3)
-
     plots = []
-    for epoch, theta in epochs.items():
+    best_epochs = {}
+    for epoch in epochs:
+        likelihoods = []
+        for key in epochs[epoch].keys():
+            likelihoods.append(float(key))
+        likelihoods.sort()
+        minLogLn = str(likelihoods[0])
+        best_epochs[epoch] = epochs[epoch][minLogLn]
+    for epoch, theta in best_epochs.items():
         groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
         x = []
         y = []
@@ -467,7 +488,7 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
     lines_fig2 = []
     lines_fig3 = []
     #plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    for epoch, theta in epochs.items():
+    for epoch, theta in best_epochs.items():
         groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
         x = []
         y = []
@@ -529,10 +550,10 @@ def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale =
     # Save the entire grid as a single figure
     plt.savefig(title+'_combined.pdf')
     plt.clf()
-    # second call for individual plots
-    plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, ax = None)
-    plot_test_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, ax = None)
-    plt.clf()
+    # # second call for individual plots
+    # plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, ax = None)
+    # plot_test_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, ax = None)
+    # plt.clf()
 
 if __name__ == "__main__":