lasp/test/test_fft.py

#!/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()