A fancier brute-force DFT implementation using NumPy features

def dft2(f):
    f = np.asarray(f)
    N = f.shape[0]
    W = np.exp(-2 * np.pi * 1.0j / N)
    k = np.arange(N)
    jk = np.outer(k, k)
    M = W**jk
    return np.dot(M, f)

f = data_piano[:2000]
F = dft2(f)
N = len(f)
dx = 1. / sampling_rate
freq = 1.0 / (N * dx) * np.arange(N)
plt.plot(freq[:100], F.real[:100])

F_numpy = np.fft.fft(f)
frequency_numpy = np.fft.fftfreq(N, dx)
print(np.allclose(F, F_numpy))
plt.plot(frequency_numpy[:100], F_numpy[:100].real)

Our algorithm dft2 also produces the same result as fft from NumPy.

Timing with size

import random

Ntest = 10
timings2 = []
dimension2 = []

for i in range(1, Ntest):
    f = [random.gauss(0, 1) for n in range(2**i)]
    timing = %timeit -o -n 1 dft2(f)
    dimension2.append(2**i)
    timings2.append(timing.average)
plt.subplot(121)    
plt.plot(np.array(dimension2), np.array(timings2), 'o', markersize=12, color='red')
plt.plot(np.array(dimension2), np.array(timings2), color='red')
plt.xlabel("$N$")
plt.ylabel("Timing / s")

plt.subplot(122)
plt.plot(np.array(dimension2)**2, np.array(timings2), 'o', markersize=12, color='red')
plt.plot(np.array(dimension2)**2, np.array(timings2), color='red')

plt.xlabel("$N^2$")
plt.show()