In [15]:

    

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import firwin, freqz, hilbert

def plot_filter(b=None, a=None, f=None, cx=False):
    if f != None:
        assert a == None
        assert b == None
        w,h = freqz(f)
    else:
        w,h = freqz(b, a)
    plt.plot(w/max(w), np.abs(h), color="steelblue")
    if cx:
        plt.plot(-w/max(w), np.abs(h), color="darkred")
    ax = plt.gca()
    ax.set_xlabel("Normalized frequency")
    ax.set_ylabel("Gain")

def notch(Wn, bandwidth):
     """
     Notch filter to kill line-noise.
     """
     f = Wn / 2.0
     R = 1.0 - 3.0 * (bandwidth / 2.0)
     num = 1.0 - 2.0 * R * np.cos(2 * np.pi * f) + R ** 2.
     denom = 2.0 - 2.0 * np.cos(2 * np.pi * f) 
     K = num / denom 
     b = np.zeros(3)
     a = np.zeros(3)
     a[0] = 1.0
     a[1] = -2.0 * R * np.cos(2 * np.pi * f)
     a[2] = R ** 2.
     b[0] = K
     b[1] = -2.0 * K * np.cos(2 * np.pi * f)
     b[2] = K
     return b, a
    
lp = firwin(11, .18, pass_zero=True)
hp = firwin(11, .23, pass_zero=False)
plot_filter(f=lp)
plot_filter(f=hp)
plt.title("OLD SCHOOL LP-HP")
plt.figure()

#Assuming 250Hz ~= fs, 50Hz = 1/5 = .2
#5 Hz bw = 1/50 = .02
bw = .05
freq = .2
b, a = notch(freq, bw)
plot_filter(b, a)
plt.title("NOTCH IIR for %s taps"%len(a))
plt.figure()

#FIR designer
f = firwin(31, [.15, .25])
plot_filter(f=f)
plt.title("FIRWIN")
plt.show()


    

In [ ]: