Skip to content
Snippets Groups Projects
save_figure12.py 7.02 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,
    customize,
    paths_simuls_regimes,
    paths_simuls_regimes_proj,
    save_fig,
)

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"


sim = get_sim(letter)
path = paths_simuls_regimes[letter]

sim_proj = get_sim(letter, proj=True)
path_proj = paths_simuls_regimes_proj[letter]

# 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=1, figsize=(10, 1.2 * 3 * 4.5 / 4), constrained_layout=True
)

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


ks_mins = [-0.02, 0.18, 0.38, 0.58, 0.78, 0.98, 1.98]
ks_maxs = [0.02, 0.22, 0.42, 0.62, 0.82, 1.02, 2.02]


# Standard Navier-Stokes
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)
mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
R2 = mean_values["R2"]
Fh = mean_values["Fh"]
epsK = mean_values["epsK"]
N = sim.params.N
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:
    # Get the data
    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"][:]
    E = Epolo + Etoro + EA
    Ee = 2 * np.minimum(EA, Epolo)
    Ed = EA + Epolo - Ee

    spectrum = Ee
    omega_emp, delta_omega_emp = compute_omega_emp_vs_kzkh(
        N, spectrum, kh, kz, omegas
    )
    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
    nlb = delta_omega_emp / omega_disp
    chi = (K**2 * epsK) ** (1 / 3) / omega_disp
DO_PLOT = [0 for i in range(len(ks_maxs))]
for nh in range(len(kh)):
    for nz in range(len(kz)):
        k = np.sqrt(kh[nh] ** 2 + kz[nz] ** 2)
        sint = kh[nh] / k
        omega_waves = N * sint
        theta = np.arcsin(sint)
        color = cm(k / (2.5 * kb))
        for i in range(len(ks_maxs)):
            if (
                k / kb > ks_mins[i]
                and k / kb < ks_maxs[i]
                and max(spectrum[nz, nh, :]) >= 1e-16
                and DO_PLOT[i] < 3
            ):
                cs = ax0.plot(
                    (omegas - omega_waves) / N,
                    spectrum[nz, nh, :] / max(E[nz, nh, :]),
                    color=color,
                    linestyle="-",
                    linewidth=1,
                    marker="o",
                    markersize=1.5,
                )
                DO_PLOT[i] += 1


# Without vortical modes
t_start, t_last = sim_proj.output.print_stdout.get_times_start_last()
tmin = t_last - 2.0
kf, kb, ko, keta, kd = compute_kf_kb_ko_keta_kd(sim_proj, tmin)
mean_values = sim_proj.output.get_mean_values(tmin=tmin, customize=customize)
R2 = mean_values["R2"]
Fh = mean_values["Fh"]
epsK = mean_values["epsK"]
N = sim_proj.params.N
path_spec = sorted(path_proj.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:
    # Get the data
    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"][:]
    E = Epolo + Etoro + EA
    Ee = 2 * np.minimum(EA, Epolo)
    Ed = EA + Epolo - Ee
    spectrum = Ee
    omega_emp, delta_omega_emp = compute_omega_emp_vs_kzkh(
        N, spectrum, kh, kz, omegas
    )
    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
    nlb = delta_omega_emp / omega_disp

DO_PLOT = [0 for i in range(len(ks_maxs))]
for nh in range(len(kh)):
    for nz in range(len(kz)):
        k = np.sqrt(kh[nh] ** 2 + kz[nz] ** 2)
        sint = kh[nh] / k
        omega_waves = N * sint
        theta = np.arcsin(sint)
        color = cm(k / (2.5 * kb))
        for i in range(len(ks_maxs)):
            if (
                k / kb > ks_mins[i]
                and k / kb < ks_maxs[i]
                and max(spectrum[nz, nh, :]) >= 1e-16
                and DO_PLOT[i] < 3
            ):
                cs = ax1.plot(
                    (omegas - omega_waves) / N,
                    spectrum[nz, nh, :] / max(E[nz, nh, :]),
                    color=color,
                    linestyle="-",
                    linewidth=1,
                    marker="o",
                    markersize=1.5,
                )
                DO_PLOT[i] += 1


# Axis
for ax in [ax0, ax1]:
    ax.set_xlim([-1, 3])
    ax.set_xticks([-1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3])
    ax.set_xticklabels(
        [
            r"$-1$",
            r"$-0.5$",
            r"$0$",
            r"$0.5$",
            r"$1$",
            r"$1.5$",
            r"$2$",
            r"$2.5$",
            r"$3$",
        ],
        fontsize=14,
    )
    ax.plot([0, 0], [1e-3, 1e0], "k-")
    ax.grid(True)

for ax in [ax0, ax1]:
    ax.set_yscale("log")
    ax.set_ylim([1e-3, 1e0])
    ax.set_yticks([1e-3, 1e-2, 1e-1, 1e0])
    ax.set_yticklabels(
        [r"$10^{-3}$", r"$10^{-2}$", r"$10^{-1}$", r"$10^{0}$"], fontsize=14
    )

for ax in [ax0, ax1]:
    ax.set_xlabel(r"$(\omega - \omega_{\boldsymbol{k}})/N$", fontsize=20)
ax1.set_yticklabels([])
ax0.set_ylabel(
    r"$E_{\rm equi} / \max\limits_{\omega} ~ E$",
    fontsize=20,
)

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)


(leg0,) = ax1.plot([-2, -2], [0, 0], color=cm(0.0))
(leg1,) = ax1.plot([-2, -2], [0, 0], color=cm(0.2 / 2.5))
(leg2,) = ax1.plot([-2, -2], [0, 0], color=cm(0.4 / 2.5))
(leg3,) = ax1.plot([-2, -2], [0, 0], color=cm(0.6 / 2.5))
(leg4,) = ax1.plot([-2, -2], [0, 0], color=cm(0.8 / 2.5))
(leg5,) = ax1.plot([-2, -2], [0, 0], color=cm(1.0 / 2.5))
(leg6,) = ax1.plot([-2, -2], [0, 0], color=cm(2.0 / 2.5))

ax1.legend(
    [leg1, leg2, leg3, leg4, leg5, leg6],
    [
        r"$k/k_{\rm b} \simeq 0.2$",
        r"$k/k_{\rm b} \simeq 0.4$",
        r"$k/k_{\rm b} \simeq 0.6$",
        r"$k/k_{\rm b} \simeq 0.8$",
        r"$k/k_{\rm b} \simeq 1.0$",
        r"$k/k_{\rm b} \simeq 2.0$",
    ],
    loc="upper right",
    fontsize=10,
fig.tight_layout()
save_fig(fig, f"figure12.png")


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