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