DSP, FFT, Numpy and more

dsp.py

import matplotlib.pyplot as plt
import numpy as np
from scipy import signal

# Our complex impulse response.
imp = signal.unit_impulse(128, dtype=complex)
x = np.linspace(0, 1, 128)
freq_domain = np.linspace(-64, 64, 128)

# Effectively, we have a bunch of 0's and 0+j's, there should be
# exactly 128 of them. The first one should have a real value of
# 1 and an imaginary value of 1+j. First, let's graph both the
# real and complex numbers in the time domain.
# plt.plot(x, imp)
#
# # This is required to display the imaginary numbers.
# # img = [y.imag for y in imp]
# # plt.plot(x, img)
#
# plt.show()

# The imaginary numbers don't really add any value for this particular 
# demonstration, so we are going to remove them from now on. Now let's plot
# the fft values of the impulse response on the frequency domain.

# fig = plt.figure(1)
#
# ax1 = fig.add_subplot(211)
# ax1.plot(x, imp)
# ax1.set_ylim((-1, 1))
# ax1.set_ylabel('Amplitude')
# ax1.set_xlabel('Time')
# ax1.set_title('Impulse Response in the Time Domain')
#
# ax2 = fig.add_subplot(212)
# freq_domain = np.linspace(-70, 70, 128)
# y = np.fft.fft(imp)
# ax2.plot(freq_domain, y)
# ax2.set_xlabel('Frequency')
# ax2.set_ylabel('Amplitude')
#
# plt.show()


# OK. Now, let's run the impulse response through a filter, plot the impulse
# response, and the filter on the same plot, then display the result of the
# the filter on the frequency domain by performing an FFT on it. First we
# are going to run the impulse response through a 4th-order butterworth
# lowpass filter

b, a = signal.butter(4, 0.2)
response = signal.lfilter(b, a, imp)

fig2 = plt.figure(2)

filter_plot = fig2.add_subplot(211)
filter_plot.set_ylabel('Amplitude')
filter_plot.set_xlabel('Time')
filter_plot.set_title('Impulse response & Butterworth lowpass filter')
filter_plot.plot(x, imp)
filter_plot.plot(x, response)

filter_freq = fig2.add_subplot(212)
filter_freq.set_title('FFT of Impulse and Filter')
filter_freq.set_ylabel('Amplitude')
filter_freq.set_xlabel('Frequency')
filter_freq.plot(freq_domain, np.fft.fft(response))
filter_freq.plot(freq_domain, np.fft.fft(imp))
filter_freq.legend()

plt.show()