dft.py

import time

from matplotlib import pyplot as plt
import numpy as np


def fold(x, base, k):
    return sum(x * base ** k)


def dft(x, N):
    assert len(x) == N

    base = np.arange(0, N)
    base = np.exp(-2 * np.pi * 1j * base / N)

    return [fold(x, base, k) for k in range(N)]


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 = dft(x, N)
    end = time.time()

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

    # plot power spectrum
    plt.plot(t, np.abs(X) ** 2)
    plt.savefig('dft_{}_{}.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.002465963363647461
started: fp=1, N=1024
elapsed_time: 0.18504905700683594
started: fp=1, N=2048
elapsed_time: 0.7435791492462158
started: fp=2, N=128
elapsed_time: 0.0024840831756591797
started: fp=2, N=1024
elapsed_time: 0.18305587768554688
started: fp=2, N=2048
elapsed_time: 0.7439591884613037
started: fp=3, N=128
elapsed_time: 0.002466917037963867
started: fp=3, N=1024
elapsed_time: 0.18359899520874023
started: fp=3, N=2048
elapsed_time: 0.7415070533752441
"""