fft.py

import time

from matplotlib import pyplot as plt
import numpy as np


def main(f, N, postfix=''):
    # compute the value of function
    t = np.arange(0, N)
    x = f(t, N)

    plt.plot(t, x)
    plt.savefig('wave_{}_{}.png'.format(postfix, N))
    plt.cla()

    # discrete fourier transition
    start = time.time()
    X = np.fft.fft(x)
    end = time.time()

    print('elapsed_time: {}'.format(end - start))

    # plot power spectrum
    plt.plot(t, np.abs(X) ** 2)
    plt.savefig('fft_{}_{}.png'.format(postfix, N))
    plt.cla()


if __name__ == '__main__':
    fs = [
        lambda t, N: np.sin(2 * np.pi * t / N * 10),
        lambda t, N: np.sin(2 * np.pi * t / N) + np.sin(2 * np.pi * t / N * 10),
        lambda t, N: np.sin(2 * np.pi * t / N) + np.sin(2 * np.pi * t / N * 10) + np.sin(2 * np.pi * t / N * 20) / 5
    ]

    for fp in [1, 2, 3]:
        f = fs[fp - 1]
        for N in [128, 1024, 2048]:
            print('started: fp={}, N={}'.format(fp, N))
            main(f, N, postfix=fp)

"""
started: fp=1, N=128
elapsed_time: 0.00033926963806152344
started: fp=1, N=1024
elapsed_time: 0.00012087821960449219
started: fp=1, N=2048
elapsed_time: 0.00020194053649902344
started: fp=2, N=128
elapsed_time: 8.702278137207031e-05
started: fp=2, N=1024
elapsed_time: 7.200241088867188e-05
started: fp=2, N=2048
elapsed_time: 0.0001506805419921875
started: fp=3, N=128
elapsed_time: 4.982948303222656e-05
started: fp=3, N=1024
elapsed_time: 6.29425048828125e-05
started: fp=3, N=2048
elapsed_time: 0.00012493133544921875
"""