FFT_tools.py

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