"""Spectra (:mod:`fluidsim.solvers.ns3d.strat.output.spectra`)
==============================================================
.. autoclass:: SpectraNS3DStrat
:members:
:private-members:
"""
import numpy as np
import h5py
from fluidsim.solvers.ns3d.output.spectra import SpectraNS3D
[docs]
class SpectraNS3DStrat(SpectraNS3D):
"""Save and plot spectra."""
[docs]
def compute(self):
"""compute the values at one time."""
get_var = self.sim.state.state_spect.get_var
b_fft = get_var("b_fft")
vx_fft = get_var("vx_fft")
vy_fft = get_var("vy_fft")
vz_fft = get_var("vz_fft")
urx_fft, ury_fft, udx_fft, udy_fft = self.sim.oper.urudfft_from_vxvyfft(
vx_fft, vy_fft
)
nrj_vx_fft = 0.5 * np.abs(vx_fft) ** 2
nrj_vy_fft = 0.5 * np.abs(vy_fft) ** 2
nrj_vz_fft = 0.5 * np.abs(vz_fft) ** 2
nrj_A_fft = 0.5 / self.sim.params.N**2 * np.abs(b_fft) ** 2
nrj_Khr_fft = 0.5 * (np.abs(urx_fft) ** 2 + np.abs(ury_fft) ** 2)
nrj_Khd_fft = 0.5 * (np.abs(udx_fft) ** 2 + np.abs(udy_fft) ** 2)
s_vx_kx, s_vx_ky, s_vx_kz = self.oper.compute_1dspectra(nrj_vx_fft)
s_vy_kx, s_vy_ky, s_vy_kz = self.oper.compute_1dspectra(nrj_vy_fft)
s_vz_kx, s_vz_ky, s_vz_kz = self.oper.compute_1dspectra(nrj_vz_fft)
s_A_kx, s_A_ky, s_A_kz = self.oper.compute_1dspectra(nrj_A_fft)
s_Khr_kx, s_Khr_ky, s_Khr_kz = self.oper.compute_1dspectra(nrj_Khr_fft)
s_Khd_kx, s_Khd_ky, s_Khd_kz = self.oper.compute_1dspectra(nrj_Khd_fft)
s_kx = s_vx_kx + s_vy_kx + s_vz_kx
s_ky = s_vx_ky + s_vy_ky + s_vz_ky
s_kz = s_vx_kz + s_vy_kz + s_vz_kz
dict_spectra1d = {
"vx_kx": s_vx_kx,
"vx_ky": s_vx_ky,
"vx_kz": s_vx_kz,
"vy_kx": s_vy_kx,
"vy_ky": s_vy_ky,
"vy_kz": s_vy_kz,
"vz_kx": s_vz_kx,
"vz_ky": s_vz_ky,
"vz_kz": s_vz_kz,
"E_kx": s_kx,
"E_ky": s_ky,
"E_kz": s_kz,
"A_kx": s_A_kx,
"A_ky": s_A_ky,
"A_kz": s_A_kz,
"Khr_kx": s_Khr_kx,
"Khr_ky": s_Khr_ky,
"Khr_kz": s_Khr_kz,
"Khd_kx": s_Khd_kx,
"Khd_ky": s_Khd_ky,
"Khd_kz": s_Khd_kz,
}
dict_spectra1d = {"spectra_" + k: v for k, v in dict_spectra1d.items()}
s_vx = self.oper.compute_3dspectrum(nrj_vx_fft)
s_vy = self.oper.compute_3dspectrum(nrj_vy_fft)
s_vz = self.oper.compute_3dspectrum(nrj_vz_fft)
s_A = self.oper.compute_3dspectrum(nrj_A_fft)
s_Khr = self.oper.compute_3dspectrum(nrj_Khr_fft)
s_Khd = self.oper.compute_3dspectrum(nrj_Khd_fft)
dict_spectra3d = {
"vx": s_vx,
"vy": s_vy,
"vz": s_vz,
"E": s_vx + s_vy + s_vz,
"A": s_A,
"Khr": s_Khr,
"Khd": s_Khd,
}
dict_spectra3d = {"spectra_" + k: v for k, v in dict_spectra3d.items()}
if self.has_to_save_kzkh():
dict_kzkh = {
"A": self.oper.compute_spectrum_kzkh(nrj_A_fft),
"Khr": self.oper.compute_spectrum_kzkh(nrj_Khr_fft),
"Khd": self.oper.compute_spectrum_kzkh(nrj_Khd_fft),
"Kz": self.oper.compute_spectrum_kzkh(nrj_vz_fft),
}
else:
dict_kzkh = None
return dict_spectra1d, dict_spectra3d, dict_kzkh
def plot_kzkh(self, tmin=0, tmax=None, key="Khd", ax=None):
super().plot_kzkh(tmin, tmax, key, ax)
def _get_key_spectrum(self, ndim, direction=None, kind=None):
if kind is None or kind == "K":
base = "E"
elif kind in set(("A", "Khd", "vx", "vy", "vz")):
base = kind
elif kind.startswith("K") and len(kind) == 2 and kind[1] in "xyz":
base = "v" + kind[1]
else:
raise ValueError(f"{kind = }")
if ndim == 1:
base += "_k" + direction
return "spectra_" + base
def _get_styleline(self, ndim, direction, kind=None):
if kind is None or kind == "K":
style_line = "r"
elif kind == "A":
style_line = "b"
elif kind == "polo":
style_line = "m"
else:
raise ValueError(f"{kind = }")
if direction == "z":
style_line += ":"
return style_line
def _ax_add_average(
self,
ax,
path_file,
ndim,
direction,
ks,
imin_plot,
imax_plot,
coef_norm,
):
def get_average_spectrum(kind):
return self._get_averaged_spectrum(
path_file, ndim, direction, imin_plot, imax_plot, kind
)
def _plot(kind, spectrum, label, linewidth=2):
ax.plot(
ks,
spectrum * coef_norm,
self._get_styleline(ndim, direction, kind),
linewidth=linewidth,
label=label,
)
spectrumK = get_average_spectrum("K")
spectrumA = get_average_spectrum("A")
_plot("K", spectrumK, f"$E_K(k_{direction})$")
_plot("A", spectrumA, f"$E_A(k_{direction})$")
if (
hasattr(self.sim.params, "projection")
and self.sim.params.projection != "toroidal"
):
spectrumKhd = get_average_spectrum("Khd")
spectrumKz = get_average_spectrum("Kz")
_plot("polo", spectrumKhd + spectrumKz, "poloidal", 1)
def plot_kzkh_cumul_diss(self, tmin=0, tmax=None):
path_file = self.path_file_kzkh
with h5py.File(path_file, "r") as h5file:
times = h5file["times"][...]
if tmax is None:
tmax = times.max()
# load kzkh spectra
spectra_K = 0
data = self.load_kzkh_mean(tmin=tmin, tmax=tmax, key_to_load="Khd")
spectra_K += data["Khd"]
data = self.load_kzkh_mean(tmin=tmin, tmax=tmax, key_to_load="Khr")
spectra_K += data["Khr"]
data = self.load_kzkh_mean(tmin=tmin, tmax=tmax, key_to_load="Kz")
spectra_K += data["Kz"]
# spectral space
kh = data["kh_spectra"]
kz = data["kz"]
deltakh = kh[1]
deltakz = kz[1]
KH, KZ = np.meshgrid(kh, kz)
K2 = KH**2 + KZ**2
K4 = K2**2
del KH, KZ
# nu_2 dissipation fluxes
fd = self.sim.params.nu_2 * K2
diss = fd * spectra_K
hflux_diss_nu2 = deltakz * diss.sum(0)
hflux_diss_nu2 = deltakh * np.cumsum(hflux_diss_nu2)
zflux_diss_nu2 = deltakh * diss.sum(1)
zflux_diss_nu2 = deltakz * np.cumsum(zflux_diss_nu2)
# nu_4 dissipation fluxes
fd = self.sim.params.nu_4 * K4
diss = fd * spectra_K
hflux_diss_nu4 = deltakz * diss.sum(0)
hflux_diss_nu4 = deltakh * np.cumsum(hflux_diss_nu4)
zflux_diss_nu4 = deltakh * diss.sum(1)
zflux_diss_nu4 = deltakz * np.cumsum(zflux_diss_nu4)
del fd, diss, K2, K4
# normalize by total dissipation
eps_tot = hflux_diss_nu2[-1] + hflux_diss_nu4[-1]
hflux_diss_nu2 /= eps_tot
zflux_diss_nu2 /= eps_tot
hflux_diss_nu4 /= eps_tot
zflux_diss_nu4 /= eps_tot
# plot
fig, ax = self.output.figure_axe()
ax.set_xlabel(r"$k$")
ax.set_ylabel(r"$D(k)/\epsilon$")
ax.plot(kh, hflux_diss_nu2, "r-", label=r"$D_2(k_h)$")
ax.plot(kz, zflux_diss_nu2, "r:", label=r"$D_2(k_z)$")
ax.plot(kh, hflux_diss_nu4, "m-", label=r"$D_4(k_h)$")
ax.plot(kz, zflux_diss_nu4, "m:", label=r"$D_4(k_z)$")
ax.set_title(
f"kzkh cumulative dissipation spectra (tmin={tmin:.2g}, tmax={tmax:.2g})\n"
+ self.output.summary_simul
)
ax.set_xscale("log")
fig.legend()