from __future__ import print_function
import sys
import scipy.signal
import numpy as np
import matplotlib
matplotlib.use('WebAgg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
def plot(title, figure_number, response, p=None):
fig = plt.figure(figure_number, dpi=78, figsize=(10,4))
w, h, lag = response
w /= np.pi * 2
index_max_h = np.argmax(h)
w_at_max = w[index_max_h]
lag_at_max = lag[index_max_h]
amp = 20 * np.log10(abs(h))
if lag[0] == 0:
lag[0] = np.nan # avoid ugly vertical line at left edge of graph
if p is not None:
critical_w = 1./p
real_critical_w = 0
prev_amp = amp[0]
for i in range(len(amp)):
if prev_amp > -3 and amp[i] < -3:
real_critical_w = w[i]
break
if prev_amp < -3 and amp[i] > -3:
real_critical_w = w[i]
break
prev_amp = amp[i]
if real_critical_w != 0:
title += ", real critical freq %.3f = %.1f bars" % (real_critical_w, 1./real_critical_w,)
plt.suptitle(title)
ax = plt.subplot2grid((1, 2),(0, 0))
ax.xaxis.set_major_locator(ticker.MultipleLocator(.1))
ax.xaxis.set_minor_locator(ticker.MultipleLocator(.01))
ax.set_ylim(max(-40, min(amp)-3), max(3, max(amp)+3))
if p is not None:
plt.axvline(critical_w)
amp_at_p = amp[min(len(amp) - 1, np.searchsorted(w, critical_w))]
plt.title("Amplitude at period %d is %.2f" % (p, amp_at_p))
elif index_max_h > 0 and index_max_h < len(h) - 1:
plt.title("Max response at frequency %.3f = %.1f bars" % (w_at_max, float(1./w_at_max)))
plt.plot(w, amp, 'b')
plt.ylabel('Amplitude [dB]')
plt.xlabel('Frequency')
ax = plt.subplot2grid((1, 2),(0, 1))
ax.set_ylim(-.9, 10.1, auto=False)
ax.xaxis.set_major_locator(ticker.MultipleLocator(.1))
ax.xaxis.set_minor_locator(ticker.MultipleLocator(.01))
plt.axhline(0)
if index_max_h > 0 and index_max_h < len(h) - 1:
plt.axvline(w_at_max)
plt.title("Lag at max response %.2f" % (lag_at_max,))
elif p is not None:
plt.axvline(critical_w)
lag_at_p = lag[min(len(lag) - 1, np.searchsorted(w, critical_w))]
plt.title("Lag at period %d is %.2f" % (p, lag_at_p,))
plt.plot(w, lag, 'r')
plt.ylabel('Lag samples')
plt.xlabel('Frequency')
plt.tight_layout(pad=2)
def ema(period, unused=None):
title = "EMA(%d)" % period
alpha = 2. / (period + 1)
numerator = [alpha] # coeffs for data[0], data[1], data[2], ...
denominator = [1, -(1 - alpha)] # first elem always 1, then -coeffs for filt[1], filt[2], ...
return title, numerator, denominator, period
def sma(period, unused=None):
title = "Simple moving average with Period = %d" % period
numerator = [1./period for _ in range(period)]
denominator = [1]
return title, numerator, denominator, period
def hp2(period, unused=None):
title = "Two-pole highpass filter with Period = %d" % period
f = 0.707106781 * 3.141592654 / period
alpha = 1 + (np.sin(f) - 1) / np.cos(f)
b = 1 - alpha / 2
c = 1 - alpha
numerator = [b*b, -2*b*b, b*b]
denominator = [1, -2*c, c*c]
return title, numerator, denominator, period
def corrected_hp2(period, unused=None):
title = "Corrected two-pole highpass filter with Period = %d" % period
f = 2 * 0.707106781 * 3.141592654 / period
alpha = 1 + (np.sin(f) - 1) / np.cos(f)
b = 1 - alpha / 2
c = 1 - alpha
numerator = [b*b, -2*b*b, b*b]
denominator = [1, -2*c, c*c]
return title, numerator, denominator, period
def butterworth_hp2(period, unused=None):
title = "Two-pole butterworth highpass filter with Period = %d" % period
t = np.tan(np.pi / period)
k = 1 / (1 + t*np.sqrt(2) + t*t)
numerator = [k, -2*k, k]
denominator = [1, k*(2*t*t - 2), k*(1 - t*np.sqrt(2) + t*t)]
return title, numerator, denominator, period
def smooth(period, unused=None):
title = "Super smoother with Period = %d" % period
f = 1.414213562 * 3.141592654 / period
a = np.exp(-f)
c2 = 2 * a * np.cos(f)
c3 = -(a**2)
c1 = 1 - c2 - c3
numerator = [c1/2, c1/2]
denominator = [1, -c2, -c3]
return title, numerator, denominator, period
def roof(hp, lp):
title = "Roofing filter with HP = %d and LP = %d" % (hp, lp)
_, numeratorHP, denominatorHP, _ = hp2(hp)
_, numeratorSM, denominatorSM, _ = smooth(lp)
numerator = np.convolve(numeratorHP, numeratorSM)
denominator = np.convolve(denominatorHP, denominatorSM)
return title, numerator, denominator, None
def corrected_roof(hp, lp):
title = "Corrected roofing filter with HP = %d and LP = %d" % (hp, lp)
_, numeratorHP, denominatorHP, _ = corrected_hp2(hp)
_, numeratorSM, denominatorSM, _ = smooth(lp)
numerator = np.convolve(numeratorHP, numeratorSM)
denominator = np.convolve(denominatorHP, denominatorSM)
return title, numerator, denominator, None
def butterworth_roof(hp, lp):
title = "Butterworth Roofing filter with HP = %d and LP = %d" % (hp, lp)
_, numeratorHP, denominatorHP, _ = butterworth_hp2(hp)
_, numeratorSM, denominatorSM, _ = smooth(lp)
numerator = np.convolve(numeratorHP, numeratorSM)
denominator = np.convolve(denominatorHP, denominatorSM)
return title, numerator, denominator, None
def calc_and_plot(filter_type, figure_number, p, q):
if p is None:
print("== Missing period for", filter_type)
return
title, numerator, denominator, p = globals()[filter_type](p, q)
w, h = scipy.signal.freqz(b=numerator, a=denominator)
w, lag = scipy.signal.group_delay((numerator, denominator))
response = (w, h, lag)
plot(title, figure_number, response, p)
return figure_number + 1
if __name__ == "__main__":
figure_number = 1
filter_type, p_, q_ = [None] * 3
for f in sys.argv[1:]:
try:
val = int(f)
if p_ is None: p_ = val
elif q_ is None: q_ = val
except:
if filter_type is not None:
figure_number = calc_and_plot(filter_type, figure_number, p_, q_)
filter_type = f
p_, q_ = None, None
figure_number = calc_and_plot(filter_type, figure_number, p_, q_)
plt.show()