Skip to content
Snippets Groups Projects
save_figure7.py 5.5 KiB
Newer Older
import matplotlib.cm
import matplotlib.pyplot as plt
import numpy as np

from util import (
    compute_kf_kb_ko_keta_kd,
    save_fig,
)

from util_simuls_regimes import get_sim

cm = matplotlib.cm.get_cmap("inferno", 100)

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


def plot_spectra(sim, ax, key="Ee", key_k="kh"):
    p_oper = sim.params.oper
    kmax = p_oper.coef_dealiasing * p_oper.nx * np.pi / p_oper.Lx
    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, key_to_load=["A", "Khd", "Kz", "Khr"]
    )
    kh = data["kh_spectra"]
    kz = data["kz"]
    delta_kh = kh[1]
    delta_kz = kz[1]
    if key == "EA":
        spectrum = data["A"]
    elif key == "EK":
        EKhd = data["Khd"]
        EKz = data["Kz"]
        Etoro = data["Khr"]
        spectrum = EKhd + EKz + Etoro
    elif key == "Epolo":
        EKhd = data["Khd"]
        EKz = data["Kz"]
        spectrum = EKhd + EKz
    elif key == "Ee":
        EA = data["A"]
        EKhd = data["Khd"]
        EKz = data["Kz"]
        Epolo = EKhd + EKz
        spectrum = 2 * np.minimum(EA, Epolo)
    elif key == "Ed":
        EA = data["A"]
        EKhd = data["Khd"]
        EKz = data["Kz"]
        Epolo = EKhd + EKz
        spectrum = EA + Epolo - 2 * np.minimum(EA, Epolo)
    elif key == "Etoro":
        spectrum = data["Khr"]
    else:
        print(f"Don't know key: {key} \n")
        exit

    if key_k == "kh":
        for iz in range(int(np.floor(np.log(len(kz)) / np.log(1.2)))):
            nz = int(np.floor(1.2**iz))
            if kz[nz] <= 10 * kb and kz[nz] >= 0.1 * kb:
                cs = ax.plot(
                    kh,
                    spectrum[nz, :],
                    color=cm(0.5 + 0.5 * np.log10((kz[nz] / kb))),
                    linestyle="-",
                )
        ks = np.array([0.1 * kf, kb])
        ax.plot(ks, 2e-1 * ks ** (-2), "k-")
        ax.text(np.sqrt(0.1 * kf * kb), 2e-3, r"$\propto k_h^{-2}$", fontsize=14)
        ks = np.array([0.1 * kf, kb])
        ax.plot(ks, 1e-10 * ks ** (1), "k-")
        ax.text(np.sqrt(0.1 * kf * kb), 3e-10, r"$\propto k_h^{1}$", fontsize=14)
        ax.set_xlim([delta_kh, kmax])
        ax.set_xlabel(r"$k_h$", fontsize=20)
    elif key_k == "kz":
        for ih in range(int(np.floor(np.log(len(kh)) / np.log(1.2)))):
            nh = int(np.floor(1.2**ih))
            if kh[nh] <= 10 * kb and kh[nh] >= 0.1 * kb:
                cs = ax.plot(
                    kz,
                    spectrum[:, nh],
                    color=cm(0.5 + 0.5 * np.log10((kh[nh] / kb))),
                    linestyle="-",
                )
        ks = np.array([0.3 * kf, ko])
        ax.plot(ks, 5e-9 * ks ** (0), "k-")
        ax.text(np.sqrt(0.3 * kf * ko), 7e-10, r"$\propto k_z^{0}$", fontsize=14)
        ks = np.array([kb, ko])
        ax.plot(ks, 1e4 * ks ** (-4), "k-")
        ax.text(np.sqrt(kb * ko), 1e-4, r"$\propto k_z^{-4}$", fontsize=14)
        ax.set_xlim([delta_kz, kmax])
        ax.set_xlabel(r"$k_z$", fontsize=20)
    ax.set_ylim([1e-10, 1e-1])
    ax.set_xticks(
        [1e1, 1e2, 1e3, kb, ko],
        [r"$10^1$", r"$10^2$", r"$10^3$", r"$k_{\rm b}$", r"$k_{\rm O}$"],
        fontsize=14,
    )
    ax.set_xscale("log")
    ax.set_yscale("log")
    # ax.grid(True)

    ax.axvline(kb, color="k", linestyle="dotted")
    ax.axvline(ko, color="k", linestyle="dashed")
    ax.axvline(kf, color="orange", linestyle="dashed")
    ax.axvline(keta, color="g", linestyle="dashed")

### Figure: Kinetic energy
nbax = 0
css = [None for i in range(4)]
fig, axes = plt.subplots(
    ncols=2, nrows=2, figsize=(10, 2 * 3 * 4.5 / 4), constrained_layout=True
)

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

axs = [ax0, ax1, ax2, ax3]
for proj in [False, True]:
    sim = get_sim("L", proj=proj)
    css[nbax] = plot_spectra(sim, axs[nbax], key="EK", key_k="kh")
    css[nbax + 2] = plot_spectra(sim, axs[nbax + 2], key="EK", key_k="kz")
    nbax += 1
for ax in [ax0, ax2]:
    ax.set_ylabel(r"$E_{\rm kin}(k_h, k_z)$", fontsize=20)
    ax.set_yticks(
        [1e-10, 1e-8, 1e-6, 1e-4, 1e-2],
        [r"$10^{-10}$", r"$10^{-8}$", r"$10^{-6}$", r"$10^{-4}$", r"$10^{-2}$"],
        fontsize=14,
    )

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

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)


norm = matplotlib.colors.Normalize(vmin=-1, vmax=1)

fig.tight_layout()
fig.subplots_adjust(right=0.85)

cbar_ax = fig.add_axes([0.88, 0.1, 0.02, 0.33])
cbar = fig.colorbar(
    matplotlib.cm.ScalarMappable(norm=norm, cmap=cm),
    cax=cbar_ax,
    cmap=cm,
    orientation="vertical",
)
cbar.set_ticks([-1, -0.5, 0.0, 0.5, 1])
cbar.set_ticklabels([r"$-1$", r"$-0.5$", r"$0.0$", r"$0.5$", r"$1$"], fontsize=14)
cbar.set_label(r"$\log_{10} (k_h/k_{\rm b})$", fontsize=20)
cbar_ax = fig.add_axes([0.88, 0.565, 0.02, 0.33])
cbar = fig.colorbar(
    matplotlib.cm.ScalarMappable(norm=norm, cmap=cm),
    cax=cbar_ax,
    cmap=cm,
    orientation="vertical",
)
cbar.set_ticks([-1, -0.5, 0.0, 0.5, 1])
cbar.set_ticklabels([r"$-1$", r"$-0.5$", r"$0.0$", r"$0.5$", r"$1$"], fontsize=14)
cbar.set_label(r"$\log_{10} (k_z/k_{\rm b})$", fontsize=20)


save_fig(fig, "figure7.png")

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