From 1848102140e0ed0e1d8459adbc899c2fe9f5f8ee Mon Sep 17 00:00:00 2001 From: tforest Date: Thu, 29 Feb 2024 00:59:25 +0100 Subject: [PATCH] Fixing some scaling for swp2 output --- swp2.py | 134 +++++++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 98 insertions(+), 36 deletions(-) diff --git a/swp2.py b/swp2.py index 766b5a2..7c6893b 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+'_'+str(breaks+1)+'_epochs.pdf') + plt.savefig(title+'_best_'+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) @@ -299,7 +299,15 @@ def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc # number of monomorphic sites S0 = L-S # print("SFS", SFS_stored) - # print("S", S, "L", L, "S0=", S0) + print("S", S, "L", L, "S0=", S0) + + my_n = len(SFS_stored)*2 + print("n=",my_n) + an = 1 + for i in range(2, my_n): + an +=1.0/i + + print("an=", an, "theta_w", S/an, "theta_w_p_site", (S/an)/L) # compute Ln Ln = log_facto(S+S0) - log_facto(S0) + np.log(float(S0)/(S+S0)) * S0 for xi in range(0, len(SFS_stored)): @@ -413,6 +421,25 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, cumul = val+cumul prop = prop_cumul + + # print("raw stairs", plots[3]) + + + # ########### + + # time = [] + # for k in plots[0][0]: + # k = int(k) + # dt = 2.0/(k*(k-1)) + # time.append(2.0/(k*(k-1))) + + # Ne = [] + # for values in plots: + # Ne.append(np.array(values[1])/(4*mu)) + # print(time) + # print(Ne[3]) + + lines_fig2 = [] for epoch, theta in best_epochs.items(): groups = np.array(list(theta.values()), dtype=object)[:, 1].tolist() @@ -423,24 +450,33 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, x += group[::-1] y += list(np.repeat(thetas[i], len(group))) if epoch == 0: - N0 = y[0] + # watterson theta + theta_w = y[0] if theta_scale : for i in range(len(y)): y[i] = y[i]/N0 + for i in range(len(y)): + y[i] = y[i]/(4*mu) x_2 = [] T = 0 for i in range(len(x)): x[i] = int(x[i]) # compute the times as: theta_k / (k*(k-1)) for i in range(0, len(x)): - T += y[i] / (x[i]*(x[i]-1)) + T += y[i]*2 / (x[i]*(x[i]-1)) x_2.append(T) # Save plotting (fig 2) - x_2 = [0]+x_2 - y = [y[0]]+y + # x_2 = [0]+x_2 + # y = [y[0]]+y # x2_plot, y2_plot = plot_straight_x_y(x_2, y) p2 = x_2, y lines_fig2.append(p2) + # print("breaks=", epoch, "scaled_theta", lines_fig2[10]) + # print(lines_fig2[3][1][0]/(4*mu)) + # print(np.array(lines_fig2[3][1])/lines_fig2[3][1][0]) + # print("size list y=", len(lines_fig2[3][1])) + #exit(0) + if input == None: saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2, "prop":prop} @@ -458,9 +494,9 @@ 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): - recent_limit_years = 500 + recent_limit_years = 100 # recent limit in coal. time - recent_limit = recent_limit_years/tgen*mu + recent_limit = recent_limit_years/tgen # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1) nb_epochs = len(plot_lines) # fig 2 & 3 @@ -480,9 +516,9 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = #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 + swp2_lines[0][k] = swp2_lines[0][k]/tgen for k in range(len(swp2_lines[1])): - swp2_lines[1][k] = swp2_lines[1][k]*4*mu + swp2_lines[1][k] = swp2_lines[1][k] # x2_plot, y2_plot = plot_straight_x_y(swp2_lines[0],swp2_lines[1]) x2_plot, y2_plot = swp2_lines[0], swp2_lines[1] p2, = ax2.plot(x2_plot, y2_plot, linestyle="-", alpha=0.75, lw=2, label = 'swp2', color="black") @@ -508,14 +544,14 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = # skip the base 0 points x_plot[0:3] t_max_below_limit = 0 - t_min_below_limit = 1 + t_min_below_limit = recent_limit recent_change = False for t in x[1:]: if t <= recent_limit: recent_change = True t_max_below_limit = max(t_max_below_limit, t) t_min_below_limit = min(t_min_below_limit, t) - Ne_max_below_limit = y[x.index(t_max_below_limit)] + Ne_max_below_limit = y[min(x.index(t_max_below_limit)+1, len(y)-1)] Ne_min_below_limit = y[x.index(t_min_below_limit)] if recent_change: print(f"\n{breaks} breaks ; This is below the recent limit of {recent_limit_years} years:\n", @@ -547,6 +583,8 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = lines_fig3.append(p3) # put the vertical line of the "recent" time limit ax3.axvline(x=recent_limit, linestyle="--") + ax3.axvline(x=recent_limit/2, linestyle="--", color="green") + if theta_scale: xlabel = "Theta scaled by N0" ylabel = "Theta scaled by N0" @@ -557,26 +595,36 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = # 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) + plt.gca().set_xlim(0, recent_limit * 3) + if recent_change: + plt.ylim(Ne_min_below_limit/3, Ne_max_below_limit *3) + else: + plt.ylim(y2_plot[0]/3, y2_plot[0]) + # plt.ylim(0, max(max_y+(max_y*0.05), max(swp2_lines[1])+(max(swp2_lines[1])*0.05))) + #plt.xlim(0, recent_limit * 3) #xlim_val = plt.gca().get_xlim() - #x_ticks = list(plt.xticks())[0] - plt.xlim(min(min_x,min(swp2_lines[0])), max(max(swp2_lines[0]), max_x)) - x_ticks = list(plt.gca().get_xticks()) - plt.gca().set_xticks(x_ticks) + x_ticks = list(plt.xticks())[0] + # plt.xlim(min(min_x,min(swp2_lines[0])), max(max(swp2_lines[0]), max_x)) + # x_ticks = list(plt.gca().get_xticks()) + # plt.gca().set_xticks(x_ticks) # plt.xticks(x_ticks) # plt.gca().set_xlim(xlim_val) - 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.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.gca().set_xticklabels([f'{k:.1f}\n{k*tgen:.1f}' for k in x_ticks], fontsize = fnt_size*0.5) + # rescale y to effective pop size # ylim_val = plt.gca().get_ylim() - plt.ylim(min(min_y,min(swp2_lines[1])), max(max_y+(max_y*0.05), max(swp2_lines[1])+(max(swp2_lines[1])*0.05))) - y_ticks = list(plt.yticks())[0] - plt.gca().set_yticks(y_ticks) + # plt.ylim(min(min_y,min(swp2_lines[1])), max(max_y+(max_y*0.05), max(swp2_lines[1])+(max(swp2_lines[1])*0.05))) + # y_ticks = list(plt.yticks())[0] + # plt.gca().set_yticks(y_ticks) # plt.gca().set_ylim(ylim_val) - plt.yticks(y_ticks) - plt.gca().set_yticklabels([f'{k/(4*mu):.0e}' for k in y_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.yticks(y_ticks) + # plt.gca().set_yticklabels([f'{k/(4*mu):.0e}' for k in y_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.text(-0.13, -0.135, 'Gen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes) + plt.subplots_adjust(bottom=0.2) # Adjust the value as needed plt.savefig(title+'_plotB_'+str(nb_epochs)+'_epochs.pdf') # close fig2 to save memory @@ -594,16 +642,30 @@ def plot_scaled_theta(plot_lines, prop, title, mu, tgen, swp2_lines = None, ax = ax3.set_xscale('log') ax3.set_yscale('log') # Scale the x-axis - x_ticks = list(ax3.get_xticks()) - ax3.set_xticks(x_ticks) - ax3.set_xlim(min(min(x_ticks), min(swp2_lines[0])), max(max_x, max(swp2_lines[0]))) - 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) + # x_ticks = list(ax3.get_xticks()) + # ax3.set_xticks(x_ticks) + # x_ticks = [i for i in range(0.1,max(max_x, max(swp2_lines[0]))), ] + # ax3.set_xticks(x_ticks) + ax3.set_xlim(0.1, max(max_x, max(swp2_lines[0]))) + x_ticks = ax3.get_xticks() + # ax3.set_xlim(min(min(x_ticks), min(swp2_lines[0])), max(max_x, max(swp2_lines[0]))) + # ax3.set_xlim(1, max(max_x, max(swp2_lines[0]))) + # 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) + # ax3.set_xticklabels([f'{k/(mu):.0e}\n{k/(mu)*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5) + ax3.set_xticklabels([f'{k:.0e}\n{k*tgen:.0e}' for k in x_ticks], fontsize = fnt_size*0.5) + # rescale y to effective pop size - y_ticks = list(ax3.get_yticks()) - ax3.set_yticks(y_ticks) - ax3.set_ylim(min(min(y_ticks), min(swp2_lines[1])), max(max_y+(max_y*0.5), max(swp2_lines[1])+(max(swp2_lines[1])*0.5))) - ax3.set_yticklabels([f'{k/(4*mu):.0e}' for k in y_ticks], fontsize = fnt_size*0.5) - plt.text(-0.13, -0.135, 'Coal. time\nGen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes) + # y_ticks = list(ax3.get_yticks()) + # ax3.set_yticks(y_ticks) + # ax3.set_ylim(min(min(y_ticks), min(swp2_lines[1])), max(max_y+(max_y*0.5), max(swp2_lines[1])+(max(swp2_lines[1])*0.5))) + # ax3.set_ylim(1, max(max_y, max(swp2_lines[1]))) + ax3.set_ylim(1, max(max_y+(max_y*0.5), max(swp2_lines[1])+(max(swp2_lines[1])*0.5))) + + # ax3.set_yticklabels([f'{k/(4*mu):.0e}' for k in y_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.text(-0.13, -0.135, 'Gen. time\nYears', ha='left', va='bottom', transform=ax3.transAxes) + plt.text(-0.13, -0.085, 'Gen. 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) @@ -638,7 +700,7 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10, rescale = 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(breaks)+' brks') - + print("breaks=", breaks, "theta0", y[0]) # add plot to the list of all plots to superimpose plots.append(p) x_ticks = x