Signal filtering utility

signalFilter.py

# Filtering utilities

import numpy as np
from numpy import pi
from numpy import fft
from scipy import signal
import scipy.signal as sig
import pylab as plt

#### Functions #################

def theta(x): 
    if x >=0: return 1
    else: return 0
vtheta = np.vectorize(theta)

def lowPass(s,tau):
    return 1./(1 + 1j*2*pi*s*tau)

def highPass(s,tau):
    return (1 + 1j*2*pi*s*tau)

def bandPass(s,f0=1.,deltaf=1.):
    return deltaf/(2*pi) / ( (deltaf/2)**2 + (s-f0)**2 )

def allPass(s,tau):
    return 1.


## Time array
N = int(1e5)
dt = 1e-6
T = N*dt
t = np.arange(N)*dt

f = fft.fftfreq(N,d=dt)
f = fft.fftshift(f)

## Signal

#x = np.sin(2*pi*5e4*t)
x = np.random.normal(0,1.0,(N))
X = fft.fft(x)/np.sqrt(T)
X = fft.fftshift(X)

## Filtered signal
#Y = X * lowPass(f,1e-5)

Y = X * bandPass(f,f0=1e5,deltaf=1e3)

y = fft.ifft( fft.ifftshift(Y) ) * np.sqrt(T)


#### Plots #####################
plt.figure(figsize=(8,8))
plt.subplot(311)                    # time traces
plt.plot(t,x, alpha=0.5, label="Raw")
plt.plot(t,y, label="Filtered")
plt.ylabel("$x$")
plt.xlabel("$t$")
plt.legend(loc='upper right')

plt.subplot(312)
plt.xlabel("$f$")
plt.ylabel("$|Y_T(f)|^2$")          # spectra
#plt.plot(f,np.abs(X)**2, alpha=0.5, label="Raw")
plt.plot(f,np.abs(Y)**2, 'g', label="Filtered")

plt.subplot(313)
plt.xlabel("$f$")
plt.ylabel("$W_{YY}(f)$")           # single sided spectral density of filtered signal
plt.loglog(f,2 * np.abs(Y)**2,'g')
plt.grid(axis='both')

plt.tight_layout()