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     print "Test numpy rfft definition"
49     # Numpy's FFT algoritm returns
50     #          n-1
51     #   X[k] = SUM x[m] exp (-j 2pi km /n)
52     #          m=0
53     # (see http://www.tramy.us/numpybook.pdf)
54     j = complex(0,1)
55     n = len(xs)
56     Xa = []
57     for k in range(n) :
58         Xa.append(sum([x*exp(-j*2*pi*k*m/n) for x,m in zip(xs,range(n))]))
59         if k < len(Xs):
60             assert (Xs[k]-Xa[k])/abs(Xa[k]) < 1e-6, \
61                 "rfft mismatch on element %d: %g != %g, relative error %g" \
62                 % (k, Xs[k], Xa[k], (Xs[k]-Xa[k])/abs(Xa[k]))
63     # Which should satisfy the discrete form of Parseval's theorem
64     #   n-1               n-1
65     #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
66     #   m=0               k=0
67     timeSum = sum([abs(x)**2 for x in xs])
68     freqSum = sum([abs(X)**2 for X in Xa])
69     assert abs(freqSum/float(n) - timeSum)/timeSum < 1e-6, \
70         "Mismatch on Parseval's, %g != 1/%d * %g" % (timeSum, n, freqSum)
72 def _test_rfft_suite() :
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     print "Test unitary rfft on Parseval's theorem"
124     # Which should satisfy the discretized integral form of Parseval's theorem
125     #   n-1              n-1
126     #   SUM |x_m|^2 dt = SUM |X_k|^2 df
127     #   m=0              k=0
128     dt = 1.0/freq
129     df = freqs-freqs
130     assert (df - 1/(len(xs)*dt))/df < 1e-6, \
131         "Mismatch in spacing, %g != 1/(%d*%g)" % (df, len(xs), dt)
132     Xa = list(Xs)
133     for k in range(len(Xs)-1,1,-1) :
134         Xa.append(Xa[k])
135     assert len(xs) == len(Xa), "Length mismatch %d != %d" % (len(xs), len(Xa))
136     lhs = sum([abs(x)**2 for x in xs]) * dt
137     rhs = sum([abs(X)**2 for X in Xa]) * df
138     assert abs(lhs - rhs)/lhs < 1e-4, "Mismatch on Parseval's, %g != %g" \
139         % (lhs, rhs)
141 def _test_unitary_rfft_parsevals_suite():
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)
165     # remove the phase due to our time shift
166     j = complex(0.0,1.0) # sqrt(-1)
167     for i in range(len(freq_axis)) :
168         f = freq_axis[i]
169         inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
170         X[i] *= inverse_phase_shift
172     expected = zeros((len(freq_axis),), dtype=float)
173     # normalized sinc(x) = sin(pi x)/(pi x)
174     # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
175     assert sinc(0.5) == 2.0/pi, "abnormal sinc()"
176     for i in range(len(freq_axis)) :
177         f = freq_axis[i]
178         expected[i] = 1.0/abs(a) * sinc(f/a)
180     if TEST_PLOTS :
181         pylab.figure()
182         pylab.subplot(211)
183         pylab.plot(arange(0, dt*samples, dt), x)
184         pylab.title('time series')
185         pylab.subplot(212)
186         pylab.plot(freq_axis, X.real, 'r.')
187         pylab.plot(freq_axis, X.imag, 'g.')
188         pylab.plot(freq_axis, expected, 'b-')
189         pylab.title('freq series')
191 def _test_unitary_rfft_rect_suite() :
192     print "Test unitary FFTs on variously shaped rectangular functions"
193     _test_unitary_rfft_rect(a=0.5)
194     _test_unitary_rfft_rect(a=2.0)
195     _test_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
196     _test_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
198 def _gaussian(a, t) :
199     return exp(-a * t**2)
201 def _test_unitary_rfft_gaussian(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) :
202     "Show fft(rect(at)) = 1/abs(a) * sinc(f/a)"
203     samp_freq = float(samp_freq)
204     a = float(a)
206     x = zeros((samples,), dtype=float)
207     dt = 1.0/samp_freq
208     for i in range(samples) :
209         t = i*dt
210         x[i] = _gaussian(a, (t-time_shift))
211     freq_axis, X = unitary_rfft(x, samp_freq)
213     # remove the phase due to our time shift
214     j = complex(0.0,1.0) # sqrt(-1)
215     for i in range(len(freq_axis)) :
216         f = freq_axis[i]
217         inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
218         X[i] *= inverse_phase_shift
220     expected = zeros((len(freq_axis),), dtype=float)
221     for i in range(len(freq_axis)) :
222         f = freq_axis[i]
223         expected[i] = sqrt(pi/a) * _gaussian(1.0/a, pi*f) # see Wikipedia, or do the integral yourself.
225     if TEST_PLOTS :
226         pylab.figure()
227         pylab.subplot(211)
228         pylab.plot(arange(0, dt*samples, dt), x)
229         pylab.title('time series')
230         pylab.subplot(212)
231         pylab.plot(freq_axis, X.real, 'r.')
232         pylab.plot(freq_axis, X.imag, 'g.')
233         pylab.plot(freq_axis, expected, 'b-')
234         pylab.title('freq series')
236 def _test_unitary_rfft_gaussian_suite() :
237     print "Test unitary FFTs on variously shaped gaussian functions"
238     _test_unitary_rfft_gaussian(a=0.5)
239     _test_unitary_rfft_gaussian(a=2.0)
240     _test_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
241     _test_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
245 def power_spectrum(data, freq=1.0) :
246     """
247     Compute the power spectrum of DATA taken with a sampling frequency FREQ.
248     DATA must be real (not complex).
249     Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
250     If the number of samples in data is not an integer power of two,
251     the FFT ignores some of the later points.
252     """
253     nsamps = floor_pow_of_two(len(data))
255     freq_axis = linspace(0, freq/2, nsamps/2+1)
256     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
257     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
258     # See Numerical Recipies for a details.
259     trans = rfft(data[0:nsamps])
260     power = trans * trans.conj() # We want the square of the amplitude.
261     return (freq_axis, power)
263 def unitary_power_spectrum(data, freq=1.0) :
264     freq_axis,power = power_spectrum(data, freq)
265     # One sided power spectral density, so 2|H(f)|**2 (see NR 2nd edition 12.0.14, p498)
266     #
267     # numpy normalizes with 1/N on the inverse transform ifft,
268     # so we should normalize the freq-space representation with 1/sqrt(N).
269     # But we're using the rfft, where N points are like N/2 complex points, so 1/sqrt(N/2)
270     # So the power gets normalized by that twice and we have 2/N
271     #
272     # On top of this, the FFT assumes a sampling freq of 1 per second,
273     # and we want to preserve area under our curves.
274     # If our total time T = len(data)/freq is smaller than 1,
275     # our df_real = freq/len(data) is bigger that the FFT expects (dt_fft = 1/len(data)),
276     # and we need to scale the powers down to conserve area.
277     # df_fft * F_fft(f) = df_real *F_real(f)
278     # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq
279     # So the power gets normalized by *that* twice and we have 2/N * freq**2
281     # power per unit time
282     # measure x(t) for time T
283     # X(f)   = int_0^T x(t) exp(-2 pi ift) dt
284     # PSD(f) = 2 |X(f)|**2 / T
286     # total_time = len(data)/float(freq)
287     # power *= 2.0 / float(freq)**2   /   total_time
288     # power *= 2.0 / freq**2   *   freq / len(data)
289     power *= 2.0 / (freq * float(len(data)))
291     return (freq_axis, power)
293 def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) :
294     x = zeros((samples,), dtype=float)
295     samp_freq = float(samp_freq)
296     for i in range(samples) :
297         x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
298     freq_axis, power = unitary_power_spectrum(x, samp_freq)
299     imax = argmax(power)
301     expected = zeros((len(freq_axis),), dtype=float)
302     df = samp_freq/float(samples) # df = 1/T, where T = total_time
303     i = int(sin_freq/df)
304     # average power per unit time is
305     #  P = <x(t)**2>
306     # average value of sin(t)**2 = 0.5    (b/c sin**2+cos**2 == 1)
307     # so average value of (int sin(t)**2 dt) per unit time is 0.5
308     #  P = 0.5
309     # we spread that power over a frequency bin of width df, sp
310     #  P(f0) = 0.5/df
311     # where f0 is the sin's frequency
312     #
313     # or :
314     # FFT of sin(2*pi*t*f0) gives
315     #  X(f) = 0.5 i (delta(f-f0) - delta(f-f0)),
316     # (area under x(t) = 0, area under X(f) = 0)
317     # so one sided power spectral density (PSD) per unit time is
318     #  P(f) = 2 |X(f)|**2 / T
319     #       = 2 * |0.5 delta(f-f0)|**2 / T
320     #       = 0.5 * |delta(f-f0)|**2 / T
321     # but we're discrete and want the integral of the 'delta' to be 1,
322     # so 'delta'*df = 1  --> 'delta' = 1/df, and
323     #  P(f) = 0.5 / (df**2 * T)
324     #       = 0.5 / df                (T = 1/df)
325     expected[i] = 0.5 / df
327     print "The power should be a peak at %g Hz of %g (%g, %g)" % \
328         (sin_freq, expected[i], freq_axis[imax], power[imax])
329     Pexp = 0
330     P    = 0
331     for i in range(len(freq_axis)) :
332         Pexp += expected[i] *df
333         P    += power[i] * df
334     print " The total power should be %g (%g)" % (Pexp, P)
336     if TEST_PLOTS :
337         pylab.figure()
338         pylab.subplot(211)
339         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
340         pylab.title('time series')
341         pylab.subplot(212)
342         pylab.plot(freq_axis, power, 'r.')
343         pylab.plot(freq_axis, expected, 'b-')
344         pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
346 def _test_unitary_power_spectrum_sin_suite() :
347     print "Test unitary power spectrums on variously shaped sin functions"
348     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
349     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
350     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
351     _test_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024)
352     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
353     # finally, with some irrational numbers, to check that I'm not getting lucky
354     _test_unitary_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)
355     # test with non-integer number of periods
356     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256)
358 def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256) :
359     x = zeros((samples,), dtype=float)
360     samp_freq = float(samp_freq)
361     x = amp
362     freq_axis, power = unitary_power_spectrum(x, samp_freq)
364     # power = <x(t)**2> = (amp)**2 * dt/T
365     # we spread that power over the entire freq_axis [0,fN], so
366     #  P(f)  = (amp)**2 dt / (T fN)
367     # where
368     #  dt = 1/samp_freq        (sample period)
369     #  T  = samples/samp_freq  (total time of data aquisition)
370     #  fN = 0.5 samp_freq      (Nyquist frequency)
371     # so
372     #  P(f) = amp**2 / (samp_freq * samples/samp_freq * 0.5 samp_freq)
373     #       = 2 amp**2 / (samp_freq*samples)
374     expected_amp = 2.0 * amp**2 / (samp_freq * samples)
375     expected = ones((len(freq_axis),), dtype=float) * expected_amp
377     print "The power should be flat at y = %g (%g)" % (expected_amp, power)
379     if TEST_PLOTS :
380         pylab.figure()
381         pylab.subplot(211)
382         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
383         pylab.title('time series')
384         pylab.subplot(212)
385         pylab.plot(freq_axis, power, 'r.')
386         pylab.plot(freq_axis, expected, 'b-')
387         pylab.title('%g samples of delta amp %g' % (samples, amp))
389 def _test_unitary_power_spectrum_delta_suite() :
390     print "Test unitary power spectrums on various delta functions"
391     _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024)
392     _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048)
393     _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)# expected = 2*computed
394     _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)# expected = 0.5*computed
395     _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024)
396     _test_unitary_power_spectrum_delta(amp=pi, samp_freq=exp(1), samples=1024)
398 def _gaussian2(area, mean, std, t) :
399     "Integral over all time = area (i.e. normalized for area=1)"
400     return area/(std*sqrt(2.0*pi)) * exp(-0.5*((t-mean)/std)**2)
402 def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10.24 ,samples=512) : #1024
403     x = zeros((samples,), dtype=float)
404     mean = float(mean)
405     for i in range(samples) :
406         t = i/float(samp_freq)
407         x[i] = _gaussian2(area, mean, std, t)
408     freq_axis, power = unitary_power_spectrum(x, samp_freq)
410     # generate the predicted curve
411     # by comparing our _gaussian2() form to _gaussian(),
412     # we see that the Fourier transform of x(t) has parameters:
413     #  std'  = 1/(2 pi std)    (references declaring std' = 1/std are converting to angular frequency, not frequency like we are)
414     #  area' = area/[std sqrt(2*pi)]   (plugging into FT of _gaussian() above)
415     #  mean' = 0               (changing the mean in the time-domain just changes the phase in the freq-domain)
416     # So our power spectral density per unit time is given by
417     #  P(f) = 2 |X(f)|**2 / T
418     # Where
419     #  T  = samples/samp_freq  (total time of data aquisition)
420     mean = 0.0
421     area = area /(std*sqrt(2.0*pi))
422     std = 1.0/(2.0*pi*std)
423     expected = zeros((len(freq_axis),), dtype=float)
424     df = float(samp_freq)/samples # 1/total_time ( = freq_axis-freq_axis = freq_axis)
425     for i in range(len(freq_axis)) :
426         f = i*df
427         gaus = _gaussian2(area, mean, std, f)
428         expected[i] = 2.0 * gaus**2 * samp_freq/samples
429     print "The power should be a half-gaussian, ",
430     print "with a peak at 0 Hz with amplitude %g (%g)" % (expected, power)
432     if TEST_PLOTS :
433         pylab.figure()
434         pylab.subplot(211)
435         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
436         pylab.title('time series')
437         pylab.subplot(212)
438         pylab.plot(freq_axis, power, 'r.')
439         pylab.plot(freq_axis, expected, 'b-')
440         pylab.title('freq series')
442 def _test_unitary_power_spectrum_gaussian_suite() :
443     print "Test unitary power spectrums on various gaussian functions"
444     _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=1024)
445     _test_unitary_power_spectrum_gaussian(area=1, std=2, samp_freq=10.0, samples=1024)
446     _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=2048)
447     _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=20.0, samples=2048)
448     _test_unitary_power_spectrum_gaussian(area=3, std=1, samp_freq=10.0, samples=1024)
449     _test_unitary_power_spectrum_gaussian(area=pi, std=sqrt(2), samp_freq=exp(1), samples=1024)
451 def window_hann(length) :
452     "Returns a Hann window array with length entries"
453     win = zeros((length,), dtype=float)
454     for i in range(length) :
455         win[i] = 0.5*(1.0-cos(2.0*pi*float(i)/(length)))
456     # avg value of cos over a period is 0
457     # so average height of Hann window is 0.5
458     return win
460 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
461                        overlap=True, window=window_hann) :
462     """
463     Compute the avg power spectrum of DATA taken with a sampling frequency FREQ.
464     DATA must be real (not complex) by breaking DATA into chunks.
465     The chunks may or may not be overlapping (by setting OVERLAP).
466     The chunks are windowed by dotting with WINDOW(CHUNK_SIZE), FFTed,
467     and the resulting spectra are averaged together.
468     See NR 13.4 for rational.
470     Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
471     CHUNK_SIZE should really be a power of 2.
472     If the number of samples in DATA is not an integer power of CHUNK_SIZE,
473     the FFT ignores some of the later points.
474     """
475     assert chunk_size == floor_pow_of_two(chunk_size), \
476         "chunk_size %d should be a power of 2" % chunk_size
478     nchunks = len(data)/chunk_size # integer division = implicit floor
479     if overlap :
480         chunk_step = chunk_size/2
481     else :
482         chunk_step = chunk_size
484     win = window(chunk_size) # generate a window of the appropriate size
485     freq_axis = linspace(0, freq/2, chunk_size/2+1)
486     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
487     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
488     # See Numerical Recipies for a details.
489     power = zeros((chunk_size/2+1,), dtype=float)
490     for i in range(nchunks) :
491         starti = i*chunk_step
492         stopi = starti+chunk_size
493         fft_chunk = rfft(data[starti:stopi]*win)
494         p_chunk = fft_chunk * fft_chunk.conj()
495         power += p_chunk.astype(float)
496     power /= float(nchunks)
497     return (freq_axis, power)
499 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
500                                overlap=True, window=window_hann) :
501     """
502     compute the average power spectrum, preserving normalization
503     """
504     freq_axis,power = avg_power_spectrum(data, freq, chunk_size,
505                                          overlap, window)
506     #        2.0 / (freq * chunk_size)          |rfft()|**2 --> unitary_power_spectrum
507     power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum()
508     #                                       * 8/3  to remove power from windowing
509     #  <[x(t)*w(t)]**2> = <x(t)**2 * w(t)**2> ~= <x(t)**2> * <w(t)**2>
510     # where the ~= is because the frequency of x(t) >> the frequency of w(t).
511     # So our calulated power has and extra <w(t)**2> in it.
512     # For the Hann window, <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
513     # For low frequency components, where the frequency of x(t) is ~= the frequency of w(t),
514     # The normalization is not perfect. ??
515     # The normalization approaches perfection as chunk_size -> infinity.
516     return (freq_axis, power)
518 def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024,
519                                          chunk_size=512, overlap=True,
520                                          window=window_hann) :
521     x = zeros((samples,), dtype=float)
522     samp_freq = float(samp_freq)
523     for i in range(samples) :
524         x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
525     freq_axis, power = unitary_avg_power_spectrum(x, samp_freq, chunk_size,
526                                                   overlap, window)
527     imax = argmax(power)
529     expected = zeros((len(freq_axis),), dtype=float)
530     df = samp_freq/float(chunk_size) # df = 1/T, where T = total_time
531     i = int(sin_freq/df)
532     expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
534     print "The power should be a peak at %g Hz of %g (%g, %g)" % \
535         (sin_freq, expected[i], freq_axis[imax], power[imax])
536     Pexp = 0
537     P    = 0
538     for i in range(len(freq_axis)) :
539         Pexp += expected[i] * df
540         P    += power[i] * df
541     print " The total power should be %g (%g)" % (Pexp, P)
543     if TEST_PLOTS :
544         pylab.figure()
545         pylab.subplot(211)
546         pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
547         pylab.title('time series')
548         pylab.subplot(212)
549         pylab.plot(freq_axis, power, 'r.')
550         pylab.plot(freq_axis, expected, 'b-')
551         pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
553 def _test_unitary_avg_power_spectrum_sin_suite() :
554     print "Test unitary avg power spectrums on variously shaped sin functions"
555     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
556     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
557     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
558     _test_unitary_avg_power_spectrum_sin(sin_freq=17, samp_freq=512, samples=1024)
559     _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
560     # test long wavelenth sin, so be closer to window frequency
561     _test_unitary_avg_power_spectrum_sin(sin_freq=1, samp_freq=1024, samples=2048)
562     # finally, with some irrational numbers, to check that I'm not getting lucky
563     _test_unitary_avg_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)
566 def test() :
567     _test_rfft_suite()
568     _test_unitary_rfft_parsevals_suite()
569     _test_unitary_rfft_rect_suite()
570     _test_unitary_rfft_gaussian_suite()
571     _test_unitary_power_spectrum_sin_suite()
572     _test_unitary_power_spectrum_delta_suite()
573     _test_unitary_power_spectrum_gaussian_suite()
574     _test_unitary_avg_power_spectrum_sin_suite()
576 if __name__ == "__main__" :
577     if TEST_PLOTS :
578         import pylab
579     test()
580     if TEST_PLOTS :
581         pylab.show()