Skip to content
Snippets Groups Projects
save_figure11.py 6.09 KiB
Newer Older
import sys

import h5py
import matplotlib.cm
import matplotlib.pyplot as plt
import numpy as np
from util_simuls_regimes import get_sim

from util import (
    compute_kf_kb_ko_keta_kd,
    compute_omega_emp_vs_kzkh,
    spectra_vs_khomega_slice,
    spectra_vs_kzomega_slice,
    paths_simuls_regimes,
    paths_simuls_regimes_proj,
    save_fig,
    customize,
)

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 "DLOWPU":
    letter = "L"


normalize = False


def plot_slices_omega_ki(sim, ax, key="Ee", key_k="kh", ik=10, PLOT_Doppler=True):
    N = sim.params.N
    t_start, t_last = sim.output.print_stdout.get_times_start_last()
    tmin = t_last - 2.0
    mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
    Uh = np.sqrt(mean_values["Uh2"])

    path_spec = sorted(path.glob(f"spatiotemporal/periodogram_[0-9]*.h5"))
    assert len(path_spec) == 1, f"Not only 1 periodogram in {path} \n"
    path_spec = path_spec[0]
    with h5py.File(path_spec, "r") as f:
        kh = f["kh_spectra"][:]
        kz = f["kz_spectra"][:]
        omegas = f["omegas"][:]
        EA = f["spectrum_A"][:]
        EKz = f["spectrum_K"][:] - f["spectrum_Khd"][:] - f["spectrum_Khr"][:]
        Epolo = f["spectrum_Khd"][:] + EKz
        Etoro = f["spectrum_Khr"][:]
        delta_kh = kh[1]
        delta_kz = kz[1]
        if key == "EA":
            spectrum = EA
        elif key == "EK":
            spectrum = Epolo + Etoro
        elif key == "Epolo":
            spectrum = Epolo
        elif key == "Ee":
            spectrum = 2 * np.minimum(EA, Epolo)
        elif key == "Ed":
            spectrum = EA + Epolo - 2 * np.minimum(EA, Epolo)
        elif key == "Etoro":
            spectrum = Etoro
        else:
            print(f"Don't know key: {key} \n")
            exit

    omega_emp, delta_omega_emp = compute_omega_emp_vs_kzkh(
        N,
        spectrum,
        kh,
        kz,
        omegas,
    )
    if key_k == "kh":
        spectra_normalized = spectra_vs_khomega_slice(
            spectrum, kh, omegas, ik, normalize=normalize
        )
        ax.set_xlabel(r"$k_h$", fontsize=20)
        ax.set_xlim([delta_kh, max(kh)])
        cs = ax.pcolormesh(
            kh,
            omegas / N,
            np.log10(spectra_normalized.transpose()),
            cmap=cm,
            vmin=-9,
            vmax=-5,
            shading="nearest",
        )
        k = (kh**2 + kz[ik] ** 2) ** 0.5
        omega_disp = kh / k
        ax.plot(kh, omega_disp, "k-")
        # ax.plot(kh, omega_emp[ik, :] / N, "g-")
        ax.plot(kh, omega_disp + 0.5 * delta_omega_emp[ik, :] / N, "b-")
        ax.plot(kh, omega_disp - 0.5 * delta_omega_emp[ik, :] / N, "b-")
        if PLOT_Doppler:
            ax.plot(kh, omega_disp + k * Uh / N, "y--")
            ax.plot(kh, omega_disp - k * Uh / N, "y--")
        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(
            spectrum, kz, omegas, ik, normalize=normalize
        )
        ax.set_xlabel(r"$k_z$", fontsize=20)
        ax.set_xlim([delta_kz, max(kz)])
        cs = ax.pcolormesh(
            kz,
            omegas / N,
            np.log10(spectra_normalized.transpose()),
            cmap=cm,
            vmin=-9,
            vmax=-5,
            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_emp[:, ik] / N, "g-")
        ax.plot(kz, omega_disp + 0.5 * delta_omega_emp[:, ik] / N, "b-")
        ax.plot(kz, omega_disp - 0.5 * delta_omega_emp[:, ik] / N, "b-")
        if PLOT_Doppler:
            ax.plot(kz, omega_disp + k * Uh / N, "y--")
            ax.plot(kz, omega_disp - k * Uh / N, "y--")
        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)
    return cs


# assert sim.params.oper.nx == sim_proj.params.oper.nx, f"Not the same resolution for simulation Without vortical modes: {sim.params.oper.nx} vs {sim_proj.params.oper.nx}"

fig, axes = plt.subplots(
    ncols=2, nrows=2, figsize=(10, 2 * 3 * 4.5 / 4), constrained_layout=True
nbax = 0
css = [None for i in range(4)]
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(letter, proj=proj)
    if proj:
        path = paths_simuls_regimes_proj[letter]
    else:
        path = paths_simuls_regimes[letter]
    css[nbax] = plot_slices_omega_ki(sim, axs[nbax], key="Ee", key_k="kh", ik=3)
    css[nbax + 2] = plot_slices_omega_ki(
        sim, axs[nbax + 2], key="Ee", key_k="kz", ik=12
    nbax += 1
for ax in [ax1, ax3]:
    ax.set_ylabel("")
    ax.set_yticks([])

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)
# Legend for delta_omega
ax0.plot([100, 100], [0.42, 1.5], "b--")
ax0.text(105, 1.1, rf"$\delta \omega$", color="b", fontsize=14)


fig.tight_layout()
# fig.subplots_adjust(right=0.85)
cbar_ax = fig.add_axes([0.88, 0.2, 0.02, 0.5])
cbar = fig.colorbar(css[0], cax=cbar_ax)
cbar.set_ticks([-5, -6, -7, -8, -9])
cbar.set_ticklabels([r"$-5$", r"$-6$", r"$-7$", r"$-8$", r"$-9$"], fontsize=14)

if normalize:
    cbar.ax.set_ylabel(
        r"$\log_{10} \left( E_{\rm equi} / \int E_{\rm equi} ~ \mathrm{d}\omega \right)$",
        fontsize=20,
    )
else:
    cbar.ax.set_ylabel(r"$\log_{10} E_{\rm equi}(k_h,k_z,\omega)$", fontsize=20)


fig.subplots_adjust(right=0.85, wspace=0.1, hspace=0.4)
save_fig(fig, f"figure11.png")


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