Grid of plots for swp2 thetas

swp2.py

     plt.title(title)
     plt.savefig(title+'_b'+str(breaks)+'.pdf')
 
-def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True):
+def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, ax = None):
     #scenari = {}
     cpt = 0
     epochs = {}
                                                                   mu = mu, relative_theta_scale = theta_scale)
     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.")
-
-    # intialize figure
     my_dpi = 300
-    plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    plt.xlim(1e-3, 1)
+    if ax is None:
+        # intialize figure
+        my_dpi = 300
+        fnt_size = 18
+        fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+    else:
+        fnt_size = 12
+        plt.rcParams['font.size'] = fnt_size
+        ax1 = ax[0,0]
+    ax1.set_xlim(1e-3, 1)
     #plt.ylim(0, 10)
-    plt.yscale('log')
-    plt.xscale('log')
-    plt.grid(True,which="both", linestyle='--', alpha = 0.3)
+    ax1.set_yscale('log')
+    ax1.set_xscale('log')
+    ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
     brkpt_lik = []
     for epoch, scenari in epochs.items():
         # sort starting by the smallest -log(Likelihood)
             # divide by N0
             y[i] = y[i]/N0
             x[i] = x[i]/N0
-        plt.plot(x, y, 'o', linestyle = "-", alpha=0.75, lw=2, label = str(epoch)+' BrkPt | Lik='+greatest_likelihood)
+        ax1.plot(x, y, 'o', linestyle = "-", alpha=0.75, lw=2, label = str(epoch)+' BrkPt | Lik='+greatest_likelihood)
         if theta_scale:
-            plt.xlabel("Coal. time")
-            plt.ylabel("Pop. size scaled by N0")
+            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)
-            plt.axvline(x=recent_scale_lower_bound)
-            plt.axvline(x=recent_scale_upper_bound)
+            ax1.axvline(x=recent_scale_lower_bound)
+            ax1.axvline(x=recent_scale_upper_bound)
         else:
             # years
-            plt.xlabel("Time (years)")
-            plt.ylabel("Individuals (N)")
-        plt.xlim(1e-5, 1)
-        plt.legend(loc='upper right')
-        plt.title(title)
-        plt.savefig(title+'_b'+str(breaks)+'.pdf')
+            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')
     # plot likelihood against nb of breakpoints
     # best possible likelihood from SFS
     # Segregating sites
     res = Ln
     # print(res)
     # basic plot likelihood
-    plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    plt.rcParams['font.size'] = '18'
-    plt.plot(np.array(brkpt_lik)[:, 0], np.array(brkpt_lik)[:, 1].astype(float), 'o', linestyle = "dotted", lw=2)
+    if ax is None:
+        fig, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+        plt.rcParams['font.size'] = '18'
+    else:
+        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)
-    # plt.axhline(y=res)
-    plt.yscale('log')
-    plt.xlabel("# breakpoints", fontsize=20)
-    plt.ylabel("$-\log\mathcal{L}$")
-    #plt.legend(loc='upper right')
-    plt.title(title)
-    plt.savefig(title+'_Breakpts_Likelihood.pdf')
+    ax2.axhline(y=-Ln)
+    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)
+    if ax is None:
+        plt.savefig(title+'_Breakpts_Likelihood.pdf')
     # AIC
-    plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-    plt.rcParams['font.size'] = '18'
+    if ax is None:
+        fig, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+        plt.rcParams['font.size'] = '18'
+    else:
+        ax3 = ax[2,1]
     AIC = 2*(len(brkpt_lik)+1)+2*np.array(brkpt_lik)[:, 1].astype(float)
-    plt.plot(np.array(brkpt_lik)[:, 0], AIC, 'o', linestyle = "dotted", lw=2)
-    # plt.axhline(y=106)
-    plt.yscale('log')
-    plt.xlabel("# breakpoints", fontsize=20)
-    plt.ylabel("AIC")
-    #plt.legend(loc='upper right')
-    plt.title(title)
-    plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
-
-def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, breaks_max = 12, ax = None):
+    ax3.plot(np.array(brkpt_lik)[:, 0], AIC, 'o', linestyle = "dotted", lw=2)
+    ax3.axhline(y=2*(len(brkpt_lik)+1)-2*Ln)
+    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)
+    if ax is None:
+        plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
+    return ax
+def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, breaks_max = 10, ax = None):
     """
     Use theta values as is to do basic plots.
     """
                 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.")
-    # intialize figure 1
-    my_dpi = 300
-    # plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
     # multiple fig
     if ax is None:
+        # intialize figure 1
+        my_dpi = 300
+        fnt_size = 18
         fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
         # Add some extra space for the second axis at the bottom
         fig.subplots_adjust(bottom=0.15)
     else:
+        fnt_size = 12
         ax1 = ax[0, 1]
         plt.subplots_adjust(wspace=0.3, hspace=0.3)
 
     twin.xaxis.set_ticks_position("bottom")
     twin.xaxis.set_label_position("bottom")
     # Offset the twin axis below the host
-    twin.spines["bottom"].set_position(("axes", -0.15))
+    if ax is None:
+        # arrange differently the second x axis if the plot is plain
+        twin.spines["bottom"].set_position(("axes", -0.15))
+    else:
+        # 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")
     # Set the x-axis locator to reduce the number of ticks to 10
     ax1.xaxis.set_major_locator(MaxNLocator(nbins=10))
     ax1.set_ylabel("theta")
-    # twin.set_ylabel("Proportion")
-    plt.legend(handles=plots, loc='upper right')
+    twin.set_ylabel("Proportion")
+    ax1.legend(handles=plots, loc='best', fontsize = fnt_size*0.5)
     if ax is None:
         plt.savefig(title+'_raw'+str(k)+'.pdf')
     # fig 2 & 3
     ax2.set_xlabel("# breaks")
     ax2.set_ylabel("theta")
     ax2.set_title("Test")
-    ax2.legend(handles=lines_fig2, loc='upper right')
+    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_ylabel("theta")
     ax3.set_title("Test")
-    ax3.legend(handles=lines_fig3, loc='upper right')
+    ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5)
     if ax is None:
         plt.savefig(title+'_plot3_'+str(k)+'_log.pdf')
     # return plots
     return ax
 
-def save_combined_pdf(output_path):
-    with PdfPages(output_path) as pdf:
-        pdf.savefig()
-
-def save_multi_image(filename):
-    pp = PdfPages(filename)
-    fig_nums = plt.get_fignums()
-    figs = [plt.figure(n) for n in fig_nums]
-    for fig in figs:
-        fig.savefig(pp, format='pdf')
-    pp.close()
-
 def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = True):
-    # plot1, plot2, plot3 = plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale)
     my_dpi = 300
     # Add some extra space for the second axis at the bottom
-    fig, axs = plt.subplots(2, 2, figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
-
+    fig, axs = plt.subplots(3, 2, figsize=(4500/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.35)  # Adjust this value based on your requirement
+    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')