Simple and stupid signal and bandstop butterworth filter

signal_and_filter-bandstop.py

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


freq1 = 60  # Hz
freq2 = 90  # Hz
sampling_rate = 1000  # samples per second
time = 1  # segundos
t = np.linspace(0, time, time * sampling_rate)

# sinal com as duas frequencias
y = np.sin(2*np.pi*freq1*t) + np.sin(2*np.pi*freq2*t)
# y + ruido
yn = y * 0.3 * np.random.rand(len(t))
dt = 250  # faixa de observacao

# plot
plt.figure(tight_layout=True)
plt.subplot(811)
plt.plot(t[:dt], y[:dt])
plt.title('y(t)')
plt.subplot(812)
plt.plot(t[:dt], yn[:dt])
plt.title('yn(t)')
plt.grid()
plt.xlabel('t(s)')

# fft
Y = np.abs(np.fft.fft(y))
Yn = np.abs(np.fft.fft(yn))
f = np.arange(0, len(Y)) * sampling_rate/len(Y)
df = int(200 * len(y) / sampling_rate)

#fft plots
plt.subplot(813)
plt.plot(f[:df], Y[:df])
plt.title('|Y(f)|')
plt.subplot(814)
plt.plot(f[:df], Yn[:df])
plt.title('|Yn(f)|')
plt.grid()
plt.xlabel('f(Hz)')

# filter bandstop - reject 60 Hztop)
lowcut = 59
highcut = 61
order = 4
nyq = 0.5 * sampling_rate
low = lowcut / nyq
high = highcut / nyq
b, a = signal.butter(order, [low, high], btype='bandstop')

y_filt = signal.filtfilt(b, a, y)
yn_filt = signal.filtfilt(b, a, yn)

plt.subplot(815)
plt.plot(t[:dt], y_filt[:dt])
plt.title('y_filt(t)')
plt.subplot(816)
plt.plot(t[:dt], yn_filt[:dt])
plt.title('yn_filt(t)')
plt.grid()
plt.xlabel('t(s)')

# fft
Y = np.abs(np.fft.fft(y_filt))
Yn = np.abs(np.fft.fft(yn_filt))
f = np.arange(0, len(Y)) * sampling_rate/len(Y)
df = int(200 * len(y_filt) / sampling_rate)

#fft plots
plt.subplot(817)
plt.plot(f[:df], Y[:df])
plt.title('|Y(f) filered|')
plt.subplot(818)
plt.plot(f[:df], Yn[:df])
plt.title('|Yn(f)| filered')
plt.grid()
plt.xlabel('f(Hz)')
plt.show()