"""Spectral multidimensional
============================
Provides:
.. autoclass:: SpectraMultiDim
:members:
:private-members:
:noindex:
:undoc-members:
"""
import os
import numpy as np
import h5py
from fluiddyn.util import mpi
from .base import SpecificOutput
[docs]
class SpectraMultiDim(SpecificOutput):
"""Saving the multidimensional spectra."""
_tag = "spectra_multidim"
[docs]
@staticmethod
def _complete_params_with_default(params):
tag = "spectra_multidim"
params.output.periods_save._set_attrib(tag, 0)
params.output._set_child(tag, attribs={"HAS_TO_PLOT_SAVED": False})
def __init__(self, output):
self.output = output
params = output.sim.params
self.nx = int(params.oper.nx)
if not params.output.HAS_TO_SAVE:
params.output.periods_save.spectra_multidim = False
super().__init__(
output,
period_save=params.output.periods_save.spectra_multidim,
has_to_plot_saved=params.output.spectra_multidim.HAS_TO_PLOT_SAVED,
)
[docs]
def _init_path_files(self):
path_run = self.output.path_run
self.path_file = path_run + "/spectra_multidim.h5"
[docs]
def _init_files(self, arrays_1st_time=None):
dict_spectra = self.compute()
if mpi.rank == 0:
if not os.path.exists(self.path_file):
arrays_1st_time = {
"kxE": self.sim.oper.kxE,
"kyE": self.sim.oper.kyE,
}
self._create_file_from_dict_arrays(
self.path_file, dict_spectra, arrays_1st_time
)
self.nb_saved_times = 1
else:
with h5py.File(self.path_file, "r") as file:
dset_times = file["times"]
self.nb_saved_times = dset_times.shape[0] + 1
# save the spectra in the file spectra_multidim.h5
self._add_dict_arrays_to_file(self.path_file, dict_spectra)
self.t_last_save = self.sim.time_stepping.t
[docs]
def _online_save(self):
"""Save the values at one time."""
tsim = self.sim.time_stepping.t
if tsim - self.t_last_save >= self.period_save:
self.t_last_save = tsim
dict_spectra = self.compute()
if mpi.rank == 0:
# save the spectra in the file spectra_multidim.h5
self._add_dict_arrays_to_file(self.path_file, dict_spectra)
self.nb_saved_times += 1
if self.has_to_plot:
self._online_plot_saving(dict_spectra)
if tsim - self.t_last_show >= self.period_show:
self.t_last_show = tsim
self.ax.get_figure().canvas.draw()
[docs]
def compute(self):
"""compute the values at one time."""
if mpi.rank == 0:
dict_results = {}
return dict_results
[docs]
def _init_online_plot(self):
if mpi.rank == 0:
fig, ax = self.output.figure_axe(numfig=1_000_000)
self.ax = ax
ax.set_xlabel("$k_x$")
ax.set_ylabel("$k_y$")
ax.set_title("Multidimensional spectra\n" + self.output.summary_simul)
[docs]
def _online_plot_saving(self, arg):
pass
[docs]
def load_mean(self, tmin=None, tmax=None):
"""Loads time averaged data between tmin and tmax."""
results = {}
# Load data
with h5py.File(self.path_file, "r") as file:
for key in file.keys():
if not key.startswith("info"):
results[key] = file[key][...]
# Time average spectra
times = results["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 of multidimensional spectra\n"
f"tmin = {tmin:8.6g} ; tmax = {tmax:8.6g}\n"
f"imin = {imin_plot:8d} ; imax = {imax_plot:8d}"
)
for key in results.keys():
if key.startswith("spectr"):
spect = results[key]
spect_averaged = spect[imin_plot : imax_plot + 1].mean(0)
results[key] = spect_averaged
return results