Commit d459d1da authored by Jason Reneuve's avatar Jason Reneuve
Browse files

save spectra kzkhomega

parent 72aea45aad3f
......@@ -393,7 +393,23 @@ class SpatioTemporalSpectra(SpecificOutput):
# stop opening files when we've reached the right one
break
series["times"] = times
# add Ki_adim arrays, times and dims order
k0_adim = np.r_[0 : ik0max + 1, ik0min:0]
k1_adim = np.r_[0 : ik1max + 1, ik1min:0]
k2_adim = np.r_[0 : ik2max + 1, ik2min:0]
K0_adim, K1_adim, K2_adim = np.meshgrid(
k0_adim, k1_adim, k2_adim, indexing="ij"
)
series.update(
{
"K0_adim": K0_adim,
"K1_adim": K1_adim,
"K2_adim": K2_adim,
"times": times,
"dims_order": order,
}
)
return series
def compute_spectra(self, tmin=0, tmax=None):
......@@ -411,14 +427,15 @@ class SpatioTemporalSpectra(SpecificOutput):
# compute spectra
print("computing temporal spectra...")
dict_spectra = {}
dict_spectra = {k: v for k, v in series.items() if k.startswith("K")}
for key, data in series.items():
if key.startswith("times"):
if not key.startswith("spect"):
continue
freq, spectra = signal.periodogram(data, fs=f_sample)
dict_spectra[key] = spectra
dict_spectra["omegas"] = 2 * pi * freq
dict_spectra["dims_order"] = series["dims_order"]
return dict_spectra
......@@ -10,9 +10,89 @@ Provides:
"""
from pathlib import Path
import numpy as np
import h5py
from fluidsim.base.output.spatiotemporal_spectra import SpatioTemporalSpectra
class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
def plot(self):
raise NotImplementedError
def loop_spectra_kzkhomega(self, spectrum_k0k1k2omega, khs, KH, kzs, KZ):
"""Compute the kz-kh-omega spectrum."""
deltakh = khs[1]
deltakz = kzs[1]
nkh = len(khs)
nkz = len(kzs)
nk0, nk1, nk2, nomega = spectrum_k0k1k2omega.shape
spectrum_kzkhomega = np.zeros((nkz, nkh, nomega))
for ik0 in range(nk0):
for ik1 in range(nk1):
for ik2 in range(nk2):
value = spectrum_k0k1k2omega[ik0, ik1, ik2, :]
kappa = KH[ik0, ik1, ik2]
ikh = int(kappa / deltakh)
kz = abs(KZ[ik0, ik1, ik2])
ikz = int(round(kz / deltakz))
if ikz >= nkz - 1:
ikz = nkz - 1
if ikh >= nkh - 1:
ikh = nkh - 1
spectrum_kzkhomega[ikz, ikh, :] += value
else:
coef_share = (kappa - khs[ikh]) / deltakh
spectrum_kzkhomega[ikz, ikh, :] += (
1 - coef_share
) * value
spectrum_kzkhomega[ikz, ikh + 1, :] += coef_share * value
return spectrum_kzkhomega
def save_spectra_kzkhomega(self, tmin=0, tmax=None):
"""save the spatiotemporal spectra, with a cylindrical average in k-space"""
if tmax is None:
tmax = self.sim.params.time_stepping.t_end
# compute spectra
dict_spectra = self.compute_spectra(tmin=tmin, tmax=tmax)
# get kz, kh
oper = self.sim.oper
order = dict_spectra["dims_order"]
KZ = oper.deltakz * dict_spectra[f"K{order[0]}_adim"]
KY = oper.deltaky * dict_spectra[f"K{order[1]}_adim"]
KX = oper.deltaky * dict_spectra[f"K{order[1]}_adim"]
KH = np.sqrt(KX ** 2 + KY ** 2)
kz_spectra = np.arange(0, KZ.max() + 1e-15, oper.deltakz)
deltakh = oper.deltakh
khmax_spectra = min(KX.max(), KY.max())
nkh_spectra = max(2, int(khmax_spectra / deltakh))
kh_spectra = deltakh * np.arange(nkh_spectra)
# perform cylindrical average
dict_spectra_kzkhomega = {
"kz_spectra": kz_spectra,
"kh_spectra": kh_spectra,
"omegas": dict_spectra["omegas"],
}
for key, data in dict_spectra.items():
if not key.startswith("spect"):
continue
dict_spectra_kzkhomega[key] = self.loop_spectra_kzkhomega(
data, kh_spectra, KH, kz_spectra, KZ
)
# save to file
path_file = Path(self.sim.output.path_run) / "spatiotemporal_spectra.h5"
with h5py.File(path_file, "w") as file:
for key, val in dict_spectra_kzkhomega.items():
file.create_dataset(key, data=val)
def plot_kzkhomega(self, tmin=0, tmax=None):
"""plot the spatiotemporal spectra, with a cylindrical average in k-space"""
if tmax is None:
tmax = self.sim.params.time_stepping.t_end
# load time
......@@ -222,7 +222,7 @@ class TestOutput(TestSimulBase):
sim3.output.temporal_spectra.save_data_as_phys_fields(
delta_index_times=2
)
sim3.output.spatiotemporal_spectra.compute_spectra()
sim3.output.spatiotemporal_spectra.save_spectra_kzkhomega()
plt.close("all")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment