Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #!/usr/bin/python3
- import math
- try:
- import matplotlib.pyplot as plt
- has_matplotlib = True
- except:
- has_matplotlib = False
- import argparse
- defaultParameters = {
- 'window_size': 1,
- 'num_points': 1024,
- }
- class Complex:
- def __init__(self, real, im):
- self.real = real
- self.im = im
- def __add__(self, n):
- if type(n) is Complex:
- return Complex(self.real+n.real, self.im+n.im)
- else:
- return Complex(self.real+n, self.im)
- def __mul__(self, n):
- if type(n) is Complex:
- return Complex(self.real*n.real - self.im*n.im, self.real*n.im + self.im*n.real)
- else:
- return Complex(self.real*n, self.im*n)
- def __truediv__(self, n):
- if type(n) is Complex:
- m2 = n.real*n.real + n.im*n.im
- return self*Complex(n.real/m2, -n.im/m2)
- else:
- return Complex(self.real/n, self.im/n)
- def __str__(self):
- return "("+str(self.real)+","+str(self.im)+")"
- def mod(self):
- return math.sqrt(self.real*self.real + self.im*self.im)
- def e(v):
- return Complex(math.cos(v), math.sin(v))
- def create_samples(window, n_samples, function):
- offset = window/n_samples
- samples = list()
- for i in range(0, n_samples):
- v = i*offset
- samples.append(function(v))
- return samples
- def perform_dft(t_samples):
- n_samples = len(t_samples)
- f_samples = [Complex(0, 0)]*n_samples
- for k in range(0, n_samples):
- for n in range(0, n_samples):
- v = -2*math.pi*k*n/n_samples
- f_samples[k] += (Complex.e(v)*t_samples[n])
- return f_samples
- sines = lambda x: 3+2*math.sin(50*x)+math.cos(90*x)
- def ramp(x):
- if x < math.pi/2:
- return x*2/math.pi
- if x < 3*math.pi/2:
- return 2-x*2/math.pi
- if x < 2*math.pi:
- return -4+x*2/math.pi
- def parseArgs():
- parser = argparse.ArgumentParser(description='Playing with DFT')
- parser.add_argument('-w', '--window', help='Window size (default = %.2f)' % defaultParameters['window_size'], default=defaultParameters['window_size'], type=float)
- parser.add_argument('-n', '--num_points', help='Number of points (default = %d)' % defaultParameters['num_points'], default=defaultParameters['num_points'], type=int)
- parser.add_argument('-p', '--plot', help='Plot time and frequency series', action='store_true')
- args = parser.parse_args()
- return args
- def plot(args, t_samples, f_samples):
- # Drop all samples for frequencies higher than Fsampling/2
- m = [round(v.mod(), 3) for v in f_samples[0:int(args.num_points/2)]]
- f, axa = plt.subplots(2)
- # Get current axes
- axa[0].set_title('DFT example')
- axa[0].set_xlim([-0.5, args.window + 0.5])
- axa[0].set_xlabel('Time')
- axa[0].plot([x*args.window/args.num_points for x in list(range(0, args.num_points))], t_samples, color = 'b')
- axa[0].grid()
- delta_freq = 2*math.pi/args.window_size
- x = [v*delta_freq for v in list(range(0, int(args.num_points/2)))]
- axa[1].set_xlim([-0.5, x[-1]+0.5])
- axa[1].set_ylabel('Power Spectral Density')
- axa[1].set_xlabel('Frequency')
- axa[1].stem(x, m, 'r-')
- axa[1].grid()
- plt.show()
- def show_dft(args, f_samples):
- delta_freq = 2*math.pi/args.window
- print("+---------+---------+")
- print("| %-7s | %-7s |" % ("Freq", "PSD"))
- print("+---------+---------+")
- for i, v in enumerate(f_samples[0:int(args.num_points/2)]):
- print("| %7.3f | %7.3f |" %(round(i*delta_freq,3), round(v.mod(),3)))
- print("+---------+---------+")
- # Main program !
- f=sines
- args = parseArgs()
- t_samples = create_samples(args.window, args.num_points, f)
- f_samples = perform_dft(t_samples)
- # Normalize value, then take only half the spectrum
- for i,v in enumerate(f_samples):
- f_samples[i] = f_samples[i]*(2/args.num_points)
- f_samples[0] = f_samples[0]/2
- if args.plot:
- if has_matplotlib:
- plot(args, t_samples, f_samples)
- else:
- print("Plotting is not possible. Please install matplotlib.")
- else:
- show_dft(args, f_samples)
Add Comment
Please, Sign In to add comment