Skip to content
Snippets Groups Projects
save_nonlinear_regimes.py 8.63 KiB
Newer Older
import glob
import os
import re
vlabarre's avatar
vlabarre committed
import sys
from pathlib import Path
vlabarre's avatar
vlabarre committed
import h5py
import matplotlib.cm
import matplotlib.pyplot as plt
import numpy as np
vlabarre's avatar
vlabarre committed
from matplotlib.collections import LineCollection
from util_simuls_regimes import get_sim
vlabarre's avatar
vlabarre committed

from fluiddyn.util import modification_date
from fluidsim import load
from fluidsim.util import load_params_simul, times_start_last_from_path
vlabarre's avatar
vlabarre committed

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

# Latex
vlabarre's avatar
vlabarre committed
plt.rcParams["text.usetex"] = True
plt.rcParams["text.latex.preamble"] = r"\usepackage{bm}"
vlabarre's avatar
vlabarre committed
cm = matplotlib.cm.get_cmap("inferno", 100)
vlabarre's avatar
vlabarre committed


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

if letter not in "DLOWPU":
vlabarre's avatar
vlabarre committed
    letter = "L"
vlabarre's avatar
vlabarre committed


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]

vlabarre's avatar
vlabarre committed
# 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}"
vlabarre's avatar
vlabarre committed
fig, axes = plt.subplots(
    ncols=2, nrows=3, figsize=(10, 3 * 3 * 4.5 / 4), constrained_layout=True
)
vlabarre's avatar
vlabarre committed

ax0 = axes[0, 0]
ax1 = axes[0, 1]
ax2 = axes[1, 0]
ax3 = axes[1, 1]
vlabarre's avatar
vlabarre committed
ax4 = axes[2, 0]
ax5 = axes[2, 1]
vlabarre's avatar
vlabarre committed
# Standard Navier-Stokes
vlabarre's avatar
vlabarre committed
t_start, t_last = times_start_last_from_path(path)
tmin = t_last - 2
vlabarre's avatar
vlabarre committed
kf, kb, ko, keta, kd = compute_kf_kb_ko_keta_kd(sim, tmin)
vlabarre's avatar
vlabarre committed
mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
vlabarre's avatar
vlabarre committed
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
vlabarre's avatar
vlabarre committed
    # print("Keys: %s" % f.keys())
    # a_group_key = list(f.keys())[0]
vlabarre's avatar
vlabarre committed
    # 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
vlabarre's avatar
vlabarre committed
    Ee = 2 * np.minimum(EA, Epolo)
vlabarre's avatar
vlabarre committed
    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)
vlabarre's avatar
vlabarre committed
    for ax in [ax2, ax4]:
vlabarre's avatar
vlabarre committed
        ax.plot(xb, np.sqrt(kb**2 - xb**2), color="k", linestyle="dotted")
vlabarre's avatar
vlabarre committed
        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="dashed",
            color="gray",
        )
        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="gray",
        )
vlabarre's avatar
vlabarre committed
    Ee = np.sum(Ee, axis=2)
    cs0 = ax0.scatter(
vlabarre's avatar
vlabarre committed
        omega_disp.flatten() / N,
        omega_emp.flatten() / N,
        c=np.log10(Ee.flatten()),
vlabarre's avatar
vlabarre committed
        cmap=cm,
        vmin=-7.0,
        vmax=-3.0,
vlabarre's avatar
vlabarre committed
        # shading="nearest",
vlabarre's avatar
vlabarre committed
    )
    cs2 = ax2.pcolormesh(
        kh,
        kz,
vlabarre's avatar
vlabarre committed
        np.abs((omega_emp - omega_disp)) / N,
        cmap=cm,
        vmin=0.0,
        vmax=1,
        shading="nearest",
    )
vlabarre's avatar
vlabarre committed
    cs4 = ax4.pcolormesh(
vlabarre's avatar
vlabarre committed
        kh,
        kz,
vlabarre's avatar
vlabarre committed
        np.log10(delta_omega_emp / omega_disp),
vlabarre's avatar
vlabarre committed
        cmap=cm,
        vmin=-0.5,
        vmax=1.5,
        shading="nearest",
    )


vlabarre's avatar
vlabarre committed
# Without vortical modes
vlabarre's avatar
vlabarre committed
t_start, t_last = times_start_last_from_path(path_proj)
tmin = t_last - 2
k_old = kb
vlabarre's avatar
vlabarre committed
kf, kb, ko, keta, kd = compute_kf_kb_ko_keta_kd(sim_proj, tmin)
vlabarre's avatar
vlabarre committed
mean_values = sim_proj.output.get_mean_values(tmin=tmin, customize=customize)
vlabarre's avatar
vlabarre committed
print(k_old, "           ", kb)
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
vlabarre's avatar
vlabarre committed
    # print("Keys: %s" % f.keys())
    # a_group_key = list(f.keys())[0]
vlabarre's avatar
vlabarre committed
    # 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
vlabarre's avatar
vlabarre committed
    Ee = 2 * np.minimum(EA, Epolo)
vlabarre's avatar
vlabarre committed
    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)
vlabarre's avatar
vlabarre committed
    for ax in [ax1, ax3]:
vlabarre's avatar
vlabarre committed
        ax.plot(xb, np.sqrt(kb**2 - xb**2), color="k", linestyle="dotted")
vlabarre's avatar
vlabarre committed
        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="dashed",
            color="gray",
        )
        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="gray",
        )
vlabarre's avatar
vlabarre committed
    Ee = np.sum(Ee, axis=2)
    cs1 = ax1.scatter(
vlabarre's avatar
vlabarre committed
        omega_disp.flatten() / N,
        omega_emp.flatten() / N,
        c=np.log10(Ee.flatten()),
vlabarre's avatar
vlabarre committed
        cmap=cm,
vlabarre's avatar
vlabarre committed
        # shading="nearest",
vlabarre's avatar
vlabarre committed
    )
    cs3 = ax3.pcolormesh(
        kh,
        kz,
vlabarre's avatar
vlabarre committed
        np.abs((omega_emp - omega_disp)) / N,
        cmap=cm,
        vmin=0.0,
        vmax=1,
        shading="nearest",
    )
vlabarre's avatar
vlabarre committed
    cs5 = ax5.pcolormesh(
vlabarre's avatar
vlabarre committed
        kh,
        kz,
vlabarre's avatar
vlabarre committed
        np.log10(delta_omega_emp / omega_disp),
vlabarre's avatar
vlabarre committed
        cmap=cm,
        vmin=-0.5,
        vmax=1.5,
        shading="nearest",
    )

for ax in [ax0, ax1]:
vlabarre's avatar
vlabarre committed
    ax.plot([0, 1], [0, 1], "k-")
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 2])
vlabarre's avatar
vlabarre committed
for ax in [ax2, ax3, ax4, ax5]:
    ax.set_xlim([kh[1], 0.8 * max(kh)])
    ax.set_ylim([kz[1], 0.8 * max(kh)])
    # ax.set_xscale('log')
    # ax.set_yscale('log')
vlabarre's avatar
vlabarre committed
    ax.plot([kh[1], max(kh)], [kh[1], max(kh)], "k-")
vlabarre's avatar
vlabarre committed
    ax.plot(
        [kh[1], kf * 0.35],
        [kf * (1 - 0.25**2) ** 0.5, kf * (1 - 0.25**2) ** 0.5],
        linestyle="--",
        color="orange",
    )
    ax.plot(
        [kf * 0.35, kf * 0.35],
        [kz[1], kf * (1 - 0.25**2) ** 0.5],
        linestyle="--",
        color="orange",
    )
    # ax.plot(xo, kb*np.sqrt(xo/delta_kh), "m--")
for ax in [ax2, ax3]:
vlabarre's avatar
vlabarre committed
    ax.set_xticklabels([])
vlabarre's avatar
vlabarre committed
for ax in [ax1, ax3, ax5]:
vlabarre's avatar
vlabarre committed
    ax.set_yticklabels([])
for ax in [ax2, ax3, ax4, ax5]:
vlabarre's avatar
vlabarre committed
    ax.set_xlabel(r"$k_h$", fontsize=16)
for ax in [ax2, ax4]:
vlabarre's avatar
vlabarre committed
    ax.set_ylabel(r"$k_z$", fontsize=16)
for ax in [ax0, ax1]:
    ax.set_xlabel(r"$\omega_{\bm{k}}/N$", fontsize=16)
ax0.set_ylabel(r"$\omega_{emp, \bm{k}}/N$", fontsize=16)
vlabarre's avatar
vlabarre committed
ax0.set_title(r"Standard Navier-Stokes" + "\n" + r"$(a)$", fontsize=16)
vlabarre's avatar
vlabarre committed
ax1.set_title(r"Without vortical modes" + "\n" + r"$(b)$", fontsize=16)
ax2.set_title(r"$(c)$", fontsize=16)
ax3.set_title(r"$(d)$", fontsize=16)
ax4.set_title(r"$(e)$", fontsize=16)
ax5.set_title(r"$(f)$", fontsize=16)
vlabarre's avatar
vlabarre committed
fig.tight_layout()
vlabarre's avatar
vlabarre committed
# fig.subplots_adjust(right=0.85)
vlabarre's avatar
vlabarre committed
cbar_ax = fig.add_axes([0.88, 0.71, 0.02, 0.23])
vlabarre's avatar
vlabarre committed
cbar = fig.colorbar(cs0, cax=cbar_ax)
cbar.set_ticks([-7, -6, -5, -4, -3])
vlabarre's avatar
vlabarre committed
cbar.ax.set_ylabel(r"$\log_{10} E_{equi}$", fontsize=16)
vlabarre's avatar
vlabarre committed
# fig.subplots_adjust(right=0.85)
vlabarre's avatar
vlabarre committed
cbar_ax = fig.add_axes([0.88, 0.385, 0.02, 0.23])
cbar = fig.colorbar(cs2, cax=cbar_ax)
vlabarre's avatar
vlabarre committed
cbar.set_ticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0])
vlabarre's avatar
vlabarre committed
cbar.ax.set_ylabel(r"$|\omega_{emp, \bm{k}}- \omega_{\bm{k}}|/N$", fontsize=16)
vlabarre's avatar
vlabarre committed
# fig.subplots_adjust(right=0.85)
vlabarre's avatar
vlabarre committed
cbar_ax = fig.add_axes([0.88, 0.06, 0.02, 0.23])
vlabarre's avatar
vlabarre committed
cbar = fig.colorbar(cs4, cax=cbar_ax)
vlabarre's avatar
vlabarre committed
cbar.set_ticks([-0.5, 0.0, 0.5, 1.0, 1.5])
vlabarre's avatar
vlabarre committed
cbar.ax.set_ylabel(
    r"$\log_{10}\left(\delta \omega_{\bm{k}}/ \omega_{\bm{k}}\right)$",
    fontsize=16,
)
fig.subplots_adjust(right=0.85, wspace=0.1, hspace=0.4)

vlabarre's avatar
vlabarre committed
save_fig(fig, f"fig_nonlinear_regimes_{letter}.png")
vlabarre's avatar
vlabarre committed


if __name__ == "__main__":
vlabarre's avatar
vlabarre committed
    plt.show()