Newer
Older
import matplotlib.cm
import matplotlib.pyplot as plt
import numpy as np
)
cm = matplotlib.cm.get_cmap("inferno", 100)
# Latex
plt.rcParams["text.usetex"] = True
plt.rcParams["text.latex.preamble"] = r"\usepackage{bm}"
def plot_spectra_vs_omegak(sim, ax, key="Ee"):
N = sim.params.N
t_start, t_last = sim.output.print_stdout.get_times_start_last()
tmin = t_last - 2.0
mean_values = sim.output.get_mean_values(tmin=tmin, customize=customize)
Uh = np.sqrt(mean_values["Uh2"])
path = sim.params.path_run
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"][:]
Epot = f["spectrum_A"][:]
EKz = f["spectrum_K"][:] - f["spectrum_Khd"][:] - f["spectrum_Khr"][:]
Epolo = f["spectrum_Khd"][:] + EKz
Etoro = f["spectrum_Khr"][:]
delta_kh = kh[1]
delta_kz = kz[1]
if key == "Epot":
spectrum = Epot
elif key == "EK":
spectrum = Epolo + Etoro
elif key == "Epolo":
spectrum = Epolo
elif key == "Ee":
spectrum = 2 * np.minimum(Epot, Epolo)
elif key == "Ed":
spectrum = Epot + Epolo - 2 * np.minimum(Epot, Epolo)
elif key == "Etoro":
spectrum = Etoro
elif key == "Etot":
spectrum = Epolo + Etoro + Epot
else:
print(f"Don't know key: {key} \n")
exit
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
omega_disp_grid = np.linspace(0, N, 25, endpoint=True)
OMEGA_DISP_GRID, OMEGA_GRID = np.meshgrid(omega_disp_grid, omegas)
SPECTRUM = np.zeros(OMEGA_DISP_GRID.shape)
for iod in range(len(omega_disp_grid) - 1):
for io in range(len(omegas)):
omega_disp_min = omega_disp_grid[iod]
omega_disp_max = omega_disp_grid[iod + 1]
IOD = np.where(
np.logical_and(
omega_disp >= omega_disp_min, omega_disp < omega_disp_max
)
)
for i in range(len(IOD[0])):
SPECTRUM[io, iod] += spectrum[IOD[0][i], IOD[1][i], io]
cs = ax.pcolormesh(
OMEGA_DISP_GRID / N,
OMEGA_GRID / N,
np.log10(SPECTRUM),
cmap=cm,
vmin=-6,
vmax=-3,
shading="nearest",
ax.set_xlabel(r"$\sin \theta_{\boldsymbol{k}}$", fontsize=20)
ax.set_xlim([0, 1])
ax.set_xticks([0, 0.25, 0.5, 0.75, 1])
ax.set_xticklabels(
[r"$0$", r"$0.25$", r"$0.5$", r"$0.75$", r"$1$"], fontsize=14
ax.set_ylabel(r"$\omega / N$", fontsize=20)
ax.set_ylim([0, 3])
ax.set_yticks([0, 1, 2, 3])
ax.set_yticklabels([r"$0$", r"$1$", r"$2$", r"$3$"], fontsize=14)
# 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}"
ncols=2, nrows=3, figsize=(10, 3 * 3 * 4.5 / 4), constrained_layout=True
css = [None for i in range(6)]
ax0 = axes[0, 0]
ax1 = axes[0, 1]
ax2 = axes[1, 0]
ax3 = axes[1, 1]
ax4 = axes[2, 0]
ax5 = axes[2, 1]
sim = get_sim("S1", proj=False)
css[0] = plot_spectra_vs_omegak(sim, ax0, key="Etot")
sim_proj = get_sim("S1", proj=True)
css[1] = plot_spectra_vs_omegak(sim_proj, ax1, key="Etot")
sim = get_sim("S2", proj=False)
css[2] = plot_spectra_vs_omegak(sim, ax2, key="Etot")
sim_proj = get_sim("S2", proj=True)
css[3] = plot_spectra_vs_omegak(sim_proj, ax3, key="Etot")
sim = get_sim("S3", proj=False)
css[4] = plot_spectra_vs_omegak(sim, ax4, key="Etot")
sim_proj = get_sim("S3", proj=True)
css[5] = plot_spectra_vs_omegak(sim_proj, ax5, key="Etot")
for ax in [ax1, ax3, ax5]:
ax.set_ylabel("")
ax.set_yticklabels([])
for ax in [ax0, ax1, ax2, ax3]:
ax.set_xlabel("")
ax.set_xticklabels([])
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)
ax2.set_title(r"$\rm (c)$", fontsize=20)
ax3.set_title(r"$\rm (d)$", fontsize=20)
ax4.set_title(r"$\rm (e)$", fontsize=20)
ax5.set_title(r"$\rm (f)$", fontsize=20)
for ax in [ax0, ax1, ax2, ax3, ax4, ax5]:
ax.plot([0, 1], [0, 1], "w--")
# fig.subplots_adjust(right=0.85)
cbar_ax = fig.add_axes([0.88, 0.19, 0.02, 0.63])
cbar = fig.colorbar(css[0], cax=cbar_ax)
cbar.set_ticks([-3, -4, -5, -6])
cbar.set_ticklabels([r"$-3$", r"$-4$", r"$-5$", r"$-6$"], fontsize=16)
cbar.ax.set_ylabel(
r"$\log_{10} E(\sin \theta_{\boldsymbol{k}},\omega)$", fontsize=20
fig.subplots_adjust(right=0.85, wspace=0.05, hspace=0.15)
save_fig(fig, f"figure14.png")