FFT_tools.py

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