47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
import numpy as np
|
|
from tests.node_tests._shared import make_field
|
|
|
|
|
|
def test_fft2d():
|
|
from backend.nodes.fft_2d import FFT2D
|
|
|
|
node = FFT2D()
|
|
|
|
N = 64
|
|
y, x = np.mgrid[0:N, 0:N] / N
|
|
freq = 5
|
|
data = np.sin(2 * np.pi * freq * x)
|
|
field = make_field(data=data, xreal=1e-6, yreal=1e-6)
|
|
|
|
spectrum, spec_mag, spec_phase, spec_psdf = node.process(field, windowing="none", level="none")
|
|
assert spectrum.data.shape == (N, N)
|
|
assert spectrum.domain == "frequency"
|
|
assert spectrum.si_unit_xy == "1/m"
|
|
centre = N // 2
|
|
row = spectrum.data[centre, :]
|
|
peak_idx = np.argmax(row[centre + 1:]) + centre + 1
|
|
assert abs(peak_idx - (centre + freq)) <= 1, f"Peak at {peak_idx}, expected ~{centre + freq}"
|
|
|
|
_, spec_mag, _, _ = node.process(field, windowing="hann", level="mean")
|
|
assert spec_mag.data.shape == (N, N)
|
|
assert np.all(spec_mag.data >= 0)
|
|
|
|
_, _, spec_phase, _ = node.process(field, windowing="none", level="none")
|
|
assert spec_phase.data.shape == (N, N)
|
|
assert spec_phase.data.min() >= -np.pi - 0.01
|
|
assert spec_phase.data.max() <= np.pi + 0.01
|
|
|
|
_, _, _, spec_psdf = node.process(field, windowing="hamming", level="plane")
|
|
assert spec_psdf.data.shape == (N, N)
|
|
assert np.all(spec_psdf.data >= 0)
|
|
assert "^2" in spec_psdf.si_unit_z
|
|
|
|
const_field = make_field(data=np.ones((32, 32)) * 3.0)
|
|
_, spec_const, _, _ = node.process(const_field, windowing="none", level="none")
|
|
centre32 = 16
|
|
dc_val = spec_const.data[centre32, centre32]
|
|
assert dc_val == spec_const.data.max()
|
|
|
|
spec_bk, _, _, _ = node.process(field, windowing="blackman", level="none")
|
|
assert spec_bk.data.shape == (N, N)
|