Simple multi-bandstop filter

bandstop.py

import pylab as plt
import numpy as np
import scipy.signal as ss
from scipy.io import wavfile
from math import pi

"""

A program that applies a multi-bandstop filter to a waveform
(eg, to filter out beeps)

-Preetum Nakkiran

"""


def bandStopImpulseResp(fStop, dF, sampRate, pts=1001):
	"""
	A manual implementation of designing a band-stop filter.
	"""

	w0 = fStop / sampRate * (2*pi)
	dW = dF / sampRate * (2*pi)	

	pts = pts/2 # only specifying half the response (by conj symmetry)
	freqResp = np.ones(pts)
	dft_i0 = w0 / pi * pts
	dft_di = dW / pi * pts
	freqResp[max(0, dft_i0 - dft_di) : dft_i0 + dft_di] = 0
	# freqResp: samples of the DTFT from 0 --> PI. (assuming impulse response real, so freq response conjugate symmetric)
	# also DFT real, even --> signal real, even

	impResp = np.fft.irfft(freqResp) # returns samples t=0 --> t=pts
	impResp = np.roll(impResp, len(impResp)/2) # will be periodic-- shift in time (still linear phase)
	# TODO: need to window?

	return impResp

# returns equivalent results as above function
def bandStopImpulseRespScipy(fStop, dF, sampRate, pts=1001):
	nyqFreq = sampRate / 2.0
	f1 = (fStop-dF) / nyqFreq
	f2 = (fStop+dF) / nyqFreq

	impResp = ss.firwin(pts, [f1, f2])
	return impResp


def multiBandStopImpulseResp(stopBands, sampRate, pts=1001):
	"""
	Returns the impulse response of a multi band-stop filter.
	stopBands: A list of (fStop, dF) pairs,
				specifying a stop band of (fStop +/- dF)
	sampRate: The sampling rate
	pts: The filter order (numtaps)
	"""
	nyqFreq = sampRate / 2.0

	stopBands.sort()

	freqs = []
	for fStop, dF in stopBands:
		fBandStart = max(0.0, (fStop-dF) / nyqFreq) # > 0
		fBandEnd = min(1.0, (fStop+dF) / nyqFreq) # < nyquest
		freqs.append(fBandStart)
		freqs.append(fBandEnd)

	impResp = ss.firwin(pts, freqs) #multi-bandstop
	return impResp


rate, data = wavfile.read('input.wav')
signal = np.array(data, dtype=float)
signal /= np.amax(np.abs(signal))

# The stop bands and bandwidths.
band1 = (3520.0, 200.0) # 3520 +/- 200 Hz
band2 = (2500.0, 100.0)
band3 = (1600.0, 100.0)

impResp = multiBandStopImpulseResp([band1, band2, band3], rate, pts=2001)

filtered = ss.fftconvolve(signal, impResp)

#renormed = filtered/np.max(np.abs(filtered))
gain = 1.5
output = np.clip(filtered * gain, -1.0, 1.0)
output = np.int16(output * (2**15 - 1))
wavfile.write('output.wav', rate, output)