lasp/test/test_fft.py

49 lines
1.1 KiB
Python
Raw Permalink Normal View History

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 15 19:45:33 2018
@author: anne
"""
import numpy as np
from lasp import Fft, getFreq
def test_forward_fft():
"""
Test that our FFT implementation equals Numpy's rfft implementation
"""
nfft = 2048
t = np.linspace(0, 1.0, nfft, endpoint=False)
freq = 10
omg = 2*np.pi*freq
sig = np.cos(omg*t)+10
sig = np.random.randn(nfft)
fft_lasp = Fft(nfft)
res_lasp = fft_lasp.fft(sig)
res_npy = np.fft.rfft(sig)
assert(np.isclose(np.linalg.norm(res_lasp- res_npy), 0))
def test_backward_fft():
"""
Test that our backward FFT implementation equals Numpy's rfft implementation
"""
nfft = 2048
freq = getFreq(nfft, nfft)
# Sig = np.zeros(nfft//2+1, dtype=complex)
Sigr = np.random.randn(nfft//2+1)
Sigi = np.random.randn(nfft//2+1)
Sig = Sigr + 1j*Sigi
fft_lasp = Fft(nfft)
sig_lasp = fft_lasp.ifft(Sig)
sig_py = np.fft.irfft(Sig)
assert(np.isclose(np.linalg.norm(sig_py- sig_lasp), 0))
if __name__ == '__main__':
test_forward_fft()
test_backward_fft()