Skip to content
Snippets Groups Projects
save_figure14.py 5.15 KiB
Newer Older
import h5py
import matplotlib.cm
import matplotlib.pyplot as plt
import numpy as np
from util_simuls_regimes import get_sim

from util import (
    save_fig,
    customize,
)

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

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


def plot_spectra_vs_omegak(sim, ax, key="Ee"):
    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 = sim.params.path_run
    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"][:]
        Epot = 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 == "Epot":
            spectrum = Epot
        elif key == "EK":
            spectrum = Epolo + Etoro
        elif key == "Epolo":
            spectrum = Epolo
        elif key == "Ee":
            spectrum = 2 * np.minimum(Epot, Epolo)
        elif key == "Ed":
            spectrum = Epot + Epolo - 2 * np.minimum(Epot, Epolo)
        elif key == "Etoro":
            spectrum = Etoro
        elif key == "Etot":
            spectrum = Epolo + Etoro + Epot
        else:
            print(f"Don't know key: {key} \n")
            exit
        KH, KZ = np.meshgrid(kh, kz)
        K = (KH**2 + KZ**2) ** 0.5
        K_NOZERO = K.copy()
        K_NOZERO[K_NOZERO == 0] = 1e-16
        omega_disp = N * KH / K_NOZERO

        omega_disp_grid = np.linspace(0, N, 25, endpoint=True)

        OMEGA_DISP_GRID, OMEGA_GRID = np.meshgrid(omega_disp_grid, omegas)
        SPECTRUM = np.zeros(OMEGA_DISP_GRID.shape)

        for iod in range(len(omega_disp_grid) - 1):
            for io in range(len(omegas)):
                omega_disp_min = omega_disp_grid[iod]
                omega_disp_max = omega_disp_grid[iod + 1]
                IOD = np.where(
                    np.logical_and(
                        omega_disp >= omega_disp_min, omega_disp < omega_disp_max
                    )
                )

                for i in range(len(IOD[0])):
                    SPECTRUM[io, iod] += spectrum[IOD[0][i], IOD[1][i], io]

        cs = ax.pcolormesh(
            OMEGA_DISP_GRID / N,
            OMEGA_GRID / N,
            np.log10(SPECTRUM),
            cmap=cm,
            vmin=-6,
            vmax=-3,
            shading="nearest",

    ax.set_xlabel(r"$\sin \theta_{\boldsymbol{k}}$", fontsize=20)
    ax.set_xlim([0, 1])
    ax.set_xticks([0, 0.25, 0.5, 0.75, 1])
    ax.set_xticklabels(
        [r"$0$", r"$0.25$", r"$0.5$", r"$0.75$", r"$1$"], fontsize=14
    ax.set_ylabel(r"$\omega / N$", fontsize=20)
    ax.set_ylim([0, 3])
    ax.set_yticks([0, 1, 2, 3])
    ax.set_yticklabels([r"$0$", r"$1$", r"$2$", r"$3$"], fontsize=14)
# 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=3, figsize=(10, 3 * 3 * 4.5 / 4), constrained_layout=True

css = [None for i in range(6)]

ax0 = axes[0, 0]
ax1 = axes[0, 1]
ax2 = axes[1, 0]
ax3 = axes[1, 1]
ax4 = axes[2, 0]
ax5 = axes[2, 1]
sim = get_sim("S1", proj=False)
css[0] = plot_spectra_vs_omegak(sim, ax0, key="Etot")

sim_proj = get_sim("S1", proj=True)
css[1] = plot_spectra_vs_omegak(sim_proj, ax1, key="Etot")

sim = get_sim("S2", proj=False)
css[2] = plot_spectra_vs_omegak(sim, ax2, key="Etot")

sim_proj = get_sim("S2", proj=True)
css[3] = plot_spectra_vs_omegak(sim_proj, ax3, key="Etot")
sim = get_sim("S3", proj=False)
css[4] = plot_spectra_vs_omegak(sim, ax4, key="Etot")

sim_proj = get_sim("S3", proj=True)
css[5] = plot_spectra_vs_omegak(sim_proj, ax5, key="Etot")
for ax in [ax1, ax3, ax5]:
    ax.set_ylabel("")
    ax.set_yticklabels([])

for ax in [ax0, ax1, ax2, ax3]:
    ax.set_xlabel("")
    ax.set_xticklabels([])


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)
for ax in [ax0, ax1, ax2, ax3, ax4, ax5]:
    ax.plot([0, 1], [0, 1], "w--")

fig.tight_layout()
# fig.subplots_adjust(right=0.85)
cbar_ax = fig.add_axes([0.88, 0.19, 0.02, 0.63])
cbar = fig.colorbar(css[0], cax=cbar_ax)
cbar.set_ticks([-3, -4, -5, -6])
cbar.set_ticklabels([r"$-3$", r"$-4$", r"$-5$", r"$-6$"], fontsize=16)
cbar.ax.set_ylabel(
    r"$\log_{10} E(\sin \theta_{\boldsymbol{k}},\omega)$", fontsize=20
fig.subplots_adjust(right=0.85, wspace=0.05, hspace=0.15)
save_fig(fig, f"figure14.png")


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