Update on swp2 plots

swp2.py

@@ -216,6 +216,18 @@ def plot_k_epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title"
     plt.title(title)
     plt.savefig(title+'_b'+str(breaks)+'.pdf')
 
+def plot_straight_x_y(x,y):
+    x_1 = [x[0]]
+    y_1 = []
+    for i in range(0, len(y)-1):
+        x_1.append(x[i])
+        x_1.append(x[i])
+        y_1.append(y[i])
+        y_1.append(y[i])
+    y_1 = y_1+[y[-1],y[-1]]
+    x_1.append(x[-1])
+    return x_1, y_1
+
 def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, ax = None):
     #scenari = {}
     cpt = 0
@@ -237,6 +249,9 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
                 x,y,likelihood,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:
+                # last break did not work, then breaks = breaks-1
+                breaks -= 1
     print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(breaks)+"\n*******\n")
     print(cpt, "theta file(s) have been scanned.")
     my_dpi = 300
@@ -249,12 +264,12 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
         fnt_size = 12
         plt.rcParams['font.size'] = fnt_size
         ax1 = ax[0,0]
-    ax1.set_xlim(1e-3, 1)
-    #plt.ylim(0, 10)
+    #ax1.set_xlim(1e-3, 1)
     ax1.set_yscale('log')
     ax1.set_xscale('log')
     ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
     brkpt_lik = []
+    top_plots = {}
     for epoch, scenari in epochs.items():
         # sort starting by the smallest -log(Likelihood)
         best10_scenari = (sorted(list(scenari.keys())))[:10]
@@ -271,30 +286,44 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
             # divide by N0
             y[i] = y[i]/N0
             x[i] = x[i]/N0
-        ax1.plot(x, y, 'o', linestyle = "-", alpha=0.75, lw=2, label = str(epoch)+' BrkPt | Lik='+greatest_likelihood)
-        if theta_scale:
-            ax1.set_xlabel("Coal. time")
-            ax1.set_ylabel("Pop. size scaled by N0")
-            recent_scale_lower_bound = 0.01
-            recent_scale_upper_bound = 0.1
-            #print(recent_scale_lower_bound, recent_scale_upper_bound)
-            ax1.axvline(x=recent_scale_lower_bound)
-            ax1.axvline(x=recent_scale_upper_bound)
-        else:
-            # years
-            plt.set_xlabel("Time (years)")
-            plt.set_ylabel("Individuals (N)")
+        top_plots[greatest_likelihood] = x,y,epoch
+    plots_likelihoods = list(top_plots.keys())
+    for i in range(len(plots_likelihoods)):
+        plots_likelihoods[i] = float(plots_likelihoods[i])
+    best10_plots = sorted(plots_likelihoods)[:10]
+    top_plot_lik = str(best10_plots[0])
+    plot_handles = []
+    p0, = ax1.plot(top_plots[top_plot_lik][0], top_plots[top_plot_lik][1], 'o', linestyle = "-",
+    alpha=1, lw=2, label = str(top_plots[top_plot_lik][2])+' epoch | Lik='+top_plot_lik)
+    plot_handles.append(p0)
+    for k, plot_Lk in enumerate(best10_plots[1:]):
+        plot_Lk = str(plot_Lk)
+        p, = ax1.plot(top_plots[plot_Lk][0], top_plots[plot_Lk][1], 'o', linestyle = "--",
+        alpha=1/(k+1), lw=1.5, label = str(top_plots[plot_Lk][2])+' epoch | Lik='+plot_Lk)
+        plot_handles.append(p)
+    if theta_scale:
+        ax1.set_xlabel("Coal. time")
+        ax1.set_ylabel("Pop. size scaled by N0")
+        # recent_scale_lower_bound = 0.01
+        # recent_scale_upper_bound = 0.1
+        # ax1.axvline(x=recent_scale_lower_bound)
+        # ax1.axvline(x=recent_scale_upper_bound)
+    else:
+        # years
+        plt.set_xlabel("Time (years)")
+        plt.set_ylabel("Individuals (N)")
         ax1.set_xlim(1e-5, 1)
-        ax1.legend(loc='best', fontsize = fnt_size*0.5)
-        ax1.set_title(title)
-        if ax is None:
-            plt.savefig(title+'_b'+str(breaks)+'.pdf')
+    ax1.legend(handles = plot_handles, loc='best', fontsize = fnt_size*0.5)
+    ax1.set_title(title)
+    if ax is None:
+        plt.savefig(title+'_b'+str(breaks)+'.pdf')
     # plot likelihood against nb of breakpoints
     # best possible likelihood from SFS
     # Segregating sites
     S = sum(SFS_stored)
-    # number of monomorphic sites
+    # Number of kept sites from which the SFS is computed
     L = L_stored
+    # number of monomorphic sites
     S0 = L-S
     # print("SFS", SFS_stored)
     # print("S", S, "L", L, "S0=", S0)
@@ -303,8 +332,6 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
     for xi in range(0, len(SFS_stored)):
         p_i = SFS_stored[xi] / float(S+S0)
         Ln += np.log(p_i) * SFS_stored[xi] - log_facto(SFS_stored[xi])
-    res = Ln
-    # print(res)
     # basic plot likelihood
     if ax is None:
         fig, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
@@ -313,13 +340,12 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
         plt.rcParams['font.size'] = fnt_size
         ax2 = ax[2,0]
     ax2.plot(np.array(brkpt_lik)[:, 0], np.array(brkpt_lik)[:, 1].astype(float), 'o', linestyle = "dotted", lw=2)
-    # plt.ylim(0,100)
-    ax2.axhline(y=-Ln)
+    ax2.axhline(y=-Ln, linestyle = "-.", color = "red", label = "$-\log\mathcal{L}$ = "+str(round(-Ln, 2)))
     ax2.set_yscale('log')
     ax2.set_xlabel("# breakpoints", fontsize=fnt_size)
     ax2.set_ylabel("$-\log\mathcal{L}$")
-    #ax2.legend(loc='best', fontsize = fnt_size*0.5)
-    ax2.set_title(title)
+    ax2.legend(loc='best', fontsize = fnt_size*0.8)
+    ax2.set_title(title+" Likelihood gain from # breakpoints")
     if ax is None:
         plt.savefig(title+'_Breakpts_Likelihood.pdf')
     # AIC
@@ -330,12 +356,14 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc
         ax3 = ax[2,1]
     AIC = 2*(len(brkpt_lik)+1)+2*np.array(brkpt_lik)[:, 1].astype(float)
     ax3.plot(np.array(brkpt_lik)[:, 0], AIC, 'o', linestyle = "dotted", lw=2)
-    ax3.axhline(y=2*(len(brkpt_lik)+1)-2*Ln)
+    AIC_ln = 2*(len(brkpt_lik)+1)-2*Ln
+    ax3.axhline(y=AIC_ln, linestyle = "-.", color = "red",
+    label = "Min. AIC = "+str(round(AIC_ln, 2)))
     ax3.set_yscale('log')
     ax3.set_xlabel("# breakpoints", fontsize=fnt_size)
     ax3.set_ylabel("AIC")
-    #ax3.legend(loc='best', fontsize = fnt_size*0.5)
-    ax3.set_title(title)
+    ax3.legend(loc='best', fontsize = fnt_size*0.8)
+    ax3.set_title(title+" AIC")
     if ax is None:
         plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
     return ax
@@ -396,8 +424,9 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         for i in range(len(y)):
             y[i] = y[i]/N0
         # plot
+        x_plot, y_plot = plot_straight_x_y(x, y)
         #plt.plot(x, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        p, = ax1.plot(x, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
+        p, = ax1.plot(x_plot, y_plot, 'o', linestyle="-.", alpha=0.75, lw=2, label = str(epoch)+' brks')
         # add plot to the list of all plots to superimpose
         plots.append(p)
     # virtual line to get the second x axis for proportions
@@ -413,11 +442,13 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
         # in a combined plot, more space between the fig and the axis
         twin.spines["bottom"].set_position(("axes", -0.35))
     #ax.legend(handles=[p0]+plots)
-    ax1.set_xlabel("# breaks")
+    ax1.set_xlabel("# bin")
     # Set the x-axis locator to reduce the number of ticks to 10
     ax1.xaxis.set_major_locator(MaxNLocator(nbins=10))
+    twin.xaxis.set_major_locator(MaxNLocator(nbins=10))
     ax1.set_ylabel("theta")
     twin.set_ylabel("Proportion")
+    ax1.set_title("Title")
     ax1.legend(handles=plots, loc='best', fontsize = fnt_size*0.5)
     if ax is None:
         plt.savefig(title+'_raw'+str(k)+'.pdf')
@@ -453,21 +484,24 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
             T += y[i] / (x[i]*(x[i]-1))
             x_2.append(T)
         # Plotting (fig 2)
-        p2, = ax2.plot(x_2, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
+        x_2 = [0]+x_2
+        y = [y[0]]+y
+        x2_plot, y2_plot = plot_straight_x_y(x_2, y)
+        p2, = ax2.plot(x2_plot, y2_plot, 'o', linestyle="-.", alpha=0.75, lw=2, label = str(epoch)+' brks')
         lines_fig2.append(p2)
         # Plotting (fig 3) which is the same but log scale for x
-        p3, = ax3.plot(x_2, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
+        p3, = ax3.plot(x2_plot, y2_plot, 'o', linestyle="-.", alpha=0.75, lw=2, label = str(epoch)+' brks')
         lines_fig3.append(p3)
-    ax2.set_xlabel("# breaks")
+    ax2.set_xlabel("Relative scale")
     ax2.set_ylabel("theta")
-    ax2.set_title("Test")
+    ax2.set_title("Title")
     ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
     if ax is None:
         plt.savefig(title+'_plot2_'+str(k)+'.pdf')
     ax3.set_xscale('log')
-    ax3.set_xlabel("log()")
+    ax3.set_xlabel("log Relative scale")
     ax3.set_ylabel("theta")
-    ax3.set_title("Test")
+    ax3.set_title("Title")
     ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
     if ax is None:
         plt.savefig(title+'_plot3_'+str(k)+'_log.pdf')
@@ -477,18 +511,19 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
 def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = True):
     my_dpi = 300
     # Add some extra space for the second axis at the bottom
-    fig, axs = plt.subplots(3, 2, figsize=(4500/my_dpi, 2970/my_dpi), dpi=my_dpi)
+    fig, axs = plt.subplots(3, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
     ax = plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, ax = axs)
     ax = plot_test_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, ax = axs)
-
     # Adjust layout to prevent clipping of titles
     plt.tight_layout()
     # Adjust absolute space between the top and bottom rows
     plt.subplots_adjust(hspace=0.7)  # Adjust this value based on your requirement
-
     # Save the entire grid as a single figure
     plt.savefig(title+'_combined.pdf')
-
+    plt.close()
+    # 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)
 
 if __name__ == "__main__":