Files
rohde-engineering-competition/fft/real_fft_test.py
2026-02-18 21:32:22 +01:00

84 lines
2.3 KiB
Python

"""This script is used to visualize the behavior of the fft if the input is
purely real."""
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
def bit_reverse(val: int, num_bits: int) -> None:
b = '{:0{num_bits}b}'.format(val, num_bits=num_bits)
return int(b[::-1], 2)
def fft_reorder(x: np.array) -> np.array:
result = np.zeros(x.size)
for i in range(x.size):
result[i] = x[bit_reverse(i, int(np.log2(x.size)))]
return result
def fft(x: np.array) -> np.array:
x = fft_reorder(x).astype('complex64')
N0 = x.size
# fig, axes = plt.subplots(2, int(np.log2(N0)))
N = 2
while N <= N0:
first_subfft_index = 0
while first_subfft_index < N0:
temp = x[first_subfft_index + N // 2]
x[first_subfft_index + N // 2] = x[first_subfft_index] - temp
x[first_subfft_index] = x[first_subfft_index] + temp
for k in range(1, N // 2):
factor = np.exp(-2j * np.pi * k / N)
x[first_subfft_index + k] \
= x[first_subfft_index + k] \
+ factor * x[first_subfft_index + k + N // 2]
for k in range(1, N // 2):
x[first_subfft_index + N - k] \
= x[first_subfft_index + k].conj()
# \
# = x[first_subfft_index + k] + temp
# for k in range(N // 2):
# factor = np.exp(- 2j * np.pi * k / N)
#
# temp = factor * x[first_subfft_index + k + N // 2]
# x[first_subfft_index + k + N // 2] \
# = x[first_subfft_index + k] - temp
# x[first_subfft_index + k] \
# = x[first_subfft_index + k] + temp
first_subfft_index += N
# axes[0][int(np.log2(N)) - 1].plot(x.real[1:])
# axes[1][int(np.log2(N)) - 1].plot(x.imag[1:])
N *= 2
plt.show()
return x
def main():
sns.set_theme()
x = np.arange(16) + 1
X = fft(x)
X_reference = np.fft.fft(x)
for X_i, X_ref_i in zip(X, X_reference):
print(
f"{X_i.real:06.2f} + j({X_i.imag:06.2f})\t\t{X_ref_i.real:06.2f} "
f"+ j({X_ref_i.imag:06.2f})")
print(np.allclose(X, X_reference, rtol=1e-05, atol=1e-08))
if __name__ == "__main__":
main()