Skip to content
Snippets Groups Projects
save_spectra_1strat.py 3.88 KiB
Newer Older
vlabarre's avatar
vlabarre committed
import matplotlib.cm
Pierre Augier's avatar
Pierre Augier committed
from math import sqrt
import matplotlib.pyplot as plt
vlabarre's avatar
vlabarre committed

plt.rcParams["text.usetex"] = True
Pierre Augier's avatar
Pierre Augier committed

from fluidsim import load

vlabarre's avatar
vlabarre committed
import numpy as np

from util import (
    save_fig,
    couples320,
    get_path_finer_resol,
    customize,
    add_letters,
)
vlabarre's avatar
vlabarre committed
cm = matplotlib.cm.get_cmap("inferno", 100)


vlabarre's avatar
vlabarre committed
N = 20
Pierre Augier's avatar
Pierre Augier committed

R_is = sorted(Ri for (Nloc, Ri) in couples320 if Nloc == N)

paths = [get_path_finer_resol(N, R_i) for R_i in R_is]
paths = [path for path in paths if path is not None]

vlabarre's avatar
vlabarre committed
fig, axes = plt.subplots(ncols=2, nrows=3, figsize=(10, 3 * 3 * 4.5 / 4))
vlabarre's avatar
vlabarre committed

vlabarre's avatar
vlabarre committed
ax0 = axes[0, 0]
ax1 = axes[0, 1]
ax2 = axes[1, 0]
ax3 = axes[1, 1]
ax4 = axes[2, 0]
ax5 = axes[2, 1]
Pierre Augier's avatar
Pierre Augier committed


for path, R_i in zip(paths, R_is):
    if R_i == 80:
        continue
    sim = load(path, hide_stdout=True)
vlabarre's avatar
vlabarre committed
    mean_values = sim.output.get_mean_values(tmin="t_last-2", customize=customize)
Pierre Augier's avatar
Pierre Augier committed
    R2 = mean_values["R2"]
vlabarre's avatar
vlabarre committed
    Uh2 = mean_values["Uh2"]
    kb = N / Uh2**0.5
    epsK = mean_values["epsK"]
vlabarre's avatar
vlabarre committed
    ko = (N**3 / epsK) ** 0.5
Pierre Augier's avatar
Pierre Augier committed
    if R2 < 10:
        R2_str = f"{R2:.1f}"
    else:
        R2_str = f"{R2:.0f}"

    Uh2 = mean_values["Uh2"]
    kb = N / sqrt(Uh2)
    epsK = mean_values["epsK"]
    ko = sqrt(N**3 / epsK)

    t_start, t_last = sim.output.print_stdout.get_times_start_last()
    tmin = max(10, min(t_start + 2, t_last - 1))
Pierre Augier's avatar
Pierre Augier committed

vlabarre's avatar
vlabarre committed
    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_kh = kh[1]
    delta_kz = kz[1]
Pierre Augier's avatar
Pierre Augier committed
    KH, KZ = np.meshgrid(kh, kz)
vlabarre's avatar
vlabarre committed

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

vlabarre's avatar
vlabarre committed
    Epolo_vs_kh = np.sum(Epolo, axis=0) * delta_kz
    Epolo_vs_kz = np.sum(Epolo, axis=1) * delta_kh

    Etoro_vs_kh = np.sum(Etoro, axis=0) * delta_kz
    Etoro_vs_kz = np.sum(Etoro, axis=1) * delta_kh
vlabarre's avatar
vlabarre committed
    EA_vs_kh = np.sum(EA, axis=0) * delta_kz
    EA_vs_kz = np.sum(EA, axis=1) * delta_kh
Pierre Augier's avatar
Pierre Augier committed

vlabarre's avatar
vlabarre committed
    color = cm((np.log10(R2) - (-1)) / (3 - (-1)))
vlabarre's avatar
vlabarre committed
    # Potential
vlabarre's avatar
vlabarre committed
    (line,) = ax0.plot(kh / ko, EA_vs_kh[:] * kh[:] ** (5 / 3), color=color)
    ax0.axvline(kb / ko, linestyle="dotted", color=color)
vlabarre's avatar
vlabarre committed
    ax1.plot(kz / ko, EA_vs_kz[:] * kz[:] ** (5 / 3), color=color)
vlabarre's avatar
vlabarre committed
    ax1.axvline(kb / ko, linestyle="dotted", color=color)
vlabarre's avatar
vlabarre committed
    # Poloidal
vlabarre's avatar
vlabarre committed
    ax2.plot(kh / ko, Epolo_vs_kh[:] * kh[:] ** (5 / 3), color=color)
vlabarre's avatar
vlabarre committed
    ax2.axvline(kb / ko, linestyle="dotted", color=color)
vlabarre's avatar
vlabarre committed
    ax3.plot(kz / ko, Epolo_vs_kz[:] * kz[:] ** (5 / 3), color=color)
vlabarre's avatar
vlabarre committed
    ax3.axvline(kb / ko, linestyle="dotted", color=color)
vlabarre's avatar
vlabarre committed

    # Toroidal
vlabarre's avatar
vlabarre committed
    ax4.plot(
        kh / ko,
        Etoro_vs_kh[:] * kh[:] ** (5 / 3),
        color=color,
        label=rf"$\mathcal{{R}}={R2_str}$",
    )
    ax4.axvline(kb / ko, linestyle="dotted", color=color)
vlabarre's avatar
vlabarre committed
    ax5.plot(kz / ko, Etoro_vs_kz[:] * kz[:] ** (5 / 3), color=color)
vlabarre's avatar
vlabarre committed
    ax5.axvline(kb / ko, linestyle="dotted", color=color)
vlabarre's avatar
vlabarre committed
ax0.set_xlabel(r"$k_h / k_{\rm O}$", fontsize=14)
ax1.set_xlabel(r"$k_z / k_{\rm O}$", fontsize=14)
ax2.set_xlabel(r"$k_h / k_{\rm O}$", fontsize=14)
ax3.set_xlabel(r"$k_z / k_{\rm O}$", fontsize=14)
ax4.set_xlabel(r"$k_h / k_{\rm O}$", fontsize=14)
ax5.set_xlabel(r"$k_z / k_{\rm O}$", fontsize=14)
ax0.set_ylabel(r"$E_{\rm A}(k_h) {k_h}^{5/3}$", fontsize=14)
ax1.set_ylabel(r"$E_{\rm A}(k_z) {k_z}^{5/3}$", fontsize=14)
ax2.set_ylabel(r"$E_{\rm polo}(k_h) {k_h}^{5/3}$", fontsize=14)
ax3.set_ylabel(r"$E_{\rm polo}(k_z) {k_z}^{5/3}$", fontsize=14)
ax4.set_ylabel(r"$E_{\rm toro}(k_h) {k_h}^{5/3}$", fontsize=14)
ax5.set_ylabel(r"$E_{\rm toro}(k_z) {k_z}^{5/3}$", fontsize=14)
vlabarre's avatar
vlabarre committed
ax4.legend(fontsize=10)
vlabarre's avatar
vlabarre committed
for ax in [ax0, ax1, ax2, ax3, ax4, ax5]:
Pierre Augier's avatar
Pierre Augier committed
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_ylim(bottom=1e-2, top=4)

vlabarre's avatar
vlabarre committed
for ax in [ax0, ax1, ax2, ax3]:
    ax.set_xlabel(None)
    ax.set_xticks([])

Pierre Augier's avatar
Pierre Augier committed
add_letters(fig, "ab")
Pierre Augier's avatar
Pierre Augier committed
fig.tight_layout()
save_fig(fig, "fig_spectra_1strat.png")

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