Save all epochs in JSON
tforest
parent
6a6d4bf6f9
commit
f90938f8d9
294
swp2.py
294
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)
|
||||
|
||||
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'],
|
||||
|
|
|
|||
Loading…
Reference in New Issue