Python script to compute discrete Fourier transform.

discrete_fourier_transform.py

#! /usr/local/bin/python3.6
"""
Discrete Fourier transform

  f(t) = 2 * sin(4 * t) + 3 * cos(2 * t)
         ( 0 <= t < 2 * pi )
"""
import math
import sys
import traceback


class DiscreteFourierTransform:
    N = 100                # Number of division
    CSV_DFT  = "DFT.csv"   # Output file (DFT)
    CSV_IDFT = "IDFT.csv"  # Output file (IDFT)

    def __init__(self):
        self.src_re, self.src_im = [], []
        self.dft_re, self.dft_im = [], []

    def make_source_data(self):
        """ Maiking source data """
        try:
            for i in range(self.N):
                val  = 2 * math.sin(4 * (2 * math.pi / self.N) * i) \
                     + 3 * math.cos(2 * (2 * math.pi / self.N) * i)
                self.src_re.append(val)
                self.src_im.append(0.0)
        except Exception as e:
            raise

    def dft(self):
        """ Discrete Fourier Transformation """
        try:
            with open(self.CSV_DFT, "w") as f:
                f.write("k,f,x_re,x_im,X_re,X_im\n")
                for k in range(self.N):
                    dft_re,dft_im = 0.0, 0.0
                    for n in range(self.N):
                        v_re = self.src_re[n] \
                             * ( math.cos((2 * math.pi / self.N) * k * n)) \
                             + self.src_im[n] \
                             * ( math.sin((2 * math.pi / self.N) * k * n))
                        v_im = self.src_re[n] \
                             * (-math.sin((2 * math.pi / self.N) * k * n)) \
                             + self.src_im[n] \
                             * ( math.cos((2 * math.pi / self.N) * k * n))
                        dft_re += v_re
                        dft_im += v_im
                    self.dft_re.append(dft_re)
                    self.dft_im.append(dft_im)
                    f.write("{:d},".format(k))
                    f.write("{:.6f},".format((2 * math.pi / self.N) * k))
                    f.write("{:.6f},".format(self.src_re[k]))
                    f.write("{:.6f},".format(self.src_im[k]))
                    f.write("{:.6f},".format(dft_re))
                    f.write("{:.6f}\n".format(dft_im))
        except Exception as e:
            raise

    def idft(self):
        """ Inverse Discrete Fourier Transformation """
        try:
            with open(self.CSV_IDFT, "w") as f:
                f.write("k,f,X_re,X_im,x_re,x_im\n")
                for n in range(self.N):
                    idft_re, idft_im  = 0.0, 0.0
                    for k in range(self.N):
                        v_re = self.dft_re[k] \
                             * (math.cos((2 * math.pi / self.N) * k * n)) \
                             - self.dft_im[k] \
                             * (math.sin((2 * math.pi / self.N) * k * n))
                        v_im = self.dft_re[k] \
                             * (math.sin((2 * math.pi / self.N) * k * n)) \
                             + self.dft_im[k] \
                             * (math.cos((2 * math.pi / self.N) * k * n))
                        idft_re += v_re
                        idft_im += v_im
                    idft_re /= self.N
                    idft_im /= self.N
                    f.write("{:d},".format(n))
                    f.write("{:.6f},".format((2 * math.pi / self.N) * n))
                    f.write("{:.6f},".format(self.dft_re[n]))
                    f.write("{:.6f},".format(self.dft_im[n]))
                    f.write("{:.6f},".format(idft_re))
                    f.write("{:.6f}\n".format(idft_im))
        except Exception as e:
            raise


if __name__ == '__main__':
    try:
        obj = DiscreteFourierTransform()
        obj.make_source_data()
        obj.dft()
        obj.idft()
    except Exception as e:
        traceback.print_exc()
        sys.exit(1)