1import cmsisdsp as dsp
2import cmsisdsp.fixedpoint as f
3
4import numpy as np
5from scipy import signal
6#import matplotlib.pyplot as plt
7import scipy.fft
8
9import colorama
10from colorama import init,Fore, Back, Style
11from numpy.testing import assert_allclose
12
13init()
14
15def printTitle(s):
16    print("\n" + Fore.GREEN + Style.BRIGHT +  s + Style.RESET_ALL)
17
18def printSubTitle(s):
19    print("\n" + Style.BRIGHT + s + Style.RESET_ALL)
20
21
22def chop(A, eps = 1e-6):
23    B = np.copy(A)
24    B[np.abs(A) < eps] = 0
25    return B
26
27# For fixed point version, compare that
28# the conjugate part is really the conjugate part
29def compareWithConjugatePart(r):
30    res = r[0::2] + 1j * r[1::2]
31    conjPart = res[nb:nb//2:-1].conj()
32    refPart = res[1:nb//2]
33    assert(np.equal(refPart , conjPart).all())
34
35nb = 32
36signal = np.cos(2 * np.pi * np.arange(nb) / nb)*np.cos(0.2*2 * np.pi * np.arange(nb) / nb)
37
38ref=scipy.fft.rfft(signal)
39invref = scipy.fft.irfft(ref)
40
41assert(len(ref) == (nb // 2) + 1)
42assert(len(invref) == nb)
43
44# Length of arrays for the float implementation
45# of the RFFT (in float so there is a factor 2
46# when the samples are complex)
47
48RFFT_F_IN_LENGTH = nb # real
49RFFT_F_OUT_LENGTH = nb # complex (so nb // 2 complex)
50
51RIFFT_F_IN_LENGTH = nb # complex
52RIFFT_F_OUT_LENGTH = nb # real
53
54# Length of arrays for the fixed point implementation
55# of the RFFT
56RFFT_Q_IN_LENGTH = nb
57RFFT_Q_OUT_LENGTH = 2*nb
58
59# Conjugate part ignored
60RIFFT_Q_IN_LENGTH = nb + 2
61RIFFT_Q_OUT_LENGTH = nb
62
63# Convert ref to CMSIS-DSP format
64referenceFloat=np.zeros(nb)
65# Replace complex datatype by real datatype
66referenceFloat[0::2] = np.real(ref)[:-1]
67referenceFloat[1::2] = np.imag(ref)[:-1]
68# Copy Nyquist frequency value into first
69# sample.This is just a storage trick so that the
70# output of the RFFT has same length as input
71# It is legacy behavior that we need to keep
72# for backward compatibility but it is not
73# very pretty
74referenceFloat[1] = np.real(ref[-1])
75
76referenceFixed=np.zeros(2*len(ref))
77referenceFixed[0::2] = np.real(ref)
78referenceFixed[1::2] = np.imag(ref)
79
80printTitle("RFFT FAST F64")
81
82printSubTitle("RFFT")
83
84
85rfftf64=dsp.arm_rfft_fast_instance_f64()
86status=dsp.arm_rfft_fast_init_f64(rfftf64,nb)
87result = dsp.arm_rfft_fast_f64(rfftf64,signal,0)
88assert(len(signal) == RFFT_F_IN_LENGTH)
89assert(len(result) == RFFT_F_OUT_LENGTH)
90
91assert_allclose(referenceFloat,result)
92
93printSubTitle("RIFFT")
94
95rifftf64=dsp.arm_rfft_fast_instance_f64()
96status=dsp.arm_rfft_fast_init_f64(rifftf64,nb)
97result = dsp.arm_rfft_fast_f64(rifftf64,referenceFloat,1)
98assert(len(referenceFloat) == RIFFT_F_IN_LENGTH)
99assert(len(result) == RIFFT_F_OUT_LENGTH)
100
101assert_allclose(invref,result,atol=1e-15)
102
103printTitle("RFFT FAST F32")
104
105printSubTitle("RFFT")
106
107
108rfftf32=dsp.arm_rfft_fast_instance_f32()
109status=dsp.arm_rfft_fast_init_f32(rfftf32,nb)
110result = dsp.arm_rfft_fast_f32(rfftf32,signal,0)
111assert(len(signal) == RFFT_F_IN_LENGTH)
112assert(len(result) == RFFT_F_OUT_LENGTH)
113
114
115assert_allclose(referenceFloat,result,rtol=3e-6)
116
117printSubTitle("RIFFT")
118
119rifftf32=dsp.arm_rfft_fast_instance_f32()
120status=dsp.arm_rfft_fast_init_f32(rifftf32,nb)
121result = dsp.arm_rfft_fast_f32(rifftf32,referenceFloat,1)
122assert(len(referenceFloat) == RIFFT_F_IN_LENGTH)
123assert(len(result) == RIFFT_F_OUT_LENGTH)
124
125assert_allclose(invref,result,atol=1e-7)
126
127# Fixed point
128
129printTitle("RFFT Q31")
130
131printSubTitle("RFFT")
132
133signalQ31 = f.toQ31(signal)
134rfftQ31=dsp.arm_rfft_instance_q31()
135status=dsp.arm_rfft_init_q31(rfftQ31,nb,0,1)
136resultQ31 = dsp.arm_rfft_q31(rfftQ31,signalQ31)
137assert(len(signalQ31) == RFFT_Q_IN_LENGTH)
138assert(len(resultQ31) == RFFT_Q_OUT_LENGTH)
139compareWithConjugatePart(resultQ31)
140
141# Drop the conjugate part which is not computed by scipy
142resultQ31 = resultQ31[:nb+2]
143assert(len(resultQ31) == RIFFT_Q_IN_LENGTH)
144
145resultF = f.Q31toF32(resultQ31) * nb
146
147assert_allclose(referenceFixed,resultF,rtol=1e-6,atol=1e-6)
148
149
150printSubTitle("RIFFT")
151
152rifftQ31=dsp.arm_rfft_instance_q31()
153status=dsp.arm_rfft_init_q31(rifftQ31,nb,1,1)
154# Apply CMSIS-DSP scaling
155referenceQ31 = f.toQ31(referenceFixed/ nb)
156resultQ31 = dsp.arm_rfft_q31(rifftQ31,referenceQ31)
157resultF = f.Q31toF32(resultQ31)
158assert(len(referenceQ31) == RIFFT_Q_IN_LENGTH)
159assert(len(resultQ31) == RIFFT_Q_OUT_LENGTH)
160
161assert_allclose(invref/nb,resultF,atol=1e-6)
162
163printTitle("RFFT Q15")
164
165printSubTitle("RFFT")
166
167signalQ15 = f.toQ15(signal)
168rfftQ15=dsp.arm_rfft_instance_q15()
169status=dsp.arm_rfft_init_q15(rfftQ15,nb,0,1)
170resultQ15 = dsp.arm_rfft_q15(rfftQ15,signalQ15)
171assert(len(signalQ15) == RFFT_Q_IN_LENGTH)
172assert(len(resultQ15) == RFFT_Q_OUT_LENGTH)
173compareWithConjugatePart(resultQ15)
174
175# Drop the conjugate part which is not computed by scipy
176resultQ15 = resultQ15[:nb+2]
177assert(len(resultQ15) == RIFFT_Q_IN_LENGTH)
178
179resultF = f.Q15toF32(resultQ15) * nb
180
181assert_allclose(referenceFixed,resultF,rtol=1e-6,atol=1e-2)
182
183
184printSubTitle("RIFFT")
185
186rifftQ15=dsp.arm_rfft_instance_q15()
187status=dsp.arm_rfft_init_q15(rifftQ15,nb,1,1)
188# Apply CMSIS-DSP scaling
189referenceQ15 = f.toQ15(referenceFixed / nb)
190resultQ15 = dsp.arm_rfft_q15(rifftQ15,referenceQ15)
191resultF = f.Q15toF32(resultQ15)
192assert(len(referenceQ15) == RIFFT_Q_IN_LENGTH)
193assert(len(resultQ15) == RIFFT_Q_OUT_LENGTH)
194
195assert_allclose(invref/nb,resultF,atol=1e-3)
196
197