Improve plotting with concordant colours

swp2.py

     ax1.set_title(title)
     breaks = len(full_dict['all_epochs']['plots'])
     if ax is None:
-        plt.savefig(title+'_b'+str(breaks)+'.pdf')
+        plt.savefig(title+'_'+str(breaks+1)+'_epochs.pdf')
     # plot likelihood against nb of breakpoints
     if ax is None:
         fig, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
     else:
         #plt.rcParams['font.size'] = fnt_size
         ax2 = ax[0][0,1]
-
-    ax2.plot(full_dict['Ln_Brks'][0], full_dict['Ln_Brks'][1], 'o', linestyle = "dotted", lw=2)
+    # Retrieve the default color cycle from rcParams
+    default_colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
+    # Create an array of colors from the default color cycle
+    colors = [default_colors[i % len(default_colors)] for i in range(len(full_dict['Ln_Brks'][0]))]
+    ax2.plot(full_dict['Ln_Brks'][0], full_dict['Ln_Brks'][1], "--", lw=1, color="black", zorder=1)
+    ax2.scatter(full_dict['Ln_Brks'][0], full_dict['Ln_Brks'][1], s=50, c=colors, marker='o', zorder=2)
     ax2.axhline(y=full_dict['best_Ln'], linestyle = "-.", color = "red", label = "$-\log\mathcal{L}$ = "+str(round(full_dict['best_Ln'], 2)))
     ax2.set_yscale('log')
     ax2.set_xlabel("# breakpoints", fontsize=fnt_size)
         #plt.rcParams['font.size'] = fnt_size
         ax3 = ax[1][0,1]
     AIC = full_dict['AIC_Brks']
-    ax3.plot(AIC[0], AIC[1], 'o', linestyle = "dotted", lw=2)
+    # ax3.plot(AIC[0], AIC[1], 'o', linestyle = "dotted", lw=2)
+    ax3.plot(AIC[0], AIC[1], "--", lw=1, color="black", zorder=1)
+    ax3.scatter(AIC[0], AIC[1], s=50, c=colors, marker='o', zorder=2)
     ax3.axhline(y=full_dict['best_AIC'], linestyle = "-.", color = "red",
     label = "Min. AIC = "+str(round(full_dict['best_AIC'], 2)))
     ax3.set_yscale('log')
     for file_name in os.listdir(folder_path):
         cpt +=1
         if os.path.isfile(os.path.join(folder_path, file_name)):
-            for k in range(breaks_max):
+            for k in range(breaks_max+1):
                 x,y,likelihood,thetas,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = k,
                                                                  tgen = tgen,
                                                                  mu = mu, relative_theta_scale = theta_scale)
     return saved_plots
 
 def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = None, n_ticks = 10, subset = None, theta_scale = False):
+    # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
+    nb_epochs = len(plot_lines)
     # fig 2 & 3
     if ax is None:
         my_dpi = 500
             swp2_lines[0][k] = swp2_lines[0][k]/tgen*mu
         for k in range(len(swp2_lines[1])):
             swp2_lines[1][k] = swp2_lines[1][k]*4*mu
-        # plot_lines = [[swp2_lines[0], swp2_lines[1]]]+plot_lines 
-
         x2_plot, y2_plot = plot_straight_x_y(swp2_lines[0],swp2_lines[1])
-        p2, = ax2.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2')
+        p2, = ax2.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2', color="black")
         lines_fig2.append(p2)
         # Plotting (fig 3) which is the same but log scale for x
-        p3, = ax3.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2')
+        p3, = ax3.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2', color="black")
         lines_fig3.append(p3)
-    nb_breaks = len(plot_lines)
     for breaks, plot in enumerate(plot_lines):
-        if subset is not None:
-            if breaks not in subset :
-                # skip if not in subset
-                if max(subset) > nb_breaks and breaks == nb_breaks:
-                    pass
-                else:
-                    continue
         x,y=plot
-        # y = [k/(4*mu) for k in y]
-        # x = [k/(mu)*tgen for k in x]
         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(breaks)+' brks')
-        lines_fig2.append(p2)
+        if subset is not None:
+            if breaks in subset:
+                masking_alpha = 0.75
+            else:
+                masking_alpha = 0
+        p2, = ax2.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=masking_alpha, lw=2, label = str(breaks)+' brks')
         # 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(breaks)+' brks')
-        lines_fig3.append(p3)
-
+        p3, = ax3.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=masking_alpha, lw=2, label = str(breaks)+' brks')
+        if subset is not None and breaks in subset:
+            # store for legend
+            lines_fig2.append(p2)
+            lines_fig3.append(p3)
     ax3.axvline(x=500/tgen*mu, linestyle="--")
     if theta_scale:
         xlabel = "Theta scaled by N0"
         ylabel = "Theta scaled by N0"
     else:
-        xlabel = "Theta scale"
-        ylabel = "Theta"
+        xlabel = "t"
+        ylabel = r"$\theta$"
     if ax is None:
         # if not ax, then use the plt syntax, not ax...
         plt.xlabel(xlabel, fontsize=fnt_size)
         plt.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5)
         plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes)
         plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
-        # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
-        plt.savefig(title+'_plot2_'+str(len(plot_lines))+'.pdf')
+        plt.savefig(title+'_plotB_'+str(nb_epochs)+'_epochs.pdf')
         # close fig2 to save memory
         plt.close(fig2)
     else:
     plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
     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')
+        plt.savefig(title+'_plotC_'+str(nb_epochs)+'_epochs_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, rescale = False, subset = None):
+def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale = False, subset = None, max_breaks = None):
+    if max_breaks:
+        nb_breaks = max_breaks
+    else:
+        nb_breaks = len(plot_lines)+1
     # multiple fig
     if ax is None:
         # intialize figure 1
-        my_dpi = 300
+        my_dpi = 500
         fnt_size = 18
         # plt.rcParams['font.size'] = fnt_size
         fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi)
+        plt.subplots_adjust(bottom=0.2)  # Adjust the value as needed
     else:
         fnt_size = 12
         # plt.rcParams['font.size'] = fnt_size
         ax1 = ax[0, 0]
         plt.subplots_adjust(wspace=0.3, hspace=0.3)
     plots = []
-    for epoch, plot in enumerate(plot_lines):
+    for breaks, plot in enumerate(plot_lines):
+        if max_breaks and breaks > max_breaks:
+            # stop plotting if it exceeds the limit
+            continue
         x,y = plot
         x_plot, y_plot = plot_straight_x_y(x,y)
-        p, = ax1.plot(x_plot, y_plot, 'o', linestyle="-", 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(breaks)+' brks')
 
         # add plot to the list of all plots to superimpose
         plots.append(p)
     #ax.legend(handles=[p0]+plots)
     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_ylabel(r'$\theta_k$', fontsize=fnt_size, rotation = 90)
     ax1.set_title(title, fontsize=fnt_size)
     ax1.legend(handles=plots, loc='best', fontsize = fnt_size*0.5)
     ax1.set_xticks(x_ticks)
     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:
         # 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.savefig(title+'_raw_'+str(nb_breaks)+'_breaks.pdf')
         plt.close(fig)
     # return plots
     return ax
     my_dpi = 300
     saved_plots_dict = save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, output = title+"_plotdata.json")
     nb_of_epochs = len(saved_plots_dict["all_epochs"]["plots"])
-    print(nb_of_epochs)
     best_epoch = saved_plots_dict["best_epoch_by_AIC"]
+    print("Best epoch based on AIC =", best_epoch)
     save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = nb_of_epochs, input = title+"_plotdata.json", output = title+"_plotdata.json")
 
     with open(title+"_plotdata.json", 'r') as json_file:
     swp2_x, swp2_y = swp2_vals[0], swp2_vals[1]
     # End of Parsing real swp2 output
     plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
-                            prop = loaded_data['prop'], title = title, ax = None)
+                            prop = loaded_data['prop'], title = title, ax = None, max_breaks = breaks)
     plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], mu = mu, tgen = tgen, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks,
     # plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], subset=list(range(0,3))+[loaded_data['best_epoch_by_AIC']]+selected_breaks,
                             prop = loaded_data['prop'], title = title, swp2_lines = [swp2_x, swp2_y], ax = None)