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 c64a7932 authored by Pierre Augier's avatar Pierre Augier
Browse files

Transonic for fluidsim/solvers/ns3d/output/spatiotemporal_spectra.py

parent a402660d0c1c
......@@ -11,31 +11,40 @@ Provides:
"""
from pathlib import Path
from math import pi
import numpy as np
from scipy import signal
import h5py
from fluidsim.base.output.spatiotemporal_spectra import SpatioTemporalSpectra
from transonic import boost, Array, Type
A4 = Array[Type(np.float32, np.float64), "4d", "C"]
A3 = "float[:,:,:]"
A1 = "float[:]"
def loop_spectra_kzkhomega(spectrum_k0k1k2omega, khs, KH, kzs, KZ):
@boost
def loop_spectra_kzkhomega(
spectrum_k0k1k2omega: A4, khs: A1, KH: A3, kzs: A1, KZ: A3
):
"""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)
spectrum_kzkhomega = np.zeros(
(nkz, nkh, nomega), dtype=spectrum_k0k1k2omega.dtype
)
for ik0 in range(nk0):
for ik1 in range(nk1):
for ik2 in range(nk2):
value = spectrum_k0k1k2omega[ik0, ik1, ik2, :]
values = spectrum_k0k1k2omega[ik0, ik1, ik2, :]
kappa = KH[ik0, ik1, ik2]
ikh = int(kappa / deltakh)
kz = abs(KZ[ik0, ik1, ik2])
......@@ -44,11 +53,15 @@ def loop_spectra_kzkhomega(spectrum_k0k1k2omega, khs, KH, kzs, KZ):
ikz = nkz - 1
if ikh >= nkh - 1:
ikh = nkh - 1
spectrum_kzkhomega[ikz, ikh, :] += value
for i, value in enumerate(values):
spectrum_kzkhomega[ikz, ikh, i] += value
else:
coef_share = (kappa - khs[ikh]) / deltakh
spectrum_kzkhomega[ikz, ikh, :] += (1 - coef_share) * value
spectrum_kzkhomega[ikz, ikh + 1, :] += coef_share * value
for i, value in enumerate(values):
spectrum_kzkhomega[ikz, ikh, i] += (
1 - coef_share
) * value
spectrum_kzkhomega[ikz, ikh + 1, i] += coef_share * value
# get one-sided spectrum in the omega dimension
nomega = nomega // 2 + 1
......@@ -105,7 +118,7 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
if not key.startswith("spectrum_"):
continue
spectra_kzkhomega[key] = loop_spectra_kzkhomega(
data, kh_spectra, KH, kz_spectra, KZ
np.ascontiguousarray(data), kh_spectra, KH, kz_spectra, KZ
)
del spectra
......@@ -144,7 +157,7 @@ class SpatioTemporalSpectraNS3D(SpatioTemporalSpectra):
if not key.startswith("spectrum_"):
continue
spectra_urud_kzkhomega[key] = loop_spectra_kzkhomega(
data, kh_spectra, KH, kz_spectra, KZ
np.ascontiguousarray(data), kh_spectra, KH, kz_spectra, KZ
)
spectra_kzkhomega[key] = spectra_urud_kzkhomega[key]
......
......@@ -133,6 +133,7 @@ def transonize():
"fluidsim/solvers/ns3d/forcing.py",
"fluidsim/util/mini_oper_modif_resol.py",
"fluidsim/base/output/spatiotemporal_spectra.py",
"fluidsim/solvers/ns3d/output/spatiotemporal_spectra.py",
]
make_backend_files([here / path for path in paths], backend=TRANSONIC_BACKEND)
......
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