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