КГ Лаба 2(3)

import matplotlib.pyplot as plot
from matplotlib import mlab
from tkinter.filedialog import askopenfilename
import pylab
import numpy as np
import scipy as sp
from PIL import Image, ImageDraw
import random, math, copy


def exp(x):
    a = 0.5
    return math.exp(-(a**2) * (x**2))

def fraction(x):
    b = 0.2
    return 1 / (1 + (b**2) * (x**2))

def sinFraction(x):
    a = 0.9
    b = 0.6
    return math.sin(a * x) / (1 + (b**2) * (x**2))

def myfft(fx):
    fx = np.asarray(fx, dtype=float) #for shape func
    N = fx.shape[0] # len of array
    result = np.zeros(N,dtype=complex) # new array with zero values
    for u in range(N):
        z = np.zeros(2)
        for k in range(N):
            p = 2 * sp.pi * k * u / N
            z[0] += fx[k] * sp.cos(p) + 0 * sp.sin(p) # real part
            z[1] += 0 * sp.cos(p) - fx[k] * sp.sin(p) # imag part
        result[u] = z[0] + z[1]*1j

    return result

# initialize x = [1, 2, 3, 4 ...]
SIZE = 256
xlist = [0] * SIZE
for i in range(SIZE):
    xlist[i] = i
mode = 2

# fx - exponenta
fx = [exp(x) for x in xlist]

#if (mode == 1):
    # My fast fourier transform
res = myfft(fx)
pylab.plot (xlist, res.real)
pylab.plot (xlist, res.imag, '--')
pylab.legend ( ("real part", "imagine part"), title = "Exponent (homemade)" )
pylab.show()
#if (mode == 2):
    # standart fft from Numpy for checking answer
res = np.fft.fft(fx)

pylab.plot (xlist, res.real)
pylab.plot (xlist, res.imag, '--')
pylab.legend ( ("real part", "imagine part"), title = "Exponent" )
pylab.show()


# fx - Fraction
fx = [fraction(x) for x in xlist]

if (mode == 1):
    res = myfft(fx)
if (mode == 2):
    res = np.fft.fft(fx)

pylab.plot (xlist, res.real)
pylab.plot (xlist, res.imag, '--')
pylab.legend ( ("real part", "imagine part"), title = "Fraction" )
pylab.show()



# fx - Fraction with sinus
fx = [sinFraction(x) for x in xlist]

if (mode == 1):
    res = myfft(fx)
if (mode == 2):
    res = np.fft.fft(fx)

pylab.plot (xlist, res.real)
pylab.plot (xlist, res.imag,'--')
pylab.legend ( ("real part", "imagine part"), title = "Fraction with sinus" )
pylab.show()