1 # Copyright (C) 2008-2012 W. Trevor King
2 # This program is free software: you can redistribute it and/or modify
3 # it under the terms of the GNU General Public License as published by
4 # the Free Software Foundation, either version 3 of the License, or
5 # (at your option) any later version.
7 # This program is distributed in the hope that it will be useful,
8 # but WITHOUT ANY WARRANTY; without even the implied warranty of
9 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
10 # GNU General Public License for more details.
12 # You should have received a copy of the GNU General Public License
13 # along with this program. If not, see <http://www.gnu.org/licenses/>.
15 """Wrap Numpy's fft module to reduce clutter.
17 Provides a unitary discrete FFT and a windowed version based on
20 Create some fake data:
23 >>> data = numpy.random.rand(10)
28 >>> rfft = unitary_rfft(data, freq=frequency)
29 >>> psd = power_spectrum(data, freq=frequency)
30 >>> upsd = unitary_power_spectrum(data, freq=frequency)
31 >>> apsd = avg_power_spectrum(data, freq=frequency, chunk_size=2048,
32 ... overlap=True, window=window_hann)
33 >>> aupsd = unitary_avg_power_spectrum(data, freq=frequency, chunk_size=2048,
34 ... overlap=True, window=window_hann)
37 import numpy as _numpy
42 # Display time- and freq-space plots of the test transforms if True
46 def floor_pow_of_two(num):
47 "Round num down to the closest exact a power of two."
48 lnum = _numpy.log2(num)
50 num = 2**_numpy.floor(lnum)
53 def round_pow_of_two(num):
54 "Round num to the closest exact a power of two on a log scale."
55 lnum = _numpy.log2(num)
57 num = 2**_numpy.round(lnum)
60 def ceil_pow_of_two(num):
61 "Round num up to the closest exact a power of two."
62 lnum = _numpy.log2(num)
64 num = 2**_numpy.ceil(lnum)
67 def _test_rfft(xs, Xs):
68 # Numpy's FFT algoritm returns
70 # X[k] = SUM x[m] exp (-j 2pi km /n)
72 # (see http://www.tramy.us/numpybook.pdf)
73 j = _numpy.complex(0,1)
77 Xa.append(sum([x*_numpy.exp(-j*2*_numpy.pi*k*m/n)
78 for x,m in zip(xs,range(n))]))
80 if (Xs[k]-Xa[k])/_numpy.abs(Xa[k]) >= 1e-6:
82 ('rfft mismatch on element {}: {} != {}, relative error {}'
84 k, Xs[k], Xa[k], (Xs[k]-Xa[k])/_numpy.abs(Xa[k])))
85 # Which should satisfy the discrete form of Parseval's theorem
87 # SUM |x_m|^2 = 1/n SUM |X_k|^2.
89 timeSum = sum([_numpy.abs(x)**2 for x in xs])
90 freqSum = sum([_numpy.abs(X)**2 for X in Xa])
91 if _numpy.abs(freqSum/_numpy.float(n) - timeSum)/timeSum >= 1e-6:
93 "Mismatch on Parseval's, {} != 1/{} * {}".format(
96 def _test_rfft_suite():
97 print('Test numpy rfft definition')
98 xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
99 _test_rfft(xs, _numpy.fft.rfft(xs))
101 def unitary_rfft(data, freq=1.0):
102 """Compute the Fourier transform of real data.
104 Unitary (preserves power [Parseval's theorem]).
106 If the units on your data are Volts,
107 and your sampling frequency is in Hz,
108 then freq_axis will be in Hz,
109 and trans will be in Volts.
111 nsamps = floor_pow_of_two(len(data))
112 # Which should satisfy the discrete form of Parseval's theorem
114 # SUM |x_m|^2 = 1/n SUM |X_k|^2.
116 # However, we want our FFT to satisfy the continuous Parseval eqn
117 # int_{-infty}^{infty} |x(t)|^2 dt = int_{-infty}^{infty} |X(f)|^2 df
118 # which has the discrete form
120 # SUM |x_m|^2 dt = SUM |X'_k|^2 df
122 # with X'_k = AX, this gives us
124 # SUM |x_m|^2 = A^2 df/dt SUM |X'_k|^2
129 # From Numerical Recipes (http://www.fizyka.umk.pl/nrbook/bookcpdf.html),
130 # Section 12.1, we see that for a sampling rate dt, the maximum frequency
131 # f_c in the transformed data is the Nyquist frequency (12.1.2)
133 # and the points are spaced out by (12.1.5)
139 # A = 1/ndf = ndt/n = dt
140 # so we can convert the Numpy transformed data to match our unitary
141 # continuous transformed data with (also NR 12.1.8)
142 # X'_k = dtX = X / <sampling freq>
143 trans = _numpy.fft.rfft(data[0:nsamps]) / _numpy.float(freq)
144 freq_axis = _numpy.linspace(0, freq/2, nsamps/2+1)
145 return (freq_axis, trans)
147 def _test_unitary_rfft_parsevals(xs, freq, freqs, Xs):
148 # Which should satisfy the discretized integral form of Parseval's theorem
150 # SUM |x_m|^2 dt = SUM |X_k|^2 df
153 df = freqs[1]-freqs[0]
154 if df - 1/(len(xs)*dt)/df >= 1e-6:
156 'Mismatch in spacing, {} != 1/({}*{})'.format(df, len(xs), dt))
158 for k in range(len(Xs)-1,1,-1):
160 if len(xs) != len(Xa):
161 raise ValueError('Length mismatch {} != {}'.format(len(xs), len(Xa)))
162 lhs = sum([_numpy.abs(x)**2 for x in xs]) * dt
163 rhs = sum([_numpy.abs(X)**2 for X in Xa]) * df
164 if _numpy.abs(lhs - rhs)/lhs >= 1e-4:
165 raise ValueError("Mismatch on Parseval's, {} != {}".format(lhs, rhs))
167 def _test_unitary_rfft_parsevals_suite():
168 print("Test unitary rfft on Parseval's theorem")
169 xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
171 freqs,Xs = unitary_rfft(xs, 1.0/dt)
172 _test_unitary_rfft_parsevals(xs, 1.0/dt, freqs, Xs)
175 if _numpy.abs(t) < 0.5:
180 def _test_unitary_rfft_rect(
181 a=1.0, time_shift=5.0, samp_freq=25.6, samples=256):
182 "Show fft(rect(at)) = 1/abs(a) * _numpy.sinc(f/a)"
183 samp_freq = _numpy.float(samp_freq)
186 x = _numpy.zeros((samples,), dtype=_numpy.float)
188 for i in range(samples):
190 x[i] = _rect(a*(t-time_shift))
191 freq_axis, X = unitary_rfft(x, samp_freq)
192 #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
194 # remove the phase due to our time shift
195 j = _numpy.complex(0.0,1.0) # sqrt(-1)
196 for i in range(len(freq_axis)):
198 inverse_phase_shift = _numpy.exp(j*2.0*_numpy.pi*time_shift*f)
199 X[i] *= inverse_phase_shift
201 expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
202 # normalized sinc(x) = sin(pi x)/(pi x)
203 # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
204 if _numpy.sinc(0.5) != 2.0/_numpy.pi:
205 raise ValueError('abnormal sinc()')
206 for i in range(len(freq_axis)):
208 expected[i] = 1.0/_numpy.abs(a) * _numpy.sinc(f/a)
211 figure = _pyplot.figure()
212 time_axes = figure.add_subplot(2, 1, 1)
213 time_axes.plot(_numpy.arange(0, dt*samples, dt), x)
214 time_axes.set_title('time series')
215 freq_axes = figure.add_subplot(2, 1, 2)
216 freq_axes.plot(freq_axis, X.real, 'r.')
217 freq_axes.plot(freq_axis, X.imag, 'g.')
218 freq_axes.plot(freq_axis, expected, 'b-')
219 freq_axes.set_title('freq series')
221 def _test_unitary_rfft_rect_suite():
222 print('Test unitary FFTs on variously shaped rectangular functions')
223 _test_unitary_rfft_rect(a=0.5)
224 _test_unitary_rfft_rect(a=2.0)
225 _test_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
226 _test_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
229 return _numpy.exp(-a * t**2)
231 def _test_unitary_rfft_gaussian(
232 a=1.0, time_shift=5.0, samp_freq=25.6, samples=256):
233 "Show fft(rect(at)) = 1/abs(a) * sinc(f/a)"
234 samp_freq = _numpy.float(samp_freq)
237 x = _numpy.zeros((samples,), dtype=_numpy.float)
239 for i in range(samples):
241 x[i] = _gaussian(a, (t-time_shift))
242 freq_axis, X = unitary_rfft(x, samp_freq)
243 #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
245 # remove the phase due to our time shift
246 j = _numpy.complex(0.0,1.0) # sqrt(-1)
247 for i in range(len(freq_axis)):
249 inverse_phase_shift = _numpy.exp(j*2.0*_numpy.pi*time_shift*f)
250 X[i] *= inverse_phase_shift
252 expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
253 for i in range(len(freq_axis)):
255 # see Wikipedia, or do the integral yourself.
256 expected[i] = _numpy.sqrt(_numpy.pi/a) * _gaussian(
260 figure = _pyplot.figure()
261 time_axes = figure.add_subplot(2, 1, 1)
262 time_axes.plot(_numpy.arange(0, dt*samples, dt), x)
263 time_axes.set_title('time series')
264 freq_axes = figure.add_subplot(2, 1, 2)
265 freq_axes.plot(freq_axis, X.real, 'r.')
266 freq_axes.plot(freq_axis, X.imag, 'g.')
267 freq_axes.plot(freq_axis, expected, 'b-')
268 freq_axes.set_title('freq series')
270 def _test_unitary_rfft_gaussian_suite():
271 print("Test unitary FFTs on variously shaped gaussian functions")
272 _test_unitary_rfft_gaussian(a=0.5)
273 _test_unitary_rfft_gaussian(a=2.0)
274 _test_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
275 _test_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
279 def power_spectrum(data, freq=1.0):
280 """Compute the power spectrum of DATA taken with a sampling frequency FREQ.
282 DATA must be real (not complex).
283 Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
284 If the number of samples in data is not an integer power of two,
285 the FFT ignores some of the later points.
287 nsamps = floor_pow_of_two(len(data))
289 freq_axis = _numpy.linspace(0, freq/2, nsamps/2+1)
290 # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
291 # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
292 # See Numerical Recipies for a details.
293 trans = _numpy.fft.rfft(data[0:nsamps])
294 power = (trans * trans.conj()).real # We want the square of the amplitude.
295 return (freq_axis, power)
297 def unitary_power_spectrum(data, freq=1.0):
298 freq_axis,power = power_spectrum(data, freq)
299 # One sided power spectral density, so 2|H(f)|**2
300 # (see NR 2nd edition 12.0.14, p498)
302 # numpy normalizes with 1/N on the inverse transform ifft,
303 # so we should normalize the freq-space representation with 1/sqrt(N).
304 # But we're using the rfft, where N points are like N/2 complex points,
306 # So the power gets normalized by that twice and we have 2/N
308 # On top of this, the FFT assumes a sampling freq of 1 per second,
309 # and we want to preserve area under our curves.
310 # If our total time T = len(data)/freq is smaller than 1,
311 # our df_real = freq/len(data) is bigger that the FFT expects
312 # (dt_fft = 1/len(data)),
313 # and we need to scale the powers down to conserve area.
314 # df_fft * F_fft(f) = df_real *F_real(f)
315 # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq
316 # So the power gets normalized by *that* twice and we have 2/N * freq**2
318 # power per unit time
319 # measure x(t) for time T
320 # X(f) = int_0^T x(t) exp(-2 pi ift) dt
321 # PSD(f) = 2 |X(f)|**2 / T
323 # total_time = len(data)/float(freq)
324 # power *= 2.0 / float(freq)**2 / total_time
325 # power *= 2.0 / freq**2 * freq / len(data)
326 power *= 2.0 / (freq * _numpy.float(len(data)))
328 return (freq_axis, power)
330 def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024):
331 x = _numpy.zeros((samples,), dtype=_numpy.float)
332 samp_freq = _numpy.float(samp_freq)
333 for i in range(samples):
334 x[i] = _numpy.sin(2.0 * _numpy.pi * (i/samp_freq) * sin_freq)
335 freq_axis, power = unitary_power_spectrum(x, samp_freq)
336 imax = _numpy.argmax(power)
338 expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
339 df = samp_freq/_numpy.float(samples) # df = 1/T, where T = total_time
341 # average power per unit time is
343 # average value of sin(t)**2 = 0.5 (b/c sin**2+cos**2 == 1)
344 # so average value of (int sin(t)**2 dt) per unit time is 0.5
346 # we spread that power over a frequency bin of width df, sp
348 # where f0 is the sin's frequency
351 # FFT of sin(2*pi*t*f0) gives
352 # X(f) = 0.5 i (delta(f-f0) - delta(f-f0)),
353 # (area under x(t) = 0, area under X(f) = 0)
354 # so one sided power spectral density (PSD) per unit time is
355 # P(f) = 2 |X(f)|**2 / T
356 # = 2 * |0.5 delta(f-f0)|**2 / T
357 # = 0.5 * |delta(f-f0)|**2 / T
358 # but we're discrete and want the integral of the 'delta' to be 1,
359 # so 'delta'*df = 1 --> 'delta' = 1/df, and
360 # P(f) = 0.5 / (df**2 * T)
361 # = 0.5 / df (T = 1/df)
362 expected[i] = 0.5 / df
364 print('The power should be a peak at {} Hz of {} ({}, {})'.format(
365 sin_freq, expected[i], freq_axis[imax], power[imax]))
368 for i in range(len(freq_axis)):
369 Pexp += expected[i] *df
371 print('The total power should be {} ({})'.format(Pexp, P))
374 figure = _pyplot.figure()
375 time_axes = figure.add_subplot(2, 1, 1)
377 _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
378 time_axes.set_title('time series')
379 freq_axes = figure.add_subplot(2, 1, 2)
380 freq_axes.plot(freq_axis, power, 'r.')
381 freq_axes.plot(freq_axis, expected, 'b-')
383 '{} samples of sin at {} Hz'.format(samples, sin_freq))
385 def _test_unitary_power_spectrum_sin_suite():
386 print('Test unitary power spectrums on variously shaped sin functions')
387 _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
388 _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
389 _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
390 _test_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024)
391 _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
392 # with some irrational numbers, to check that I'm not getting lucky
393 _test_unitary_power_spectrum_sin(
394 sin_freq=_numpy.pi, samp_freq=100*_numpy.exp(1), samples=1024)
395 # test with non-integer number of periods
396 _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256)
398 def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256):
399 x = _numpy.zeros((samples,), dtype=_numpy.float)
400 samp_freq = _numpy.float(samp_freq)
402 freq_axis, power = unitary_power_spectrum(x, samp_freq)
404 # power = <x(t)**2> = (amp)**2 * dt/T
405 # we spread that power over the entire freq_axis [0,fN], so
406 # P(f) = (amp)**2 dt / (T fN)
408 # dt = 1/samp_freq (sample period)
409 # T = samples/samp_freq (total time of data aquisition)
410 # fN = 0.5 samp_freq (Nyquist frequency)
412 # P(f) = amp**2 / (samp_freq * samples/samp_freq * 0.5 samp_freq)
413 # = 2 amp**2 / (samp_freq*samples)
414 expected_amp = 2.0 * amp**2 / (samp_freq * samples)
415 expected = _numpy.ones(
416 (len(freq_axis),), dtype=_numpy.float) * expected_amp
418 print('The power should be flat at y = {} ({})'.format(
419 expected_amp, power[0]))
422 figure = _pyplot.figure()
423 time_axes = figure.add_subplot(2, 1, 1)
425 _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
426 time_axes.set_title('time series')
427 freq_axes = figure.add_subplot(2, 1, 2)
428 freq_axes.plot(freq_axis, power, 'r.')
429 freq_axes.plot(freq_axis, expected, 'b-')
430 freq_axes.set_title('{} samples of delta amp {}'.format(samples, amp))
432 def _test_unitary_power_spectrum_delta_suite():
433 print('Test unitary power spectrums on various delta functions')
434 _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024)
435 _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048)
436 # expected = 2*computed
437 _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)
438 # expected = 0.5*computed
439 _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)
440 _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024)
441 _test_unitary_power_spectrum_delta(
442 amp=_numpy.pi, samp_freq=_numpy.exp(1), samples=1024)
444 def _gaussian2(area, mean, std, t):
445 "Integral over all time = area (i.e. normalized for area=1)"
446 return area/(std*_numpy.sqrt(2.0*_numpy.pi)) * _numpy.exp(
447 -0.5*((t-mean)/std)**2)
449 def _test_unitary_power_spectrum_gaussian(
450 area=2.5, mean=5, std=1, samp_freq=10.24 ,samples=512):
451 x = _numpy.zeros((samples,), dtype=_numpy.float)
452 mean = _numpy.float(mean)
453 for i in range(samples):
454 t = i/_numpy.float(samp_freq)
455 x[i] = _gaussian2(area, mean, std, t)
456 freq_axis, power = unitary_power_spectrum(x, samp_freq)
458 # generate the predicted curve
459 # by comparing our _gaussian2() form to _gaussian(),
460 # we see that the Fourier transform of x(t) has parameters:
461 # std' = 1/(2 pi std) (references declaring std' = 1/std are
462 # converting to angular frequency,
463 # not frequency like we are)
464 # area' = area/[std sqrt(2*pi)] (plugging into FT of _gaussian() above)
465 # mean' = 0 (changing the mean in the time-domain just
466 # changes the phase in the freq-domain)
467 # So our power spectral density per unit time is given by
468 # P(f) = 2 |X(f)|**2 / T
470 # T = samples/samp_freq (total time of data aquisition)
472 area = area /(std*_numpy.sqrt(2.0*_numpy.pi))
473 std = 1.0/(2.0*_numpy.pi*std)
474 expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
475 # 1/total_time ( = freq_axis[1]-freq_axis[0] = freq_axis[1])
476 df = _numpy.float(samp_freq)/samples
477 for i in range(len(freq_axis)):
479 gaus = _gaussian2(area, mean, std, f)
480 expected[i] = 2.0 * gaus**2 * samp_freq/samples
481 print(('The power should be a half-gaussian, '
482 'with a peak at 0 Hz with amplitude {} ({})').format(
483 expected[0], power[0]))
486 figure = _pyplot.figure()
487 time_axes = figure.add_subplot(2, 1, 1)
489 _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
490 time_axes.set_title('time series')
491 freq_axes = figure.add_subplot(2, 1, 2)
492 freq_axes.plot(freq_axis, power, 'r.')
493 freq_axes.plot(freq_axis, expected, 'b-')
494 freq_axes.set_title('freq series')
496 def _test_unitary_power_spectrum_gaussian_suite():
497 print('Test unitary power spectrums on various gaussian functions')
498 _test_unitary_power_spectrum_gaussian(
499 area=1, std=1, samp_freq=10.0, samples=1024)
500 _test_unitary_power_spectrum_gaussian(
501 area=1, std=2, samp_freq=10.0, samples=1024)
502 _test_unitary_power_spectrum_gaussian(
503 area=1, std=1, samp_freq=10.0, samples=2048)
504 _test_unitary_power_spectrum_gaussian(
505 area=1, std=1, samp_freq=20.0, samples=2048)
506 _test_unitary_power_spectrum_gaussian(
507 area=3, std=1, samp_freq=10.0, samples=1024)
508 _test_unitary_power_spectrum_gaussian(
509 area=_numpy.pi, std=_numpy.sqrt(2), samp_freq=_numpy.exp(1),
512 def window_hann(length):
513 "Returns a Hann window array with length entries"
514 win = _numpy.zeros((length,), dtype=_numpy.float)
515 for i in range(length):
516 win[i] = 0.5*(1.0-_numpy.cos(2.0*_numpy.pi*_numpy.float(i)/(length)))
517 # avg value of cos over a period is 0
518 # so average height of Hann window is 0.5
521 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
522 overlap=True, window=window_hann):
523 """Compute the avgerage power spectrum of DATA.
525 Taken with a sampling frequency FREQ.
527 DATA must be real (not complex) by breaking DATA into chunks.
528 The chunks may or may not be overlapping (by setting OVERLAP).
529 The chunks are windowed by dotting with WINDOW(CHUNK_SIZE), FFTed,
530 and the resulting spectra are averaged together.
531 See NR 13.4 for rational.
533 Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
534 CHUNK_SIZE should really be a power of 2.
535 If the number of samples in DATA is not an integer power of CHUNK_SIZE,
536 the FFT ignores some of the later points.
538 if chunk_size != floor_pow_of_two(chunk_size):
540 'chunk_size {} should be a power of 2'.format(chunk_size))
542 nchunks = len(data)/chunk_size # integer division = implicit floor
544 chunk_step = chunk_size/2
546 chunk_step = chunk_size
548 win = window(chunk_size) # generate a window of the appropriate size
549 freq_axis = _numpy.linspace(0, freq/2, chunk_size/2+1)
550 # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
551 # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
552 # See Numerical Recipies for a details.
553 power = _numpy.zeros((chunk_size/2+1,), dtype=_numpy.float)
554 for i in range(nchunks):
555 starti = i*chunk_step
556 stopi = starti+chunk_size
557 fft_chunk = _numpy.fft.rfft(data[starti:stopi]*win)
558 p_chunk = (fft_chunk * fft_chunk.conj()).real
559 power += p_chunk.astype(_numpy.float)
560 power /= _numpy.float(nchunks)
561 return (freq_axis, power)
563 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
564 overlap=True, window=window_hann):
565 """Compute the average power spectrum, preserving normalization
567 freq_axis,power = avg_power_spectrum(
568 data, freq, chunk_size, overlap, window)
569 # 2.0 / (freq * chunk_size) |rfft()|**2 --> unitary_power_spectrum
570 # see unitary_power_spectrum()
571 power *= 2.0 / (freq*_numpy.float(chunk_size)) * 8/3
572 # * 8/3 to remove power from windowing
573 # <[x(t)*w(t)]**2> = <x(t)**2 * w(t)**2> ~= <x(t)**2> * <w(t)**2>
574 # where the ~= is because the frequency of x(t) >> the frequency of w(t).
575 # So our calulated power has and extra <w(t)**2> in it.
576 # For the Hann window,
577 # <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
578 # For low frequency components,
579 # where the frequency of x(t) is ~= the frequency of w(t),
580 # the normalization is not perfect. ??
581 # The normalization approaches perfection as chunk_size -> infinity.
582 return (freq_axis, power)
584 def _test_unitary_avg_power_spectrum_sin(
585 sin_freq=10, samp_freq=512, samples=1024, chunk_size=512, overlap=True,
587 x = _numpy.zeros((samples,), dtype=_numpy.float)
588 samp_freq = _numpy.float(samp_freq)
589 for i in range(samples):
590 x[i] = _numpy.sin(2.0 * _numpy.pi * (i/samp_freq) * sin_freq)
591 freq_axis, power = unitary_avg_power_spectrum(
592 x, samp_freq, chunk_size, overlap, window)
593 imax = _numpy.argmax(power)
595 expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
596 df = samp_freq/_numpy.float(chunk_size) # df = 1/T, where T = total_time
598 expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
600 print('The power should peak at {} Hz of {} ({}, {})'.format(
601 sin_freq, expected[i], freq_axis[imax], power[imax]))
604 for i in range(len(freq_axis)):
605 Pexp += expected[i] * df
607 print('The total power should be {} ({})'.format(Pexp, P))
610 figure = _pyplot.figure()
611 time_axes = figure.add_subplot(2, 1, 1)
613 _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
614 time_axes.set_title('time series')
615 freq_axes = figure.add_subplot(2, 1, 2)
616 freq_axes.plot(freq_axis, power, 'r.')
617 freq_axes.plot(freq_axis, expected, 'b-')
619 '{} samples of sin at {} Hz'.format(samples, sin_freq))
621 def _test_unitary_avg_power_spectrum_sin_suite():
622 print('Test unitary avg power spectrums on variously shaped sin functions')
623 _test_unitary_avg_power_spectrum_sin(
624 sin_freq=5, samp_freq=512, samples=1024)
625 _test_unitary_avg_power_spectrum_sin(
626 sin_freq=5, samp_freq=512, samples=2048)
627 _test_unitary_avg_power_spectrum_sin(
628 sin_freq=5, samp_freq=512, samples=4098)
629 _test_unitary_avg_power_spectrum_sin(
630 sin_freq=17, samp_freq=512, samples=1024)
631 _test_unitary_avg_power_spectrum_sin(
632 sin_freq=5, samp_freq=1024, samples=2048)
633 # test long wavelenth sin, so be closer to window frequency
634 _test_unitary_avg_power_spectrum_sin(
635 sin_freq=1, samp_freq=1024, samples=2048)
636 # finally, with some irrational numbers, to check that I'm not
638 _test_unitary_avg_power_spectrum_sin(
639 sin_freq=_numpy.pi, samp_freq=100*_numpy.exp(1), samples=1024)
644 _test_unitary_rfft_parsevals_suite()
645 _test_unitary_rfft_rect_suite()
646 _test_unitary_rfft_gaussian_suite()
647 _test_unitary_power_spectrum_sin_suite()
648 _test_unitary_power_spectrum_delta_suite()
649 _test_unitary_power_spectrum_gaussian_suite()
650 _test_unitary_avg_power_spectrum_sin_suite()
652 if __name__ == '__main__':
653 from optparse import OptionParser
655 p = OptionParser('%prog [options]', epilog='Run FFT-tools unit tests.')
656 p.add_option('-p', '--plot', dest='plot', action='store_true',
657 help='Display time- and freq-space plots of test transforms.')
659 options,args = p.parse_args()
662 import matplotlib.pyplot as _pyplot