Fourier Transform

dft.py

import math
from matplotlib import pyplot as plt

# Let's initialize two lists with the real
# and imaginary numbers.
real_list = [36, 22, 45, 15]
imag_list = [0, 0, 0, 0]

results_f = []  # Forward
results_i = []  # Inverse

N = len(real_list)

# Forward DFT
for u in range(N):
    real = imag = 0
    for x in range(N):
        cosine =  math.cos(2. * math.pi * u * x/N)
        sine   = -math.sin(2. * math.pi * u * x/N)

        real += (real_list[x] * cosine) - (imag_list[x] * sine)
        imag += (real_list[x] * sine) + (imag_list[x] * cosine)

    result = [real / N, imag / N]
    results_f.append(result)

# Coefficients
print('Coefficients:')
for u in range(N):
    print('({:.2f}, {:.2f}),'.format(results_f[u][0], results_f[u][1]))

# Magnitude
mag = []
print('\nMagnitude:')
for u in range(N):
    mag_val = math.sqrt(results_f[u][0] * results_f[u][0] +
                        results_f[u][1] * results_f[u][1])
    print('{:.2f},'.format(mag_val))
    mag.append(mag_val)

# Phase
pha = []
print('\nPhase:')
for u in range(N):
    pha_val = math.atan2(results_f[u][1], results_f[u][0])
    print('{:.2f},'.format(pha_val))
    pha.append(pha_val)

# Now let's plot these two
plt.subplot(121), plt.plot(mag, 'bo')
plt.title('Magnitude')
plt.subplot(122), plt.plot(pha, 'g*')
plt.title('Phase')
plt.savefig('mag-phase.png')

# Now let's go back

# Inverse
for x in range(N):
    real = imag = 0
    for u in range(N):
        cosine = math.cos(2. * math.pi * u * x/N)
        sine   = math.sin(2. * math.pi * u * x/N)

        real += (results_f[u][0] * cosine) - (results_f[u][1] * sine)
        imag += (results_f[u][0] * sine) + (results_f[u][1] * cosine)

    result = [real, imag]
    results_i.append(result)

# Our List
print('\nOur List:')
for u in range(N):
    print('({:.2f}, {:.2f}),'.format(results_i[u][0], abs(results_i[u][1])))