1 # Copyright (C) 2008-2010 W. Trevor King <wking@drexel.edu>
3 # This file is part of Hooke.
5 # Hooke is free software: you can redistribute it and/or
6 # modify it under the terms of the GNU Lesser General Public
7 # License as published by the Free Software Foundation, either
8 # version 3 of the License, or (at your option) any later version.
10 # Hooke is distributed in the hope that it will be useful,
11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 # GNU Lesser General Public License for more details.
15 # You should have received a copy of the GNU Lesser General Public
16 # License along with Hooke. If not, see
17 # <http://www.gnu.org/licenses/>.
19 """Wrap :mod:`numpy.fft` to produce 1D unitary transforms and power spectra.
21 Define some FFT wrappers to reduce clutter.
22 Provides a unitary discrete FFT and a windowed version.
23 Based on numpy.fft.rfft.
27 * :func:`unitary_rfft`
28 * :func:`power_spectrum`
29 * :func:`unitary_power_spectrum`
30 * :func:`avg_power_spectrum`
31 * :func:`unitary_avg_power_spectrum`
36 from numpy import log2, floor, round, ceil, abs, pi, exp, cos, sin, sqrt, \
37 sinc, arctan2, array, ones, arange, linspace, zeros, \
38 uint16, float, concatenate, fromfile, argmax, complex
39 from numpy.fft import rfft
44 def floor_pow_of_two(num):
45 """Round `num` down to the closest exact a power of two.
50 >>> floor_pow_of_two(3)
52 >>> floor_pow_of_two(11)
54 >>> floor_pow_of_two(15)
62 def round_pow_of_two(num):
63 """Round `num` to the closest exact a power of two on a log scale.
68 >>> round_pow_of_two(2.9) # Note rounding on *log scale*
70 >>> round_pow_of_two(11)
72 >>> round_pow_of_two(15)
80 def ceil_pow_of_two(num):
81 """Round `num` up to the closest exact a power of two.
86 >>> ceil_pow_of_two(3)
88 >>> ceil_pow_of_two(11)
90 >>> ceil_pow_of_two(15)
98 def unitary_rfft(data, freq=1.0):
99 """Compute the unitary Fourier transform of real data.
101 Unitary = preserves power [Parseval's theorem].
106 Real (not complex) data taken with a sampling frequency `freq`.
112 freq_axis,trans : numpy.ndarray
113 Arrays ready for plotting.
118 If the units on your data are Volts,
119 and your sampling frequency is in Hz,
120 then `freq_axis` will be in Hz,
121 and `trans` will be in Volts.
123 nsamps = floor_pow_of_two(len(data))
124 # Which should satisfy the discrete form of Parseval's theorem
126 # SUM |x_m|^2 = 1/n SUM |X_k|^2.
128 # However, we want our FFT to satisfy the continuous Parseval eqn
129 # int_{-infty}^{infty} |x(t)|^2 dt = int_{-infty}^{infty} |X(f)|^2 df
130 # which has the discrete form
132 # SUM |x_m|^2 dt = SUM |X'_k|^2 df
134 # with X'_k = AX, this gives us
136 # SUM |x_m|^2 = A^2 df/dt SUM |X'_k|^2
141 # From Numerical Recipes (http://www.fizyka.umk.pl/nrbook/bookcpdf.html),
142 # Section 12.1, we see that for a sampling rate dt, the maximum frequency
143 # f_c in the transformed data is the Nyquist frequency (12.1.2)
145 # and the points are spaced out by (12.1.5)
151 # A = 1/ndf = ndt/n = dt
152 # so we can convert the Numpy transformed data to match our unitary
153 # continuous transformed data with (also NR 12.1.8)
154 # X'_k = dtX = X / <sampling freq>
155 trans = rfft(data[0:nsamps]) / float(freq)
156 freq_axis = linspace(0, freq/2, nsamps/2+1)
157 return (freq_axis, trans)
159 def power_spectrum(data, freq=1.0):
160 """Compute the power spectrum of the time series `data`.
165 Real (not complex) data taken with a sampling frequency `freq`.
171 freq_axis,power : numpy.ndarray
172 Arrays ready for plotting.
176 If the number of samples in `data` is not an integer power of two,
177 the FFT ignores some of the later points.
181 unitary_power_spectrum,avg_power_spectrum
183 nsamps = floor_pow_of_two(len(data))
185 freq_axis = linspace(0, freq/2, nsamps/2+1)
186 # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
187 # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
188 # See Numerical Recipies for a details.
189 trans = rfft(data[0:nsamps])
190 power = trans * trans.conj() # We want the square of the amplitude.
191 return (freq_axis, power)
193 def unitary_power_spectrum(data, freq=1.0):
194 """Compute the unitary power spectrum of the time series `data`.
198 power_spectrum,unitary_avg_power_spectrum
200 freq_axis,power = power_spectrum(data, freq)
201 # One sided power spectral density, so 2|H(f)|**2 (see NR 2nd edition 12.0.14, p498)
203 # numpy normalizes with 1/N on the inverse transform ifft,
204 # so we should normalize the freq-space representation with 1/sqrt(N).
205 # But we're using the rfft, where N points are like N/2 complex points, so 1/sqrt(N/2)
206 # So the power gets normalized by that twice and we have 2/N
208 # On top of this, the FFT assumes a sampling freq of 1 per second,
209 # and we want to preserve area under our curves.
210 # If our total time T = len(data)/freq is smaller than 1,
211 # our df_real = freq/len(data) is bigger that the FFT expects (dt_fft = 1/len(data)),
212 # and we need to scale the powers down to conserve area.
213 # df_fft * F_fft(f) = df_real *F_real(f)
214 # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq
215 # So the power gets normalized by *that* twice and we have 2/N * freq**2
217 # power per unit time
218 # measure x(t) for time T
219 # X(f) = int_0^T x(t) exp(-2 pi ift) dt
220 # PSD(f) = 2 |X(f)|**2 / T
222 # total_time = len(data)/float(freq)
223 # power *= 2.0 / float(freq)**2 / total_time
224 # power *= 2.0 / freq**2 * freq / len(data)
225 power *= 2.0 / (freq * float(len(data)))
227 return (freq_axis, power)
229 def window_hann(length):
230 r"""Returns a Hann window array with length entries
234 The Hann window with length :math:`L` is defined as
236 .. math:: w_i = \frac{1}{2} (1-\cos(2\pi i/L))
238 win = zeros((length,), dtype=float)
239 for i in range(length):
240 win[i] = 0.5*(1.0-cos(2.0*pi*float(i)/(length)))
241 # avg value of cos over a period is 0
242 # so average height of Hann window is 0.5
245 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
246 overlap=True, window=window_hann):
247 """Compute the avgerage power spectrum of `data`.
252 Real (not complex) data taken with a sampling frequency `freq`.
256 Number of samples per chunk. Use a power of two.
257 overlap: {True,False}
258 If `True`, each chunk overlaps the previous chunk by half its
259 length. Otherwise, the chunks are end-to-end, and not
262 Weights used to "smooth" the chunks, there is a whole science
263 behind windowing, but if you're not trying to squeeze every
264 drop of information out of your data, you'll be OK with the
269 freq_axis,power : numpy.ndarray
270 Arrays ready for plotting.
274 The average power spectrum is computed by breaking `data` into
275 chunks of length `chunk_size`. These chunks are transformed
276 individually into frequency space and then averaged together.
278 See Numerical Recipes 2 section 13.4 for a good introduction to
281 If the number of samples in `data` is not a multiple of
282 `chunk_size`, we ignore the extra points.
284 assert chunk_size == floor_pow_of_two(chunk_size), \
285 "chunk_size %d should be a power of 2" % chunk_size
287 nchunks = len(data)/chunk_size # integer division = implicit floor
289 chunk_step = chunk_size/2
291 chunk_step = chunk_size
293 win = window(chunk_size) # generate a window of the appropriate size
294 freq_axis = linspace(0, freq/2, chunk_size/2+1)
295 # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
296 # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
297 # See Numerical Recipies for a details.
298 power = zeros((chunk_size/2+1,), dtype=float)
299 for i in range(nchunks):
300 starti = i*chunk_step
301 stopi = starti+chunk_size
302 fft_chunk = rfft(data[starti:stopi]*win)
303 p_chunk = fft_chunk * fft_chunk.conj()
304 power += p_chunk.astype(float)
305 power /= float(nchunks)
306 return (freq_axis, power)
308 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
309 overlap=True, window=window_hann):
310 """Compute the unitary avgerage power spectrum of `data`.
314 avg_power_spectrum,unitary_power_spectrum
316 freq_axis,power = avg_power_spectrum(data, freq, chunk_size,
318 # 2.0 / (freq * chunk_size) |rfft()|**2 --> unitary_power_spectrum
319 power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum()
320 # * 8/3 to remove power from windowing
321 # <[x(t)*w(t)]**2> = <x(t)**2 * w(t)**2> ~= <x(t)**2> * <w(t)**2>
322 # where the ~= is because the frequency of x(t) >> the frequency of w(t).
323 # So our calulated power has and extra <w(t)**2> in it.
324 # For the Hann window, <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
325 # For low frequency components, where the frequency of x(t) is ~= the frequency of w(t),
326 # The normalization is not perfect. ??
327 # The normalization approaches perfection as chunk_size -> infinity.
328 return (freq_axis, power)
332 class TestRFFT (unittest.TestCase):
333 r"""Ensure Numpy's FFT algorithm acts as expected.
338 The expected return values are [#numpybook]_:
340 .. math:: X_k = \sum_{m=0}^{n-1} x_m \exp^{-j 2\pi k_m/n}
342 .. [#numpybook] http://www.tramy.us/numpybook.pdf
344 def run_rfft(self, xs, Xs):
349 Xa.append(sum([x*exp(-j*2*pi*k*m/n) for x,m in zip(xs,range(n))]))
351 assert (Xs[k]-Xa[k])/abs(Xa[k]) < 1e-6, \
352 "rfft mismatch on element %d: %g != %g, relative error %g" \
353 % (k, Xs[k], Xa[k], (Xs[k]-Xa[k])/abs(Xa[k]))
354 # Which should satisfy the discrete form of Parseval's theorem
356 # SUM |x_m|^2 = 1/n SUM |X_k|^2.
358 timeSum = sum([abs(x)**2 for x in xs])
359 freqSum = sum([abs(X)**2 for X in Xa])
360 assert abs(freqSum/float(n) - timeSum)/timeSum < 1e-6, \
361 "Mismatch on Parseval's, %g != 1/%d * %g" % (timeSum, n, freqSum)
364 xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
365 self.run_rfft(xs, rfft(xs))
367 class TestUnitaryRFFT (unittest.TestCase):
368 """Verify `unitary_rfft`.
370 def run_unitary_rfft_parsevals(self, xs, freq, freqs, Xs):
371 """Check the discretized integral form of Parseval's theorem
378 .. math:: \sum_{m=0}^{n-1} |x_m|^2 dt = \sum_{k=0}^{n-1} |X_k|^2 df
381 df = freqs[1]-freqs[0]
382 assert (df - 1/(len(xs)*dt))/df < 1e-6, \
383 "Mismatch in spacing, %g != 1/(%d*%g)" % (df, len(xs), dt)
385 for k in range(len(Xs)-1,1,-1):
387 assert len(xs) == len(Xa), "Length mismatch %d != %d" % (len(xs), len(Xa))
388 lhs = sum([abs(x)**2 for x in xs]) * dt
389 rhs = sum([abs(X)**2 for X in Xa]) * df
390 assert abs(lhs - rhs)/lhs < 1e-4, "Mismatch on Parseval's, %g != %g" \
393 def test_unitary_rfft_parsevals(self):
394 "Test unitary rfft on Parseval's theorem"
395 xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
397 freqs,Xs = unitary_rfft(xs, 1.0/dt)
398 self.run_unitary_rfft_parsevals(xs, 1.0/dt, freqs, Xs)
401 r"""Rectangle function.
408 \rect(t) = \begin{cases}
409 1& \text{if $|t| < 0.5$}, \\
410 0& \text{if $|t| \ge 0.5$}.
418 def run_unitary_rfft_rect(self, a=1.0, time_shift=5.0, samp_freq=25.6,
420 r"""Test `unitary_rttf` on known function `rect(at)`.
427 .. math:: \rfft(\rect(at)) = 1/|a|\cdot\sinc(f/a)
429 samp_freq = float(samp_freq)
432 x = zeros((samples,), dtype=float)
434 for i in range(samples):
436 x[i] = self.rect(a*(t-time_shift))
437 freq_axis, X = unitary_rfft(x, samp_freq)
438 #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
440 # remove the phase due to our time shift
441 j = complex(0.0,1.0) # sqrt(-1)
442 for i in range(len(freq_axis)):
444 inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
445 X[i] *= inverse_phase_shift
447 expected = zeros((len(freq_axis),), dtype=float)
448 # normalized sinc(x) = sin(pi x)/(pi x)
449 # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
450 assert sinc(0.5) == 2.0/pi, "abnormal sinc()"
451 for i in range(len(freq_axis)):
453 expected[i] = 1.0/abs(a) * sinc(f/a)
458 pylab.plot(arange(0, dt*samples, dt), x)
459 pylab.title('time series')
461 pylab.plot(freq_axis, X.real, 'r.')
462 pylab.plot(freq_axis, X.imag, 'g.')
463 pylab.plot(freq_axis, expected, 'b-')
464 pylab.title('freq series')
466 def test_unitary_rfft_rect(self):
467 "Test unitary FFTs on variously shaped rectangular functions."
468 self.run_unitary_rfft_rect(a=0.5)
469 self.run_unitary_rfft_rect(a=2.0)
470 self.run_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
471 self.run_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
473 def gaussian(self, a, t):
474 r"""Gaussian function.
479 .. math:: \gaussian(a,t) = \exp^{-at^2}
481 return exp(-a * t**2)
483 def run_unitary_rfft_gaussian(self, a=1.0, time_shift=5.0, samp_freq=25.6,
485 r"""Test `unitary_rttf` on known function `gaussian(a,t)`.
494 \rfft(\gaussian(a,t)) = \sqrt{\pi/a} \cdot \gaussian(1/a,\pi f)
496 samp_freq = float(samp_freq)
499 x = zeros((samples,), dtype=float)
501 for i in range(samples):
503 x[i] = self.gaussian(a, (t-time_shift))
504 freq_axis, X = unitary_rfft(x, samp_freq)
505 #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
507 # remove the phase due to our time shift
508 j = complex(0.0,1.0) # sqrt(-1)
509 for i in range(len(freq_axis)):
511 inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
512 X[i] *= inverse_phase_shift
514 expected = zeros((len(freq_axis),), dtype=float)
515 for i in range(len(freq_axis)):
517 expected[i] = sqrt(pi/a) * self.gaussian(1.0/a, pi*f) # see Wikipedia, or do the integral yourself.
522 pylab.plot(arange(0, dt*samples, dt), x)
523 pylab.title('time series')
525 pylab.plot(freq_axis, X.real, 'r.')
526 pylab.plot(freq_axis, X.imag, 'g.')
527 pylab.plot(freq_axis, expected, 'b-')
528 pylab.title('freq series')
530 def test_unitary_rfft_gaussian(self):
531 "Test unitary FFTs on variously shaped gaussian functions."
532 self.run_unitary_rfft_gaussian(a=0.5)
533 self.run_unitary_rfft_gaussian(a=2.0)
534 self.run_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
535 self.run_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
537 class TestUnitaryPowerSpectrum (unittest.TestCase):
538 def run_unitary_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
540 x = zeros((samples,), dtype=float)
541 samp_freq = float(samp_freq)
542 for i in range(samples):
543 x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
544 freq_axis, power = unitary_power_spectrum(x, samp_freq)
547 expected = zeros((len(freq_axis),), dtype=float)
548 df = samp_freq/float(samples) # df = 1/T, where T = total_time
550 # average power per unit time is
552 # average value of sin(t)**2 = 0.5 (b/c sin**2+cos**2 == 1)
553 # so average value of (int sin(t)**2 dt) per unit time is 0.5
555 # we spread that power over a frequency bin of width df, sp
557 # where f0 is the sin's frequency
560 # FFT of sin(2*pi*t*f0) gives
561 # X(f) = 0.5 i (delta(f-f0) - delta(f-f0)),
562 # (area under x(t) = 0, area under X(f) = 0)
563 # so one sided power spectral density (PSD) per unit time is
564 # P(f) = 2 |X(f)|**2 / T
565 # = 2 * |0.5 delta(f-f0)|**2 / T
566 # = 0.5 * |delta(f-f0)|**2 / T
567 # but we're discrete and want the integral of the 'delta' to be 1,
568 # so 'delta'*df = 1 --> 'delta' = 1/df, and
569 # P(f) = 0.5 / (df**2 * T)
570 # = 0.5 / df (T = 1/df)
571 expected[i] = 0.5 / df
573 print "The power should be a peak at %g Hz of %g (%g, %g)" % \
574 (sin_freq, expected[i], freq_axis[imax], power[imax])
577 for i in range(len(freq_axis)):
578 Pexp += expected[i] *df
580 print " The total power should be %g (%g)" % (Pexp, P)
585 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
586 pylab.title('time series')
588 pylab.plot(freq_axis, power, 'r.')
589 pylab.plot(freq_axis, expected, 'b-')
590 pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
592 def test_unitary_power_spectrum_sin(self):
593 "Test unitary power spectrums on variously shaped sin functions"
594 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
595 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
596 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
597 self.run_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024)
598 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
599 # finally, with some irrational numbers, to check that I'm not getting lucky
600 self.run_unitary_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)
601 # test with non-integer number of periods
602 self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256)
604 def run_unitary_power_spectrum_delta(self, amp=1, samp_freq=1,
608 x = zeros((samples,), dtype=float)
609 samp_freq = float(samp_freq)
611 freq_axis, power = unitary_power_spectrum(x, samp_freq)
613 # power = <x(t)**2> = (amp)**2 * dt/T
614 # we spread that power over the entire freq_axis [0,fN], so
615 # P(f) = (amp)**2 dt / (T fN)
617 # dt = 1/samp_freq (sample period)
618 # T = samples/samp_freq (total time of data aquisition)
619 # fN = 0.5 samp_freq (Nyquist frequency)
621 # P(f) = amp**2 / (samp_freq * samples/samp_freq * 0.5 samp_freq)
622 # = 2 amp**2 / (samp_freq*samples)
623 expected_amp = 2.0 * amp**2 / (samp_freq * samples)
624 expected = ones((len(freq_axis),), dtype=float) * expected_amp
626 print "The power should be flat at y = %g (%g)" % (expected_amp, power[0])
631 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
632 pylab.title('time series')
634 pylab.plot(freq_axis, power, 'r.')
635 pylab.plot(freq_axis, expected, 'b-')
636 pylab.title('%g samples of delta amp %g' % (samples, amp))
638 def _test_unitary_power_spectrum_delta(self):
639 "Test unitary power spectrums on various delta functions"
640 _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024)
641 _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048)
642 _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)# expected = 2*computed
643 _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)# expected = 0.5*computed
644 _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024)
645 _test_unitary_power_spectrum_delta(amp=pi, samp_freq=exp(1), samples=1024)
647 def gaussian(self, area, mean, std, t):
648 "Integral over all time = area (i.e. normalized for area=1)"
649 return area/(std*sqrt(2.0*pi)) * exp(-0.5*((t-mean)/std)**2)
651 def run_unitary_power_spectrum_gaussian(self, area=2.5, mean=5, std=1,
652 samp_freq=10.24 ,samples=512):
655 x = zeros((samples,), dtype=float)
657 for i in range(samples):
658 t = i/float(samp_freq)
659 x[i] = self.gaussian(area, mean, std, t)
660 freq_axis, power = unitary_power_spectrum(x, samp_freq)
662 # generate the predicted curve
663 # by comparing our self.gaussian() form to _gaussian(),
664 # we see that the Fourier transform of x(t) has parameters:
665 # std' = 1/(2 pi std) (references declaring std' = 1/std are converting to angular frequency, not frequency like we are)
666 # area' = area/[std sqrt(2*pi)] (plugging into FT of _gaussian() above)
667 # mean' = 0 (changing the mean in the time-domain just changes the phase in the freq-domain)
668 # So our power spectral density per unit time is given by
669 # P(f) = 2 |X(f)|**2 / T
671 # T = samples/samp_freq (total time of data aquisition)
673 area = area /(std*sqrt(2.0*pi))
674 std = 1.0/(2.0*pi*std)
675 expected = zeros((len(freq_axis),), dtype=float)
676 df = float(samp_freq)/samples # 1/total_time ( = freq_axis[1]-freq_axis[0] = freq_axis[1])
677 for i in range(len(freq_axis)):
679 gaus = self.gaussian(area, mean, std, f)
680 expected[i] = 2.0 * gaus**2 * samp_freq/samples
681 print "The power should be a half-gaussian, ",
682 print "with a peak at 0 Hz with amplitude %g (%g)" % (expected[0], power[0])
687 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
688 pylab.title('time series')
690 pylab.plot(freq_axis, power, 'r.')
691 pylab.plot(freq_axis, expected, 'b-')
692 pylab.title('freq series')
694 def test_unitary_power_spectrum_gaussian(self):
695 "Test unitary power spectrums on various gaussian functions"
697 for std in [1,sqrt(2)]:
698 for samp_freq in [10.0, exp(1)]:
699 for samples in [1024,2048]:
700 self.run_unitary_power_spectrum_gaussian(
701 area=area, std=std, samp_freq=samp_freq,
704 class TestUnitaryAvgPowerSpectrum (unittest.TestCase):
705 def run_unitary_avg_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
706 samples=1024, chunk_size=512,
707 overlap=True, window=window_hann):
710 x = zeros((samples,), dtype=float)
711 samp_freq = float(samp_freq)
712 for i in range(samples):
713 x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
714 freq_axis, power = unitary_avg_power_spectrum(x, samp_freq, chunk_size,
718 expected = zeros((len(freq_axis),), dtype=float)
719 df = samp_freq/float(chunk_size) # df = 1/T, where T = total_time
721 expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
723 print "The power should be a peak at %g Hz of %g (%g, %g)" % \
724 (sin_freq, expected[i], freq_axis[imax], power[imax])
727 for i in range(len(freq_axis)):
728 Pexp += expected[i] * df
730 print " The total power should be %g (%g)" % (Pexp, P)
735 pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
736 pylab.title('time series')
738 pylab.plot(freq_axis, power, 'r.')
739 pylab.plot(freq_axis, expected, 'b-')
740 pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
742 def test_unitary_avg_power_spectrum_sin(self):
743 "Test unitary avg power spectrums on variously shaped sin functions."
744 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
745 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
746 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
747 self.run_unitary_avg_power_spectrum_sin(sin_freq=17, samp_freq=512, samples=1024)
748 self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
749 # test long wavelenth sin, so be closer to window frequency
750 self.run_unitary_avg_power_spectrum_sin(sin_freq=1, samp_freq=1024, samples=2048)
751 # finally, with some irrational numbers, to check that I'm not getting lucky
752 self.run_unitary_avg_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)