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