Skip to content
Snippets Groups Projects
save_spectra_slices_regimes.py 7.84 KiB
Newer Older
vlabarre's avatar
vlabarre committed
import sys

import matplotlib.pyplot as plt
vlabarre's avatar
vlabarre committed
import numpy as np
from util_simuls_regimes import get_sim

from fluidsim import load

from util import (
vlabarre's avatar
vlabarre committed
    compute_kf_kb_ko_keta_kd,
    customize,
    get_paths,
    pos_closest_value,
    save_fig,
)
vlabarre's avatar
vlabarre committed

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

if letter not in "DLOWP":
    letter = "P"


sim = get_sim(letter)

vlabarre's avatar
vlabarre committed
coef_compensate = 5 / 3
vlabarre's avatar
vlabarre committed

p_oper = sim.params.oper
kmax = p_oper.coef_dealiasing * p_oper.nx * np.pi / p_oper.Lx


t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = t_last - 2.0


if letter == "P":
    fig, ax = plt.subplots()

    temp = sim.output.spectra.load3d_mean(tmin)
    print(temp)
    EK = temp["spectra_E"]
    EA = temp["spectra_A"]
vlabarre's avatar
vlabarre committed
    EKhd = temp["spectra_Khd"]
    EKhr = temp["spectra_Khr"]
vlabarre's avatar
vlabarre committed
    EKh = EKhr + EKhd
    EKz = EK - EKh
    Epolo = EKhd + EKz
    Etoro = EKhr
vlabarre's avatar
vlabarre committed

vlabarre's avatar
vlabarre committed
    k = temp["k"]
    ax.plot(
        k,
vlabarre's avatar
vlabarre committed
        EA * k**coef_compensate,
vlabarre's avatar
vlabarre committed
        "b-",
vlabarre's avatar
vlabarre committed
        label=(r"$E_{A}$"),
vlabarre's avatar
vlabarre committed
    )
    ax.plot(
        k,
vlabarre's avatar
vlabarre committed
        Epolo * k**coef_compensate,
vlabarre's avatar
vlabarre committed
        "m-",
vlabarre's avatar
vlabarre committed
        label=(r"$E_{polo}$"),
vlabarre's avatar
vlabarre committed
    )
    ax.plot(
        k,
vlabarre's avatar
vlabarre committed
        Etoro * k**coef_compensate,
vlabarre's avatar
vlabarre committed
        "r-",
vlabarre's avatar
vlabarre committed
        label=(r"$E_{toro}$"),
vlabarre's avatar
vlabarre committed
    ax.legend(fontsize=10, loc="upper left")
vlabarre's avatar
vlabarre committed

    ax.set_xscale("log")
    ax.set_yscale("log")

vlabarre's avatar
vlabarre committed
    ax.set_xlabel(r"$k$", fontsize=16)
    ax.set_ylabel(r"$E(k) ~ k^{5/3}$", fontsize=16)
vlabarre's avatar
vlabarre committed
    # ax.text(1.1*kb, 1e-10, r"$k_b$", fontsize=16)
    # ax.text(1.1*ko, 1e-10, r"$k_O$", fontsize=16)
vlabarre's avatar
vlabarre committed
    ax.set_xlim([k[1], 0.8 * max(k)])
vlabarre's avatar
vlabarre committed
    ax.set_ylim(bottom=1e-4, top=1e1)

    fig.tight_layout()
    save_fig(fig, f"fig_spectra_slices_regime_{letter}_3D.png")


if True:
vlabarre's avatar
vlabarre committed
    kf, kb, ko, keta, kd = compute_kf_kb_ko_keta_kd(sim, tmin)
    data = sim.output.spectra.load_kzkh_mean(
        tmin, key_to_load=["A", "Khd", "Kz", "Khr"]
    )
vlabarre's avatar
vlabarre committed
    kh = data["kh_spectra"]
    kz = data["kz"]
    delta_kz = kz[1]
    delta_kh = kh[1]

    EA = data["A"]
    EKhd = data["Khd"]
    EKz = data["Kz"]
    EKhr = data["Khr"]
vlabarre's avatar
vlabarre committed
    Epolo = EKhd + EKz
    Etoro = EKhr
vlabarre's avatar
vlabarre committed
    EKh = EKhr + EKhd
vlabarre's avatar
vlabarre committed

    kstxts = ["kfkb", "kbko", "koketa"]
vlabarre's avatar
vlabarre committed
    kss = [0.5 * kb, (kb * ko) ** 0.5, (ko * keta) ** 0.5]
vlabarre's avatar
vlabarre committed

    for n in range(3):

        ks = kss[n]
        kstxt = kstxts[n]

        ikz = pos_closest_value(kz, ks)
        ikh = pos_closest_value(kh, ks)
vlabarre's avatar
vlabarre committed
        print(ks, kz[ikz], ikz)
vlabarre's avatar
vlabarre committed
        print(ks, kh[ikh], ikh)
vlabarre's avatar
vlabarre committed
        EA_kz = EA[:, ikh]
        Epolo_kz = Epolo[:, ikh]
        Etoro_kz = Etoro[:, ikh]
        EKz_kz = EKz[:, ikh]
        EKh_kz = EKh[:, ikh]
vlabarre's avatar
vlabarre committed
        EA_kh = EA[ikz, :]
        Epolo_kh = Epolo[ikz, :]
        Etoro_kh = Etoro[ikz, :]
        EKz_kh = EKz[ikz, :]
        EKh_kh = EKh[ikz, :]
vlabarre's avatar
vlabarre committed
        fig, ax = plt.subplots()  # figsize=(10, 3 * 4.5 / 2))
vlabarre's avatar
vlabarre committed

        if kstxt == "kfkb":
vlabarre's avatar
vlabarre committed
            kstitle = (
                r"$k_h=$"
                + f"{kh[ikh]/kb:.1f}"
                + r"$k_b ~ (--)$"
                + " or "
                + r"$k_z=$"
                + f"{kz[ikz]/kb:.1f}"
                + r"$k_b ~ (-)$"
            )
vlabarre's avatar
vlabarre committed
        elif kstxt == "kbko":
vlabarre's avatar
vlabarre committed
            kstitle = (
                r"$k_h=$"
                + f"{kh[ikh]/ko:.1f}"
                + r"$k_O ~ (--)$"
                + " or "
                + r"$k_z=$"
                + f"{kz[ikz]/ko:.1f}"
                + r"$k_O ~ (-)$"
            )
vlabarre's avatar
vlabarre committed
        else:
vlabarre's avatar
vlabarre committed
            kstitle = (
                r"$k_h=$"
                + f"{kh[ikh]/ko:.1f}"
                + r"$k_O ~ (--)$"
                + " or "
                + r"$k_z=$"
                + f"{kz[ikz]/ko:.1f}"
                + r"$k_O ~ (-)$"
            )
vlabarre's avatar
vlabarre committed

        ax.plot(
            kz,
vlabarre's avatar
vlabarre committed
            EA_kz * kz**coef_compensate,
vlabarre's avatar
vlabarre committed
            "b--",
            label=None,
        )
        ax.plot(
            kz,
vlabarre's avatar
vlabarre committed
            Epolo_kz * kz**coef_compensate,
vlabarre's avatar
vlabarre committed
            "m--",
            label=None,
        )
        ax.plot(
            kz,
vlabarre's avatar
vlabarre committed
            Etoro_kz * kz**coef_compensate,
vlabarre's avatar
vlabarre committed
            "r--",
            label=None,
        )
        """
        ax.plot(
            kz,
            EKz_kz * kz ** coef_compensate,
            "g--",
        )
        ax.plot(
            kz,
            EKh_kz * kz ** coef_compensate,
            "y--",
        )
        """

        ax.plot(
            kh,
vlabarre's avatar
vlabarre committed
            EA_kh * kh**coef_compensate,
vlabarre's avatar
vlabarre committed
            "b-",
vlabarre's avatar
vlabarre committed
            label=(r"$E_A$"),
vlabarre's avatar
vlabarre committed
        )
        ax.plot(
            kh,
vlabarre's avatar
vlabarre committed
            Epolo_kh * kh**coef_compensate,
vlabarre's avatar
vlabarre committed
            "m-",
vlabarre's avatar
vlabarre committed
            label=(r"$E_{polo}$"),
vlabarre's avatar
vlabarre committed
        )
        ax.plot(
            kh,
vlabarre's avatar
vlabarre committed
            Etoro_kh * kh**coef_compensate,
vlabarre's avatar
vlabarre committed
            "r-",
vlabarre's avatar
vlabarre committed
            label=(r"$E_{toro}$"),
vlabarre's avatar
vlabarre committed
        )
        """
        ax.plot(
            kh,
            EKz_kh * kh ** coef_compensate,
            "g-",
            label=(
                r"$E_{z}$"
            ),
        )
        ax.plot(
            kh,
            EKh_kh * kh ** coef_compensate,
            "y-",
            label=(
                r"$E_{h}$"
            ),
        )
        """
vlabarre's avatar
vlabarre committed

vlabarre's avatar
vlabarre committed
        if letter != "P":
            if kstxt == "kfkb":
                x = [delta_kz, kb]
vlabarre's avatar
vlabarre committed
                y = [
                    5e-3 * x[0] ** (-2 + coef_compensate),
                    5e-3 * x[1] ** (-2 + coef_compensate),
                ]
                ax.plot(x, y, "k-")
                ax.text(
                    (x[0] * x[1]) ** 0.5,
                    0.1 * 5e-3 * ((x[0] * x[1]) ** 0.5) ** (-2 + coef_compensate),
                    r"$k_h^{-2}$",
                    fontsize=16,
vlabarre's avatar
vlabarre committed
                )
vlabarre's avatar
vlabarre committed
                ax.legend(fontsize=10, loc="lower left")
vlabarre's avatar
vlabarre committed

            if kstxt == "kbko":
                x = [kb, ko]
vlabarre's avatar
vlabarre committed
                y = [
                    2e-4 * kb ** (-5 / 3 + coef_compensate),
                    2e-4 * ko ** (-5 / 3 + coef_compensate),
                ]
                ax.plot(x, y, "k-")
                ax.text(
                    ks,
                    0.1 * 2e-4 * ks ** (-5 / 3 + coef_compensate),
                    r"$k_h^{-5/3}$",
                    fontsize=16,
vlabarre's avatar
vlabarre committed
                )
vlabarre's avatar
vlabarre committed
                # ax.legend(fontsize=10, loc = "lower left")
vlabarre's avatar
vlabarre committed

            if kstxt == "koketa":
                x = [delta_kz, kb]
vlabarre's avatar
vlabarre committed
                y = [
                    1e-11 * delta_kz ** (1 + coef_compensate),
                    1e-11 * kb ** (1 + coef_compensate),
                ]
                ax.plot(x, y, "k-")
                ax.text(
                    (delta_kz * kb) ** 0.5,
                    0.2
                    * 1e-11
                    * ((delta_kz * kb) ** (0.5 * (1 + coef_compensate))),
                    r"$k_h^1$",
                    fontsize=16,
vlabarre's avatar
vlabarre committed
                )
vlabarre's avatar
vlabarre committed
                # ax.legend(fontsize=10, loc = "upper left")
vlabarre's avatar
vlabarre committed

            x = [delta_kz, ko]
vlabarre's avatar
vlabarre committed
            y = [1e-9 * delta_kz**coef_compensate, 1e-9 * ko**coef_compensate]
            ax.plot(x, y, "k--")
            ax.text(
                (x[0] * x[1]) ** 0.5,
                0.07 * 1e-9 * ((kb * ko) ** (0.5 * coef_compensate)),
                r"$k_z^0$",
                fontsize=16,
vlabarre's avatar
vlabarre committed
            ax.axvline(kb, color="k", linestyle="dotted")
            ax.text(1.1 * kb, 2e-10, r"$k_b$", fontsize=16)
            ax.axvline(ko, color="k", linestyle="dashed")
            ax.text(1.1 * ko, 2e-10, r"$k_O$", fontsize=16)
vlabarre's avatar
vlabarre committed
        # ax.legend(fontsize=10, loc = "lower right")
vlabarre's avatar
vlabarre committed

        ax.set_xscale("log")
        ax.set_yscale("log")

vlabarre's avatar
vlabarre committed
        ax.set_xlabel(r"$k_h, k_z$", fontsize=16)
vlabarre's avatar
vlabarre committed
        ax.set_ylabel(
            r"$E_i ~ k_h^{5/3} ~ (-)$  or $E_i ~ k_z^{5/3} (--)$", fontsize=16
        )
        # ax.text(1.1*kb, 1e-10, r"$k_b$", fontsize=16)
        # ax.text(1.1*ko, 1e-10, r"$k_O$", fontsize=16)
vlabarre's avatar
vlabarre committed
        ax.set_xlim([kh[1], 0.8 * max(kh)])
vlabarre's avatar
vlabarre committed
        ax.set_ylim(bottom=1e-10, top=1e0)

vlabarre's avatar
vlabarre committed
        # plt.legend(loc=2, prop={'size': 10})
vlabarre's avatar
vlabarre committed
        ax.set_title(kstitle, fontsize=12)
vlabarre's avatar
vlabarre committed

        fig.tight_layout()
        save_fig(fig, f"fig_spectra_slices_regime_{letter}_{kstxt}.png")

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