diff --git a/swp2.py b/swp2.py index e4300c6..edabe79 100644 --- a/swp2.py +++ b/swp2.py @@ -176,7 +176,7 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", 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) @@ -184,8 +184,12 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", 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) @@ -202,7 +206,9 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", #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') @@ -341,7 +347,7 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, 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) @@ -443,6 +449,8 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, 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 @@ -463,40 +471,34 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = 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) @@ -509,8 +511,7 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = 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: @@ -533,29 +534,37 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = 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) @@ -565,7 +574,7 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale = #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) @@ -579,7 +588,7 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale = 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 @@ -588,8 +597,8 @@ def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = 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: @@ -628,7 +637,7 @@ def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = 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)