def get_stat(alpha_ref, stat):
data = simul_data[simul_data.alpha == alpha_ref].copy()
if stat == "mean_nucleotide_div":
def get_val(Nb_ref):
h1 = data[data.Nb ==
Nb_ref].loc[:, "mean_nucleotide_div_pop_a"].values
h2 = data[data.Nb ==
Nb_ref].loc[:, "mean_nucleotide_div_pop_c"].values
return h2 / h1
h_simul_list = [
list(zip(t_coal, get_val(Nb_ref))) for Nb_ref in Nb_list
]
h_theory_list = [
list(
zip(
t_coal,
ctu.admix_coal_time_ratio(t_coal, alpha_ref,
Nb_ref / Na)))
for Nb_ref in Nb_list
]
y_label = r"$\frac{\pi_A}{\pi_0}$"
elif stat == "mean_num_seg_sites":
n = data.num_samples.unique()
def get_val(Nb_ref):
# h = data[data.Nb == Nb_ref].loc[:, "mean_num_seg_sites_pop_b"].values
h2 = data[data.Nb ==
Nb_ref].loc[:, "mean_num_seg_sites_pop_a"].values
h1 = data[data.Nb ==
Nb_ref].loc[:, "mean_num_seg_sites_pop_c"].values
h = h1 / h2
return h
h_simul_list = [
list(zip(t_coal, get_val(Nb_ref))) for Nb_ref in Nb_list
]
h_theory_list = [
list(
zip(
t_coal,
ctu.s_admix_ratio((2 * Na) * t_coal, n, 2 * Na,
2 * Nb_ref, alpha_ref),
)) for Nb_ref in Nb_list
]
y_label = r"$\frac{S_A}{S_0}$"
return (h_simul_list, h_theory_list, y_label)
def get_stat(Nb_ref, stat):
data = simul_data[simul_data.Nb == Nb_ref]
x = data.t_div.values.reshape(psize) / (2 * Na)
y = data.alpha.values.reshape(psize)
if stat == "mean_nucleotide_div":
h1 = data.loc[:, "mean_nucleotide_div_pop_a"].values
h2 = data.loc[:, "mean_nucleotide_div_pop_c"].values
z = (h2 / h1).reshape(psize)
z_th = ctu.admix_coal_time_ratio(x, y, Nb_ref / Na)
s_label = r" $\frac{\pi_A}{\pi_0}$"
elif stat == "mean_num_seg_sites":
h1 = data.loc[:, "mean_num_seg_sites_pop_a"].values
h2 = data.loc[:, "mean_num_seg_sites_pop_c"].values
z = (h2 / h1).reshape(psize)
n = data.num_samples.unique()
z_th = ctu.s_admix_ratio((2 * Na) * x, n, 2 * Na, 2 * Nb_ref, y)
s_label = r" $\frac{S_A}{S_0}$"
return (x, y, z, z_th, s_label)
def contour_stats(simul_data,
alpha_ref,
stat,
digits=2,
savefig=True,
showfig=True):
Nb_list = simul_data.Nb.unique()
t_div_list = simul_data.t_div.unique()
Na = simul_data.Na.unique()
psize = len(t_div_list), len(Nb_list)
data = simul_data[simul_data.alpha == alpha_ref]
x = data.t_div.values.reshape(psize) / (2 * Na)
y = data.Nb.values.reshape(psize) / Na
if stat == "mean_nucleotide_div":
H1 = data.loc[:, "mean_nucleotide_div_pop_a"].values
H2 = data.loc[:, "mean_nucleotide_div_pop_c"].values
res = H2 / H1
z = res.reshape(psize)
z_th = ctu.admix_coal_time_ratio(x, alpha_ref, y)
s_label = r"$\frac{\pi_A}{\pi_0}$"
figname = "../figures/contour_htz_alpha-{}.pdf".format(alpha_ref)
lr = 0
midpoint = 1
nticks = 15
digits = 1
elif stat == "mean_num_seg_sites":
n = data.num_samples.unique()
H1 = data.loc[:, "mean_num_seg_sites_pop_a"].values
H2 = data.loc[:, "mean_num_seg_sites_pop_c"].values
res = H2 / H1
z = res.reshape(psize)
z_th = ctu.s_admix_ratio((2 * Na) * x, n, 2 * Na, 2 * Na * y,
alpha_ref)
s_label = r"$\frac{S_A}{S_0}$"
figname = "../figures/contour_num_seg_sites_alpha-{}.pdf".format(
alpha_ref)
lr = 0
midpoint = 1
nticks = 15
digits = 1
elif stat == "tajima_d":
n = 2 * data.num_samples.unique()
res = data.loc[:, "tajima_d_pop_c"].values
z = res.reshape(psize)
lr = 0
z_th = ctu.tajima_d_admix((2 * Na) * x, n, Na, Na * y, alpha_ref)
s_label = r""
figname = "../figures/contour_tajimas_d_admix_alpha-{}.png".format(
alpha_ref)
midpoint = 0
nticks = 15
digits = 3
cmap = plt.get_cmap("bwr")
norm = MidPointNorm(midpoint=midpoint)
z_max = z.max()
z_min = z.min() if z.min() < midpoint else midpoint - 0.05
cmap_levels, cmap_ticks = set_cmap_levels(z_max,
z_min,
midpoint=midpoint,
digits=digits,
nticks=nticks)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_title(r"$\alpha={{{}}}$".format(alpha_ref), size=16)
ax.set_xlabel(r"split time ($2 N_0$ coalescent units)", size=16)
ax.set_ylabel(r"$\frac{N_1}{N_0}$", size=20, rotation=0, labelpad=10)
ax.set_xlim((x.min(), x.max()))
ax.set_ylim((y.min(), y.max()))
ct = ax.contourf(x, y, z, levels=cmap_levels, cmap=cmap, norm=norm)
if stat != "tajima_d":
ct_th = ax.contour(x,
y,
z_th,
levels=cmap_levels,
colors="black",
linestyles="dashed")
ax.clabel(ct_th, cmap_ticks, inline=True, fmt=f"%.1f", fontsize=10)
axcb = fig.colorbar(ct)
axcb.set_label(s_label, size=20, rotation=lr, labelpad=12)
axcb.set_ticks(cmap_ticks)
fig.tight_layout()
if savefig:
fig.savefig(figname)
if showfig:
fig.show()
pass
def lines_stats(simul_data,
alpha_ref,
stat,
digits=2,
savefig=True,
showfig=True):
data = simul_data[simul_data.alpha == alpha_ref]
Nb_list = simul_data.Nb.unique()[::2]
t_div_list = simul_data.t_div.unique()
Na = simul_data.Na.unique()
t_coal = t_div_list / (2 * Na)
if stat == "mean_nucleotide_div":
def get_val(Nb_ref):
h1 = data[data.Nb ==
Nb_ref].loc[:, "mean_nucleotide_div_pop_a"].values
h2 = data[data.Nb ==
Nb_ref].loc[:, "mean_nucleotide_div_pop_c"].values
return h2 / h1
h_simul_list = [
list(zip(t_coal, get_val(Nb_ref))) for Nb_ref in Nb_list
]
h_theory_list = [
list(
zip(t_coal,
ctu.admix_coal_time_ratio(t_coal, alpha_ref, Nb_ref / Na)))
for Nb_ref in Nb_list
]
y_label = r"$\frac{\pi_A}{\pi_0}$"
lr = 0
figname = "../figures/lines_htz_alpha-{}.pdf".format(alpha_ref)
elif stat == "mean_num_seg_sites":
n = data.num_samples.unique()
def get_val(Nb_ref):
# h = data[data.Nb == Nb_ref].loc[:, "mean_num_seg_sites_pop_b"].values
h2 = data[data.Nb == Nb_ref].loc[:,
"mean_num_seg_sites_pop_a"].values
h1 = data[data.Nb == Nb_ref].loc[:,
"mean_num_seg_sites_pop_c"].values
h = h1 / h2
return h
h_simul_list = [
list(zip(t_coal, get_val(Nb_ref))) for Nb_ref in Nb_list
]
h_theory_list = [
list(
zip(
t_coal,
ctu.s_admix_ratio((2 * Na) * t_coal, n, 2 * Na, 2 * Nb_ref,
alpha_ref))) for Nb_ref in Nb_list
]
y_label = r"$\frac{S_A}{S_0}$"
lr = 0
figname = "../figures/lines_num_seg_sites_alpha-{}.pdf".format(
alpha_ref)
elif stat == "tajima_d":
n = 2 * data.num_samples.unique()
def get_val(Nb_ref):
h = data[data.Nb == Nb_ref].loc[:, "tajima_d_pop_c"].values
return h
h_simul_list = [
list(zip(t_coal, get_val(Nb_ref))) for Nb_ref in Nb_list
]
h_theory_list = [
list(
zip(
t_coal,
ctu.tajima_d_admix((2 * Na) * t_coal, n, Na, Nb_ref,
alpha_ref))) for Nb_ref in Nb_list
]
y_label = r"$\hat{\theta}_{\pi_A} - \hat{\theta}_{S_A}}$"
lr = 90
figname = "../figures/lines_tajimas_d_admix_alpha-{}.png".format(
alpha_ref)
elif stat == "prop_diff":
n = data.num_samples.unique()
def get_val(Nb_ref):
h1 = data[data.Nb ==
Nb_ref].loc[:, "mean_nucleotide_div_pop_a"].values
h2 = data[data.Nb ==
Nb_ref].loc[:, "mean_nucleotide_div_pop_c"].values
s1 = data[data.Nb == Nb_ref].loc[:,
"mean_num_seg_sites_pop_a"].values
s2 = data[data.Nb == Nb_ref].loc[:,
"mean_num_seg_sites_pop_c"].values
s = h2 / h1
h = s2 / s1
return s - h
h_simul_list = [
list(zip(t_coal, get_val(Nb_ref))) for Nb_ref in Nb_list
]
h_theory_list = [
list(
zip(
t_coal,
ctu.admix_coal_time_ratio(t_coal, alpha_ref, Nb_ref / Na) -
ctu.s_admix_ratio(
(2 * Na) * t_coal, n, 2 * Na, 2 * Nb_ref, alpha_ref),
)) for Nb_ref in Nb_list
]
y_label = r"$\frac{\pi_A}{\pi_0} - \frac{S_A}{S_0}$"
lr = 90
figname = "../figures/lines_prop_diff_alpha-{}.png".format(alpha_ref)
cmap = plt.get_cmap("Spectral")
hs = np.array(h_simul_list)
ht = np.array(h_theory_list)
y_min, y_max = (
min(hs[:, :, 1].min(), ht[:, :, 1].min()),
max(hs[:, :, 1].max(), ht[:, :, 1].max()),
)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim((0, 1))
ax.set_ylim((0.9 * y_min, 1.1 * y_max))
ax.set_title(r"$\alpha={{{}}}$".format(alpha_ref), size=16)
ax.set_xlabel(
r"$\mathrm{time} \, (2 N_0 \, \mathrm{coalescent \, units})$", size=16)
ax.set_ylabel(y_label, rotation=lr, size=20, labelpad=10)
ax.plot(np.linspace(0, 1, 20),
np.ones(20),
color="black",
linestyle="dotted",
linewidth=1.5)
line_segments_simul = LineCollection(h_simul_list,
linewidths=1.5,
linestyles="solid",
cmap=cmap)
line_segments_simul.set_array(Nb_list / Na)
ax.add_collection(line_segments_simul)
line_segments_theory = LineCollection(h_theory_list,
linewidths=1.3,
linestyles="dashed",
cmap=cmap)
ax.add_collection(line_segments_theory)
line_segments_theory.set_array(Nb_list / Na)
axcb = fig.colorbar(line_segments_simul)
axcb.set_label(r"$\frac{N_1}{N_0}$", rotation=0, size=20, labelpad=12)
fig.tight_layout()
if savefig:
fig.savefig(figname)
if showfig:
fig.show()