Skip to content
Snippets Groups Projects
save_nonlinear_regimes.py 8.44 KiB
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"


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


sim_proj = get_sim(letter, proj=True)
print(sim_proj.params.projection)
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]


# Standard Navier-Stokes
t_start, t_last = times_start_last_from_path(path)
tmin = t_last - 2
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"]
params = load_params_simul(path)
nh = nx = params.oper.nx
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:
    # List all groups
    # print("Keys: %s" % f.keys())
    # a_group_key = list(f.keys())[0]
    # 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
    xb = np.linspace(kh[1], kb, 50, endpoint=True)
    for ax in [ax0]:
        ax.plot(xb, np.sqrt(kb**2 - xb**2), color="c", linestyle="dashed")
        a = 3
        xa = np.linspace(kh[1], a**1.5 * ko, 50, endpoint=True)
        ax.plot(
            xa,
            xa * np.sqrt((a**1.5 * ko / xa) ** 0.8 - 1),
            linestyle="dotted",
            color="c",
        )
        a = 1 / 3
        xa = np.linspace(kh[1], a**1.5 * ko, 50, endpoint=True)
        ax.plot(
            xa,
            xa * np.sqrt((a**1.5 * ko / xa) ** 0.8 - 1),
            linestyle="dotted",
            color="c",
        )

    Ee = np.sum(Ee, axis=2)
    E = np.sum(E, axis=2)
    cs0 = ax0.pcolormesh(
        kh,
        kz,
        np.log10(delta_omega_emp / omega_disp),
        cmap=cm,
        vmin=-0.5,
        vmax=1.5,
        shading="nearest",
    )


# Without vortical modes
t_start, t_last = times_start_last_from_path(path_proj)
tmin = t_last - 2
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"]
params = load_params_simul(path)
nh = nx = params.oper.nx
proj = params.projection
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:
    # List all groups
    # print("Keys: %s" % f.keys())
    # a_group_key = list(f.keys())[0]
    # 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
    xb = np.linspace(kh[1], kb, 50, endpoint=True)
    for ax in [ax1]:
        ax.plot(xb, np.sqrt(kb**2 - xb**2), color="c", linestyle="dashed")
        a = 3
        xa = np.linspace(kh[1], a**1.5 * ko, 50, endpoint=True)
        ax.plot(
            xa,
            xa * np.sqrt((a**1.5 * ko / xa) ** 0.8 - 1),
            linestyle="dotted",
            color="c",
        )
        a = 1 / 3
        xa = np.linspace(kh[1], a**1.5 * ko, 50, endpoint=True)
        ax.plot(
            xa,
            xa * np.sqrt((a**1.5 * ko / xa) ** 0.8 - 1),
            linestyle="dotted",
            color="c",
        )

    Ee = np.sum(Ee, axis=2)
    E = np.sum(E, axis=2)
    cs1 = ax1.pcolormesh(
        kh,
        kz,
        np.log10(delta_omega_emp / omega_disp),
        cmap=cm,
        vmin=-0.5,
        vmax=1.5,
        shading="nearest",
    )

for ax in [ax0, ax1]:
    ax.plot([0, 1], [1, 1], "k-")
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 3])


# Forcing
delta_kh = kh[1]
delta_kz = kz[1]
kf_min = sim.params.forcing.nkmin_forcing * delta_kz
kf_max = sim.params.forcing.nkmax_forcing * delta_kz
angle = sim.params.forcing.tcrandom_anisotropic.angle
delta_angle = sim.params.forcing.tcrandom_anisotropic.delta_angle
for ax in [ax0, ax1]:
    ax.set_xlim([kh[1], 0.8 * max(kh)])
    ax.set_ylim([kz[1], 0.8 * max(kh)])
    # ax.plot([kh[1], max(kh)], [kh[1], max(kh)], "k-")
    # Forcing
    ax.add_patch(
        patches.Arc(
            xy=(0, 0),
            width=2 * kf_max,
            height=2 * kf_max,
            angle=0,
            theta1=90.0 - degrees(angle) - 0.5 * degrees(delta_angle),
            theta2=90.0 - degrees(angle) + 0.5 * degrees(delta_angle),
            linestyle="-",
            color="orange",
            linewidth=1,
        )
    )
    ax.add_patch(
        patches.Arc(
            xy=(0, 0),
            width=2 * kf_min,
            height=2 * kf_min,
            angle=0,
            theta1=90.0 - degrees(angle) - 0.5 * degrees(delta_angle),
            theta2=90.0 - degrees(angle) + 0.5 * degrees(delta_angle),
            linestyle="-",
            color="orange",
            linewidth=1,
        )
    )
    ax.plot(
        [
            kf_min * np.sin(angle - 0.5 * delta_angle),
            kf_max * np.sin(angle - 0.5 * delta_angle),
        ],
        [
            kf_min * np.cos(angle - 0.5 * delta_angle),
            kf_max * np.cos(angle - 0.5 * delta_angle),
        ],
        linestyle="-",
        color="orange",
        linewidth=1,
    )
    ax.plot(
        [
            kf_min * np.sin(angle + 0.5 * delta_angle),
            kf_max * np.sin(angle + 0.5 * delta_angle),
        ],
        [
            kf_min * np.cos(angle + 0.5 * delta_angle),
            kf_max * np.cos(angle + 0.5 * delta_angle),
        ],
        linestyle="-",
        color="orange",
        linewidth=1,
    )


for ax in [ax0, ax1]:
    ax.set_xlabel(r"$k_h$", fontsize=20)
    ax.set_xticks([20, 40, 60, 80, 100])
    ax.set_xticklabels(
        [r"$20$", r"$40$", r"$60$", r"$80$", r"$100$"], fontsize=14
    )
    ax.set_yticks([20, 40, 60, 80, 100])
    ax.set_yticklabels(
        [r"$20$", r"$40$", r"$60$", r"$80$", r"$100$"], fontsize=14
    )

for ax in [ax1]:
    ax.set_yticklabels([])

ax0.set_ylabel(r"$k_z$", fontsize=20)


ax0.set_title(r"Standard Navier-Stokes" + "\n" + r"$(a)$", fontsize=20)
ax1.set_title(r"Without vortical modes" + "\n" + r"$(b)$", fontsize=20)


fig.tight_layout()

# fig.subplots_adjust(right=0.85)
cbar_ax = fig.add_axes([0.88, 0.17, 0.02, 0.65])
cbar = fig.colorbar(cs1, cax=cbar_ax)
cbar.set_ticks([-0.5, 0.0, 0.5, 1.0, 1.5])
cbar.set_ticklabels([r"$-0.5$", r"$0$", r"$0.5$", r"$1$", r"$1.5$"], fontsize=14)
cbar.ax.set_ylabel(
    r"$\log_{10}\left(\delta \omega/ \omega_{\boldsymbol{k}}\right)$",
    fontsize=20,
)


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


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