Newer
Older
from math import sqrt
import matplotlib.pyplot as plt
from util import (
save_fig,
couples320,
get_path_finer_resol,
customize,
add_letters,
)
R_is = sorted(Ri for (Nloc, Ri) in couples320 if Nloc == N)
paths = [get_path_finer_resol(N, R_i) for R_i in R_is]
paths = [path for path in paths if path is not None]
fig, axes = plt.subplots(ncols=2, nrows=3, figsize=(10, 3 * 3 * 4.5 / 4))
ax0 = axes[0, 0]
ax1 = axes[0, 1]
ax2 = axes[1, 0]
ax3 = axes[1, 1]
ax4 = axes[2, 0]
ax5 = axes[2, 1]
for path, R_i in zip(paths, R_is):
if R_i == 80:
continue
sim = load(path, hide_stdout=True)
mean_values = sim.output.get_mean_values(tmin="t_last-2", customize=customize)
Uh2 = mean_values["Uh2"]
kb = N / Uh2**0.5
epsK = mean_values["epsK"]
if R2 < 10:
R2_str = f"{R2:.1f}"
else:
R2_str = f"{R2:.0f}"
Uh2 = mean_values["Uh2"]
kb = N / sqrt(Uh2)
epsK = mean_values["epsK"]
ko = sqrt(N**3 / epsK)
t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = max(10, min(t_start + 2, t_last - 1))
data = sim.output.spectra.load_kzkh_mean(
tmin, key_to_load=["A", "Khd", "Kz", "Khr"]
)
kh = data["kh_spectra"]
kz = data["kz"]
delta_kh = kh[1]
delta_kz = kz[1]
EA = data["A"]
EKhd = data["Khd"]
EKz = data["Kz"]
EKhr = data["Khr"]
Epolo_vs_kh = np.sum(Epolo, axis=0) * delta_kz
Epolo_vs_kz = np.sum(Epolo, axis=1) * delta_kh
Etoro_vs_kh = np.sum(Etoro, axis=0) * delta_kz
Etoro_vs_kz = np.sum(Etoro, axis=1) * delta_kh
EA_vs_kh = np.sum(EA, axis=0) * delta_kz
EA_vs_kz = np.sum(EA, axis=1) * delta_kh
(line,) = ax0.plot(kh / ko, EA_vs_kh[:] * kh[:] ** (5 / 3), color=color)
ax0.axvline(kb / ko, linestyle="dotted", color=color)
ax1.plot(kz / ko, EA_vs_kz[:] * kz[:] ** (5 / 3), color=color)
ax2.plot(kh / ko, Epolo_vs_kh[:] * kh[:] ** (5 / 3), color=color)
ax3.plot(kz / ko, Epolo_vs_kz[:] * kz[:] ** (5 / 3), color=color)
ax4.plot(
kh / ko,
Etoro_vs_kh[:] * kh[:] ** (5 / 3),
color=color,
label=rf"$\mathcal{{R}}={R2_str}$",
)
ax4.axvline(kb / ko, linestyle="dotted", color=color)
ax5.plot(kz / ko, Etoro_vs_kz[:] * kz[:] ** (5 / 3), color=color)
ax0.set_xlabel(r"$k_h / k_{\rm O}$", fontsize=14)
ax1.set_xlabel(r"$k_z / k_{\rm O}$", fontsize=14)
ax2.set_xlabel(r"$k_h / k_{\rm O}$", fontsize=14)
ax3.set_xlabel(r"$k_z / k_{\rm O}$", fontsize=14)
ax4.set_xlabel(r"$k_h / k_{\rm O}$", fontsize=14)
ax5.set_xlabel(r"$k_z / k_{\rm O}$", fontsize=14)
ax0.set_ylabel(r"$E_{\rm A}(k_h) {k_h}^{5/3}$", fontsize=14)
ax1.set_ylabel(r"$E_{\rm A}(k_z) {k_z}^{5/3}$", fontsize=14)
ax2.set_ylabel(r"$E_{\rm polo}(k_h) {k_h}^{5/3}$", fontsize=14)
ax3.set_ylabel(r"$E_{\rm polo}(k_z) {k_z}^{5/3}$", fontsize=14)
ax4.set_ylabel(r"$E_{\rm toro}(k_h) {k_h}^{5/3}$", fontsize=14)
ax5.set_ylabel(r"$E_{\rm toro}(k_z) {k_z}^{5/3}$", fontsize=14)
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_ylim(bottom=1e-2, top=4)
for ax in [ax0, ax1, ax2, ax3]:
ax.set_xlabel(None)
ax.set_xticks([])
fig.tight_layout()
save_fig(fig, "fig_spectra_1strat.png")
if __name__ == "__main__":
plt.show()