"""Spectra output (:mod:`fluidsim.solvers.ns2d.strat.output.spectra`)
=====================================================================
.. autoclass:: SpectraNS2DStrat
:members:
:private-members:
"""
import h5py
import numpy as np
from fluidsim.base.output.spectra import Spectra
[docs]
class SpectraNS2DStrat(Spectra):
"""Save and plot spectra."""
[docs]
def compute(self):
"""compute the values at one time."""
# energy_fft = self.output.compute_energy_fft()
energyK_fft, energyA_fft = self.output.compute_energies_fft()
energy_fft = energyK_fft + energyA_fft
energyK_ux_fft, energyK_uy_fft = self.output.compute_energies2_fft()
energyK, energyA, energyK_ux = self.output.compute_energies()
# Compute the kinetic energy spectra 1D for the two velocity components
# and two directions
spectrum1Dkx_EK_ux, spectrum1Dky_EK_ux = self.spectra1D_from_fft(
energyK_ux_fft
)
spectrum1Dkx_EK_uy, spectrum1Dky_EK_uy = self.spectra1D_from_fft(
energyK_uy_fft
)
spectrum1Dkx_EK, spectrum1Dky_EK = self.spectra1D_from_fft(energyK_fft)
# Compute the potential energy spectra 1D two directions
spectrum1Dkx_EA, spectrum1Dky_EA = self.spectra1D_from_fft(energyA_fft)
# Compute the total energy spectra 1D
spectrum1Dkx_E, spectrum1Dky_E = self.spectra1D_from_fft(energy_fft)
# Dictionary with the 1D kinetic energy spectra
dict_spectra1D = {
"spectrum1Dkx_EK_ux": spectrum1Dkx_EK_ux,
"spectrum1Dky_EK_ux": spectrum1Dky_EK_ux,
"spectrum1Dkx_EK_uy": spectrum1Dkx_EK_uy,
"spectrum1Dky_EK_uy": spectrum1Dky_EK_uy,
"spectrum1Dkx_EK": spectrum1Dkx_EK,
"spectrum1Dky_EK": spectrum1Dky_EK,
"spectrum1Dkx_EA": spectrum1Dkx_EA,
"spectrum1Dky_EA": spectrum1Dky_EA,
"spectrum1Dkx_E": spectrum1Dkx_E,
"spectrum1Dky_E": spectrum1Dky_E,
}
# compute the kinetic energy spectra 2D
spectrum2D_E = self.spectrum2D_from_fft(energy_fft)
spectrum2D_EK = self.spectrum2D_from_fft(energyK_fft)
spectrum2D_EK_ux = self.spectrum2D_from_fft(energyK_ux_fft)
spectrum2D_EK_uy = self.spectrum2D_from_fft(energyK_uy_fft)
spectrum2D_EA = self.spectrum2D_from_fft(energyA_fft)
dict_spectra2D = {
"spectrum2D_EK_ux": spectrum2D_EK_ux,
"spectrum2D_EK_uy": spectrum2D_EK_uy,
"spectrum2D_EK": spectrum2D_EK,
"spectrum2D_EA": spectrum2D_EA,
"spectrum2D_E": spectrum2D_E,
}
return dict_spectra1D, dict_spectra2D
def _online_plot_saving(self, dict_spectra1D, dict_spectra2D):
if (
self.nx == self.params.oper.ny
and self.params.oper.Lx == self.params.oper.Ly
):
spectrum2D_EK = dict_spectra2D["spectrum2D_EK"]
spectrum2D_EA = dict_spectra2D["spectrum2D_EA"]
spectrum2D = spectrum2D_EK + spectrum2D_EA
khE = self.oper.khE
coef_norm = khE ** (3.0)
self.ax.loglog(khE, spectrum2D * coef_norm, "k")
lin_inf, lin_sup = self.ax.get_ylim()
if lin_inf < 10e-6:
lin_inf = 10e-6
self.ax.set_ylim([lin_inf, lin_sup])
else:
print(
"you need to implement the ploting "
"of the spectra for this case"
)
def load1d_means(
self, tmin=0, tmax=None, delta_t=2, versus_kx=True, versus_ky=True
):
means = {}
with h5py.File(self.path_file1D, "r") as h5file:
# Open data from file
dset_times = h5file["times"]
dset_kxE = h5file["kxE"]
dset_kyE = h5file["kyE"]
times = dset_times[...]
means["times"] = times
means["kx"] = dset_kxE[...]
means["ky"] = dset_kyE[...]
if tmin is None:
tmin = 0
if tmax is None:
tmax = np.max(times)
# Compute average from tmin and tmax for plot
delta_t_save = np.mean(times[1:] - times[0:-1])
delta_i_plot = int(np.round(delta_t / delta_t_save))
delta_t = delta_t_save * delta_i_plot
if delta_i_plot == 0:
delta_i_plot = 1
imin_plot = np.argmin(abs(times - tmin))
imax_plot = np.argmin(abs(times - tmax))
tmin_plot = times[imin_plot]
tmax_plot = times[imax_plot]
if versus_kx:
# Open data set 1D spectra
dset_spectrum1Dkx_EA = h5file["spectrum1Dkx_EA"]
dset_spectrum1Dkx_EK = h5file["spectrum1Dkx_EK"]
dset_spectrum1Dkx_E = h5file["spectrum1Dkx_E"]
# Average in time between tmin and tmax
means["E_kx"] = (
dset_spectrum1Dkx_E[imin_plot : imax_plot + 1]
).mean(0)
means["EK_kx"] = (
dset_spectrum1Dkx_EK[imin_plot : imax_plot + 1]
).mean(0)
means["EA_kx"] = (
dset_spectrum1Dkx_EA[imin_plot : imax_plot + 1]
).mean(0)
if versus_ky:
# Open data set 1D spectra
dset_spectrum1Dky_EA = h5file["spectrum1Dky_EA"]
dset_spectrum1Dky_EK = h5file["spectrum1Dky_EK"]
dset_spectrum1Dky_E = h5file["spectrum1Dky_E"]
# Average in time between tmin and tmax
means["E_ky"] = (
dset_spectrum1Dky_E[imin_plot : imax_plot + 1]
).mean(0)
means["EK_ky"] = (
dset_spectrum1Dky_EK[imin_plot : imax_plot + 1]
).mean(0)
means["EA_ky"] = (
dset_spectrum1Dky_EA[imin_plot : imax_plot + 1]
).mean(0)
print(
"""load spectra
tmin = {:8.6g} ; tmax = {:8.6g} ; delta_t = {:8.6g}
imin = {:8d} ; imax = {:8d} ; delta_i = {:8d}""".format(
tmin_plot, tmax_plot, delta_t, imin_plot, imax_plot, delta_i_plot
)
)
return means
[docs]
def plot1d(
self,
tmin=0,
tmax=None,
delta_t=2,
coef_compensate=5 / 3,
level2=1.0,
level3=1.0,
yrange=5,
):
"""Plot spectra one-dimensional."""
print(
f"plot1d(tmin={tmin}, tmax={tmax}, delta_t={delta_t:.2f},"
+ f" coef_compensate={coef_compensate:.3f})"
)
means = self.load1d_means(
tmin, tmax, delta_t, versus_kx=True, versus_ky=True
)
kx = means["kx"]
ky = means["ky"]
# Parameters figure E(k_x)
fig, ax = self.output.figure_axe()
ax.set_xlabel("$k_x$, $k_z$")
ax.set_ylabel(r"$E(k)k^{{{}}}$".format(round(coef_compensate, 2)))
ax.set_title("1D spectra\n" + self.output.summary_simul)
ax.set_xscale("log")
ax.set_yscale("log")
# ax.set_ylim(ymin=1e-6, ymax=1e3)
id_kx_dealiasing = (
np.argmin(abs(kx - (kx.max() * self.sim.oper.coef_dealiasing))) - 1
)
id_ky_dealiasing = (
np.argmin(abs(ky - (ky.max() * self.sim.oper.coef_dealiasing))) - 1
)
# Remove modes dealiased.
E_kx_plot = means["E_kx"][:id_kx_dealiasing]
EK_kx_plot = means["EK_kx"][:id_kx_dealiasing]
EA_kx_plot = means["EA_kx"][:id_kx_dealiasing]
kx_plot = kx[:id_kx_dealiasing]
# Remove shear modes if there is no energy on them.
if self.sim.params.oper.NO_SHEAR_MODES:
E_kx_plot = E_kx_plot[1:]
EK_kx_plot = EK_kx_plot[1:]
EA_kx_plot = EA_kx_plot[1:]
kx_plot = kx_plot[1:]
# Compute k_b: L_b = U / N
U = np.sqrt(np.mean(abs(self.sim.state.get_var("ux")) ** 2))
k_b = self.sim.params.N / U
ax.axvline(x=k_b, color="y", linestyle="--", label="$k_b$")
# Plot ozmidov scale
k_o = (self.sim.params.N**3 / self.sim.params.forcing.forcing_rate) ** (
1 / 2
)
ax.axvline(x=k_o, color="y", linestyle=":", label="$k_o$")
# ax.plot(
# kx_plot, E_kx_plot * kx_plot ** coef_compensate, "k", label="$E(k_x)$"
# )
ax.plot(
kx_plot,
EK_kx_plot * kx_plot**coef_compensate,
"r",
label="$E_K(k_x)$",
)
ax.plot(
kx_plot,
EA_kx_plot * kx_plot**coef_compensate,
"b",
label="$E_A(k_x)$",
)
# Remove modes dealiased.
E_ky_plot = means["E_ky"][:id_ky_dealiasing]
EK_ky_plot = means["EK_ky"][:id_ky_dealiasing]
EA_ky_plot = means["EA_ky"][:id_ky_dealiasing]
ky_plot = ky[:id_ky_dealiasing]
# Remove shear modes if there is no energy on them.
if self.sim.params.oper.NO_SHEAR_MODES:
E_ky_plot = E_ky_plot[1:]
EK_ky_plot = EK_ky_plot[1:]
EA_ky_plot = EA_ky_plot[1:]
ky_plot = ky_plot[1:]
# ax.plot(
# ky_plot,
# E_ky_plot * ky_plot ** coef_compensate,
# "k--",
# label="$E(k_z)$",
# )
ax.plot(
ky_plot,
EK_ky_plot * ky_plot**coef_compensate,
"r--",
label="$E_K(k_z)$",
)
ax.plot(
ky_plot,
EA_ky_plot * ky_plot**coef_compensate,
"b--",
label="$E_A(k_z)$",
)
# Plot scaling lines
kx = kx_plot[40:id_kx_dealiasing]
if level3:
ax.plot(
kx,
level3 * kx ** (-3) * kx**coef_compensate,
"k:",
label=r"$k^{-3}$",
)
if level2:
ax.plot(
kx,
level2 * kx ** (-2) * kx**coef_compensate,
"k-.",
label=r"$k^{-2}$",
)
# Plot forcing wave-number k_f
nkmax = self.sim.params.forcing.nkmax_forcing
nkmin = self.sim.params.forcing.nkmin_forcing
pforcing = self.sim.params.forcing
if pforcing.enable and pforcing.type == "tcrandom_anisotropic":
angle = pforcing.tcrandom_anisotropic.angle
try:
if angle.endswith("°"):
angle = np.pi / 180 * float(angle[:-1])
except AttributeError:
pass
# Band forcing region kx
k_fxmin = nkmin * self.sim.oper.deltak * np.sin(angle)
k_fxmin = max(k_fxmin, self.sim.oper.deltakx)
k_fxmax = nkmax * self.sim.oper.deltak * np.sin(angle)
# Band forcing region ky
k_fymin = nkmin * self.sim.oper.deltak * np.cos(angle)
k_fymin = max(k_fymin, self.sim.oper.deltaky)
k_fymax = nkmax * self.sim.oper.deltak * np.cos(angle)
# Plot forcing band
ax.axvspan(k_fxmin, k_fxmax, alpha=0.15, color="black")
ax.axvspan(k_fymin, k_fymax, alpha=0.15, color="black")
# Set limits axis y
if yrange is not None:
ymax = 2 * max(
(E_kx_plot * kx_plot**coef_compensate).max(),
(E_ky_plot * ky_plot**coef_compensate).max(),
)
ax.set_ylim((10 ** (-yrange) * ymax, ymax))
ax.legend()
# from ipdb import set_trace
# set_trace()
def load2d_means(self, tmin=0, tmax=1000, delta_t=2):
means = {}
# Load data from file
with h5py.File(self.path_file2D, "r") as h5file:
dset_times = h5file["times"]
means["kh"] = h5file["khE"][...]
times = dset_times[...]
means["times"] = times
# Compute average from tmin and tmax for plot
delta_t_save = np.mean(times[1:] - times[0:-1])
delta_i_plot = int(np.round(delta_t / delta_t_save))
if delta_i_plot == 0 and delta_t != 0.0:
delta_i_plot = 1
delta_t = delta_i_plot * delta_t_save
imin_plot = np.argmin(abs(times - tmin))
imax_plot = np.argmin(abs(times - tmax))
tmin_plot = times[imin_plot]
tmax_plot = times[imax_plot]
print(
"""plot 2D spectra
tmin = {:8.6g} ; tmax = {:8.6g} ; delta_t = {:8.6g}
imin = {:8d} ; imax = {:8d} ; delta_i = {:8d}""".format(
tmin_plot,
tmax_plot,
delta_t,
imin_plot,
imax_plot,
delta_i_plot,
)
)
dset_spectrum_E = h5file["spectrum2D_E"]
dset_spectrum_EK = h5file["spectrum2D_EK"]
dset_spectrum_EA = h5file["spectrum2D_EA"]
means["E"] = dset_spectrum_E[imin_plot : imax_plot + 1].mean(0)
means["EK"] = dset_spectrum_EK[imin_plot : imax_plot + 1].mean(0)
means["EA"] = dset_spectrum_EA[imin_plot : imax_plot + 1].mean(0)
return means
[docs]
def plot2d(self, tmin=0, tmax=1000, delta_t=2, coef_compensate=3):
"""Plot 2D spectra."""
print(
f"plot2s(tmin={tmin}, tmax={tmax}, delta_t={delta_t:.2f},"
f" coef_compensate={coef_compensate:.3f})"
)
means = self.load2d_means(tmin, tmax, delta_t)
kh = means["kh"]
E = means["E"]
EK = means["EK"]
EA = means["EA"]
# Parameters figure
fig, ax = self.output.figure_axe()
ax.set_xlabel(r"$k$")
ax.set_ylabel(r"$E(k)$")
ax.set_title("2D spectra\n" + self.output.summary_simul)
ax.set_xscale("log")
ax.set_yscale("log")
ax.plot(kh, E, "k", label="E")
ax.plot(kh, EK, "r", label="EK")
ax.plot(kh, EA, "b", label="EA")
ax.legend()