Commit 355f8a2f authored by Jason Reneuve's avatar Jason Reneuve
Browse files

temporal spectra from spatiotemporal_spectra

parent 0d5168e838d6
Pipeline #20656 failed with stage
in 11 minutes and 39 seconds
......@@ -553,7 +553,7 @@ class SpatioTemporalSpectra(SpecificOutput):
)
return freq, spectrum / self.domega
def compute_temporal_spectra(self, tmin=0, tmax=None, dtype=None):
def compute_spectra(self, tmin=0, tmax=None, dtype=None):
"""compute spatiotemporal spectra from files"""
if tmax is None:
tmax = self.sim.params.time_stepping.t_end
......@@ -562,7 +562,7 @@ class SpatioTemporalSpectra(SpecificOutput):
series = self.load_time_series(tmin=tmin, tmax=tmax, dtype=dtype)
# compute spectra
print("computing temporal spectra...")
print("performing time fft...")
spectra = {k: v for k, v in series.items() if k.startswith("K")}
......
......@@ -540,9 +540,7 @@ class TemporalSpectra3D(SpecificOutput):
)
return freq, spectrum / self.domega
def compute_temporal_spectra(
self, region=None, tmin=0, tmax=None, dtype=None
):
def compute_spectra(self, region=None, tmin=0, tmax=None, dtype=None):
"""compute temporal spectra from files"""
if region is None:
region = self._get_default_region()
......@@ -583,7 +581,7 @@ class TemporalSpectra3D(SpecificOutput):
# TODO: check if spectra are saved before computing everything
# compute spectra
spectra = self.compute_temporal_spectra(
spectra = self.compute_spectra(
region=region, tmin=tmin, tmax=tmax, dtype=dtype
)
......@@ -711,9 +709,7 @@ class TemporalSpectra3D(SpecificOutput):
if tmax is None:
tmax = self.sim.params.time_stepping.t_end
spectra = self.compute_temporal_spectra(
region=region, tmin=tmin, tmax=tmax
)
spectra = self.compute_spectra(region=region, tmin=tmin, tmax=tmax)
path_file = Path(self.sim.output.path_run) / "temporal_spectra.h5"
with h5py.File(path_file, "w") as file:
......
......@@ -98,7 +98,7 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
# compute spectra
print("Computing spectra...")
spectra = self.compute_temporal_spectra(tmin=tmin, tmax=tmax, dtype=dtype)
spectra = self.compute_spectra(tmin=tmin, tmax=tmax, dtype=dtype)
# get kz, kh
params_oper = self.sim.params.oper
......@@ -467,3 +467,46 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
spectra["dims_order"] = order
return spectra
def compute_temporal_spectra(
self, tmin=0, tmax=None, dtype=None, compute_urud=False
):
"""compute the temporal spectra by averaging over Fourier space"""
if tmax is None:
tmax = self.sim.params.time_stepping.t_end
tspectra = {}
# compute kxkykzomega spectra
spectra = self.compute_spectra(tmin=tmin, tmax=tmax, dtype=dtype)
if compute_urud:
spectra.update(
self.compute_spectra_urud(tmin=tmin, tmax=tmax, dtype=dtype)
)
order = spectra["dims_order"]
KX = spectra[f"K{order[2]}_adim"]
kx_max = self.sim.params.oper.nx // 2 * 2 * pi / self.sim.params.oper.Lx
def _sum_wavenumber(field):
n0, n1, n2 = field.shape[:3]
result = 0.0
for i0 in range(n0):
for i1 in range(n1):
for i2 in range(n2):
value = field[i0, i1, i2]
kx = KX[i0, i1, i2]
if kx != 0.0 and kx != kx_max:
value *= 2
result += value
return result
# average over Fourier space (kx,ky,kz)
for key, spectrum in spectra.items():
if not key.startswith("spectrum_"):
continue
tspectra[key] = _sum_wavenumber(spectrum)
tspectra["omegas"] = spectra["omegas"]
return tspectra
......@@ -244,18 +244,19 @@ class TestOutput(TestSimulBase):
sim3.output.temporal_spectra.save_spectra()
spatiotemporal_spectra = sim3.output.spatiotemporal_spectra
series = spatiotemporal_spectra.load_time_series()
series_kxkykz = spatiotemporal_spectra.load_time_series()
tspectra_kxyz = spatiotemporal_spectra.compute_temporal_spectra()
tspectra_mean = (
sim3.output.temporal_spectra.compute_temporal_spectra()
spectra_kxkykzomega = spatiotemporal_spectra.compute_spectra()
spectra_omega = sim3.output.temporal_spectra.compute_spectra()
spectra_omega_from_spatiotemp = (
spatiotemporal_spectra.compute_temporal_spectra()
)
means = sim3.output.spatial_means.load()
deltakx = 2 * pi / self.params.oper.Lx
order = tspectra_kxyz["dims_order"]
KX = deltakx * tspectra_kxyz[f"K{order[2]}_adim"]
order = spectra_kxkykzomega["dims_order"]
KX = deltakx * spectra_kxkykzomega[f"K{order[2]}_adim"]
kx_max = self.params.oper.nx // 2 * deltakx
def sum_wavenumber(field):
......@@ -281,44 +282,56 @@ class TestOutput(TestSimulBase):
coef = delta_kz * delta_kh * delta_omega
for letter in "xyz":
vi_fft = series[f"v{letter}_Fourier"]
tspectrum_kxyz = tspectra_kxyz["spectrum_v" + letter]
tspectrum_mean = tspectra_mean["spectrum_v" + letter]
vi_fft = series_kxkykz[f"v{letter}_Fourier"]
spectrum_kxkykzomega = spectra_kxkykzomega["spectrum_v" + letter]
spectrum_omega = spectra_omega["spectrum_v" + letter]
spectrum_omega_from_spatiotemp = spectra_omega_from_spatiotemp[
"spectrum_v" + letter
]
spectrum_kzkhomega = spectra_kzkhomega["spectrum_v" + letter]
energy_series = 0.5 * sum_wavenumber(
Etot_series_kxkykz = 0.5 * sum_wavenumber(
(abs(vi_fft) ** 2).mean(axis=-1)
)
assert energy_series > 0, (letter, vi_fft)
assert Etot_series_kxkykz > 0, (letter, vi_fft)
energy_tspect_kxyz = (
Etot_kxkykzomega = (
0.5
* delta_omega
* sum_wavenumber(tspectrum_kxyz.sum(axis=-1))
* sum_wavenumber(spectrum_kxkykzomega.sum(axis=-1))
)
energy_tspect_mean = 0.5 * delta_omega * tspectrum_mean.sum()
energy_kzkhomega = 0.5 * coef * spectrum_kzkhomega.sum()
Etot_omega = 0.5 * delta_omega * spectrum_omega.sum()
Etot_kzkhomega = 0.5 * coef * spectrum_kzkhomega.sum()
# `:-1` because the last time is saved twice in spatial_means
energy_mean = means["E" + letter][:-1].mean()
Etot_mean = means["E" + letter][:-1].mean()
assert np.allclose(energy_tspect_mean, energy_tspect_kxyz), (
assert np.allclose(Etot_omega, Etot_kxkykzomega), (
letter,
energy_tspect_kxyz / energy_mean,
Etot_kxkykzomega / Etot_mean,
)
assert np.allclose(energy_series, energy_tspect_kxyz), (
assert np.allclose(Etot_series_kxkykz, Etot_kxkykzomega), (
letter,
energy_tspect_kxyz / energy_series,
Etot_kxkykzomega / Etot_series_kxkykz,
)
assert np.allclose(energy_series, energy_kzkhomega), (
assert np.allclose(Etot_series_kxkykz, Etot_kzkhomega), (
letter,
energy_kzkhomega / energy_series,
Etot_kzkhomega / Etot_series_kxkykz,
)
assert np.allclose(energy_mean, energy_series), (
assert np.allclose(Etot_mean, Etot_series_kxkykz), (
letter,
energy_series / energy_mean,
Etot_series_kxkykz / Etot_mean,
)
assert np.allclose(
spectrum_omega, spectrum_omega_from_spatiotemp
), (
letter,
spectrum_omega,
spectrum_omega_from_spatiotemp,
spectrum_omega / spectrum_omega_from_spatiotemp,
)
spectrum_Khd = spectra_kzkhomega["spectrum_Khd"]
......
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