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