Skip to content
Snippets Groups Projects
save_figure8.py 7.04 KiB
Newer Older
import sys
import matplotlib.cm
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

from math import degrees

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}"
print(sys.argv)
letter = sys.argv[-1]
if letter not in "DLOWP":
    letter = "L"
def plot_spectra(sim, ax, key="Ee"):
    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 == "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
    cs = ax.pcolormesh(
        kh,
        kz,
        np.log10(spectrum),
        cmap=cm,
        vmin=-9.0,
        vmax=-1.0,
        shading="nearest",
    th = np.linspace(0, np.pi / 2, 50)
    # Forcing
    kf_min = sim.params.forcing.nkmin_forcing * delta_kz
    kf_max = sim.params.forcing.nkmax_forcing * delta_kz
    angle = sim.params.forcing.tcrandom_anisotropic.angle
    delta_angle = sim.params.forcing.tcrandom_anisotropic.delta_angle
    # Forcing
    ax.add_patch(
        patches.Arc(
            xy=(0, 0),
            width=2 * kf_max,
            height=2 * kf_max,
            angle=0,
            theta1=90.0 - degrees(angle) - 0.5 * degrees(delta_angle),
            theta2=90.0 - degrees(angle) + 0.5 * degrees(delta_angle),
            linestyle="-",
            color="orange",
            linewidth=1,
        )
    )
    ax.add_patch(
        patches.Arc(
            xy=(0, 0),
            width=2 * kf_min,
            height=2 * kf_min,
            angle=0,
            theta1=90.0 - degrees(angle) - 0.5 * degrees(delta_angle),
            theta2=90.0 - degrees(angle) + 0.5 * degrees(delta_angle),
            linestyle="-",
            color="orange",
            linewidth=1,
        )
    )
    ax.plot(
        [
            kf_min * np.sin(angle - 0.5 * delta_angle),
            kf_max * np.sin(angle - 0.5 * delta_angle),
        ],
        [
            kf_min * np.cos(angle - 0.5 * delta_angle),
            kf_max * np.cos(angle - 0.5 * delta_angle),
        ],
        linestyle="-",
        color="orange",
        linewidth=1,
    )
    ax.plot(
        [
            kf_min * np.sin(angle + 0.5 * delta_angle),
            kf_max * np.sin(angle + 0.5 * delta_angle),
        ],
        [
            kf_min * np.cos(angle + 0.5 * delta_angle),
            kf_max * np.cos(angle + 0.5 * delta_angle),
        ],
        linestyle="-",
        color="orange",
        linewidth=1,
    )

    # kb
    ax.plot(kb * np.sin(th), kb * np.cos(th), color="cyan", linestyle="dashed")

    # Chi_L = 1/3, 3
    a = 1 / 3
    xa = np.linspace(kh[1], a**1.5 * ko, 50, endpoint=True)
    ax.plot(
        xa,
        xa * np.sqrt((a**1.5 * ko / xa) ** 0.8 - 1),
        linestyle="dotted",
        color="cyan",
    )
    a = 3
    xa = np.linspace(kh[1], a**1.5 * ko, 50, endpoint=True)
    ax.plot(
        xa,
        xa * np.sqrt((a**1.5 * ko / xa) ** 0.8 - 1),
        linestyle="dotted",
        color="cyan",
    )
    # Chi_d = 1
    a = 1.0
    xa = np.linspace(kh[1], kd, 50, endpoint=True)
    ax.plot(
        xa,
        xa * np.sqrt((kd / xa) ** (4 / 3) - 1),
        linestyle="dotted",
        color="magenta",
    )
    # Chi_nu = 1 (Kolmogorov scale)
    ax.plot(keta * np.sin(th), keta * np.cos(th), linestyle="dotted", color="g")

    ax.set_xticks(
        [1e1, kb, 1e2, ko, 1e3],
        [r"$10^1$", r"$k_{\rm b}$", r"$10^2$", r"$k_{\rm O}$", r"$10^3$"],
        fontsize=14,
    )
    ax.set_yticks(
        [1e1, kb, 1e2, ko, 1e3],
        [r"$10^1$", r"$k_{\rm b}$", r"$10^2$", r"$k_{\rm O}$", r"$10^3$"],
        fontsize=14,
    )

    ax.set_xlim([delta_kh, kmax])
    ax.set_ylim([delta_kz, kmax])
    ax.set_xscale("log")
    ax.set_yscale("log")
    # ax.grid(True)

    return cs


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

ax0 = axes[0, 0]
ax1 = axes[0, 1]
ax2 = axes[1, 0]
ax3 = axes[1, 1]
ax4 = axes[2, 0]
ax5 = axes[2, 1]

axs = [ax0, ax1, ax2, ax3, ax4, ax5]

sim = get_sim(letter)
cs0 = plot_spectra(sim, ax0, key="Etoro")
cs2 = plot_spectra(sim, ax2, key="Epolo")
cs4 = plot_spectra(sim, ax4, key="EA")

sim = get_sim(letter, proj=True)
cs1 = plot_spectra(sim, ax1, key="Etoro")
cs3 = plot_spectra(sim, ax3, key="Epolo")
cs5 = plot_spectra(sim, ax5, key="EA")


for ax in [ax0, ax2, ax4]:
    ax.set_ylabel(r"$k_z$", fontsize=20)
for ax in [ax4, ax5]:
    ax.set_xlabel(r"$k_h$", fontsize=20)


for ax in [ax0, ax1, ax2, ax3]:
    ax.set_xticklabels([])
for ax in [ax1, ax3, ax5]:
    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)
ax4.set_title(r"$\rm (e)$", fontsize=20)
ax5.set_title(r"$\rm (f)$", fontsize=20)


# Add annotations for the lines and forcing region
ax1.text(
    4.5,
    1.5e1,
    r"Forcing",
    color="orange",
    rotation=0.7 * 180 / np.pi,
    fontsize=10,
)
ax1.text(1e1, 1.5e1, r"$\chi_{\mathbf{k}} < 1/3$", color="cyan", fontsize=14)
ax1.text(1e2, 1.5e2, r"$\chi_{\mathbf{k}} < 3$", color="cyan", fontsize=14)
ax1.text(7e1, 5e1, r"$\gamma_{\mathbf{k}} < 1$", color="m", fontsize=14)


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

cbar_ax0 = fig.add_axes([0.88, 0.675, 0.02, 0.25])
cbar0 = fig.colorbar(cs0, cax=cbar_ax0, cmap=cm, orientation="vertical")
cbar0.set_label(r"$\log_{10} E_{\rm toro}$", fontsize=20)
cbar_ax2 = fig.add_axes([0.88, 0.3725, 0.02, 0.25])
cbar2 = fig.colorbar(cs2, cax=cbar_ax2, cmap=cm, orientation="vertical")
cbar2.set_label(r"$\log_{10} E_{\rm polo}$", fontsize=20)

cbar_ax4 = fig.add_axes([0.88, 0.07, 0.02, 0.25])
cbar4 = fig.colorbar(cs4, cax=cbar_ax4, cmap=cm, orientation="vertical")
cbar4.set_label(r"$\log_{10} E_{\rm pot}$", fontsize=20)
for cbar in [cbar0, cbar2, cbar4]:
    cbar.set_ticks([-9, -7, -5, -3, -1])
    cbar.set_ticklabels(
        [r"$-9$", r"$-7$", r"$-5$", r"$-3$", r"$-1$"], fontsize=14
    )


save_fig(fig, "figure8.png")

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