Skip to content
Snippets Groups Projects
save_figure4.py 4.77 KiB
Newer Older
import sys

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

from fluidsim.util import load_params_simul, times_start_last_from_path

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

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

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


print(sys.argv)
letter = sys.argv[-1]

if letter not in "DLOWPU":
    letter = "L"


def plot_omega_spectra(sim, ax):
    path = sim.params.path_run
    t_start, t_last = times_start_last_from_path(path)
    tmin = t_last - 2
    mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
    params = load_params_simul(path)
    proj = params.projection
    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:
        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)

        EA = np.sum(EA, axis=0)
        Epolo = np.sum(Epolo, axis=0)
        Etoro = np.sum(Etoro, axis=0)
        Ee = np.sum(Ee, axis=0)
        E = np.sum(E, axis=0)

        EA = np.sum(EA, axis=0)
        Epolo = np.sum(Epolo, axis=0)
        Etoro = np.sum(Etoro, axis=0)
        Ee = np.sum(Ee, axis=0)
        E = np.sum(E, axis=0)

        coef_compensate = 0
        cs = ax.plot(
            omegas / N,
            EA * omegas**coef_compensate,
            color="b",
            label=r"$E_{\rm pot}(\omega)$",
        )
        cs = ax.plot(
            omegas / N,
            Epolo * omegas**coef_compensate,
            color="g",
            label=r"$E_{\rm polo}(\omega)$",
        )
        if proj == None:
            cs = ax.plot(
                omegas / N,
                Etoro * omegas**coef_compensate,
                color="r",
                label=r"$E_{\rm toro}(\omega)$",
            )
        om = np.array([0.1 * N, N])
        ax.plot(
            om / N,
            1e-4 * (om / N) ** (-2 + coef_compensate),
            "--",
            color="gray",
            label=None,
        )
        ax.text(0.5, 1e-3, r"$\omega^{-2}$", color="gray", fontsize=14)

        ax.plot(
            om / N,
            1e-5 * (om / N) ** (-3 / 2 + coef_compensate),
            "-.",
            color="gray",
            label=None,
        )
        ax.text(0.5, 5e-6, r"$\omega^{-3/2}$", color="gray", fontsize=14)

        # Forcing
        angle = sim.params.forcing.tcrandom_anisotropic.angle
        delta_angle = sim.params.forcing.tcrandom_anisotropic.delta_angle
        omega_fmin = N * np.sin(angle - 0.5 * delta_angle)
        omega_fmax = N * np.sin(angle + 0.5 * delta_angle)
        ax.axvline(omega_fmin / N, color="orange", linestyle="dashed")
        ax.axvline(omega_fmax / N, color="orange", linestyle="dashed")

        ax.set_xlabel(r"$\omega/N$", fontsize=20)
        ax.set_xscale("log")
        ax.set_xticks([1e-1, 1e0])
        ax.set_xticklabels([r"$10^{-1}$", r"$10^{0}$"], fontsize=14)
        ax.set_xlim([min(omegas) / N, max(omegas) / N])
        ax.set_ylabel(r"$E(\omega)$", fontsize=20)
        ax.set_yscale("log")
        ax.set_yticks([1e-6, 1e-5, 1e-4, 1e-3, 1e-2])
        ax.set_yticklabels(
            [
                r"$10^{-6}$",
                r"$10^{-5}$",
                r"$10^{-4}$",
                r"$10^{-3}$",
                r"$10^{-2}$",
            ],
            fontsize=14,
        )
        ax.legend(loc="upper right", fontsize=14)
        return cs


sim = get_sim(letter)
sim_proj = get_sim(letter, proj=True)


# TODO: uncomment this assert
# 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]


cs0 = plot_omega_spectra(sim, ax0)
cs1 = plot_omega_spectra(sim_proj, 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()


save_fig(fig, f"figure4.png")


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