richard
/
geinf
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Improve plotting with swp2 control curve

master
tforest 2024-02-25 00:27:26 +01:00
parent 8c8545a9ff
commit 4aedf5e280
1 changed files with 178 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

236
swp2.py
View File

@ -126,10 +126,11 @@ def plot_straight_x_y(x,y):
def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title", def plot_all_epochs_thetafolder(full_dict, mu, tgen, title = "Title",
theta_scale = True, ax = None, input = None, output = None): theta_scale = True, ax = None, input = None, output = None):
my_dpi = 300 my_dpi = 500
L = full_dict["L"]
if ax is None: if ax is None:
# intialize figure # intialize figure
my_dpi = 300 #my_dpi = 300
fnt_size = 18 fnt_size = 18
# plt.rcParams['font.size'] = fnt_size # plt.rcParams['font.size'] = fnt_size
fig, ax1 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi) 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 = ax[1][0,0]
ax1.set_yscale('log') ax1.set_yscale('log')
ax1.set_xscale('log') ax1.set_xscale('log')
ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
plot_handles = [] plot_handles = []
best_plot = full_dict['all_epochs']['best'] 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]) alpha=1, lw=2, label = str(best_plot[2])+' brks | Lik='+best_plot[3])
plot_handles.append(p0) plot_handles.append(p0)
#ax1.grid(True,which="both", linestyle='--', alpha = 0.3)
for k, plot_Lk in enumerate(full_dict['all_epochs']['plots']): for k, plot_Lk in enumerate(full_dict['all_epochs']['plots']):
plot_Lk = str(full_dict['all_epochs']['plots'][k][3]) plot_Lk = str(full_dict['all_epochs']['plots'][k][3])
# plt.rcParams['font.size'] = fnt_size # 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) alpha=1/(k+1), lw=1.5, label = str(full_dict['all_epochs']['plots'][k][2])+' brks | Lik='+plot_Lk)
plot_handles.append(p) plot_handles.append(p)
if theta_scale: 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) # ax1.axvline(x=recent_scale_upper_bound)
else: else:
# years # years
plt.set_xlabel("Time (years)", fontsize=fnt_size) if ax is not None:
plt.set_ylabel("Individuals (N)", fontsize=fnt_size) 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 # plt.rcParams['font.size'] = fnt_size
# print(fnt_size, "rcParam font.size=", plt.rcParams['font.size']) # print(fnt_size, "rcParam font.size=", plt.rcParams['font.size'])
ax1.legend(handles = plot_handles, loc='best', fontsize = fnt_size*0.5) ax1.legend(handles = plot_handles, loc='best', fontsize = fnt_size*0.5)
ax1.set_title(title) ax1.set_title(title)
breaks = len(full_dict['all_epochs']['plots'])
if ax is None: if ax is None:
plt.savefig(title+'_b'+str(breaks)+'.pdf') plt.savefig(title+'_b'+str(breaks)+'.pdf')
# plot likelihood against nb of breakpoints # 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") ax3.set_title(title+" AIC")
if ax is None: if ax is None:
plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf') plt.savefig(title+'_Breakpts_Likelihood_AIC.pdf')
# return plots else:
return ax[0], ax[1] # 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): def save_all_epochs_thetafolder(folder_path, mu, tgen, title = "Title", theta_scale = True, input = None, output = None):
#scenari = {} #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 # do something with the theta without bp and skip the plotting
N0 = y[0] N0 = y[0]
#continue #continue
for i in range(len(y)): if theta_scale:
# divide by N0 for i in range(len(y)):
y[i] = y[i]/N0 # divide by N0
x[i] = x[i]/N0 y[i] = y[i]/N0
x[i] = x[i]/N0
top_plots[greatest_likelihood] = x,y,epoch top_plots[greatest_likelihood] = x,y,epoch
plots_likelihoods = list(top_plots.keys()) plots_likelihoods = list(top_plots.keys())
for i in range(len(plots_likelihoods)): 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 # to return : plots ; Ln_Brks ; AIC_Brks ; best_Ln ; best_AIC
# 'plots' dict keys: 'best', {epochs}('0', '1',...) # 'plots' dict keys: 'best', {epochs}('0', '1',...)
if input == None: if input == None:
saved_plots = {"S":S, "S0":S0, "L":L, "all_epochs":plots, "Ln_Brks":Ln_Brks, saved_plots = {"S":S, "S0":S0, "L":L, "mu":mu, "tgen":tgen,
"AIC_Brks":AIC_Brks, "best_Ln":best_Ln, "all_epochs":plots, "Ln_Brks":Ln_Brks,
"best_AIC":best_AIC, "best_epoch_by_AIC":selected_brks_nb} "AIC_Brks":AIC_Brks, "best_Ln":best_Ln,
"best_AIC":best_AIC, "best_epoch_by_AIC":selected_brks_nb}
else: else:
# if the dict has to be loaded from input # if the dict has to be loaded from input
with open(input, 'r') as json_file: 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["S"] = S
saved_plots["S0"] = S0 saved_plots["S0"] = S0
saved_plots["L"] = L saved_plots["L"] = L
saved_plots["mu"] = mu
saved_plots["tgen"] = tgen
saved_plots["all_epochs"] = plots saved_plots["all_epochs"] = plots
saved_plots["Ln_Brks"] = Ln_Brks saved_plots["Ln_Brks"] = Ln_Brks
saved_plots["AIC_Brks"] = AIC_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): for k in range(2, len(y)+2):
prop.append(y[k-2] / (k - 1) / sum_theta_i) prop.append(y[k-2] / (k - 1) / sum_theta_i)
prop = prop[::-1] prop = prop[::-1]
# normalise to N0 (N0 of epoch1) if theta_scale :
for i in range(len(y)): # normalise to N0 (N0 of epoch1)
y[i] = y[i]/N0 for i in range(len(y)):
y[i] = y[i]/N0
# x_plot, y_plot = plot_straight_x_y(x, y) # x_plot, y_plot = plot_straight_x_y(x, y)
p = x, y p = x, y
# add plot to the list of all plots to superimpose # 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))) y += list(np.repeat(thetas[i], len(group)))
if epoch == 0: if epoch == 0:
N0 = y[0] N0 = y[0]
for i in range(len(y)): if theta_scale :
y[i] = y[i]/N0 for i in range(len(y)):
y[i] = y[i]/N0
x_2 = [] x_2 = []
T = 0 T = 0
for i in range(len(x)): 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) json.dump(saved_plots, json_file)
return saved_plots 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 # fig 2 & 3
if ax is None: if ax is None:
my_dpi = 300 my_dpi = 500
fnt_size = 18 fnt_size = 18
fig2, ax2 = plt.subplots(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi) 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) 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_fig2 = []
lines_fig3 = [] lines_fig3 = []
#plt.figure(figsize=(5000/my_dpi, 2800/my_dpi), dpi=my_dpi) #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) nb_breaks = len(plot_lines)
for breaks, plot in enumerate(plot_lines): for breaks, plot in enumerate(plot_lines):
if subset is not None: if subset is not None:
if breaks not in subset : if breaks not in subset :
# skip if 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 pass
else: else:
continue continue
x,y=plot 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) 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') p2, = ax2.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(breaks)+' brks')
lines_fig2.append(p2) lines_fig2.append(p2)
# Plotting (fig 3) which is the same but log scale for x # 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') p3, = ax3.plot(x2_plot, y2_plot, 'o', linestyle="-", alpha=0.75, lw=2, label = str(breaks)+' brks')
lines_fig3.append(p3) lines_fig3.append(p3)
ax2.set_xlabel("Relative scale", fontsize=fnt_size)
ax2.set_ylabel("theta", fontsize=fnt_size) ax3.axvline(x=500/tgen*mu, linestyle="--")
ax2.set_title(title, fontsize=fnt_size) if theta_scale:
ax2.legend(handles=lines_fig2, loc='best', fontsize = fnt_size*0.5) xlabel = "Theta scaled by N0"
ylabel = "Theta scaled by N0"
else:
xlabel = "Theta scale"
ylabel = "Theta"
if ax is None: 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) # nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1)
plt.savefig(title+'_plot2_'+str(len(plot_lines))+'.pdf') plt.savefig(title+'_plot2_'+str(len(plot_lines))+'.pdf')
# close fig2 to save memory # close fig2 to save memory
plt.close(fig2) 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_xscale('log')
ax3.set_yscale('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_ylabel("theta", fontsize=fnt_size)
ax3.set_title(title, fontsize=fnt_size) ax3.set_title(title, fontsize=fnt_size)
ax3.legend(handles=lines_fig3, loc='best', fontsize = fnt_size*0.5) 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: if ax is None:
# nb of plot_lines represent the number of epochs stored (len(plot_lines) = #breaks+1) # 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') 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) plt.close(fig3)
return ax 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 # multiple fig
if ax is None: if ax is None:
# intialize figure 1 # intialize figure 1
@ -494,7 +552,6 @@ def plot_raw_stairs(plot_lines, prop, title, ax = None, n_ticks = 10):
ax1 = ax[0, 0] ax1 = ax[0, 0]
plt.subplots_adjust(wspace=0.3, hspace=0.3) plt.subplots_adjust(wspace=0.3, hspace=0.3)
plots = [] plots = []
for epoch, plot in enumerate(plot_lines): for epoch, plot in enumerate(plot_lines):
x,y = plot x,y = plot
x_plot, y_plot = plot_straight_x_y(x,y) 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 plots
return ax 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 my_dpi = 300
saved_plots_dict = save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, output = title+"_plotdata.json")
save_k_theta(folder_path, mu, tgen, title, theta_scale, breaks_max = breaks, output = title+"_plotdata.json") nb_of_epochs = len(saved_plots_dict["all_epochs"]["plots"])
save_all_epochs_thetafolder(folder_path, mu, tgen, title, theta_scale, input = title+"_plotdata.json", output = title+"_plotdata.json") 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: with open(title+"_plotdata.json", 'r') as json_file:
loaded_data = json.load(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) # START OF COMBINED PLOT CODE
# fig1.tight_layout()
# Adjust absolute space between the top and bottom rows # # plot page 1 of summary
fig1.subplots_adjust(hspace=0.35) # Adjust this value based on your requirement # fig1, ax1 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
# plot page 2 of summary # # fig1.tight_layout()
fig2, ax2 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi) # # Adjust absolute space between the top and bottom rows
# fig2.tight_layout() # fig1.subplots_adjust(hspace=0.35) # Adjust this value based on your requirement
ax1 = plot_raw_stairs(plot_lines = loaded_data['raw_stairs'], # # plot page 2 of summary
prop = loaded_data['prop'], title = title, ax = ax1) # fig2, ax2 = plt.subplots(2, 2, figsize=(5000/my_dpi, 2970/my_dpi), dpi=my_dpi)
ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], # # fig2.tight_layout()
prop = loaded_data['prop'], title = title, ax = ax1, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks) # ax1 = plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
ax2 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], # prop = loaded_data['prop'], title = title, ax = ax1)
prop = loaded_data['prop'], title = title, ax = ax2) # ax1 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2]) # prop = loaded_data['prop'], title = title, ax = ax1, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks)
fig1.savefig(title+'_combined_p1.pdf') # ax2 = plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'],
print("Wrote", title+'_combined_p1.pdf') # prop = loaded_data['prop'], title = title, ax = ax2)
fig2.savefig(title+'_combined_p2.pdf') # ax1, ax2 = plot_all_epochs_thetafolder(loaded_data, mu, tgen, title, theta_scale, ax = [ax1, ax2])
print("Wrote", title+'_combined_p2.pdf')
# 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'], plot_raw_stairs(plot_lines = loaded_data['raw_stairs'],
prop = loaded_data['prop'], title = title, ax = None) prop = loaded_data['prop'], title = title, ax = None)
plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], plot_scaled_theta(plot_lines = loaded_data['scaled_stairs'], mu = mu, tgen = tgen, subset=[loaded_data['best_epoch_by_AIC']]+selected_breaks,
prop = loaded_data['prop'], title = title, ax = None) # 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(fig1)
plt.close(fig2) # 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__": if __name__ == "__main__":