From f90938f8d9058c11c5cd915a037f55ded5d7c5cb Mon Sep 17 00:00:00 2001 From: tforest Date: Tue, 19 Dec 2023 01:37:00 +0100 Subject: [PATCH] Save all epochs in JSON --- swp2.py | 298 ++++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 279 insertions(+), 19 deletions(-) diff --git a/swp2.py b/swp2.py index 499c481..a66b9f2 100644 --- a/swp2.py +++ b/swp2.py @@ -16,7 +16,7 @@ def log_facto(k): val += np.log(i) return val -def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_theta_scale = False): +def parse_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_theta_scale = False): with open(stwp_theta_file, "r") as swp_file: # Read the first line line = swp_file.readline() @@ -109,15 +109,8 @@ def return_x_y_from_stwp_theta_file(stwp_theta_file, breaks, mu, tgen, relative_ x.append(list(t.values())[time]) x.append(list(t.values())[time]) x.append(list(t.values())[len(t.values())-1]) - # if relative_theta_scale: - # # rescale - # #N0 = y[0] - # # for i in range(len(y)): - # # # divide by N0 - # # y[i] = y[i]/N0 - # # x[i] = x[i]/N0 - return x,y,likelihood,thetas,sfs,L + return x,y,likelihood,thetas,sfs,L def plot_k_epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title", theta_scale = True): scenari = {} @@ -125,7 +118,7 @@ def plot_k_epochs_thetafolder(folder_path, mu, tgen, breaks = 2, title = "Title" for file_name in os.listdir(folder_path): if os.path.isfile(os.path.join(folder_path, file_name)): # Perform actions on each file - x, y, likelihood, theta, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks, + x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks, tgen = tgen, mu = mu, relative_theta_scale = theta_scale) if x == 0 or y == 0: @@ -178,7 +171,8 @@ def plot_straight_x_y(x,y): x_1.append(x[-1]) return x_1, y_1 -def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, ax = None): +def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", + theta_scale = True, ax = None, input = None, output = None): #scenari = {} cpt = 0 epochs = {} @@ -186,7 +180,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc breaks = 0 cpt +=1 if os.path.isfile(os.path.join(folder_path, file_name)): - x, y, likelihood, theta, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks, + x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks, tgen = tgen, mu = mu, relative_theta_scale = theta_scale) SFS_stored = sfs @@ -196,7 +190,265 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc epochs[breaks] = {} epochs[breaks][likelihood] = x,y breaks += 1 - x,y,likelihood,theta,sfs,L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = breaks, + x,y,likelihood,theta,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks, + tgen = tgen, + mu = mu, relative_theta_scale = theta_scale) + if x == 0: + # last break did not work, then breaks = breaks-1 + breaks -= 1 + print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(breaks)+"\n*******\n") + print(cpt, "theta file(s) have been scanned.") + my_dpi = 300 + if ax is None: + # intialize figure + 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) + else: + fnt_size = 12 + # plt.rcParams['font.size'] = fnt_size + ax1 = ax[1][0,0] + ax1.set_yscale('log') + ax1.set_xscale('log') + ax1.grid(True,which="both", linestyle='--', alpha = 0.3) + brkpt_lik = [] + top_plots = {} + for epoch, scenari in epochs.items(): + # sort starting by the smallest -log(Likelihood) + best10_scenari = (sorted(list(scenari.keys())))[:10] + greatest_likelihood = best10_scenari[0] + # store the tuple breakpoints and likelihood for later plot + brkpt_lik.append((epoch, greatest_likelihood)) + x, y = scenari[greatest_likelihood] + #without breakpoint + if epoch == 0: + # 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 + top_plots[greatest_likelihood] = x,y,epoch + plots_likelihoods = list(top_plots.keys()) + for i in range(len(plots_likelihoods)): + plots_likelihoods[i] = float(plots_likelihoods[i]) + best10_plots = sorted(plots_likelihoods)[:10] + top_plot_lik = str(best10_plots[0]) + plot_handles = [] + # plt.rcParams['font.size'] = fnt_size + p0, = ax1.plot(top_plots[top_plot_lik][0], top_plots[top_plot_lik][1], 'o', linestyle = "-", + alpha=1, lw=2, label = str(top_plots[top_plot_lik][2])+' brks | Lik='+top_plot_lik) + plot_handles.append(p0) + for k, plot_Lk in enumerate(best10_plots[1:]): + plot_Lk = str(plot_Lk) + # plt.rcParams['font.size'] = fnt_size + p, = ax1.plot(top_plots[plot_Lk][0], top_plots[plot_Lk][1], 'o', linestyle = "--", + alpha=1/(k+1), lw=1.5, label = str(top_plots[plot_Lk][2])+' brks | Lik='+plot_Lk) + plot_handles.append(p) + if theta_scale: + ax1.set_xlabel("Coal. time", fontsize=fnt_size) + ax1.set_ylabel("Pop. size scaled by N0", fontsize=fnt_size) + # recent_scale_lower_bound = 0.01 + # recent_scale_upper_bound = 0.1 + # ax1.axvline(x=recent_scale_lower_bound) + # 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) + # 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) + if ax is None: + plt.savefig(title+'_b'+str(breaks)+'.pdf') + # plot likelihood against nb of breakpoints + # best possible likelihood from SFS + # Segregating sites + S = sum(SFS_stored) + # Number of kept sites from which the SFS is computed + L = L_stored + # number of monomorphic sites + S0 = L-S + # print("SFS", SFS_stored) + # print("S", S, "L", L, "S0=", S0) + # 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)): + p_i = SFS_stored[xi] / float(S+S0) + Ln += np.log(p_i) * SFS_stored[xi] - log_facto(SFS_stored[xi]) + # basic plot likelihood + if ax is None: + fig, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi) + # plt.rcParams['font.size'] = fnt_size + else: + #plt.rcParams['font.size'] = fnt_size + ax2 = ax[0][0,1] + ax2.plot(np.array(brkpt_lik)[:, 0], np.array(brkpt_lik)[:, 1].astype(float), 'o', linestyle = "dotted", lw=2) + ax2.axhline(y=-Ln, linestyle = "-.", color = "red", label = "$-\log\mathcal{L}$ = "+str(round(-Ln, 2))) + ax2.set_yscale('log') + ax2.set_xlabel("# breakpoints", fontsize=fnt_size) + ax2.set_ylabel("$-\log\mathcal{L}$", fontsize=fnt_size) + ax2.legend(loc='best', fontsize = fnt_size*0.5) + ax2.set_title(title+" Likelihood gain from # breakpoints") + if ax is None: + plt.savefig(title+'_Breakpts_Likelihood.pdf') + # AIC + if ax is None: + fig, ax3 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi) + # plt.rcParams['font.size'] = '18' + else: + #plt.rcParams['font.size'] = fnt_size + ax3 = ax[1][0,1] + AIC = [] + for brk in np.array(brkpt_lik)[:, 0]: + brk = int(brk) + AIC.append((2*brk+1)+2*np.array(brkpt_lik)[brk, 1].astype(float)) + ax3.plot(np.array(brkpt_lik)[:, 0], AIC, 'o', linestyle = "dotted", lw=2) + # AIC = 2*k - 2ln(L) ; where k is the number of parameters, here brks+1 + AIC_ln = 2*(len(brkpt_lik)+1) - 2*Ln + ax3.axhline(y=AIC_ln, linestyle = "-.", color = "red", + label = "Min. AIC = "+str(round(AIC_ln, 2))) + selected_brks_nb = AIC.index(min(AIC)) + ax3.set_yscale('log') + ax3.set_xlabel("# breakpoints", fontsize=fnt_size) + ax3.set_ylabel("AIC") + ax3.legend(loc='best', fontsize = fnt_size*0.5) + ax3.set_title(title+" AIC") + if ax is None: + plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf') + print("S", S) + # 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 = {} + cpt = 0 + epochs = {} + plots = {} + # store ['best'], and [0] for epoch 0 etc... + for file_name in os.listdir(folder_path): + breaks = 0 + cpt +=1 + if os.path.isfile(os.path.join(folder_path, file_name)): + x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks, + tgen = tgen, + mu = mu, relative_theta_scale = theta_scale) + SFS_stored = sfs + L_stored = L + while not (x == 0 and y == 0): + if breaks not in epochs.keys(): + epochs[breaks] = {} + epochs[breaks][likelihood] = x,y + breaks += 1 + x,y,likelihood,theta,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks, + tgen = tgen, + mu = mu, relative_theta_scale = theta_scale) + if x == 0: + # last break did not work, then breaks = breaks-1 + breaks -= 1 + print("\n*******\n"+title+"\n--------\n"+"mu="+str(mu)+"\ntgen="+str(tgen)+"\nbreaks="+str(breaks)+"\n*******\n") + print(cpt, "theta file(s) have been scanned.") + + brkpt_lik = [] + top_plots = {} + for epoch, scenari in epochs.items(): + # sort starting by the smallest -log(Likelihood) + best10_scenari = (sorted(list(scenari.keys())))[:10] + greatest_likelihood = best10_scenari[0] + # store the tuple breakpoints and likelihood for later plot + brkpt_lik.append((epoch, greatest_likelihood)) + x, y = scenari[greatest_likelihood] + #without breakpoint + if epoch == 0: + # 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 + top_plots[greatest_likelihood] = x,y,epoch + plots_likelihoods = list(top_plots.keys()) + for i in range(len(plots_likelihoods)): + plots_likelihoods[i] = float(plots_likelihoods[i]) + best10_plots = sorted(plots_likelihoods)[:10] + top_plot_lik = str(best10_plots[0]) + # store x,y,brks,likelihood + plots['best'] = (top_plots[top_plot_lik][0], top_plots[top_plot_lik][1], str(top_plots[top_plot_lik][2]), top_plot_lik) + for k, plot_Lk in enumerate(best10_plots[1:]): + plot_Lk = str(plot_Lk) + plots[str(top_plots[plot_Lk][2])] = (top_plots[plot_Lk][0], top_plots[plot_Lk][1], str(top_plots[plot_Lk][2]), plot_Lk) + + # plot likelihood against nb of breakpoints + # best possible likelihood from SFS + # Segregating sites + S = sum(SFS_stored) + # Number of kept sites from which the SFS is computed + L = L_stored + # number of monomorphic sites + S0 = L-S + # print("SFS", SFS_stored) + # print("S", S, "L", L, "S0=", S0) + # 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)): + p_i = SFS_stored[xi] / float(S+S0) + Ln += np.log(p_i) * SFS_stored[xi] - log_facto(SFS_stored[xi]) + # basic plot likelihood + Ln_Brks = [list(np.array(brkpt_lik)[:, 0]), list(np.array(brkpt_lik)[:, 1].astype(float))] + best_Ln = -Ln + AIC = [] + for brk in np.array(brkpt_lik)[:, 0]: + brk = int(brk) + AIC.append((2*brk+1)+2*np.array(brkpt_lik)[brk, 1].astype(float)) + AIC_Brks = [list(np.array(brkpt_lik)[:, 0]), AIC] + # AIC = 2*k - 2ln(L) ; where k is the number of parameters, here brks+1 + AIC_ln = 2*(len(brkpt_lik)+1) - 2*Ln + best_AIC = AIC_ln + + # to return : plots ; Ln_Brks ; AIC_Brks ; best_Ln ; best_AIC + # 'plots' dict keys: 'best', {epochs}('0', '1',...) + if input == None: + saved_plots = {"all_epochs":plots, "Ln_Brks":Ln_Brks, + "AIC_Brks":AIC_Brks, "best_Ln":best_Ln, + "best_AIC":best_AIC} + else: + # if the dict has to be loaded from input + with open(input, 'r') as json_file: + saved_plots = json.load(json_file) + saved_plots["all_epochs"] = plots + saved_plots["Ln_Brks"] = Ln_Brks + saved_plots["AIC_Brks"] = AIC_Brks + saved_plots["best_Ln"] = best_Ln + saved_plots["best_AIC"] = best_AIC + if output == None: + output = title+"_plotdata.json" + with open(output, 'w') as json_file: + json.dump(saved_plots, json_file) + return saved_plots + +def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, ax = None): + #scenari = {} + cpt = 0 + epochs = {} + for file_name in os.listdir(folder_path): + breaks = 0 + cpt +=1 + if os.path.isfile(os.path.join(folder_path, file_name)): + x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks, + tgen = tgen, + mu = mu, relative_theta_scale = theta_scale) + SFS_stored = sfs + L_stored = L + while not (x == 0 and y == 0): + if breaks not in epochs.keys(): + epochs[breaks] = {} + epochs[breaks][likelihood] = x,y + breaks += 1 + x,y,likelihood,theta,sfs,L = parse_stwp_theta_file(folder_path+file_name, breaks = breaks, tgen = tgen, mu = mu, relative_theta_scale = theta_scale) if x == 0: @@ -330,7 +582,7 @@ def plot_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_sc return ax[0], ax[1] def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, - breaks_max = 10, output = None): + breaks_max = 10, input = None, output = None): """ Save theta values as is to do basic plots. """ @@ -341,7 +593,7 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, cpt +=1 if os.path.isfile(os.path.join(folder_path, file_name)): for k in range(breaks_max): - x,y,likelihood,thetas,sfs,L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = k, + 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) if thetas == 0: @@ -423,9 +675,16 @@ def save_k_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, # x2_plot, y2_plot = plot_straight_x_y(x_2, y) p2 = x_2, y lines_fig2.append(p2) - - saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2, - "prop":prop} + if input == None: + saved_plots = {"raw_stairs":plots, "scaled_stairs":lines_fig2, + "prop":prop} + else: + # if the dict has to be loaded from input + with open(input, 'r') as json_file: + saved_plots = json.load(json_file) + saved_plots["raw_stairs"] = plots + saved_plots["scaled_stairs"] = lines_fig2 + saved_plots["prop"] = prop if output == None: output = title+"_plotdata.json" with open(output, 'w') as json_file: @@ -536,7 +795,7 @@ def plot_test_theta(folder_path, mu, tgen, title = "Title", theta_scale = True, cpt +=1 if os.path.isfile(os.path.join(folder_path, file_name)): for k in range(breaks_max): - x, y, likelihood, theta, sfs, L = return_x_y_from_stwp_theta_file(folder_path+file_name, breaks = k, + x, y, likelihood, theta, sfs, L = parse_stwp_theta_file(folder_path+file_name, breaks = k, tgen = tgen, mu = mu, relative_theta_scale = theta_scale) if thetas == 0: @@ -713,6 +972,7 @@ def combined_plot(folder_path, mu, tgen, breaks, title = "Title", theta_scale = ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], prop = loaded_data['prop'], title = title, ax = ax1) ax1, ax2 = plot_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, ax = [ax1, ax2]) + save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, input = title+"_plotdata.json", output = title+"_plotdata.json") fig1.savefig(title+'_combined_p1.pdf') fig2.savefig(title+'_combined_p2.pdf') plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],