Skip to content
Snippets Groups Projects
save_figure6.py 7.32 KiB
Newer Older
import sys
import numpy as np
import matplotlib.pyplot as plt

from util_simuls_regimes import get_sim

from util import save_fig, compute_kf_kb_ko_keta_kd, customize

# Latex
plt.rcParams["text.usetex"] = True
plt.rcParams["text.latex.preamble"] = r"\usepackage{bm}"

print(sys.argv)
letter = sys.argv[-1]
if letter not in "DLOWPU":
    letter = "L"


fig, axes = plt.subplots(
    ncols=2, nrows=2, figsize=(10, 1.2 * 2 * 3 * 4.5 / 4), constrained_layout=True
)

ax0 = axes[0, 0]
ax1 = axes[0, 1]
ax2 = axes[1, 0]
ax3 = axes[1, 1]


coef_compensate = 5 / 3

# Standard Navier-Stokes
sim = get_sim(letter)
t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = t_last - 2.0
kf, kb, ko, keta, kd = compute_kf_kb_ko_keta_kd(sim, tmin)
data = sim.output.spectra.load_kzkh_mean(tmin)
kh = data["kh_spectra"]
kz = data["kz"]

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]
KH, KZ = np.meshgrid(kh, kz)

EA = data["A"]
EKhd = data["Khd"]
EKz = data["Kz"]
EKhr = data["Khr"]

Epolo = EKhd + EKz
Etoro = EKhr
EKh = EKhr + EKhd

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

Epot_vs_kh = np.sum(EA, axis=0) * delta_kz
Epot_vs_kz = np.sum(EA, axis=1) * delta_kh

ax0.plot(
    kh,
    Epolo_vs_kh * kh ** (coef_compensate),
    "-g",
    label=r"$E_{\rm polo}(k_h)$",
)
ax0.plot(
    kh,
    Etoro_vs_kh * kh ** (coef_compensate),
    "-r",
    label=r"$E_{\rm toro}(k_h)$",
)
ax0.plot(
    kh,
    Epot_vs_kh * kh ** (coef_compensate),
    "-b",
    label=r"$E_{\rm pot}(k_h)$",
)
ax0.plot(
    kz,
    Epolo_vs_kz * kz ** (coef_compensate),
    "--g",
    label=r"$E_{\rm polo}(k_z)$",
)
ax0.plot(
    kz,
    Etoro_vs_kz * kz ** (coef_compensate),
    "--r",
    label=r"$E_{\rm toro}(k_z)$",
)
ax0.plot(
    kz,
    Epot_vs_kz * kz ** (coef_compensate),
    "--b",
    label=r"$E_{\rm pot}(k_z)$",
)

sim.output.spect_energy_budg.plot_fluxes(
    tmin=tmin, key_k="kh", ax=ax2, plot_conversion=False
)
for ax in [ax0, ax2]:
    ax.axvline(kb, color="k", linestyle="dotted")
    ax.axvline(ko, color="k", linestyle="dashed")
    ax.axvline(kf, color="orange", linestyle="dashed")
    ax.set_xlim([kh[1], max(kh)])
    ax.set_xticks([1e1, 1e2, 1e3, kb, ko])
    ax.set_xticklabels(
        [r"$10^1$", r"$10^2$", r"$10^3$", r"$k_{\rm b}$", r"$k_{\rm O}$"],
        fontsize=14,
    )
# sim.output.spect_energy_budg.plot_fluxes(tmin=tmin, key_k="kz", ax=ax4)
# Without vortical modes
sim = get_sim(letter, proj=True)
t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = t_last - 2.0
data = sim.output.spectra.load1d_mean(tmin)
kf, kb, ko, keta, kd = compute_kf_kb_ko_keta_kd(sim, tmin)
data = sim.output.spectra.load_kzkh_mean(tmin)
kh = data["kh_spectra"]
kz = data["kz"]

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]
KH, KZ = np.meshgrid(kh, kz)

EA = data["A"]
EKhd = data["Khd"]
EKz = data["Kz"]
EKhr = data["Khr"]

Epolo = EKhd + EKz
Etoro = EKhr
EKh = EKhr + EKhd

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

Epot_vs_kh = np.sum(EA, axis=0) * delta_kz
Epot_vs_kz = np.sum(EA, axis=1) * delta_kh

ax1.plot(
    kh,
    Epolo_vs_kh * kh ** (coef_compensate),
    "-g",
    label=r"$E_{\rm polo}(k_h)$",
)
ax1.plot(
    kh,
    Epot_vs_kh * kh ** (coef_compensate),
    "-b",
    label=r"$E_{\rm pot}(k_h)$",
)
ax1.plot(
    kz,
    Epolo_vs_kz * kz ** (coef_compensate),
    "--g",
    label=r"$E_{\rm polo}(k_z)$",
)
ax1.plot(
    kz,
    Epot_vs_kz * kz ** (coef_compensate),
    "--b",
    label=r"$E_{\rm pot}(k_z)$",
)

tmp = sim.output.spect_energy_budg.plot_fluxes(
    tmin=tmin, key_k="kh", ax=ax3, plot_conversion=False
)
for ax in [ax1, ax3]:
    ax.axvline(kb, color="k", linestyle="dotted")
    ax.axvline(ko, color="k", linestyle="dashed")
    ax.axvline(kf, color="orange", linestyle="dashed")
    ax.set_xlim([kh[1], max(kh)])
    ax.set_xticks([1e1, 1e2, 1e3, kb, ko])
    ax.set_xticklabels(
        [r"$10^1$", r"$10^2$", r"$10^3$", r"$k_{\rm b}$", r"$k_{\rm O}$"],
        fontsize=14,
    )
# tmp = sim.output.spect_energy_budg.plot_fluxes(tmin=tmin, key_k="kz", ax=ax5)
for ax in [ax0, ax1]:
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_ylim([1e-3, 1e1])
    ax.set_xlabel(r"$k_h, k_z$", fontsize=20)

    N = sim.params.N
    mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
    Uh2 = mean_values["Uh2"]
    Uh = np.sqrt(Uh2)
    epsK = mean_values["epsK"]
    # Plot k^-2 and k^-3
    k = np.array([2e1, 8e2])
    ax.plot(
        k,
        4.0 * (k / k[0]) ** (-2.0 + coef_compensate),
        "--",
        color="gray",
        label=None,
    )
    ax.text(5e2, 3.0, r"$\propto k_i^{-2}$", fontsize=14, color="gray")
    ax.plot(
        k,
        4.0 * (k / k[0]) ** (-3.0 + coef_compensate),
        "-.",
        color="gray",
        label=None,
    )
    ax.text(5e2, 1e-2, r"$\propto k_i^{-3}$", fontsize=14, color="gray")

    """
    ax.plot(
        k,
        0.5 * (epsK**(2/3)) * k ** (-5/3 + coef_compensate),
        "-.",
        color="gray",
        label=None,
    )
    ax.plot(
        k,
        0.02 * N * Uh * k ** (-5/3 + coef_compensate),
        "-.",
        color="gray",
        label=None,
    )
    """

for ax in [ax2, ax3]:
    ax.set_ylim([-0.2, 1.2])
    lines = [line for line in ax.get_lines()]
    for l in range(3):
        lines[l].remove()
    labels = [
        r"$\Pi_{\rm kin}(k_h)/\varepsilon$",
        r"$\Pi_{\rm pot}(k_h)/\varepsilon$",
        r"$\varepsilon_{\rm kin}(k_h)/\varepsilon$",
        r"$\varepsilon_{\rm pot}(k_h)/\varepsilon$",
    ]
    ax.legend(loc="upper left", fontsize=10, labels=labels)

for ax in [ax2]:
    ax.set_yticks([-0.2, 0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.2])
    ax.set_yticklabels(
        [
            r"$-0.2$",
            r"$0$",
            r"$0.2$",
            r"$0.4$",
            r"$0.6$",
            r"$0.8$",
            r"$1.0$",
            r"$1.2$",
        ],
        fontsize=14,
    )

for ax in [ax1, ax3]:
    ax.set_yticklabels([])
    ax.set_ylabel("")

ax0.set_yticks([1e-3, 1e-2, 1e-1, 1e0, 1e1])
ax0.set_yticklabels(
    [r"$10^{-3}$", r"$10^{-2}$", r"$10^{-1}$", r"$10^{0}$", r"$10^{1}$"],
    fontsize=14,
)
ax0.set_ylabel(r"$E_{\rm 1D} \times k_i^{5/3}$", fontsize=20)

for ax in [ax0, ax1]:
    ax.legend(loc="lower center", fontsize=10)
    # ax.grid(True)


ax2.set_xlabel(r"$k_h$", fontsize=20)
ax3.set_xlabel(r"$k_h$", fontsize=20)
# ax4.set_xlabel(r"$k_z$", fontsize=20)
# ax5.set_xlabel(r"$k_z$", fontsize=20)

ax2.set_ylabel(r"$\Pi(k_h)/ \varepsilon$", fontsize=20)
# ax4.set_ylabel(r"$\Pi(k_z)/ \varepsilon$", fontsize=20)

ax0.set_title(r"Standard Navier-Stokes" + "\n" + r"$\rm (a)$", fontsize=20)
ax1.set_title(r"Without vortical modes" + "\n" + r"$\rm (b)$", fontsize=20)
ax2.set_title(r"$\rm (c)$", fontsize=20)
ax3.set_title(r"$\rm (d)$", fontsize=20)
# ax4.set_title(r"$\rm (e)$", fontsize=20)
# ax5.set_title(r"$\rm (f)$", fontsize=20)


fig.tight_layout()
save_fig(fig, f"figure6.png")

if __name__ == "__main__":
    plt.show()