Some progress on subplots

swp2.py

 from scipy.special import gammaln
 from matplotlib.backends.backend_pdf import PdfPages
 from matplotlib.ticker import MaxNLocator
+import io
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes
 
 def log_facto(k):
     k = int(k)
     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):
+def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, breaks_max = 12, ax = None):
     """
     Use theta values as is to do basic plots.
     """
     my_dpi = 300
     # plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
     # multiple fig
-    fig, ax = 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)
-    twin = ax.twiny()
-    # Offset the right spine of twin2
-    # twin.spines.right.set_position(("axes", 1.2))
+    if ax is None:
+        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:
+        ax1 = ax[0, 1]
+        plt.subplots_adjust(wspace=0.3, hspace=0.3)
+
+    twin = ax1.twiny()
+
     plots = []
     for epoch, theta in epochs.items():
         groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
             y[i] = y[i]/N0
         # plot
         #plt.plot(x, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        p, = ax.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')
         # add plot to the list of all plots to superimpose
         plots.append(p)
     # virtual line to get the second x axis for proportions
     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.1))
+    twin.spines["bottom"].set_position(("axes", -0.15))
     #ax.legend(handles=[p0]+plots)
-    ax.set_xlabel("# breaks")
-    # Set the x-axis locator to reduce the number of ticks = 10
-    ax.xaxis.set_major_locator(MaxNLocator(nbins=10))
-    ax.set_ylabel("theta")
+    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')
-    plt.savefig(title+'_raw'+str(k)+'.pdf')
+    if ax is None:
+        plt.savefig(title+'_raw'+str(k)+'.pdf')
     # fig 2 & 3
-    plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+    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:
+        # 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, theta in epochs.items():
         groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist()
         x = []
             T += y[i] / (x[i]*(x[i]-1))
             x_2.append(T)
         # Plotting (fig 2)
-        plt.plot(x_2, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        plt.xlabel("# breaks")
-        plt.ylabel("theta")
-        plt.legend(loc='upper right')
-        plt.savefig(title+'_plot2_'+str(k)+'.pdf')
+        p2, = ax2.plot(x_2, y, 'o', linestyle="dotted", 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
-        plt.plot(x_2, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
-        plt.xscale('log')
-        plt.xlabel("# breaks")
-        plt.ylabel("theta")
-        plt.legend(loc='upper right')
+        p3, = ax3.plot(x_2, y, 'o', linestyle="dotted", alpha=0.75, lw=2, label = str(epoch)+' brks')
+        lines_fig3.append(p3)
+    ax2.set_xlabel("# breaks")
+    ax2.set_ylabel("theta")
+    ax2.set_title("Test")
+    ax2.legend(handles=lines_fig2, loc='upper right')
+    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')
+    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)
     pp.close()
 
 def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = True):
-    plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale)
-    plot_test_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks)
-    save_multi_image(title+"_combined.pdf")
+    # 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)
+
+    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
+
+    # Save the entire grid as a single figure
+    plt.savefig(title+'_combined.pdf')
+
+
 if __name__ == "__main__":
 
     if len(sys.argv) != 4: