Fixing some scaling for swp2 output
parent
25d7ef0858
commit
1848102140
134
swp2.py
134
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
|
||||
|
|
|
|||
Loading…
Reference in New Issue