1 #!/usr/bin/python
3 """
4 Define some FFT wrappers to reduce clutter.
5 Provides a unitary discrete FFT and a windowed version.
6 Based on numpy.fft.rfft.
8 Main entry functions:
9   unitary_rfft(data, freq=1.0)
10   power_spectrum(data, freq=1.0)
11   unitary_power_spectrum(data, freq=1.0)
12   avg_power_spectrum(data, freq=1.0, chunk_size=2048, overlap=True, window=window_hann)
13   unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048, overlap=True, window=window_hann)
14 """
16 from numpy import log2, floor, round, ceil, abs, pi, exp, cos, sin, sqrt, \
17     sinc, arctan2, array, ones, arange, linspace, zeros, \
18     uint16, float, concatenate, fromfile, argmax, complex
19 from numpy.fft import rfft
22 # print time- and freq- space plots of the test transforms if True
23 TEST_PLOTS = False
24 #TEST_PLOTS = True
26 def floor_pow_of_two(num) :
27     "Round num down to the closest exact a power of two."
28     lnum = log2(num)
29     if int(lnum) != lnum :
30         num = 2**floor(lnum)
31     return num
33 def round_pow_of_two(num) :
34     "Round num to the closest exact a power of two on a log scale."
35     lnum = log2(num)
36     if int(lnum) != lnum :
37         num = 2**round(lnum)
38     return num
40 def ceil_pow_of_two(num) :
41     "Round num up to the closest exact a power of two."
42     lnum = log2(num)
43     if int(lnum) != lnum :
44         num = 2**ceil(lnum)
45     return num
47 def _test_rfft(xs, Xs) :
48     # Numpy's FFT algoritm returns
49     #          n-1
50     #   X[k] = SUM x[m] exp (-j 2pi km /n)
51     #          m=0
52     # (see http://www.tramy.us/numpybook.pdf)
53     j = complex(0,1)
54     n = len(xs)
55     Xa = []
56     for k in range(n) :
57         Xa.append(sum([x*exp(-j*2*pi*k*m/n) for x,m in zip(xs,range(n))]))
58         if k < len(Xs):
59             assert (Xs[k]-Xa[k])/abs(Xa[k]) < 1e-6, \
60                 "rfft mismatch on element %d: %g != %g, relative error %g" \
61                 % (k, Xs[k], Xa[k], (Xs[k]-Xa[k])/abs(Xa[k]))
62     # Which should satisfy the discrete form of Parseval's theorem
63     #   n-1               n-1
64     #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
65     #   m=0               k=0
66     timeSum = sum([abs(x)**2 for x in xs])
67     freqSum = sum([abs(X)**2 for X in Xa])
68     assert abs(freqSum/float(n) - timeSum)/timeSum < 1e-6, \
69         "Mismatch on Parseval's, %g != 1/%d * %g" % (timeSum, n, freqSum)
71 def _test_rfft_suite() :
72     print "Test numpy rfft definition"
73     xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
74     _test_rfft(xs, rfft(xs))
76 def unitary_rfft(data, freq=1.0) :
77     """
78     Compute the Fourier transform of real data.
79     Unitary (preserves power [Parseval's theorem]).
81     If the units on your data are Volts,
82     and your sampling frequency is in Hz,
83     then freq_axis will be in Hz,
84     and trans will be in Volts.
85     """
86     nsamps = floor_pow_of_two(len(data))
87     # Which should satisfy the discrete form of Parseval's theorem
88     #   n-1               n-1
89     #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
90     #   m=0               k=0
91     # However, we want our FFT to satisfy the continuous Parseval eqn
92     #   int_{-infty}^{infty} |x(t)|^2 dt = int_{-infty}^{infty} |X(f)|^2 df
93     # which has the discrete form
94     #   n-1              n-1
95     #   SUM |x_m|^2 dt = SUM |X'_k|^2 df
96     #   m=0              k=0
97     # with X'_k = AX, this gives us
98     #   n-1                     n-1
99     #   SUM |x_m|^2 = A^2 df/dt SUM |X'_k|^2
100     #   m=0                     k=0
101     # so we see
102     #   A^2 df/dt = 1/n
103     #   A^2 = 1/n dt/df
104     # From Numerical Recipes (http://www.fizyka.umk.pl/nrbook/bookcpdf.html),
105     # Section 12.1, we see that for a sampling rate dt, the maximum frequency
106     # f_c in the transformed data is the Nyquist frequency (12.1.2)
107     #   f_c = 1/2dt
108     # and the points are spaced out by (12.1.5)
109     #   df = 1/ndt
110     # so
111     #   dt = 1/ndf
112     #   dt/df = 1/ndf^2
113     #   A^2 = 1/n^2df^2
114     #   A = 1/ndf = ndt/n = dt
115     # so we can convert the Numpy transformed data to match our unitary
116     # continuous transformed data with (also NR 12.1.8)
117     #   X'_k = dtX = X / <sampling freq>
118     trans = rfft(data[0:nsamps]) / float(freq)
119     freq_axis = linspace(0, freq/2, nsamps/2+1)
120     return (freq_axis, trans)
122 def _test_unitary_rfft_parsevals(xs, freq, freqs, Xs):
123     # Which should satisfy the discretized integral form of Parseval's theorem
124     #   n-1              n-1
125     #   SUM |x_m|^2 dt = SUM |X_k|^2 df
126     #   m=0              k=0
127     dt = 1.0/freq
128     df = freqs-freqs
129     assert (df - 1/(len(xs)*dt))/df < 1e-6, \
130         "Mismatch in spacing, %g != 1/(%d*%g)" % (df, len(xs), dt)
131     Xa = list(Xs)
132     for k in range(len(Xs)-1,1,-1) :
133         Xa.append(Xa[k])
134     assert len(xs) == len(Xa), "Length mismatch %d != %d" % (len(xs), len(Xa))
135     lhs = sum([abs(x)**2 for x in xs]) * dt
136     rhs = sum([abs(X)**2 for X in Xa]) * df
137     assert abs(lhs - rhs)/lhs < 1e-4, "Mismatch on Parseval's, %g != %g" \
138         % (lhs, rhs)
140 def _test_unitary_rfft_parsevals_suite():
141     print "Test unitary rfft on Parseval's theorem"
142     xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
143     dt = pi
144     freqs,Xs = unitary_rfft(xs, 1.0/dt)
145     _test_unitary_rfft_parsevals(xs, 1.0/dt, freqs, Xs)
147 def _rect(t) :
148     if abs(t) < 0.5 :
149         return 1
150     else :
151         return 0
153 def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) :
154     "Show fft(rect(at)) = 1/abs(a) * sinc(f/a)"
155     samp_freq = float(samp_freq)
156     a = float(a)
158     x = zeros((samples,), dtype=float)
159     dt = 1.0/samp_freq
160     for i in range(samples) :
161         t = i*dt
162         x[i] = _rect(a*(t-time_shift))
163     freq_axis, X = unitary_rfft(x, samp_freq)
164     #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
166     # remove the phase due to our time shift
167     j = complex(0.0,1.0) # sqrt(-1)
168     for i in range(len(freq_axis)) :
169         f = freq_axis[i]
170         inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
171         X[i] *= inverse_phase_shift
173     expected = zeros((len(freq_axis),), dtype=float)
174     # normalized sinc(x) = sin(pi x)/(pi x)
175     # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
176     assert sinc(0.5) == 2.0/pi, "abnormal sinc()"
177     for i in range(len(freq_axis)) :
178         f = freq_axis[i]
179         expected[i] = 1.0/abs(a) * sinc(f/a)
181     if TEST_PLOTS :
182         pylab.figure()
183         pylab.subplot(211)
184         pylab.plot(arange(0, dt*samples, dt), x)
185         pylab.title('time series')
186         pylab.subplot(212)
187         pylab.plot(freq_axis, X.real, 'r.')
188         pylab.plot(freq_axis, X.imag, 'g.')
189         pylab.plot(freq_axis, expected, 'b-')
190         pylab.title('freq series')
192 def _test_unitary_rfft_rect_suite() :
193     print "Test unitary FFTs on variously shaped rectangular functions"
194     _test_unitary_rfft_rect(a=0.5)
195     _test_unitary_rfft_rect(a=2.0)
196     _test_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
197     _test_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
199 def _gaussian(a, t) :
200     return exp(-a * t**2)
202 def _test_unitary_rfft_gaussian(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) :
203     "Show fft(rect(at)) = 1/abs(a) * sinc(f/a)"
204     samp_freq = float(samp_freq)
205     a = float(a)
207     x = zeros((samples,), dtype=float)
208     dt = 1.0/samp_freq
209     for i in range(samples) :
210         t = i*dt
211         x[i] = _gaussian(a, (t-time_shift))
212     freq_axis, X = unitary_rfft(x, samp_freq)
213     #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
215     # remove the phase due to our time shift
216     j = complex(0.0,1.0) # sqrt(-1)
217     for i in range(len(freq_axis)) :
218         f = freq_axis[i]
219         inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
220         X[i] *= inverse_phase_shift
222     expected = zeros((len(freq_axis),), dtype=float)
223     for i in range(len(freq_axis)) :
224         f = freq_axis[i]
225         expected[i] = sqrt(pi/a) * _gaussian(1.0/a, pi*f) # see Wikipedia, or do the integral yourself.
227     if TEST_PLOTS :
228         pylab.figure()
229         pylab.subplot(211)
230         pylab.plot(arange(0, dt*samples, dt), x)
231         pylab.title('time series')
232         pylab.subplot(212)
233         pylab.plot(freq_axis, X.real, 'r.')
234         pylab.plot(freq_axis, X.imag, 'g.')
235         pylab.plot(freq_axis, expected, 'b-')
236         pylab.title('freq series')
238 def _test_unitary_rfft_gaussian_suite() :
239     print "Test unitary FFTs on variously shaped gaussian functions"
240     _test_unitary_rfft_gaussian(a=0.5)
241     _test_unitary_rfft_gaussian(a=2.0)
242     _test_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
243     _test_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
247 def power_spectrum(data, freq=1.0) :
248     """
249     Compute the power spectrum of DATA taken with a sampling frequency FREQ.
250     DATA must be real (not complex).
251     Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
252     If the number of samples in data is not an integer power of two,
253     the FFT ignores some of the later points.
254     """
255     nsamps = floor_pow_of_two(len(data))
257     freq_axis = linspace(0, freq/2, nsamps/2+1)
258     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
259     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
260     # See Numerical Recipies for a details.
261     trans = rfft(data[0:nsamps])
262     power = trans * trans.conj() # We want the square of the amplitude.
263     return (freq_axis, power)
265 def unitary_power_spectrum(data, freq=1.0) :
266     freq_axis,power = power_spectrum(data, freq)
267     # One sided power spectral density, so 2|H(f)|**2 (see NR 2nd edition 12.0.14, p498)
268     #
269     # numpy normalizes with 1/N on the inverse transform ifft,
270     # so we should normalize the freq-space representation with 1/sqrt(N).
271     # But we're using the rfft, where N points are like N/2 complex points, so 1/sqrt(N/2)
272     # So the power gets normalized by that twice and we have 2/N
273     #
274     # On top of this, the FFT assumes a sampling freq of 1 per second,
275     # and we want to preserve area under our curves.
276     # If our total time T = len(data)/freq is smaller than 1,
277     # our df_real = freq/len(data) is bigger that the FFT expects (dt_fft = 1/len(data)),
278     # and we need to scale the powers down to conserve area.
279     # df_fft * F_fft(f) = df_real *F_real(f)
280     # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq
281     # So the power gets normalized by *that* twice and we have 2/N * freq**2
283     # power per unit time
284     # measure x(t) for time T
285     # X(f)   = int_0^T x(t) exp(-2 pi ift) dt
286     # PSD(f) = 2 |X(f)|**2 / T
288     # total_time = len(data)/float(freq)
289     # power *= 2.0 / float(freq)**2   /   total_time
290     # power *= 2.0 / freq**2   *   freq / len(data)
291     power *= 2.0 / (freq * float(len(data)))
293     return (freq_axis, power)
295 def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) :
296     x = zeros((samples,), dtype=float)
297     samp_freq = float(samp_freq)
298     for i in range(samples) :
299         x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
300     freq_axis, power = unitary_power_spectrum(x, samp_freq)
301     imax = argmax(power)
303     expected = zeros((len(freq_axis),), dtype=float)
304     df = samp_freq/float(samples) # df = 1/T, where T = total_time
305     i = int(sin_freq/df)
306     # average power per unit time is
307     #  P = <x(t)**2>
308     # average value of sin(t)**2 = 0.5    (b/c sin**2+cos**2 == 1)
309     # so average value of (int sin(t)**2 dt) per unit time is 0.5
310     #  P = 0.5
311     # we spread that power over a frequency bin of width df, sp
312     #  P(f0) = 0.5/df
313     # where f0 is the sin's frequency
314     #
315     # or :
316     # FFT of sin(2*pi*t*f0) gives
317     #  X(f) = 0.5 i (delta(f-f0) - delta(f-f0)),
318     # (area under x(t) = 0, area under X(f) = 0)
319     # so one sided power spectral density (PSD) per unit time is
320     #  P(f) = 2 |X(f)|**2 / T
321     #       = 2 * |0.5 delta(f-f0)|**2 / T
322     #       = 0.5 * |delta(f-f0)|**2 / T
323     # but we're discrete and want the integral of the 'delta' to be 1,
324     # so 'delta'*df = 1  --> 'delta' = 1/df, and
325     #  P(f) = 0.5 / (df**2 * T)
326     #       = 0.5 / df                (T = 1/df)
327     expected[i] = 0.5 / df
329     print "The power should be a peak at %g Hz of %g (%g, %g)" % \
330         (sin_freq, expected[i], freq_axis[imax], power[imax])
331     Pexp = 0
332     P    = 0
333     for i in range(len(freq_axis)) :
334         Pexp += expected[i] *df
335         P    += power[i] * df
336     print " The total power should be %g (%g)" % (Pexp, P)
338     if TEST_PLOTS :
339         pylab.figure()
340         pylab.subplot(211)
341         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
342         pylab.title('time series')
343         pylab.subplot(212)
344         pylab.plot(freq_axis, power, 'r.')
345         pylab.plot(freq_axis, expected, 'b-')
346         pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
348 def _test_unitary_power_spectrum_sin_suite() :
349     print "Test unitary power spectrums on variously shaped sin functions"
350     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
351     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
352     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
353     _test_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024)
354     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
355     # finally, with some irrational numbers, to check that I'm not getting lucky
356     _test_unitary_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)
357     # test with non-integer number of periods
358     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256)
360 def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256) :
361     x = zeros((samples,), dtype=float)
362     samp_freq = float(samp_freq)
363     x = amp
364     freq_axis, power = unitary_power_spectrum(x, samp_freq)
366     # power = <x(t)**2> = (amp)**2 * dt/T
367     # we spread that power over the entire freq_axis [0,fN], so
368     #  P(f)  = (amp)**2 dt / (T fN)
369     # where
370     #  dt = 1/samp_freq        (sample period)
371     #  T  = samples/samp_freq  (total time of data aquisition)
372     #  fN = 0.5 samp_freq      (Nyquist frequency)
373     # so
374     #  P(f) = amp**2 / (samp_freq * samples/samp_freq * 0.5 samp_freq)
375     #       = 2 amp**2 / (samp_freq*samples)
376     expected_amp = 2.0 * amp**2 / (samp_freq * samples)
377     expected = ones((len(freq_axis),), dtype=float) * expected_amp
379     print "The power should be flat at y = %g (%g)" % (expected_amp, power)
381     if TEST_PLOTS :
382         pylab.figure()
383         pylab.subplot(211)
384         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
385         pylab.title('time series')
386         pylab.subplot(212)
387         pylab.plot(freq_axis, power, 'r.')
388         pylab.plot(freq_axis, expected, 'b-')
389         pylab.title('%g samples of delta amp %g' % (samples, amp))
391 def _test_unitary_power_spectrum_delta_suite() :
392     print "Test unitary power spectrums on various delta functions"
393     _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024)
394     _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048)
395     _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)# expected = 2*computed
396     _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)# expected = 0.5*computed
397     _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024)
398     _test_unitary_power_spectrum_delta(amp=pi, samp_freq=exp(1), samples=1024)
400 def _gaussian2(area, mean, std, t) :
401     "Integral over all time = area (i.e. normalized for area=1)"
402     return area/(std*sqrt(2.0*pi)) * exp(-0.5*((t-mean)/std)**2)
404 def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10.24 ,samples=512) : #1024
405     x = zeros((samples,), dtype=float)
406     mean = float(mean)
407     for i in range(samples) :
408         t = i/float(samp_freq)
409         x[i] = _gaussian2(area, mean, std, t)
410     freq_axis, power = unitary_power_spectrum(x, samp_freq)
412     # generate the predicted curve
413     # by comparing our _gaussian2() form to _gaussian(),
414     # we see that the Fourier transform of x(t) has parameters:
415     #  std'  = 1/(2 pi std)    (references declaring std' = 1/std are converting to angular frequency, not frequency like we are)
416     #  area' = area/[std sqrt(2*pi)]   (plugging into FT of _gaussian() above)
417     #  mean' = 0               (changing the mean in the time-domain just changes the phase in the freq-domain)
418     # So our power spectral density per unit time is given by
419     #  P(f) = 2 |X(f)|**2 / T
420     # Where
421     #  T  = samples/samp_freq  (total time of data aquisition)
422     mean = 0.0
423     area = area /(std*sqrt(2.0*pi))
424     std = 1.0/(2.0*pi*std)
425     expected = zeros((len(freq_axis),), dtype=float)
426     df = float(samp_freq)/samples # 1/total_time ( = freq_axis-freq_axis = freq_axis)
427     for i in range(len(freq_axis)) :
428         f = i*df
429         gaus = _gaussian2(area, mean, std, f)
430         expected[i] = 2.0 * gaus**2 * samp_freq/samples
431     print "The power should be a half-gaussian, ",
432     print "with a peak at 0 Hz with amplitude %g (%g)" % (expected, power)
434     if TEST_PLOTS :
435         pylab.figure()
436         pylab.subplot(211)
437         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
438         pylab.title('time series')
439         pylab.subplot(212)
440         pylab.plot(freq_axis, power, 'r.')
441         pylab.plot(freq_axis, expected, 'b-')
442         pylab.title('freq series')
444 def _test_unitary_power_spectrum_gaussian_suite() :
445     print "Test unitary power spectrums on various gaussian functions"
446     _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=1024)
447     _test_unitary_power_spectrum_gaussian(area=1, std=2, samp_freq=10.0, samples=1024)
448     _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=2048)
449     _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=20.0, samples=2048)
450     _test_unitary_power_spectrum_gaussian(area=3, std=1, samp_freq=10.0, samples=1024)
451     _test_unitary_power_spectrum_gaussian(area=pi, std=sqrt(2), samp_freq=exp(1), samples=1024)
453 def window_hann(length) :
454     "Returns a Hann window array with length entries"
455     win = zeros((length,), dtype=float)
456     for i in range(length) :
457         win[i] = 0.5*(1.0-cos(2.0*pi*float(i)/(length)))
458     # avg value of cos over a period is 0
459     # so average height of Hann window is 0.5
460     return win
462 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
463                        overlap=True, window=window_hann) :
464     """
465     Compute the avg power spectrum of DATA taken with a sampling frequency FREQ.
466     DATA must be real (not complex) by breaking DATA into chunks.
467     The chunks may or may not be overlapping (by setting OVERLAP).
468     The chunks are windowed by dotting with WINDOW(CHUNK_SIZE), FFTed,
469     and the resulting spectra are averaged together.
470     See NR 13.4 for rational.
472     Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
473     CHUNK_SIZE should really be a power of 2.
474     If the number of samples in DATA is not an integer power of CHUNK_SIZE,
475     the FFT ignores some of the later points.
476     """
477     assert chunk_size == floor_pow_of_two(chunk_size), \
478         "chunk_size %d should be a power of 2" % chunk_size
480     nchunks = len(data)/chunk_size # integer division = implicit floor
481     if overlap :
482         chunk_step = chunk_size/2
483     else :
484         chunk_step = chunk_size
486     win = window(chunk_size) # generate a window of the appropriate size
487     freq_axis = linspace(0, freq/2, chunk_size/2+1)
488     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
489     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
490     # See Numerical Recipies for a details.
491     power = zeros((chunk_size/2+1,), dtype=float)
492     for i in range(nchunks) :
493         starti = i*chunk_step
494         stopi = starti+chunk_size
495         fft_chunk = rfft(data[starti:stopi]*win)
496         p_chunk = fft_chunk * fft_chunk.conj()
497         power += p_chunk.astype(float)
498     power /= float(nchunks)
499     return (freq_axis, power)
501 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
502                                overlap=True, window=window_hann) :
503     """
504     compute the average power spectrum, preserving normalization
505     """
506     freq_axis,power = avg_power_spectrum(data, freq, chunk_size,
507                                          overlap, window)
508     #        2.0 / (freq * chunk_size)          |rfft()|**2 --> unitary_power_spectrum
509     power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum()
510     #                                       * 8/3  to remove power from windowing
511     #  <[x(t)*w(t)]**2> = <x(t)**2 * w(t)**2> ~= <x(t)**2> * <w(t)**2>
512     # where the ~= is because the frequency of x(t) >> the frequency of w(t).
513     # So our calulated power has and extra <w(t)**2> in it.
514     # For the Hann window, <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
515     # For low frequency components, where the frequency of x(t) is ~= the frequency of w(t),
516     # The normalization is not perfect. ??
517     # The normalization approaches perfection as chunk_size -> infinity.
518     return (freq_axis, power)
520 def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024,
521                                          chunk_size=512, overlap=True,
522                                          window=window_hann) :
523     x = zeros((samples,), dtype=float)
524     samp_freq = float(samp_freq)
525     for i in range(samples) :
526         x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
527     freq_axis, power = unitary_avg_power_spectrum(x, samp_freq, chunk_size,
528                                                   overlap, window)
529     imax = argmax(power)
531     expected = zeros((len(freq_axis),), dtype=float)
532     df = samp_freq/float(chunk_size) # df = 1/T, where T = total_time
533     i = int(sin_freq/df)
534     expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
536     print "The power should be a peak at %g Hz of %g (%g, %g)" % \
537         (sin_freq, expected[i], freq_axis[imax], power[imax])
538     Pexp = 0
539     P    = 0
540     for i in range(len(freq_axis)) :
541         Pexp += expected[i] * df
542         P    += power[i] * df
543     print " The total power should be %g (%g)" % (Pexp, P)
545     if TEST_PLOTS :
546         pylab.figure()
547         pylab.subplot(211)
548         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
549         pylab.title('time series')
550         pylab.subplot(212)
551         pylab.plot(freq_axis, power, 'r.')
552         pylab.plot(freq_axis, expected, 'b-')
553         pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
555 def _test_unitary_avg_power_spectrum_sin_suite() :
556     print "Test unitary avg power spectrums on variously shaped sin functions"
557     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
558     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
559     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
560     _test_unitary_avg_power_spectrum_sin(sin_freq=17, samp_freq=512, samples=1024)
561     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
562     # test long wavelenth sin, so be closer to window frequency
563     _test_unitary_avg_power_spectrum_sin(sin_freq=1, samp_freq=1024, samples=2048)
564     # finally, with some irrational numbers, to check that I'm not getting lucky
565     _test_unitary_avg_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)
568 def test() :
569     _test_rfft_suite()
570     _test_unitary_rfft_parsevals_suite()
571     _test_unitary_rfft_rect_suite()
572     _test_unitary_rfft_gaussian_suite()
573     _test_unitary_power_spectrum_sin_suite()
574     _test_unitary_power_spectrum_delta_suite()
575     _test_unitary_power_spectrum_gaussian_suite()
576     _test_unitary_avg_power_spectrum_sin_suite()
578 if __name__ == "__main__" :
579     if TEST_PLOTS :
580         import pylab
581     test()
582     if TEST_PLOTS :
583         pylab.show()