Read about our upcoming Code of Conduct on this issue

This instance will be upgraded to Heptapod 0.28.1 on 2022-01-26 at 16:00 UTC+1 (a few minutes of down time)

Commit b4f056ba authored by Pierre Augier's avatar Pierre Augier
Browse files

spatiotemporal_spectra: more testing + cleanup

parent 3fad3d217485
......@@ -333,7 +333,7 @@ class SpatioTemporalSpectra(SpecificOutput):
create_ds("probes_ik2_loc", data=self.probes_ik2_loc)
for key in self.keys_fields:
create_ds(
f"spect_{key}_loc",
key + "_Fourier_loc",
(self.probes_nb_loc, 1),
maxshape=(self.probes_nb_loc, None),
dtype=self.datatype,
......@@ -352,9 +352,9 @@ class SpatioTemporalSpectra(SpecificOutput):
dset.resize((self.probes_nb_loc, self.number_times_in_file))
dset[:, -1] = v
def _add_probes_data_to_dict(self, data_dict, key):
def _add_probes_data_to_dict(self, data, key):
"""Probes fields in Fourier space and append data to a dict object"""
data_dict[f"spect_{key}_loc"] = self.sim.state.get_var(f"{key}_fft")[
data[key + "_Fourier_loc"] = self.sim.state.get_var(f"{key}_fft")[
self.probes_ik0_loc, self.probes_ik1_loc, self.probes_ik2_loc
]
......@@ -402,7 +402,7 @@ class SpatioTemporalSpectra(SpecificOutput):
order = file.attrs["dims_order"]
region = file.attrs["probes_region"]
if dtype is None:
dtype = file[f"spect_{keys[0]}_loc"].dtype
dtype = file[keys[0] + "_Fourier_loc"].dtype
# get list of useful files, from tmin
tmins_files = np.array([float(p.name[14:-3]) for p in paths_1st_rank])
......@@ -439,7 +439,7 @@ class SpatioTemporalSpectra(SpecificOutput):
iksmin = np.array([1 - ikzmax, 1 - ikymax, 0])
ik0max, ik1max, ik2max = iksmax[order]
ik0min, ik1min, ik2min = iksmin[order]
spect_shape = (
shape_series = (
ik0max + 1 - ik0min,
ik1max + 1 - ik1min,
ik2max + 1 - ik2min,
......@@ -447,7 +447,9 @@ class SpatioTemporalSpectra(SpecificOutput):
)
# load series, rebuild as state_spect arrays + time
series = {f"spect_{k}": np.empty(spect_shape, dtype=dtype) for k in keys}
series = {
f"{k}_Fourier": np.empty(shape_series, dtype=dtype) for k in keys
}
with Progress() as progress:
task_ranks = progress.add_task("Rearranging...", total=len(ranks))
task_files = progress.add_task("Rank 00000...", total=npaths)
......@@ -493,7 +495,7 @@ class SpatioTemporalSpectra(SpecificOutput):
# load data at desired times for all keys_fields
for key in keys:
skey = f"spect_{key}"
skey = key + "_Fourier"
data = file[skey + "_loc"][:, its_file]
for i, it in enumerate(its):
series[skey][ik0, ik1, ik2, it] = data[:, i]
......@@ -541,10 +543,11 @@ class SpatioTemporalSpectra(SpecificOutput):
spectra = {k: v for k, v in series.items() if k.startswith("K")}
for key, data in series.items():
if not key.startswith("spect"):
if "_Fourier" not in key:
continue
key_spectrum = "spectrum_" + key.split("_Fourier")[0]
freq, spectrum = signal.periodogram(data, fs=f_sample)
spectra[key] = spectrum
spectra[key_spectrum] = spectrum
spectra["omegas"] = 2 * pi * freq
spectra["dims_order"] = series["dims_order"]
......
......@@ -47,7 +47,7 @@ class PrintStdOutNS3D(PrintStdOutBase):
if nt > 1:
nt -= 1
it = np.zeros(nt, dtype=np.int)
it = np.zeros(nt, dtype=int)
t = np.zeros(nt)
deltat = np.zeros(nt)
......
......@@ -20,49 +20,48 @@ import h5py
from fluidsim.base.output.spatiotemporal_spectra import SpatioTemporalSpectra
class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
def loop_spectra_kzkhomega(self, spectrum_k0k1k2omega, khs, KH, kzs, KZ):
"""Compute the kz-kh-omega spectrum."""
if spectrum_k0k1k2omega.dtype == "complex64":
dtype = "float32"
else:
dtype = "float64"
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), dtype=dtype)
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
# get one-sided spectrum in the omega dimension
nomega = nomega // 2 + 1
spectrum_onesided = np.zeros((nkz, nkh, nomega))
spectrum_onesided[:, :, 0] = spectrum_kzkhomega[:, :, 0]
spectrum_onesided[:, :, 1:] = 0.5 * (
spectrum_kzkhomega[:, :, 1:nomega]
+ spectrum_kzkhomega[:, :, -1:-nomega:-1]
)
return spectrum_onesided / (deltakz * deltakh)
def loop_spectra_kzkhomega(spectrum_k0k1k2omega, khs, KH, kzs, KZ):
"""Compute the kz-kh-omega spectrum."""
if spectrum_k0k1k2omega.dtype == "complex64":
dtype = "float32"
else:
dtype = "float64"
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), dtype=dtype)
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
# get one-sided spectrum in the omega dimension
nomega = nomega // 2 + 1
spectrum_onesided = np.zeros((nkz, nkh, nomega))
spectrum_onesided[:, :, 0] = spectrum_kzkhomega[:, :, 0]
spectrum_onesided[:, :, 1:] = 0.5 * (
spectrum_kzkhomega[:, :, 1:nomega]
+ spectrum_kzkhomega[:, :, -1:-nomega:-1]
)
return spectrum_onesided / (deltakz * deltakh)
class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
def save_spectra_kzkhomega(
self, tmin=0, tmax=None, dtype=None, save_urud=False
):
......@@ -103,26 +102,26 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
"omegas": omegas_onesided,
}
for key, data in spectra.items():
if not key.startswith("spect"):
if not key.startswith("spectrum_"):
continue
spectra_kzkhomega[key] = self.loop_spectra_kzkhomega(
spectra_kzkhomega[key] = loop_spectra_kzkhomega(
data, kh_spectra, KH, kz_spectra, KZ
)
del spectra
# total kinetic energy
spectra_kzkhomega["spect_K"] = 0.5 * (
spectra_kzkhomega["spect_vx"]
+ spectra_kzkhomega["spect_vy"]
+ spectra_kzkhomega["spect_vz"]
spectra_kzkhomega["spectrum_K"] = 0.5 * (
spectra_kzkhomega["spectrum_vx"]
+ spectra_kzkhomega["spectrum_vy"]
+ spectra_kzkhomega["spectrum_vz"]
)
# potential energy
try:
N = self.sim.params.N
spectra_kzkhomega["spect_A"] = (
0.5 / N ** 2 * spectra_kzkhomega["spect_b"]
spectra_kzkhomega["spectrum_A"] = (
0.5 / N ** 2 * spectra_kzkhomega["spectrum_b"]
)
except AttributeError:
pass
......@@ -138,20 +137,23 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
# toroidal/poloidal decomposition
if save_urud:
print("Computing ur, ud spectra...")
spectra_kzkhomega = {}
spectra_urud_kzkhomega = {}
spectra = self.compute_spectra_urud(tmin=tmin, tmax=tmax, dtype=dtype)
for key, data in spectra.items():
if not key.startswith("spect"):
if not key.startswith("spectrum_"):
continue
spectra_kzkhomega[key] = self.loop_spectra_kzkhomega(
spectra_urud_kzkhomega[key] = loop_spectra_kzkhomega(
data, kh_spectra, KH, kz_spectra, KZ
)
spectra_kzkhomega[key] = spectra_urud_kzkhomega[key]
with h5py.File(path_file, "a") as file:
for key, val in spectra_kzkhomega.items():
for key, val in spectra_urud_kzkhomega.items():
file.create_dataset(key, data=val)
return spectra_kzkhomega
def plot_kzkhomega(
self,
key_field=None,
......@@ -159,7 +161,7 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
tmax=None,
dtype=None,
equation=None,
cmap=None,
cmap="viridis",
vmin=None,
vmax=None,
):
......@@ -168,8 +170,6 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
key_field = self.keys_fields[0]
if tmax is None:
tmax = self.sim.params.time_stepping.t_end
if cmap is None:
cmap = "viridis"
path_file = path_file = (
Path(self.sim.output.path_run) / "spatiotemporal_spectra.h5"
......@@ -177,10 +177,9 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
spectra_kzkhomega = {}
key_spect = "spect_" + key_field
key_spect = "spectrum_" + key_field
# load spectra from file if it exists
if path_file.is_file():
if path_file.exists():
# we should check if times match?
print("loading spectra from file...")
with h5py.File(path_file, "r") as file:
......@@ -197,7 +196,7 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
spectra_kzkhomega["omegas"] = file["omegas"][...]
else:
# compute spectra and save to file, then load
if key_spect.startswith("spect_Kh"):
if key_spect.startswith("spectrum_Kh"):
save_urud = True
else:
save_urud = False
......@@ -282,7 +281,13 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
vmax = np.log10(spect[np.isfinite(spect)].max())
im = ax.pcolormesh(
xaxis, yaxis, np.log10(spect), cmap=cmap, vmin=vmin, vmax=vmax
xaxis,
yaxis,
np.log10(spect),
cmap=cmap,
vmin=vmin,
vmax=vmax,
shading="nearest",
)
fig.colorbar(im)
......@@ -328,8 +333,8 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
# toroidal/poloidal decomposition
# urx_fft, ury_fft contain shear modes!
vx_fft = series["spect_vx"]
vy_fft = series["spect_vy"]
vx_fft = series["vx_Fourier"]
vy_fft = series["vy_Fourier"]
if vx_fft.dtype == "complex64":
float_dtype = "float32"
elif vx_fft.dtype == "complex64":
......@@ -365,18 +370,18 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
del series
# ud
spectra["spect_Khd"] = np.zeros(udx_fft.shape, dtype=dtype)
spectra["spectrum_Khd"] = np.zeros(udx_fft.shape, dtype=dtype)
freq, spectrum = signal.periodogram(udx_fft, fs=f_sample)
spectra["spect_Khd"] += 0.5 * spectrum
spectra["spectrum_Khd"] += 0.5 * spectrum
freq, spectrum = signal.periodogram(udy_fft, fs=f_sample)
spectra["spect_Khd"] += 0.5 * spectrum
spectra["spectrum_Khd"] += 0.5 * spectrum
# ur
spectra["spect_Khr"] = np.zeros(udx_fft.shape, dtype=dtype)
spectra["spectrum_Khr"] = np.zeros(udx_fft.shape, dtype=dtype)
freq, spectrum = signal.periodogram(urx_fft, fs=f_sample)
spectra["spect_Khr"] += 0.5 * spectrum
spectra["spectrum_Khr"] += 0.5 * spectrum
freq, spectrum = signal.periodogram(ury_fft, fs=f_sample)
spectra["spect_Khr"] += 0.5 * spectrum
spectra["spectrum_Khr"] += 0.5 * spectrum
spectra["omegas"] = 2 * pi * freq
spectra["dims_order"] = order
......
......@@ -241,7 +241,7 @@ class TestNoShearModes(TestSimulBase):
EKhs = sim.output.spatial_means.load()["EKhs"]
E = sim.output.spatial_means.load()["E"]
ratio = EKhs[-1] / E[-1]
self.assertGreater(1e-15, ratio)
assert ratio < 1e-15
# check energy in shear modes
spectrum = data["Khr"] + data["Khd"] + data["Kz"] + data["A"]
......
......@@ -120,6 +120,13 @@ class TestOutput(TestSimulBase):
def test_output(self):
sim = self.sim
# put energy in vz
vz = sim.state.state_phys.get_var("vz")
X, Y, Z = sim.oper.get_XYZ_loc()
vz += 0.05 * np.cos(2 * pi * X / sim.oper.Lx)
sim.state.statespect_from_statephys()
sim.time_stepping.start()
# testing phaseshift
......@@ -223,6 +230,40 @@ class TestOutput(TestSimulBase):
delta_index_times=2
)
sim3.output.temporal_spectra.save_spectra()
spatiotemporal_spectra = sim3.output.spatiotemporal_spectra
series = spatiotemporal_spectra.load_time_series()
spectra = spatiotemporal_spectra.save_spectra_kzkhomega(
save_urud=True
)
delta_kz = spectra["kz_spectra"][1]
delta_kh = spectra["kh_spectra"][1]
delta_omega = spectra["omegas"][1]
coef = delta_kz * delta_kh * delta_omega
for letter in "xyz":
vi_fft = series[f"v{letter}_Fourier"]
spectrum_vi = spectra["spectrum_v" + letter]
# 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()
# TODO: fix this and plug this condition
# assert np.allclose(energy_fft, energy_spe)
spectrum_Khd = spectra["spectrum_Khd"]
spectrum_vz = spectra["spectrum_vz"]
# 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:`
# because k \cdot \hat v = 0, for kh = 0, Kz = 0
assert np.allclose(spectrum_vz[1:, 0, 1:].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