FFT_tools.py

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