"""Energy budget (:mod:`fluidsim.solvers.ns2d.strat.output.spect_energy_budget`)
================================================================================
.. autoclass:: SpectralEnergyBudgetNS2DStrat
:members:
:private-members:
"""
import numpy as np
import h5py
from fluiddyn.util import mpi
from fluidsim.base.output.spect_energy_budget import (
SpectralEnergyBudgetBase,
cumsum_inv,
)
[docs]
class SpectralEnergyBudgetNS2DStrat(SpectralEnergyBudgetBase):
"""Save and plot energy budget in spectral space."""
[docs]
def compute(self):
"""compute the spectral energy budget at one time."""
oper = self.sim.oper
ux = self.sim.state.state_phys.get_var("ux")
uy = self.sim.state.state_phys.get_var("uy")
rot_fft = self.sim.state.state_spect.get_var("rot_fft")
b_fft = self.sim.state.state_spect.get_var("b_fft")
ux_fft, uy_fft = oper.vecfft_from_rotfft(rot_fft)
px_b_fft, py_b_fft = oper.gradfft_from_fft(b_fft)
px_b = oper.ifft2(px_b_fft)
py_b = oper.ifft2(py_b_fft)
Fb = -ux * px_b - uy * py_b
Fb_fft = oper.fft2(Fb)
oper.dealiasing(Fb_fft)
px_rot_fft, py_rot_fft = oper.gradfft_from_fft(rot_fft)
px_rot = oper.ifft2(px_rot_fft)
py_rot = oper.ifft2(py_rot_fft)
px_ux_fft, py_ux_fft = oper.gradfft_from_fft(ux_fft)
px_ux = oper.ifft2(px_ux_fft)
py_ux = oper.ifft2(py_ux_fft)
px_uy_fft, py_uy_fft = oper.gradfft_from_fft(uy_fft)
px_uy = oper.ifft2(px_uy_fft)
py_uy = oper.ifft2(py_uy_fft)
Frot = -ux * px_rot - uy * (py_rot + self.params.beta)
Frot_fft = oper.fft2(Frot)
oper.dealiasing(Frot_fft)
Fx = -ux * px_ux - uy * (py_ux)
Fx_fft = oper.fft2(Fx)
oper.dealiasing(Fx_fft)
Fy = -ux * px_uy - uy * (py_uy)
Fy_fft = oper.fft2(Fy)
oper.dealiasing(Fy_fft)
# Frequency dissipation viscosity
f_d, f_d_hypo = self.sim.compute_freq_diss()
freq_diss_EK = f_d + f_d_hypo
# Energy budget terms. Nonlinear transfer terms, exchange kinetic and
# potential energy B, dissipation terms.
transferZ_fft = (
np.real(rot_fft.conj() * Frot_fft + rot_fft * Frot_fft.conj()) / 2.0
)
transferEKu_fft = np.real(ux_fft.conj() * Fx_fft)
transferEKv_fft = np.real(uy_fft.conj() * Fy_fft)
B_fft = np.real(uy_fft.conj() * b_fft)
if self.params.N == 0:
transferEA_fft = np.zeros_like(transferZ_fft)
else:
transferEA_fft = (1 / self.params.N**2) * np.real(
b_fft.conj() * Fb_fft
)
dissEKu_fft = np.real(freq_diss_EK * (ux_fft.conj() * ux_fft))
dissEKv_fft = np.real(freq_diss_EK * (uy_fft.conj() * uy_fft))
dissEK_fft = np.real(
freq_diss_EK
* (
ux_fft.conj() * ux_fft
+ ux_fft * ux_fft.conj()
+ uy_fft.conj() * uy_fft
+ uy_fft * uy_fft.conj()
)
/ 2.0
)
if self.params.N == 0:
dissEA_fft = np.zeros_like(dissEK_fft)
else:
dissEA_fft = (1 / self.params.N**2) * np.real(
freq_diss_EK * (b_fft.conj() * b_fft)
)
transferEK_fft = (
np.real(
ux_fft.conj() * Fx_fft
+ ux_fft * Fx_fft.conj()
+ uy_fft.conj() * Fy_fft
+ uy_fft * Fy_fft.conj()
)
/ 2.0
)
# Transfer spectrum 1D Kinetic energy, potential energy and exchange
# energy
transferEK_kx, transferEK_ky = self.spectra1D_from_fft(transferEK_fft)
transferEKu_kx, transferEKu_ky = self.spectra1D_from_fft(transferEKu_fft)
transferEKv_kx, transferEKv_ky = self.spectra1D_from_fft(transferEKv_fft)
transferEA_kx, transferEA_ky = self.spectra1D_from_fft(transferEA_fft)
B_kx, B_ky = self.spectra1D_from_fft(B_fft)
dissEK_kx, dissEK_ky = self.spectra1D_from_fft(dissEK_fft)
dissEKu_kx, dissEKu_ky = self.spectra1D_from_fft(dissEKu_fft)
dissEKv_kx, dissEKv_ky = self.spectra1D_from_fft(dissEKv_fft)
dissEA_kx, dissEA_ky = self.spectra1D_from_fft(dissEA_fft)
# Transfer spectrum shell mean
transferEK_2d = self.spectrum2D_from_fft(transferEK_fft)
transferEKu_2d = self.spectrum2D_from_fft(transferEKu_fft)
transferEKv_2d = self.spectrum2D_from_fft(transferEKv_fft)
transferEA_2d = self.spectrum2D_from_fft(transferEA_fft)
B_2d = self.spectrum2D_from_fft(B_fft)
dissEKu_2d = self.spectrum2D_from_fft(dissEKu_fft)
dissEKv_2d = self.spectrum2D_from_fft(dissEKv_fft)
dissEA_2d = self.spectrum2D_from_fft(dissEA_fft)
transferZ_2d = self.spectrum2D_from_fft(transferZ_fft)
# Dissipation rate at one time
epsilon_kx = dissEKu_kx.sum() + dissEKv_kx.sum() + dissEA_kx.sum()
epsilon_ky = dissEKu_ky.sum() + dissEKv_ky.sum() + dissEA_ky.sum()
# Variables saved in a dictionary
dict_results = {
"transferEK_kx": transferEK_kx,
"transferEK_ky": transferEK_ky,
"transferEKu_kx": transferEKu_kx,
"transferEKu_ky": transferEKu_ky,
"transferEKv_kx": transferEKv_kx,
"transferEKv_ky": transferEKv_ky,
"transferEKu_2d": transferEKu_2d,
"transferEKv_2d": transferEKv_2d,
"transferEK_2d": transferEK_2d,
"transferEA_kx": transferEA_kx,
"transferEA_ky": transferEA_ky,
"transferEA_2d": transferEA_2d,
"transferZ_2d": transferZ_2d,
"B_kx": B_kx,
"B_ky": B_ky,
"B_2d": B_2d,
"dissEK_kx": dissEK_kx,
"dissEK_ky": dissEK_ky,
"dissEKu_kx": dissEKu_kx,
"dissEKu_ky": dissEKu_ky,
"dissEKu_2d": dissEKu_2d,
"dissEKv_kx": dissEKv_kx,
"dissEKv_ky": dissEKv_ky,
"dissEKv_2d": dissEKv_2d,
"dissEA_kx": dissEA_kx,
"dissEA_ky": dissEA_ky,
"dissEA_2d": dissEA_2d,
"epsilon_kx": epsilon_kx,
"epsilon_ky": epsilon_ky,
}
if mpi.rank == 0:
small_value = 3e-10
for k, v in dict_results.items():
if k.startswith("transfer"):
if abs(v.sum()) > small_value:
print("warning: (abs(v.sum()) > small_value) for " + k)
print("k = ", k)
print("abs(v.sum()) = ", abs(v.sum()))
return dict_results
def _online_plot_saving(self, dict_results):
transfer2D_EA = dict_results["transferEA_2d"]
transfer2D_EK = dict_results["transferEK_2d"]
transfer2D_E = transfer2D_EA + transfer2D_EK
transfer2D_Z = dict_results["transferZ_2d"]
khE = self.oper.khE
PiE = cumsum_inv(transfer2D_E) * self.oper.deltak
PiZ = cumsum_inv(transfer2D_Z) * self.oper.deltak
self.axe_a.plot(khE + khE[1], PiE, "k")
self.axe_b.plot(khE + khE[1], PiZ, "g")
[docs]
def load_mean(self, tmin=None, tmax=None, keys_to_load=None):
"""Loads data spect_energy_budget.
It computes the mean between tmin and tmax."""
means = {}
with h5py.File(self.path_file, "r") as file:
times = file["times"][...]
nt = len(times)
if tmin is None:
imin_plot = 0
else:
imin_plot = np.argmin(abs(times - tmin))
if tmax is None:
imax_plot = nt - 1
else:
imax_plot = np.argmin(abs(times - tmax))
tmin = times[imin_plot]
tmax = times[imax_plot]
print(
"compute mean spectral energy budget\n"
f"tmin = {tmin:8.6g} ; tmax = {tmax:8.6g}\n"
f"imin = {imin_plot:8d} ; imax = {imax_plot:8d}"
)
for key in list(file.keys()):
if key.startswith("k"):
means[key] = file[key][...]
if keys_to_load is not None:
if isinstance(keys_to_load, str):
keys_to_load = [keys_to_load]
for key in keys_to_load:
if key not in file.keys():
print(key, file.keys())
raise ValueError
else:
keys_to_load = [
key
for key in file.keys()
if key != "times" and not key.startswith(("k", "info"))
]
for key in keys_to_load:
spect = file[key][imin_plot : imax_plot + 1].mean(0)
means[key] = spect
return means
[docs]
def compute_fluxes_mean(self, tmin=None, tmax=None):
"""compute the fluxes mean."""
# Load data from file
with h5py.File(self.path_file, "r") as file:
times = file["times"][...]
nt = len(times)
kxE = file["kxE"][...]
kyE = file["kyE"][...]
transferEK_kx = file["transferEK_kx"][...]
transferEK_ky = file["transferEK_ky"][...]
transferEA_kx = file["transferEA_kx"][...]
transferEA_ky = file["transferEA_ky"][...]
dissEK_kx = file["dissEK_kx"][...]
dissEA_kx = file["dissEA_kx"][...]
dissEK_ky = file["dissEK_ky"][...]
dissEA_ky = file["dissEA_ky"][...]
conv_kx = file["B_kx"][...]
conv_ky = file["B_ky"][...]
if tmin is None:
imin_plot = 0
else:
imin_plot = np.argmin(abs(times - tmin))
if tmax is None:
imax_plot = nt - 1
else:
imax_plot = np.argmin(abs(times - tmax))
tmin = times[imin_plot]
tmax = times[imax_plot]
# compute means kx
transferEK_kx = transferEK_kx[imin_plot : imax_plot + 1].mean(0)
transferEA_kx = transferEA_kx[imin_plot : imax_plot + 1].mean(0)
dissEK_kx = dissEK_kx[imin_plot : imax_plot + 1].mean(0)
dissEA_kx = dissEA_kx[imin_plot : imax_plot + 1].mean(0)
id_kx_dealiasing = np.argmin(kxE - self.sim.oper.kxmax_dealiasing) - 1
id_ky_dealiasing = np.argmin(kyE - self.sim.oper.kymax_dealiasing) - 1
conv_kx = conv_kx[imin_plot : imax_plot + 1].mean(0)[:id_kx_dealiasing]
conv_kx = cumsum_inv(conv_kx[1:]) * self.oper.deltakx
transferEK_kx = transferEK_kx[:id_kx_dealiasing]
transferEA_kx = transferEA_kx[:id_kx_dealiasing]
dissEK_kx = dissEK_kx[:id_kx_dealiasing]
dissEA_kx = dissEA_kx[:id_kx_dealiasing]
kxE = kxE[:id_kx_dealiasing]
# Transfer and dissipation kx
PiEK_kx = cumsum_inv(transferEK_kx[1:]) * self.oper.deltakx
PiEA_kx = cumsum_inv(transferEA_kx[1:]) * self.oper.deltakx
DissEK_kx = dissEK_kx[1:].cumsum() * self.oper.deltakx
DissEA_kx = dissEA_kx[1:].cumsum() * self.oper.deltakx
# compute means ky
transferEK_ky = transferEK_ky[imin_plot : imax_plot + 1].mean(0)
transferEA_ky = transferEA_ky[imin_plot : imax_plot + 1].mean(0)
dissEK_ky = dissEK_ky[imin_plot : imax_plot + 1].mean(0)
dissEA_ky = dissEA_ky[imin_plot : imax_plot + 1].mean(0)
conv_ky = conv_ky[imin_plot : imax_plot + 1].mean(0)[:id_ky_dealiasing]
conv_ky = cumsum_inv(conv_ky[1:]) * self.oper.deltaky
transferEK_ky = transferEK_ky[:id_ky_dealiasing]
transferEA_ky = transferEA_ky[:id_ky_dealiasing]
dissEK_ky = dissEK_ky[:id_ky_dealiasing]
dissEA_ky = dissEA_ky[:id_ky_dealiasing]
kyE = kyE[:id_ky_dealiasing]
# Transfer and dissipation ky
PiEK_ky = cumsum_inv(transferEK_ky[1:]) * self.oper.deltaky
PiEA_ky = cumsum_inv(transferEA_ky[1:]) * self.oper.deltaky
DissEK_ky = dissEK_ky[1:].cumsum() * self.oper.deltaky
DissEA_ky = dissEA_ky[1:].cumsum() * self.oper.deltaky
# save data in dictionary
fluxes = {
"kxE": kxE,
"kyE": kyE,
"PiEAmean_x": PiEA_kx,
"PiEKmean_x": PiEK_kx,
"DissEAmean_x": DissEA_kx,
"DissEKmean_x": DissEK_kx,
"PiEAmean_y": PiEA_ky,
"PiEKmean_y": PiEK_ky,
"DissEAmean_y": DissEA_ky,
"DissEKmean_y": DissEK_ky,
"Conv_x": conv_kx,
"Conv_y": conv_ky,
}
return fluxes
[docs]
def plot(
self, tmin=None, tmax=None, delta_t=2, plot_diss_EK=False, plot_conv=False
):
"""Plot the energy budget."""
# load data from file
with h5py.File(self.path_file, "r") as file:
times = file["times"][...]
nt = len(times)
if tmin is None:
imin_plot = 0
else:
imin_plot = np.argmin(abs(times - tmin))
if tmax is None:
imax_plot = nt - 1
else:
imax_plot = np.argmin(abs(times - tmax))
# 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
print(f"plot(tmin={tmin}, tmax={tmax}, delta_t={delta_t:.2f})")
tmin = times[imin_plot]
tmax = times[imax_plot]
print(
"plot spectral energy budget\n"
f"tmin = {tmin:8.6g} ; tmax = {tmax:8.6g} ;\n"
f"delta_t = {delta_t:8.6g} ; imin = {imin_plot:8d} ;\n"
f"imax = {imax_plot:8d} ; delta_i = {delta_i_plot:8d}"
)
fluxes = self.compute_fluxes_mean(tmin, tmax)
# Parameters of figure 1
fig, ax1 = self.output.figure_axe()
ax1.set_xlabel("$k_x/k_{f,x}$")
ax1.set_ylabel(r"$\Pi(k_x)/\epsilon$")
ax1.set_xscale("log")
ax1.set_yscale("linear")
ax1.set_title("Spectral energy budget\n" + self.output.summary_simul)
# Parameters of figure 2
fig, ax2 = self.output.figure_axe()
ax2.set_xlabel("$k_z/k_{f,z}$")
ax2.set_ylabel(r"$\Pi(k_z)/\epsilon$")
ax2.set_xscale("log")
ax2.set_yscale("linear")
ax2.set_title("Spectral energy budget\n" + self.output.summary_simul)
# Compute time average dissipation to normalize fluxes
spatial_mean_results = self.output.spatial_means.load()
t = spatial_mean_results["t"]
epsK_tot = spatial_mean_results["epsK_tot"]
epsA_tot = spatial_mean_results["epsA_tot"]
eps_tot = epsK_tot + epsA_tot
imin_plot_spatial_times = np.argmin(abs(t - tmin))
imax_plot_spatial_times = np.argmin(abs(t - tmax))
eps_tot_time_average = eps_tot[
imin_plot_spatial_times : imax_plot_spatial_times + 1
].mean(0)
# calculate forcing wave-number k_f to normalize k_x and k_z
pforcing = self.sim.params.forcing
if pforcing.enable and pforcing.type == "tcrandom_anisotropic":
nkmax = pforcing.nkmax_forcing
nkmin = pforcing.nkmin_forcing
k_f = ((nkmax + nkmin) / 2) * self.sim.oper.deltak
angle = pforcing.tcrandom_anisotropic.angle
try:
if angle.endswith("°"):
angle = np.pi / 180 * float(angle[:-1])
except AttributeError:
pass
k_fx = np.sin(angle) * k_f
k_fy = np.cos(angle) * k_f
# 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 range
ax1.axvspan(k_fxmin / k_fx, k_fxmax / k_fx, alpha=0.15, color="black")
ax2.axvspan(k_fymin / k_fy, k_fymax / k_fy, alpha=0.15, color="black")
# Plot ozmidov scale
k_o = (self.params.N**3 / self.params.forcing.forcing_rate) ** (1 / 2)
ax1.axvline(x=k_o / k_fx, color="black", linestyle="--")
ax2.axvline(x=k_o / k_fy, color="black", linestyle="--")
ax1.text((k_o / k_fx) + 2, 1, r"$k_o$")
ax2.text((k_o / k_fy) + 2, 1.7, r"$k_o$")
else:
print("Forcing is not enabled k_x and k_z are not normalized")
k_fx = 1
k_fy = 1
# extract variables from the dictionary to plot
kxE = fluxes["kxE"]
kyE = fluxes["kyE"]
PiEK_kx = fluxes["PiEKmean_x"]
PiEA_kx = fluxes["PiEAmean_x"]
DissEK_kx = fluxes["DissEKmean_x"]
DissEA_kx = fluxes["DissEAmean_x"]
PiEK_ky = fluxes["PiEKmean_y"]
PiEA_ky = fluxes["PiEAmean_y"]
DissEK_ky = fluxes["DissEKmean_y"]
DissEA_ky = fluxes["DissEAmean_y"]
conv_kx = fluxes["Conv_x"]
conv_ky = fluxes["Conv_y"]
def plot_x(x, y, style, label):
ax1.plot(x[1:] / k_fx, y / eps_tot_time_average, style, label=label)
plot_x(kxE, (PiEK_kx + PiEA_kx), "k--", r"$\Pi/\epsilon$")
plot_x(kxE, PiEK_kx, "r", r"$\Pi_K/\epsilon$")
plot_x(kxE, PiEA_kx, "b", r"$\Pi_A/\epsilon$")
plot_x(kxE, (DissEK_kx + DissEA_kx), "g", r"$D/\epsilon$")
if plot_diss_EK:
plot_x(kxE, DissEK_kx, "r--", r"$D_K$")
if plot_conv:
plot_x(kxE, conv_kx, "c", r"$C$")
ax1.axhline(y=0, color="k", linestyle=":")
def plot_y(x, y, style, label):
ax2.plot(x[1:] / k_fy, y / eps_tot_time_average, style, label=label)
plot_y(kyE, (PiEK_ky + PiEA_ky), "k--", r"$\Pi/\epsilon$")
plot_y(kyE, PiEK_ky, "r", r"$\Pi_K/\epsilon$")
plot_y(kyE, PiEA_ky, "b", r"$\Pi_A/\epsilon$")
plot_y(kyE, (DissEK_ky + DissEA_ky), "g", r"$D/\epsilon$")
if plot_diss_EK:
plot_y(kyE, DissEK_ky, "r--", r"$D_K$")
if plot_conv:
plot_y(kyE, conv_ky, "c", r"$C$")
ax2.axhline(y=0, color="k", linestyle=":")
# 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
ax1.axvline(x=k_b / k_fx, color="k", linestyle="--")
ax2.axvline(x=k_b / k_fy, color="k", linestyle="--")
ax1.text((k_b / k_fx) + 0.5, 1, r"$k_b$")
ax2.text((k_b / k_fy) + 0.5, 1.7, r"$k_b$")
ax1.legend()
ax2.legend()