rtfir/example.py at main · vfiksdal/rtfir · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
#!/usr/bin/env python3
#
# Script to test realtime fir filters
#
import matplotlib.pyplot as plt
import numpy as np
import sys,rtfir


def chirp(fs_Hz, rep_Hz, f0_Hz, f1_Hz, phase_rad=0):
    # Estimate timing
    T_s=1/rep_Hz
    c=(f1_Hz-f0_Hz)/T_s
    n=int(fs_Hz/rep_Hz)
    t_s=np.linspace(0,T_s,n)

    # Phase, phi_Hz, is integral of frequency, f(t) = ct + f0.
    phi_Hz=(c*t_s**2)/2+(f0_Hz*t_s)
    phi_rad=2*np.pi*phi_Hz
    phi_rad+=phase_rad
    return np.sin(phi_rad)

def fft(unfiltered,filtered,fs,maxf):
    # Calculate fft of unfiltered and filtered data
    N=len(unfiltered)
    T=1.0/fs
    length=int(maxf/(fs/float(N)))
    xf=np.linspace(0.0, 1.0/(2.0*T),int(N/2)+1)[:length]
    unfilteredf=np.abs(np.fft.rfft(unfiltered)[:length])
    filteredf=np.abs(np.fft.rfft(filtered)[:length])
    return unfilteredf,filteredf,xf

# Test setup
span=100.0
fs=1000.0
t=60
taps=64
i=1
while i<len(sys.argv):
    if sys.argv[i]=='--samplerate':
        i+=1
        fs=float(sys.argv[i])
    elif sys.argv[i]=='--length':
        i+=1
        t=float(sys.argv[i])
    elif sys.argv[i]=='--taps':
        i+=1
        taps=int(sys.argv[i])
    elif sys.argv[i]=='--span':
        i+=1
        span=float(sys.argv[i])
    elif sys.argv[i]=='--help':
        print('RTFIR Python Example 1.0')
        print('Tests RTFIR implementation for C++ and python')
        print('Fiksdal(C)2021')
        print('')
        print('Usage: '+sys.argv[0]+'')
        print('')
        print('Options:')
        print('\t--samplerate HZ\tSet sampling frequency in hertz')
        print('\t--length S\tLength of synthesized recording in seconds')
        print('\t--taps N\tNumber of taps for the FIR filter')
        print('\t--span HZ\tLength of frequency sweep to test with in Hz')
        print('')
        exit()
    else:
        print('Invalid parameter: '+sys.argv[i])
        exit()
    i+=1


# Create a frequency sweep of 1-100 Hz
sweep=chirp(fs,1/t,1.0,span)
flow=span/8.0*3.0
fhigh=span/8.0*5.0

# Create lowpass filter
lpf=rtfir.RTFIR_lowpass(taps,flow/fs)

# Create highpass filter
hpf=rtfir.RTFIR_highpass(taps,fhigh/fs)

# Create bandpass filter
bpf=rtfir.RTFIR_bandpass(taps,flow/fs,fhigh/fs)

# Create bandstop filter
bsf=rtfir.RTFIR_bandstop(taps,flow/fs,fhigh/fs)


# Filter chirp
lowpass=np.zeros(len(sweep))
highpass=np.zeros(len(sweep))
bandpass=np.zeros(len(sweep))
bandstop=np.zeros(len(sweep))
for i in range(0,len(sweep)):
    lowpass[i]=lpf.Filter(sweep[i])
    highpass[i]=hpf.Filter(sweep[i])
    bandpass[i]=bpf.Filter(sweep[i])
    bandstop[i]=bsf.Filter(sweep[i])


# Plot FFT data
fig, axs = plt.subplots(2, 2)
unfilteredf,filteredf,xf=fft(sweep,lowpass,fs,span*1.1)
axs[0, 0].plot(xf,unfilteredf)
axs[0, 0].plot(xf,filteredf)
axs[0, 0].set_title('Lowpass')
unfilteredf,filteredf,xf=fft(sweep,highpass,fs,span*1.1)
axs[0, 1].plot(xf,unfilteredf)
axs[0, 1].plot(xf,filteredf)
axs[0, 1].set_title('Highpass')
unfilteredf,filteredf,xf=fft(sweep,bandpass,fs,span*1.1)
axs[1, 0].plot(xf,unfilteredf)
axs[1, 0].plot(xf,filteredf)
axs[1, 0].set_title('Bandpass')
unfilteredf,filteredf,xf=fft(sweep,bandstop,fs,span*1.1)
axs[1, 1].plot(xf,unfilteredf)
axs[1, 1].plot(xf,filteredf)
axs[1, 1].set_title('Bandstop')
for ax in axs.flat: ax.set(xlabel='Frequency [Hz]',ylabel='Magnitude')
for ax in axs.flat: ax.label_outer()
fig.suptitle('Spectrum of filtered chirp\nFlow='+str(flow)+'Hz Fhigh='+str(fhigh)+'Hz Taps='+str(taps)+'\nFs='+str(fs)+'Hz t='+str(t)+'s')
fig.set_size_inches(11.7,8.3)
fig.tight_layout()
plt.savefig('test_fft.png')