First version of the new layout of combined plots

swp2.py

     else:
         fnt_size = 12
         # plt.rcParams['font.size'] = fnt_size
-        ax1 = ax[0,0]
+        ax1 = ax[1][0,0]
     ax1.set_yscale('log')
     ax1.set_xscale('log')
     ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
         # plt.rcParams['font.size'] = fnt_size
     else:
         #plt.rcParams['font.size'] = fnt_size
-        ax2 = ax[2,0]
+        ax2 = ax[0][0,1]
     ax2.plot(np.array(brkpt_lik)[:, 0], np.array(brkpt_lik)[:, 1].astype(float), 'o', linestyle = "dotted", lw=2)
     ax2.axhline(y=-Ln, linestyle = "-.", color = "red", label = "$-\log\mathcal{L}$ = "+str(round(-Ln, 2)))
     ax2.set_yscale('log')
         # plt.rcParams['font.size'] = '18'
     else:
         #plt.rcParams['font.size'] = fnt_size
-        ax3 = ax[2,1]
-    AIC = 2*(len(brkpt_lik)+1)+2*np.array(brkpt_lik)[:, 1].astype(float)
+        ax3 = ax[1][0,1]
+    AIC = []
+    for brk in np.array(brkpt_lik)[:, 0]:
+        brk = int(brk)
+        AIC.append((2*brk+1)+2*np.array(brkpt_lik)[brk, 1].astype(float))
     ax3.plot(np.array(brkpt_lik)[:, 0], AIC, 'o', linestyle = "dotted", lw=2)
-    AIC_ln = 2*(len(brkpt_lik)+1)-2*Ln
+    # AIC = 2*k - 2ln(L) ; where k is the number of parameters, here brks+1
+    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)))
+    selected_brks_nb = AIC.index(min(AIC))
     ax3.set_yscale('log')
     ax3.set_xlabel("# breakpoints", fontsize=fnt_size)
     ax3.set_ylabel("AIC")
         plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
     print("S", S)
     # return plots
-    return ax
+    return ax[0], ax[1]
 
 def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True,
     breaks_max = 10, output = None):
         json.dump(saved_plots, json_file)
     return saved_plots
 
-def plot_raw_stairs(plot_lines, plot_lines2, prop, title, ax = None, n_ticks = 10):
+def plot_scaled_theta(plot_lines, prop, title, ax = None, n_ticks = 10):
+    # fig 2 & 3
+    if ax is None:
+        my_dpi = 300
+        fnt_size = 18
+        fig2, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+        fig3, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+    else:
+        # plt.rcParams['font.size'] = fnt_size
+        fnt_size = 12
+        # place of plots on the grid
+        ax2 = ax[1,0]
+        ax3 = ax[1,1]
+    lines_fig2 = []
+    lines_fig3 = []
+    #plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+    for epoch, plot in enumerate(plot_lines):
+        x,y=plot
+        x2_plot, y2_plot = plot_straight_x_y(x,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(x2_plot, y2_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(epoch)+' brks')
+        lines_fig3.append(p3)
+    ax2.set_xlabel("Relative scale", fontsize=fnt_size)
+    ax2.set_ylabel("theta", fontsize=fnt_size)
+    ax2.set_title(title, fontsize=fnt_size)
+    ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
+    if ax is None:
+        # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
+        plt.savefig(title+'_plot2_'+str(len(plot_lines))+'.pdf')
+        # close fig2 to save memory
+        plt.close(fig2)
+    ax3.set_xscale('log')
+    ax3.set_yscale('log')
+    ax3.set_xlabel("log Relative scale", fontsize=fnt_size)
+    ax3.set_ylabel("theta", fontsize=fnt_size)
+    ax3.set_title(title, fontsize=fnt_size)
+    ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
+    if ax is None:
+        # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
+        plt.savefig(title+'_plot3_'+str(len(plot_lines))+'_log.pdf')
+        # close fig3 to save memory
+        plt.close(fig3)
+    return ax
+
+def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10):
     # multiple fig
     if ax is None:
         # intialize figure 1
     else:
         fnt_size = 12
         # plt.rcParams['font.size'] = fnt_size
-        ax1 = ax[0, 1]
+        ax1 = ax[0, 0]
         plt.subplots_adjust(wspace=0.3, hspace=0.3)
     plots = []
 
     # print(x_ticks)
     #print(prop, "\n", sum(prop))
     #ax.legend(handles=[p0]+plots)
-    ax1.set_xlabel("# bin", fontsize=fnt_size)
+    ax1.set_xlabel("# bin & cumul. prop. of sites", fontsize=fnt_size)
     # Set the x-axis locator to reduce the number of ticks to 10
     ax1.set_ylabel("theta", fontsize=fnt_size)
-    ax1.set_title("Title", fontsize=fnt_size)
+    ax1.set_title(title, fontsize=fnt_size)
     ax1.legend(handles=plots, loc='best', fontsize = fnt_size*0.5)
     ax1.set_xticks(x_ticks)
     step = len(x_ticks)//(n_ticks-1)
     ax1.set_xticks(values)
     ax1.set_xticklabels([f'{values[k]}\n{val:.2f}' for k, val in enumerate(new_prop)], fontsize = fnt_size*0.8)
     if ax is None:
-        plt.savefig(title+'_raw'+str(k)+'.pdf')
-    # fig 2 & 3
-    if ax is None:
-        fig2, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-        fig3, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    else:
-        # plt.rcParams['font.size'] = fnt_size
-        # place of plots on the grid
-        ax2 = ax[1,0]
-        ax3 = ax[1,1]
-    lines_fig2 = []
-    lines_fig3 = []
-    #plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    for epoch, plot in enumerate(plot_lines2):
-        x,y=plot
-        x2_plot, y2_plot = plot_straight_x_y(x,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(x2_plot, y2_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        lines_fig3.append(p3)
-    ax2.set_xlabel("Relative scale", fontsize=fnt_size)
-    ax2.set_ylabel("theta", fontsize=fnt_size)
-    ax2.set_title("Title", fontsize=fnt_size)
-    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_yscale('log')
-    ax3.set_xlabel("log Relative scale", fontsize=fnt_size)
-    ax3.set_ylabel("theta", fontsize=fnt_size)
-    ax3.set_title("Title", fontsize=fnt_size)
-    ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
-    if ax is None:
-        plt.savefig(title+'_plot3_'+str(k)+'_log.pdf')
-        plt.clf()
+        # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
+        plt.savefig(title+'_raw'+str(len(plot_lines))+'.pdf')
+        plt.close(fig)
     # return plots
     return ax
 
     ax1.set_xlabel("# bin", fontsize=fnt_size)
     # Set the x-axis locator to reduce the number of ticks to 10
     ax1.set_ylabel("theta", fontsize=fnt_size)
-    ax1.set_title("Title", fontsize=fnt_size)
+    ax1.set_title(title, fontsize=fnt_size)
     ax1.legend(handles=plots, loc='best', fontsize = fnt_size*0.5)
     ax1.set_xticks(x_ticks)
     if len(prop) >= 18:
         lines_fig3.append(p3)
     ax2.set_xlabel("Relative scale", fontsize=fnt_size)
     ax2.set_ylabel("theta", fontsize=fnt_size)
-    ax2.set_title("Title", fontsize=fnt_size)
+    ax2.set_title(title, fontsize=fnt_size)
     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_yscale('log')
     ax3.set_xlabel("log Relative scale", fontsize=fnt_size)
     ax3.set_ylabel("theta", fontsize=fnt_size)
-    ax3.set_title("Title", fontsize=fnt_size)
+    ax3.set_title(title, fontsize=fnt_size)
     ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
     if ax is None:
         plt.savefig(title+'_plot3_'+str(k)+'_log.pdf')
     # 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.clf()
     # # plot_test_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, ax = None)
     # # plt.clf()
     # save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, output = title+"_plotdata.json")
-
     with open(title+"_plotdata.json", 'r') as json_file:
         loaded_data = json.load(json_file)
-
+    # plot page 1 of summary
     fig1, ax1 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
-    # plot_test_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, ax = ax1)
-    ax1 = plot_raw_stairs(plot_lines = loaded_data['raw_stairs'], plot_lines2 = loaded_data['scaled_stairs'],
+    # fig1.tight_layout()
+    # Adjust absolute space between the top and bottom rows
+    fig1.subplots_adjust(hspace=0.35)  # Adjust this value based on your requirement
+    # plot page 2 of summary
+    fig2, ax2 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
+    # fig2.tight_layout()
+    ax1 = plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
+                            prop = loaded_data['prop'], title = title, ax = ax1)
+
+    ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
                             prop = loaded_data['prop'], title = title, ax = ax1)
+    ax1, ax2 = plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, ax = [ax1, ax2])
+    fig1.savefig(title+'_combined_p1.pdf')
+    fig2.savefig(title+'_combined_p2.pdf')
+    plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
+                            prop = loaded_data['prop'], title = title, ax = None)
+    plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
+                            prop = loaded_data['prop'], title = title, ax = None)
 
-    plt.savefig(title+'_raw_scaled.pdf')
-    fig1.clf()
+    plt.close(fig1)
+    plt.close(fig2)
 if __name__ == "__main__":
 
     if len(sys.argv) != 4: