3 # Wrap Numpy's fft module to produce 1D unitary transforms and power spectra.
5 # Copyright (C) 2008 W. Trevor King
8 # Redistribution and use in source and binary forms, with or without
9 # modification, are permitted provided that the following conditions
12 # * Redistributions of source code must retain the above copyright
13 # notice, this list of conditions and the following disclaimer.
15 # * Redistributions in binary form must reproduce the above copyright
16 # notice, this list of conditions and the following disclaimer in
17 # the documentation and/or other materials provided with the
20 # * Neither the name of the copyright holders nor the names of the
21 # contributors may be used to endorse or promote products derived
22 # from this software without specific prior written permission.
24 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
25 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
26 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
27 # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
28 # COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
29 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
30 # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
31 # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
32 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
33 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
34 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
35 # POSSIBILITY OF SUCH DAMAGE.
38 Define some FFT wrappers to reduce clutter.
39 Provides a unitary discrete FFT and a windowed version.
40 Based on numpy.fft.rfft.
44 * :func:`unitary_rfft`
45 * :func:`power_spectrum`
46 * :func:`unitary_power_spectrum`
47 * :func:`avg_power_spectrum`
48 * :func:`unitary_avg_power_spectrum`
53 from numpy import log2, floor, round, ceil, abs, pi, exp, cos, sin, sqrt, \
54 sinc, arctan2, array, ones, arange, linspace, zeros, \
55 uint16, float, concatenate, fromfile, argmax, complex
56 from numpy.fft import rfft
61 def floor_pow_of_two(num):
62 """Round `num` down to the closest exact a power of two.
67 >>> floor_pow_of_two(3)
69 >>> floor_pow_of_two(11)
71 >>> floor_pow_of_two(15)
79 def round_pow_of_two(num):
80 """Round `num` to the closest exact a power of two on a log scale.
85 >>> round_pow_of_two(2.9) # Note rounding on *log scale*
87 >>> round_pow_of_two(11)
89 >>> round_pow_of_two(15)
97 def ceil_pow_of_two(num):
98 """Round `num` up to the closest exact a power of two.
103 >>> ceil_pow_of_two(3)
105 >>> ceil_pow_of_two(11)
107 >>> ceil_pow_of_two(15)
111 if int(lnum) != lnum:
115 def unitary_rfft(data, freq=1.0):
116 """Compute the unitary Fourier transform of real data.
118 Unitary = preserves power [Parseval's theorem].
123 Real (not complex) data taken with a sampling frequency `freq`.
129 freq_axis,trans : numpy.ndarray
130 Arrays ready for plotting.
135 If the units on your data are Volts,
136 and your sampling frequency is in Hz,
137 then `freq_axis` will be in Hz,
138 and `trans` will be in Volts.
140 nsamps = floor_pow_of_two(len(data))
141 # Which should satisfy the discrete form of Parseval's theorem
143 # SUM |x_m|^2 = 1/n SUM |X_k|^2.
145 # However, we want our FFT to satisfy the continuous Parseval eqn
146 # int_{-infty}^{infty} |x(t)|^2 dt = int_{-infty}^{infty} |X(f)|^2 df
147 # which has the discrete form
149 # SUM |x_m|^2 dt = SUM |X'_k|^2 df
151 # with X'_k = AX, this gives us
153 # SUM |x_m|^2 = A^2 df/dt SUM |X'_k|^2
158 # From Numerical Recipes (http://www.fizyka.umk.pl/nrbook/bookcpdf.html),
159 # Section 12.1, we see that for a sampling rate dt, the maximum frequency
160 # f_c in the transformed data is the Nyquist frequency (12.1.2)
162 # and the points are spaced out by (12.1.5)
168 # A = 1/ndf = ndt/n = dt
169 # so we can convert the Numpy transformed data to match our unitary
170 # continuous transformed data with (also NR 12.1.8)
171 # X'_k = dtX = X / <sampling freq>
172 trans = rfft(data[0:nsamps]) / float(freq)
173 freq_axis = linspace(0, freq/2, nsamps/2+1)
174 return (freq_axis, trans)
176 def power_spectrum(data, freq=1.0):
177 """Compute the power spectrum of the time series `data`.
182 Real (not complex) data taken with a sampling frequency `freq`.
188 freq_axis,power : numpy.ndarray
189 Arrays ready for plotting.
193 If the number of samples in `data` is not an integer power of two,
194 the FFT ignores some of the later points.
198 unitary_power_spectrum,avg_power_spectrum
200 nsamps = floor_pow_of_two(len(data))
202 freq_axis = linspace(0, freq/2, nsamps/2+1)
203 # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
204 # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
205 # See Numerical Recipies for a details.
206 trans = rfft(data[0:nsamps])
207 power = trans * trans.conj() # We want the square of the amplitude.
208 return (freq_axis, power)
210 def unitary_power_spectrum(data, freq=1.0):
211 """Compute the unitary power spectrum of the time series `data`.
215 power_spectrum,unitary_avg_power_spectrum
217 freq_axis,power = power_spectrum(data, freq)
218 # One sided power spectral density, so 2|H(f)|**2 (see NR 2nd edition 12.0.14, p498)
220 # numpy normalizes with 1/N on the inverse transform ifft,
221 # so we should normalize the freq-space representation with 1/sqrt(N).
222 # But we're using the rfft, where N points are like N/2 complex points, so 1/sqrt(N/2)
223 # So the power gets normalized by that twice and we have 2/N
225 # On top of this, the FFT assumes a sampling freq of 1 per second,
226 # and we want to preserve area under our curves.
227 # If our total time T = len(data)/freq is smaller than 1,
228 # our df_real = freq/len(data) is bigger that the FFT expects (dt_fft = 1/len(data)),
229 # and we need to scale the powers down to conserve area.
230 # df_fft * F_fft(f) = df_real *F_real(f)
231 # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq
232 # So the power gets normalized by *that* twice and we have 2/N * freq**2
234 # power per unit time
235 # measure x(t) for time T
236 # X(f) = int_0^T x(t) exp(-2 pi ift) dt
237 # PSD(f) = 2 |X(f)|**2 / T
239 # total_time = len(data)/float(freq)
240 # power *= 2.0 / float(freq)**2 / total_time
241 # power *= 2.0 / freq**2 * freq / len(data)
242 power *= 2.0 / (freq * float(len(data)))
244 return (freq_axis, power)
246 def window_hann(length):
247 r"""Returns a Hann window array with length entries
251 The Hann window with length :math:`L` is defined as
253 .. math:: w_i = \frac{1}{2} (1-\cos(2\pi i/L))
255 win = zeros((length,), dtype=float)
256 for i in range(length):
257 win[i] = 0.5*(1.0-cos(2.0*pi*float(i)/(length)))
258 # avg value of cos over a period is 0
259 # so average height of Hann window is 0.5
262 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
263 overlap=True, window=window_hann):
264 """Compute the avgerage power spectrum of `data`.
269 Real (not complex) data taken with a sampling frequency `freq`.
273 Number of samples per chunk. Use a power of two.
274 overlap: {True,False}
275 If `True`, each chunk overlaps the previous chunk by half its
276 length. Otherwise, the chunks are end-to-end, and not
279 Weights used to "smooth" the chunks, there is a whole science
280 behind windowing, but if you're not trying to squeeze every
281 drop of information out of your data, you'll be OK with the
286 freq_axis,power : numpy.ndarray
287 Arrays ready for plotting.
291 The average power spectrum is computed by breaking `data` into
292 chunks of length `chunk_size`. These chunks are transformed
293 individually into frequency space and then averaged together.
295 See Numerical Recipes 2 section 13.4 for a good introduction to
298 If the number of samples in `data` is not a multiple of
299 `chunk_size`, we ignore the extra points.
301 assert chunk_size == floor_pow_of_two(chunk_size), \
302 "chunk_size %d should be a power of 2" % chunk_size
304 nchunks = len(data)/chunk_size # integer division = implicit floor
306 chunk_step = chunk_size/2
308 chunk_step = chunk_size
310 win = window(chunk_size) # generate a window of the appropriate size
311 freq_axis = linspace(0, freq/2, chunk_size/2+1)
312 # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
313 # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
314 # See Numerical Recipies for a details.
315 power = zeros((chunk_size/2+1,), dtype=float)
316 for i in range(nchunks):
317 starti = i*chunk_step
318 stopi = starti+chunk_size
319 fft_chunk = rfft(data[starti:stopi]*win)
320 p_chunk = fft_chunk * fft_chunk.conj()
321 power += p_chunk.astype(float)
322 power /= float(nchunks)
323 return (freq_axis, power)
325 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
326 overlap=True, window=window_hann):
327 """Compute the unitary avgerage power spectrum of `data`.
331 avg_power_spectrum,unitary_power_spectrum
333 freq_axis,power = avg_power_spectrum(data, freq, chunk_size,
335 # 2.0 / (freq * chunk_size) |rfft()|**2 --> unitary_power_spectrum
336 power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum()
337 # * 8/3 to remove power from windowing
338 # <[x(t)*w(t)]**2> = <x(t)**2 * w(t)**2> ~= <x(t)**2> * <w(t)**2>
339 # where the ~= is because the frequency of x(t) >> the frequency of w(t).
340 # So our calulated power has and extra <w(t)**2> in it.
341 # For the Hann window, <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
342 # For low frequency components, where the frequency of x(t) is ~= the frequency of w(t),
343 # The normalization is not perfect. ??
344 # The normalization approaches perfection as chunk_size -> infinity.
345 return (freq_axis, power)
349 class TestRFFT (unittest.TestCase):
350 r"""Ensure Numpy's FFT algorithm acts as expected.
355 The expected return values are [#numpybook]_:
357 .. math:: X_k = \sum_{m=0}^{n-1} x_m \exp^{-j 2\pi k_m/n}
359 .. [#numpybook] http://www.tramy.us/numpybook.pdf
361 def run_rfft(self, xs, Xs):
366 Xa.append(sum([x*exp(-j*2*pi*k*m/n) for x,m in zip(xs,range(n))]))
368 assert (Xs[k]-Xa[k])/abs(Xa[k]) < 1e-6, \
369 "rfft mismatch on element %d: %g != %g, relative error %g" \
370 % (k, Xs[k], Xa[k], (Xs[k]-Xa[k])/abs(Xa[k]))
371 # Which should satisfy the discrete form of Parseval's theorem
373 # SUM |x_m|^2 = 1/n SUM |X_k|^2.
375 timeSum = sum([abs(x)**2 for x in xs])
376 freqSum = sum([abs(X)**2 for X in Xa])
377 assert abs(freqSum/float(n) - timeSum)/timeSum < 1e-6, \
378 "Mismatch on Parseval's, %g != 1/%d * %g" % (timeSum, n, freqSum)
381 xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
382 self.run_rfft(xs, rfft(xs))
384 class TestUnitaryRFFT (unittest.TestCase):
385 """Verify `unitary_rfft`.
387 def run_unitary_rfft_parsevals(self, xs, freq, freqs, Xs):
388 """Check the discretized integral form of Parseval's theorem
395 .. math:: \sum_{m=0}^{n-1} |x_m|^2 dt = \sum_{k=0}^{n-1} |X_k|^2 df
398 df = freqs[1]-freqs[0]
399 assert (df - 1/(len(xs)*dt))/df < 1e-6, \
400 "Mismatch in spacing, %g != 1/(%d*%g)" % (df, len(xs), dt)
402 for k in range(len(Xs)-1,1,-1):
404 assert len(xs) == len(Xa), "Length mismatch %d != %d" % (len(xs), len(Xa))
405 lhs = sum([abs(x)**2 for x in xs]) * dt
406 rhs = sum([abs(X)**2 for X in Xa]) * df
407 assert abs(lhs - rhs)/lhs < 1e-4, "Mismatch on Parseval's, %g != %g" \
410 def test_unitary_rfft_parsevals(self):
411 "Test unitary rfft on Parseval's theorem"
412 xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
414 freqs,Xs = unitary_rfft(xs, 1.0/dt)
415 self.run_unitary_rfft_parsevals(xs, 1.0/dt, freqs, Xs)
418 r"""Rectangle function.
425 \rect(t) = \begin{cases}
426 1& \text{if $|t| < 0.5$}, \\
427 0& \text{if $|t| \ge 0.5$}.
435 def run_unitary_rfft_rect(self, a=1.0, time_shift=5.0, samp_freq=25.6,
437 r"""Test `unitary_rttf` on known function `rect(at)`.
444 .. math:: \rfft(\rect(at)) = 1/|a|\cdot\sinc(f/a)
446 samp_freq = float(samp_freq)
449 x = zeros((samples,), dtype=float)
451 for i in range(samples):
453 x[i] = self.rect(a*(t-time_shift))
454 freq_axis, X = unitary_rfft(x, samp_freq)
455 #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
457 # remove the phase due to our time shift
458 j = complex(0.0,1.0) # sqrt(-1)
459 for i in range(len(freq_axis)):
461 inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
462 X[i] *= inverse_phase_shift
464 expected = zeros((len(freq_axis),), dtype=float)
465 # normalized sinc(x) = sin(pi x)/(pi x)
466 # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
467 assert sinc(0.5) == 2.0/pi, "abnormal sinc()"
468 for i in range(len(freq_axis)):
470 expected[i] = 1.0/abs(a) * sinc(f/a)
475 pylab.plot(arange(0, dt*samples, dt), x)
476 pylab.title('time series')
478 pylab.plot(freq_axis, X.real, 'r.')
479 pylab.plot(freq_axis, X.imag, 'g.')
480 pylab.plot(freq_axis, expected, 'b-')
481 pylab.title('freq series')
483 def test_unitary_rfft_rect(self):
484 "Test unitary FFTs on variously shaped rectangular functions."
485 self.run_unitary_rfft_rect(a=0.5)
486 self.run_unitary_rfft_rect(a=2.0)
487 self.run_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
488 self.run_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
490 def gaussian(self, a, t):
491 r"""Gaussian function.
496 .. math:: \gaussian(a,t) = \exp^{-at^2}
498 return exp(-a * t**2)
500 def run_unitary_rfft_gaussian(self, a=1.0, time_shift=5.0, samp_freq=25.6,
502 r"""Test `unitary_rttf` on known function `gaussian(a,t)`.
511 \rfft(\gaussian(a,t)) = \sqrt{\pi/a} \cdot \gaussian(1/a,\pi f)
513 samp_freq = float(samp_freq)
516 x = zeros((samples,), dtype=float)
518 for i in range(samples):
520 x[i] = self.gaussian(a, (t-time_shift))
521 freq_axis, X = unitary_rfft(x, samp_freq)
522 #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
524 # remove the phase due to our time shift
525 j = complex(0.0,1.0) # sqrt(-1)
526 for i in range(len(freq_axis)):
528 inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
529 X[i] *= inverse_phase_shift
531 expected = zeros((len(freq_axis),), dtype=float)
532 for i in range(len(freq_axis)):
534 expected[i] = sqrt(pi/a) * self.gaussian(1.0/a, pi*f) # see Wikipedia, or do the integral yourself.
539 pylab.plot(arange(0, dt*samples, dt), x)
540 pylab.title('time series')
542 pylab.plot(freq_axis, X.real, 'r.')
543 pylab.plot(freq_axis, X.imag, 'g.')
544 pylab.plot(freq_axis, expected, 'b-')
545 pylab.title('freq series')
547 def test_unitary_rfft_gaussian(self):
548 "Test unitary FFTs on variously shaped gaussian functions."
549 self.run_unitary_rfft_gaussian(a=0.5)
550 self.run_unitary_rfft_gaussian(a=2.0)
551 self.run_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
552 self.run_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
554 class TestUnitaryPowerSpectrum (unittest.TestCase):
555 def run_unitary_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
557 x = zeros((samples,), dtype=float)
558 samp_freq = float(samp_freq)
559 for i in range(samples):
560 x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
561 freq_axis, power = unitary_power_spectrum(x, samp_freq)
564 expected = zeros((len(freq_axis),), dtype=float)
565 df = samp_freq/float(samples) # df = 1/T, where T = total_time
567 # average power per unit time is
569 # average value of sin(t)**2 = 0.5 (b/c sin**2+cos**2 == 1)
570 # so average value of (int sin(t)**2 dt) per unit time is 0.5
572 # we spread that power over a frequency bin of width df, sp
574 # where f0 is the sin's frequency
577 # FFT of sin(2*pi*t*f0) gives
578 # X(f) = 0.5 i (delta(f-f0) - delta(f-f0)),
579 # (area under x(t) = 0, area under X(f) = 0)
580 # so one sided power spectral density (PSD) per unit time is
581 # P(f) = 2 |X(f)|**2 / T
582 # = 2 * |0.5 delta(f-f0)|**2 / T
583 # = 0.5 * |delta(f-f0)|**2 / T
584 # but we're discrete and want the integral of the 'delta' to be 1,
585 # so 'delta'*df = 1 --> 'delta' = 1/df, and
586 # P(f) = 0.5 / (df**2 * T)
587 # = 0.5 / df (T = 1/df)
588 expected[i] = 0.5 / df
590 print "The power should be a peak at %g Hz of %g (%g, %g)" % \
591 (sin_freq, expected[i], freq_axis[imax], power[imax])
594 for i in range(len(freq_axis)):
595 Pexp += expected[i] *df
597 print " The total power should be %g (%g)" % (Pexp, P)
602 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
603 pylab.title('time series')
605 pylab.plot(freq_axis, power, 'r.')
606 pylab.plot(freq_axis, expected, 'b-')
607 pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
609 def test_unitary_power_spectrum_sin(self):
610 "Test unitary power spectrums on variously shaped sin functions"
611 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
612 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
613 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
614 self.run_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024)
615 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
616 # finally, with some irrational numbers, to check that I'm not getting lucky
617 self.run_unitary_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)
618 # test with non-integer number of periods
619 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256)
621 def run_unitary_power_spectrum_delta(self, amp=1, samp_freq=1,
625 x = zeros((samples,), dtype=float)
626 samp_freq = float(samp_freq)
628 freq_axis, power = unitary_power_spectrum(x, samp_freq)
630 # power = <x(t)**2> = (amp)**2 * dt/T
631 # we spread that power over the entire freq_axis [0,fN], so
632 # P(f) = (amp)**2 dt / (T fN)
634 # dt = 1/samp_freq (sample period)
635 # T = samples/samp_freq (total time of data aquisition)
636 # fN = 0.5 samp_freq (Nyquist frequency)
638 # P(f) = amp**2 / (samp_freq * samples/samp_freq * 0.5 samp_freq)
639 # = 2 amp**2 / (samp_freq*samples)
640 expected_amp = 2.0 * amp**2 / (samp_freq * samples)
641 expected = ones((len(freq_axis),), dtype=float) * expected_amp
643 print "The power should be flat at y = %g (%g)" % (expected_amp, power[0])
648 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
649 pylab.title('time series')
651 pylab.plot(freq_axis, power, 'r.')
652 pylab.plot(freq_axis, expected, 'b-')
653 pylab.title('%g samples of delta amp %g' % (samples, amp))
655 def _test_unitary_power_spectrum_delta(self):
656 "Test unitary power spectrums on various delta functions"
657 _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024)
658 _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048)
659 _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)# expected = 2*computed
660 _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)# expected = 0.5*computed
661 _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024)
662 _test_unitary_power_spectrum_delta(amp=pi, samp_freq=exp(1), samples=1024)
664 def gaussian(self, area, mean, std, t):
665 "Integral over all time = area (i.e. normalized for area=1)"
666 return area/(std*sqrt(2.0*pi)) * exp(-0.5*((t-mean)/std)**2)
668 def run_unitary_power_spectrum_gaussian(self, area=2.5, mean=5, std=1,
669 samp_freq=10.24 ,samples=512):
672 x = zeros((samples,), dtype=float)
674 for i in range(samples):
675 t = i/float(samp_freq)
676 x[i] = self.gaussian(area, mean, std, t)
677 freq_axis, power = unitary_power_spectrum(x, samp_freq)
679 # generate the predicted curve
680 # by comparing our self.gaussian() form to _gaussian(),
681 # we see that the Fourier transform of x(t) has parameters:
682 # std' = 1/(2 pi std) (references declaring std' = 1/std are converting to angular frequency, not frequency like we are)
683 # area' = area/[std sqrt(2*pi)] (plugging into FT of _gaussian() above)
684 # mean' = 0 (changing the mean in the time-domain just changes the phase in the freq-domain)
685 # So our power spectral density per unit time is given by
686 # P(f) = 2 |X(f)|**2 / T
688 # T = samples/samp_freq (total time of data aquisition)
690 area = area /(std*sqrt(2.0*pi))
691 std = 1.0/(2.0*pi*std)
692 expected = zeros((len(freq_axis),), dtype=float)
693 df = float(samp_freq)/samples # 1/total_time ( = freq_axis[1]-freq_axis[0] = freq_axis[1])
694 for i in range(len(freq_axis)):
696 gaus = self.gaussian(area, mean, std, f)
697 expected[i] = 2.0 * gaus**2 * samp_freq/samples
698 print "The power should be a half-gaussian, ",
699 print "with a peak at 0 Hz with amplitude %g (%g)" % (expected[0], power[0])
704 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
705 pylab.title('time series')
707 pylab.plot(freq_axis, power, 'r.')
708 pylab.plot(freq_axis, expected, 'b-')
709 pylab.title('freq series')
711 def test_unitary_power_spectrum_gaussian(self):
712 "Test unitary power spectrums on various gaussian functions"
714 for std in [1,sqrt(2)]:
715 for samp_freq in [10.0, exp(1)]:
716 for samples in [1024,2048]:
717 self.run_unitary_power_spectrum_gaussian(
718 area=area, std=std, samp_freq=samp_freq,
721 class TestUnitaryAvgPowerSpectrum (unittest.TestCase):
722 def run_unitary_avg_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
723 samples=1024, chunk_size=512,
724 overlap=True, window=window_hann):
727 x = zeros((samples,), dtype=float)
728 samp_freq = float(samp_freq)
729 for i in range(samples):
730 x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
731 freq_axis, power = unitary_avg_power_spectrum(x, samp_freq, chunk_size,
735 expected = zeros((len(freq_axis),), dtype=float)
736 df = samp_freq/float(chunk_size) # df = 1/T, where T = total_time
738 expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
740 print "The power should be a peak at %g Hz of %g (%g, %g)" % \
741 (sin_freq, expected[i], freq_axis[imax], power[imax])
744 for i in range(len(freq_axis)):
745 Pexp += expected[i] * df
747 print " The total power should be %g (%g)" % (Pexp, P)
752 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
753 pylab.title('time series')
755 pylab.plot(freq_axis, power, 'r.')
756 pylab.plot(freq_axis, expected, 'b-')
757 pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
759 def test_unitary_avg_power_spectrum_sin(self):
760 "Test unitary avg power spectrums on variously shaped sin functions."
761 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
762 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
763 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
764 self.run_unitary_avg_power_spectrum_sin(sin_freq=17, samp_freq=512, samples=1024)
765 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
766 # test long wavelenth sin, so be closer to window frequency
767 self.run_unitary_avg_power_spectrum_sin(sin_freq=1, samp_freq=1024, samples=2048)
768 # finally, with some irrational numbers, to check that I'm not getting lucky
769 self.run_unitary_avg_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)