Intuitive impementation of discrete Fourier transform (and inverse DFT) in Python (without numpy)

mydft.py

#!/usr/bin/env python3

import math

# >> Discrete Fourier transform for sampled signals
# x [in]: sampled signals, a list of magnitudes (real numbers)
# yr [out]: real parts of the sinusoids
# yi [out]: imaginary parts of the sinusoids
def dft(x):
  N, yr, yi = len(x), [], []
  for k in range(N):
    real, imag = 0, 0
    for n in range(N):
      theta = -k * (2 * math.pi) * (float(n) / N)
      real += x[n] * math.cos(theta)
      imag += x[n] * math.sin(theta)
    yr.append(real / N)
    yi.append(imag / N)
  return yr, yi


# >> Inverse discrete Fourier transform
# yr [in]: real parts of the sinusoids
# yi [in]: imaginary parts of the sinusoids
# x [out]: sampled signals (real parts only), a list of magnitude
def idft(yr, yi):
  N, x = len(yr), []
  for n in range(N):
    real, imag = 0, 0
    for k in range(N):
      theta = k * (2 * math.pi) * (float(n) / N)
      real += (yr[k] * math.cos(theta)) - (yi[k] * math.sin(theta))
      # imag += (yr[k] * math.sin(theta)) + (yi[k] * math.cos(theta))
    x.append(real)
  return x


if __name__ == '__main__':
  x = [1, 2, 3, 5]
  yr, yi = dft(x) # y=[2.75+0j, -0.5+0.75j, 0.75+0j, -0.5-0.75j]
  x2 = idft(yr, yi) # x2=[1, 2, 3, 5]