From 4aedf5e2804d9a29e7af48e1449886cb1817f2d2 Mon Sep 17 00:00:00 2001 From: tforest Date: Sun, 25 Feb 2024 00:27:26 +0100 Subject: [PATCH] Improve plotting with swp2 control curve --- swp2.py | 236 ++++++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 178 insertions(+), 58 deletions(-) diff --git a/swp2.py b/swp2.py index ee5b848..e4300c6 100644 --- a/swp2.py +++ b/swp2.py @@ -126,10 +126,11 @@ def plot_straight_x_y(x,y): def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", theta_scale = True, ax = None, input = None, output = None): - my_dpi = 300 + my_dpi = 500 + L = full_dict["L"] if ax is None: # intialize figure - my_dpi = 300 + #my_dpi = 300 fnt_size = 18 # plt.rcParams['font.size'] = fnt_size fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi) @@ -139,16 +140,16 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", ax1 = ax[1][0,0] ax1.set_yscale('log') ax1.set_xscale('log') - ax1.grid(True,which="both", linestyle='--', alpha = 0.3) plot_handles = [] best_plot = full_dict['all_epochs']['best'] - p0, = ax1.plot(best_plot[0], best_plot[1], 'o', linestyle = "-", + p0, = ax1.plot(best_plot[0], best_plot[1], linestyle = "-", alpha=1, lw=2, label = str(best_plot[2])+' brks | Lik='+best_plot[3]) plot_handles.append(p0) + #ax1.grid(True,which="both", linestyle='--', alpha = 0.3) for k, plot_Lk in enumerate(full_dict['all_epochs']['plots']): plot_Lk = str(full_dict['all_epochs']['plots'][k][3]) # plt.rcParams['font.size'] = fnt_size - p, = ax1.plot(full_dict['all_epochs']['plots'][k][0], full_dict['all_epochs']['plots'][k][1], 'o', linestyle = "--", + p, = ax1.plot(full_dict['all_epochs']['plots'][k][0], full_dict['all_epochs']['plots'][k][1], linestyle = "-", alpha=1/(k+1), lw=1.5, label = str(full_dict['all_epochs']['plots'][k][2])+' brks | Lik='+plot_Lk) plot_handles.append(p) if theta_scale: @@ -160,12 +161,20 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", # ax1.axvline(x=recent_scale_upper_bound) else: # years - plt.set_xlabel("Time (years)", fontsize=fnt_size) - plt.set_ylabel("Individuals (N)", fontsize=fnt_size) + if ax is not None: + plt.set_xlabel("Time (years)", fontsize=fnt_size) + plt.set_ylabel("Effective pop. size (Ne)", fontsize=fnt_size) + else: + plt.xlabel("Time (years)", fontsize=fnt_size) + plt.ylabel("Effective pop. size (Ne)", fontsize=fnt_size) + # x_ticks = ax1.get_xticks() + # ax1.set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5) + # ax1.set_xticklabels([f'{k}\n{k/(mu)}\n{k/(mu)*tgen}' for k in x_ticks], fontsize = fnt_size*0.8) # plt.rcParams['font.size'] = fnt_size # print(fnt_size, "rcParam font.size=", plt.rcParams['font.size']) ax1.legend(handles = plot_handles, loc='best', fontsize = fnt_size*0.5) ax1.set_title(title) + breaks = len(full_dict['all_epochs']['plots']) if ax is None: plt.savefig(title+'_b'+str(breaks)+'.pdf') # plot likelihood against nb of breakpoints @@ -203,8 +212,9 @@ def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", ax3.set_title(title+" AIC") if ax is None: plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf') - # return plots - return ax[0], ax[1] + else: + # return plots + return ax[0], ax[1] def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, input = None, output = None): #scenari = {} @@ -248,10 +258,11 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc # do something with the theta without bp and skip the plotting N0 = y[0] #continue - for i in range(len(y)): - # divide by N0 - y[i] = y[i]/N0 - x[i] = x[i]/N0 + if theta_scale: + for i in range(len(y)): + # divide by N0 + y[i] = y[i]/N0 + x[i] = x[i]/N0 top_plots[greatest_likelihood] = x,y,epoch plots_likelihoods = list(top_plots.keys()) for i in range(len(plots_likelihoods)): @@ -294,9 +305,10 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc # to return : plots ; Ln_Brks ; AIC_Brks ; best_Ln ; best_AIC # 'plots' dict keys: 'best', {epochs}('0', '1',...) if input == None: - saved_plots = {"S":S, "S0":S0, "L":L, "all_epochs":plots, "Ln_Brks":Ln_Brks, - "AIC_Brks":AIC_Brks, "best_Ln":best_Ln, - "best_AIC":best_AIC, "best_epoch_by_AIC":selected_brks_nb} + saved_plots = {"S":S, "S0":S0, "L":L, "mu":mu, "tgen":tgen, + "all_epochs":plots, "Ln_Brks":Ln_Brks, + "AIC_Brks":AIC_Brks, "best_Ln":best_Ln, + "best_AIC":best_AIC, "best_epoch_by_AIC":selected_brks_nb} else: # if the dict has to be loaded from input with open(input, 'r') as json_file: @@ -304,6 +316,8 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc saved_plots["S"] = S saved_plots["S0"] = S0 saved_plots["L"] = L + saved_plots["mu"] = mu + saved_plots["tgen"] = tgen saved_plots["all_epochs"] = plots saved_plots["Ln_Brks"] = Ln_Brks saved_plots["AIC_Brks"] = AIC_Brks @@ -369,9 +383,10 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, for k in range(2, len(y)+2): prop.append(y[k-2] / (k - 1) / sum_theta_i) prop = prop[::-1] - # normalise to N0 (N0 of epoch1) - for i in range(len(y)): - y[i] = y[i]/N0 + if theta_scale : + # normalise to N0 (N0 of epoch1) + for i in range(len(y)): + y[i] = y[i]/N0 # x_plot, y_plot = plot_straight_x_y(x, y) p = x, y # add plot to the list of all plots to superimpose @@ -394,8 +409,9 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, y += list(np.repeat(thetas[i], len(group))) if epoch == 0: N0 = y[0] - for i in range(len(y)): - y[i] = y[i]/N0 + if theta_scale : + for i in range(len(y)): + y[i] = y[i]/N0 x_2 = [] T = 0 for i in range(len(x)): @@ -426,10 +442,10 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, json.dump(saved_plots, json_file) return saved_plots -def plot_scaled_theta(plot_lines, prop, title, ax = None, n_ticks = 10, subset = None): +def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = None, n_ticks = 10, subset = None, theta_scale = False): # fig 2 & 3 if ax is None: - my_dpi = 300 + my_dpi = 500 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) @@ -442,37 +458,79 @@ def plot_scaled_theta(plot_lines, prop, title, ax = None, n_ticks = 10, subset = lines_fig2 = [] lines_fig3 = [] #plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi) + if swp2_lines: + for k in range(len(swp2_lines[0])): + 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') + 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') + 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-1: + 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) # 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) - 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) + + 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" if ax is None: + # if not ax, then use the plt syntax, not ax... + plt.xlabel(xlabel, fontsize=fnt_size) + plt.ylabel(ylabel, fontsize=fnt_size) + plt.xlim(left=0) + x_ticks = list(plt.xticks())[0] + plt.gca().set_xticks(x_ticks) + plt.gca().set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5) + plt.title(title, 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') # close fig2 to save memory plt.close(fig2) + else: + # when ax subplotting is used + ax2.set_xlabel(xlabel, fontsize=fnt_size) + ax2.set_ylabel(ylabel, fontsize=fnt_size) + ax2.set_title(title, fontsize=fnt_size) + ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5) ax3.set_xscale('log') ax3.set_yscale('log') - ax3.set_xlabel("log Relative scale", fontsize=fnt_size) + ax3.set_xlabel("time log 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) + x_ticks = list(ax3.get_xticks()) + ax3.set_xlim(left=min(x_ticks)) + ax3.set_xticks(x_ticks) + ax3.set_xticklabels([f'{k:.0e}\n{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], 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 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') @@ -480,7 +538,7 @@ def plot_scaled_theta(plot_lines, prop, title, ax = None, n_ticks = 10, subset = plt.close(fig3) return ax -def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10): +def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale = False, subset = None): # multiple fig if ax is None: # intialize figure 1 @@ -494,7 +552,6 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10): ax1 = ax[0, 0] plt.subplots_adjust(wspace=0.3, hspace=0.3) plots = [] - for epoch, plot in enumerate(plot_lines): x,y = plot x_plot, y_plot = plot_straight_x_y(x,y) @@ -527,40 +584,103 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10): # return plots return ax -def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = True, selected_breaks = []): +def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = False, selected_breaks = []): my_dpi = 300 - - save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, output = title+"_plotdata.json") - save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, input = title+"_plotdata.json", output = title+"_plotdata.json") + 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"] + 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: 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) - # 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, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks) - ax2 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], - prop = loaded_data['prop'], title = title, ax = ax2) - ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2]) - fig1.savefig(title+'_combined_p1.pdf') - print("Wrote", title+'_combined_p1.pdf') - fig2.savefig(title+'_combined_p2.pdf') - print("Wrote", title+'_combined_p2.pdf') + + # START OF COMBINED PLOT CODE + + # # plot page 1 of summary + # fig1, ax1 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi) + # # 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, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks) + # ax2 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], + # prop = loaded_data['prop'], title = title, ax = ax2) + # ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2]) + + # fig1.savefig(title+'_combined_p1.pdf') + # print("Wrote", title+'_combined_p1.pdf') + # fig2.savefig(title+'_combined_p2.pdf') + # print("Wrote", title+'_combined_p2.pdf') + + # END OF COMBINED PLOT CODE + + + # Start of Parsing real swp2 output + folder_splitted = folder_path.split("/") + swp2_summary = "/".join(folder_splitted[:-2])+'/'+folder_splitted[-3]+".final.summary" + swp2_vals = parse_stairwayplot_output_summary(stwplt_out = swp2_summary) + 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) - plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], - prop = loaded_data['prop'], title = title, ax = None) + 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) + plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = None) - plt.close(fig1) - plt.close(fig2) + # plt.close(fig1) + # plt.close(fig2) +def parse_stairwayplot_output_summary(stwplt_out, xlim = None, ylim = None, title = "default title", plot = False): + #col 5 + year = [] + # col 6 + ne_median = [] + ne_2_5 = [] + ne_97_5 = [] + ne_12_5 = [] + # col 10 + ne_87_5 = [] + with open(stwplt_out, "r") as stwplt_stream: + for line in stwplt_stream: + ## Line format + # mutation_per_site n_estimation theta_per_site_median theta_per_site_2.5% theta_per_site_97.5% year Ne_median Ne_2.5% Ne_97.5% Ne_12.5% Ne_87.5% + if not line.startswith("mutation_per_site"): + #not header + values = line.strip().split() + year.append(float(values[5])) + ne_median.append(float(values[6])) + ne_2_5.append(float(values[7])) + ne_97_5.append(float(values[8])) + ne_12_5.append(float(values[9])) + ne_87_5.append(float(values[10])) + + vals = [year, ne_median, ne_2_5, ne_97_5, ne_12_5, ne_87_5] + if plot : + # plot parsed data + label = ["Ne median", "Ne 2.5%", "Ne 97.5%", "Ne 12.5%", "Ne 87.5%"] + for i in range(1, 5): + fig, = plt.plot(year, vals[i], '--', alpha = 0.4) + fig.set_label(label[i]) + # # last plot is median + fig, = plt.plot(year, ne_median, 'r-', lw=2) + fig.set_label(label[0]) + plt.legend() + plt.ylabel("Individuals (Ne)") + plt.xlabel("Time (years)") + if xlim: + plt.xlim(xlim) + if ylim: + plt.ylim(ylim) + plt.title(title) + plt.show() + plt.close() + return vals if __name__ == "__main__":