Read about our upcoming Code of Conduct on this issue

Commit 94be6a62 authored by Jason Reneuve's avatar Jason Reneuve
Browse files

fix tests incompressibility + _comput_spectrum

parent 2a5dedd6df4a
......@@ -169,10 +169,16 @@ class SpatioTemporalSpectra(SpecificOutput):
return
if params_st_spec.probes_region is not None:
self.probes_region = params_st_spec.probes_region
ikxmax, ikymax, ikzmax = self.probes_region
ikxmax, ikymax, ikzmax = params_st_spec.probes_region
ikxmax = min(ikxmax, params.oper.nx // 2)
ikymax = min(ikymax, params.oper.ny // 2)
ikzmax = min(ikzmax, params.oper.nz // 2)
self.probes_region = ikxmax, ikymax, ikzmax
else:
ikxmax = ikymax = ikzmax = 4
ikxmax = min(ikxmax, params.oper.nx // 2)
ikymax = min(ikymax, params.oper.ny // 2)
ikzmax = min(ikzmax, params.oper.nz // 2)
self.probes_region = ikxmax, ikymax, ikzmax
self.file_max_size = params_st_spec.file_max_size
......@@ -532,14 +538,18 @@ class SpatioTemporalSpectra(SpecificOutput):
def _compute_spectrum(self, data):
if not hasattr(self, "f_sample"):
# TODO: Fix this value (can be read from a data file)
self.f_sample = 1
paths = sorted(self.path_dir.glob("rank*.h5"))
with h5py.File(paths[0], "r") as file:
self.f_sample = 1.0 / file.attrs["period_save"]
self.domega = 2 * pi * self.f_sample / data.shape[-1]
# TODO: I'm not sure if detrend=False is good in prod, but it's much
# better for testing
freq, spectrum = signal.periodogram(
data, fs=self.f_sample, scaling="spectrum", detrend=False
data,
fs=self.f_sample,
scaling="spectrum",
detrend=False,
)
return freq, spectrum / self.domega
......
......@@ -94,16 +94,19 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
spectra = self.compute_spectra(tmin=tmin, tmax=tmax, dtype=dtype)
# get kz, kh
oper = self.sim.oper
params_oper = self.sim.params.oper
deltakz = 2 * pi / params_oper.Lz
deltaky = 2 * pi / params_oper.Ly
deltakx = 2 * pi / params_oper.Lx
order = spectra["dims_order"]
KZ = oper.deltakz * spectra[f"K{order[0]}_adim"]
KY = oper.deltaky * spectra[f"K{order[1]}_adim"]
KX = oper.deltakx * spectra[f"K{order[2]}_adim"]
KZ = deltakz * spectra[f"K{order[0]}_adim"]
KY = deltaky * spectra[f"K{order[1]}_adim"]
KX = deltakx * spectra[f"K{order[2]}_adim"]
KH = np.sqrt(KX ** 2 + KY ** 2)
kz_spectra = np.arange(0, KZ.max() + 1e-15, oper.deltakz)
kz_spectra = np.arange(0, KZ.max() + 1e-15, deltakz)
deltakh = oper.deltakh
deltakh = max(deltakx, deltaky)
khmax_spectra = min(KX.max(), KY.max())
nkh_spectra = max(2, int(khmax_spectra / deltakh))
kh_spectra = deltakh * np.arange(nkh_spectra)
......@@ -355,18 +358,20 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
vy_fft = series["vy_Fourier"]
if vx_fft.dtype == "complex64":
float_dtype = "float32"
elif vx_fft.dtype == "complex64":
float_dtype = "float128"
elif vx_fft.dtype == "complex128":
float_dtype = "float64"
oper = self.sim.oper
params_oper = self.sim.params.oper
deltaky = 2 * pi / params_oper.Ly
deltakx = 2 * pi / params_oper.Lx
order = series["dims_order"]
shapeK = series[f"K{order[1]}_adim"].shape
KY = np.zeros(shapeK + (1,), dtype=float_dtype)
KX = np.zeros(shapeK + (1,), dtype=float_dtype)
KY[..., 0] = oper.deltaky * series[f"K{order[1]}_adim"]
KX[..., 0] = oper.deltakx * series[f"K{order[2]}_adim"]
KY[..., 0] = deltaky * series[f"K{order[1]}_adim"]
KX[..., 0] = deltakx * series[f"K{order[2]}_adim"]
inv_Kh_square_nozero = KX ** 2 + KY ** 2
inv_Kh_square_nozero[inv_Kh_square_nozero == 0] = 1e-14
......
......@@ -114,7 +114,12 @@ class TestOutput(TestSimulBase):
params.output.temporal_spectra.probes_region = probes_region
params.output.temporal_spectra.SAVE_AS_FLOAT32 = True
params.output.spatiotemporal_spectra.probes_region = (6, 8, 10)
nx, ny, nz = params.oper.nx, params.oper.ny, params.oper.nz
params.output.spatiotemporal_spectra.probes_region = (
nx // 2,
ny // 2,
nz // 2,
)
params.output.spatiotemporal_spectra.SAVE_AS_COMPLEX64 = True
def test_output(self):
......@@ -243,6 +248,9 @@ class TestOutput(TestSimulBase):
delta_omega = spectra["omegas"][1]
coef = delta_kz * delta_kh * delta_omega
print(spectra["kz_spectra"])
print(spectra["kh_spectra"])
for letter in "xyz":
vi_fft = series[f"v{letter}_Fourier"]
spectrum_vi = spectra["spectrum_v" + letter]
......@@ -250,9 +258,12 @@ class TestOutput(TestSimulBase):
# TODO: compute energy from vi_fft and spectrum_vi
energy_fft = (0.5 * abs(vi_fft) ** 2).mean(axis=-1).sum()
assert energy_fft > 0, (letter, vi_fft)
energy_spe = coef * spectrum_vi.sum()
energy_spe = 0.5 * coef * spectrum_vi.sum()
# TODO: fix this and plug this condition
# assert np.allclose(energy_fft, energy_spe)
assert np.allclose(energy_fft, energy_spe), (
letter,
energy_spe / energy_fft,
)
spectrum_Khd = spectra["spectrum_Khd"]
spectrum_vz = spectra["spectrum_vz"]
......@@ -260,9 +271,10 @@ class TestOutput(TestSimulBase):
# because k \cdot \hat v = 0, for kz = 0, Khd = 0
assert np.allclose(spectrum_Khd[0].sum(), 0.0)
# TODO: understand why we need the `1:`
# DONE: understand why we need the `1:`
# energy in the mode kx=ky=kz=0 is not zero, for any field in state_spect.
# because k \cdot \hat v = 0, for kh = 0, Kz = 0
assert np.allclose(spectrum_vz[1:, 0, 1:].sum(), 0.0)
assert np.allclose(spectrum_vz[1:, 0, :].sum(), 0.0)
sim3.output.spatiotemporal_spectra.plot_kzkhomega(
key_field="Khr", equation="kh=1"
......
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