Skip to content
Snippets Groups Projects
save_figure15.py 7.28 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,
    compute_E_waves_vs_kh_kz,
    spectra_vs_khomega_slice,
    spectra_vs_kzomega_slice,
)

from util_simuls_regimes import get_sim

from util_dataframe import df, df_proj, df_ratio_one, df_proj_ratio_one

cm = matplotlib.cm.get_cmap("inferno", 100)
cmbin = matplotlib.cm.get_cmap("binary", 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"


normalize = False

# def plot_spectra(sim, ax, key="Ee"):


def plot_slices_omega_ki_waves(sim, ax, key="Ee", key_k="kh", ik=10, delta=1e-1):
    N = sim.params.N
    E_waves_omega, E_waves, E_tot, kh, kz, omegas = compute_E_waves_vs_kh_kz(
        sim, delta=delta

    if key_k == "kh":
        spectra_normalized = spectra_vs_khomega_slice(
            E_waves_omega, kh, omegas, ik, normalize=normalize
        )
        ax.set_xlabel(r"$k_h$", fontsize=20)
        ax.set_xlim([kh[1], max(kh)])
        cs = ax.pcolormesh(
            kh,
            omegas / N,
            np.log10(spectra_normalized.transpose()),
            cmap=cm,
            vmin=-9,
            vmax=-4,
            shading="nearest",
        )
        k = (kh**2 + kz[ik] ** 2) ** 0.5
        omega_disp = kh / k
        ax.plot(kh, omega_disp, "k-")
        ax.plot(kh, omega_disp + 0.5 * delta * omega_disp, "b-")
        ax.plot(kh, omega_disp - 0.5 * delta * omega_disp, "b-")
        ax.text(10, 2.5, rf"$k_z={kz[ik]:.1f}$", color="w", fontsize=14)

    elif key_k == "kz":
        spectra_normalized = spectra_vs_kzomega_slice(
            E_waves_omega, kz, omegas, ik, normalize=normalize
        )
        ax.set_xlabel(r"$k_z$", fontsize=20)
        ax.set_xlim([kz[1], max(kz)])
        cs = ax.pcolormesh(
            kz,
            omegas / N,
            np.log10(spectra_normalized.transpose()),
            cmap=cm,
            vmin=-9,
            vmax=-4,
            shading="nearest",
        )
        k = (kh[ik] ** 2 + kz**2) ** 0.5
        omega_disp = kh[ik] / k
        ax.plot(kz, omega_disp, "k-")
        ax.plot(kz, omega_disp + 0.5 * delta * omega_disp, "b-")
        ax.plot(kz, omega_disp - 0.5 * delta * omega_disp, "b-")
        ax.text(15, 2.5, rf"$k_h={kh[ik]:.1f}$", color="w", fontsize=14)

    ax.set_ylabel(r"$\omega / N$", fontsize=20)
    ax.set_ylim([0, 3])
    ax.set_xticks([25, 50, 75, 100, 125])
    ax.set_xticklabels(
        [r"$25$", r"$50$", r"$75$", r"$100$", r"$125$"], fontsize=14
    ax.set_yticks([0, 1, 2, 3])
    ax.set_yticklabels([r"$0$", r"$1$", r"$2$", r"$3$"], fontsize=14)
    cs.cmap.set_bad("k")
    cs.cmap.set_under("k")
def plot_ratio_E_waves_vs_kh_kz(sim, ax):
    N = sim.params.N
    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)
    E_waves_omega, E_waves, E_tot, kh, kz, omegas = compute_E_waves_vs_kh_kz(sim)
    ratio_waves = E_waves / E_tot

    ax.set_xlabel(r"$k_h$", fontsize=20)
    ax.set_xlim([kh[1], max(kh)])
    ax.set_xticks([25, 50, 75, 100, 125])
    ax.set_xticklabels(
        [r"$25$", r"$50$", r"$75$", r"$100$", r"$125$"], fontsize=14
    )
    ax.set_ylabel(r"$k_z$", fontsize=20)
    ax.set_ylim([kz[1], max(kz)])
    ax.set_yticks([25, 50, 75, 100, 125])
    ax.set_yticklabels(
        [r"$25$", r"$50$", r"$75$", r"$100$", r"$125$"], fontsize=14
    )
    cs = ax.pcolormesh(
        kh,
        kz,
        ratio_waves,
        cmap=cmbin,
        vmin=0,
        vmax=1,
        shading="nearest",
    )
    f = kh / np.tan(0.3)
    ax.plot(kh, f, linestyle="-", linewidth=1, color="orange")
    f = kh / np.tan(0.15)
    ax.plot(kh, f, linestyle="--", linewidth=1, color="orange")
    f = kh / np.tan(0.075)
    ax.plot(kh, f, linestyle="--", linewidth=1, color="orange")
    f = kh / np.tan(0.225)
    ax.plot(kh, f, linestyle="--", linewidth=1, color="orange")

    # kb
    th = np.linspace(0, np.pi / 2, 50)
    ax.plot(kb * np.sin(th), kb * np.cos(th), linestyle="dashed", color="c")
    # 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="c",

    delta_kh = kh[1]
    delta_kz = kz[1]
    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,
    )

    return cs


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

ax0 = axes[0]
ax1 = axes[1]


delta = 0.1
sim = get_sim(letter, proj=False)
# cs0 = plot_slices_omega_ki_waves(sim, ax0, key="Ee", key_k="kh", ik=3, delta=delta)
cs0 = plot_ratio_E_waves_vs_kh_kz(sim, ax0)
sim = get_sim(letter, proj=True)
# cs1 = plot_slices_omega_ki_waves(sim, ax1, key="Ee", key_k="kh", ik=3, delta=delta)
cs1 = plot_ratio_E_waves_vs_kh_kz(sim, ax1)

ax1.set_yticklabels([])
ax1.set_ylabel("")

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)
fig.tight_layout()
fig.subplots_adjust(right=0.85)
cbar_ax2 = fig.add_axes([0.88, 0.17, 0.02, 0.65])
cbar2 = fig.colorbar(cs1, cax=cbar_ax2, cmap=cmbin, orientation="vertical")
cbar2.set_label(r"$\tilde{E}_{\rm wave}(k_h,k_z)$", fontsize=20)
cbar2.set_ticks([0, 0.2, 0.4, 0.6, 0.8, 1])
cbar2.set_ticklabels(
    [r"$0$", r"$0.2$", r"$0.4$", r"$0.6$", r"$0.8$", r"$1$"], fontsize=14
)


save_fig(fig, "figure15.png")

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