FFT_tools.py

   1 #!/usr/bin/python
   2 #
   3 # Wrap Numpy's fft module to produce 1D unitary transforms and power spectra.
   4 #
   5 # Copyright (C) 2008  W. Trevor King
   6 # All rights reserved.
   7 #
   8 # Redistribution and use in source and binary forms, with or without
   9 # modification, are permitted provided that the following conditions
  10 # are met:
  11 #
  12 #  * Redistributions of source code must retain the above copyright
  13 #    notice, this list of conditions and the following disclaimer.
  14 #
  15 #  * Redistributions in binary form must reproduce the above copyright
  16 #    notice, this list of conditions and the following disclaimer in
  17 #    the documentation and/or other materials provided with the
  18 #    distribution.
  19 #
  20 #  * Neither the name of the copyright holders nor the names of the
  21 #    contributors may be used to endorse or promote products derived
  22 #    from this software without specific prior written permission.
  23 #
  24 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  25 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  26 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
  27 # FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
  28 # COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
  29 # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
  30 # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  31 # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  32 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  33 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  34 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  35 # POSSIBILITY OF SUCH DAMAGE.
  36
  37 """
  38 Define some FFT wrappers to reduce clutter.
  39 Provides a unitary discrete FFT and a windowed version.
  40 Based on numpy.fft.rfft.
  41
  42 Main entry functions:
  43
  44 * :func:`unitary_rfft`
  45 * :func:`power_spectrum`
  46 * :func:`unitary_power_spectrum`
  47 * :func:`avg_power_spectrum`
  48 * :func:`unitary_avg_power_spectrum`
  49 """
  50
  51 import unittest
  52
  53 from numpy import log2, floor, round, ceil, abs, pi, exp, cos, sin, sqrt, \
  54     sinc, arctan2, array, ones, arange, linspace, zeros, \
  55     uint16, float, concatenate, fromfile, argmax, complex
  56 from numpy.fft import rfft
  57
  58
  59 TEST_PLOTS = False
  60
  61 def floor_pow_of_two(num):
  62     """Round `num` down to the closest exact a power of two.
  63
  64     Examples
  65     --------
  66
  67     >>> floor_pow_of_two(3)
  68     2
  69     >>> floor_pow_of_two(11)
  70     8
  71     >>> floor_pow_of_two(15)
  72     8
  73     """
  74     lnum = log2(num)
  75     if int(lnum) != lnum:
  76         num = 2**floor(lnum)
  77     return int(num)
  78
  79 def round_pow_of_two(num):
  80     """Round `num` to the closest exact a power of two on a log scale.
  81
  82     Examples
  83     --------
  84
  85     >>> round_pow_of_two(2.9) # Note rounding on *log scale*
  86     4
  87     >>> round_pow_of_two(11)
  88     8
  89     >>> round_pow_of_two(15)
  90     16
  91     """
  92     lnum = log2(num)
  93     if int(lnum) != lnum:
  94         num = 2**round(lnum)
  95     return int(num)
  96
  97 def ceil_pow_of_two(num):
  98     """Round `num` up to the closest exact a power of two.
  99
 100     Examples
 101     --------
 102
 103     >>> ceil_pow_of_two(3)
 104     4
 105     >>> ceil_pow_of_two(11)
 106     16
 107     >>> ceil_pow_of_two(15)
 108     16
 109     """
 110     lnum = log2(num)
 111     if int(lnum) != lnum:
 112         num = 2**ceil(lnum)
 113     return int(num)
 114
 115 def unitary_rfft(data, freq=1.0):
 116     """Compute the unitary Fourier transform of real data.
 117
 118     Unitary = preserves power [Parseval's theorem].
 119
 120     Parameters
 121     ----------
 122     data : iterable
 123         Real (not complex) data taken with a sampling frequency `freq`.
 124     freq : real
 125         Sampling frequency.
 126
 127     Returns
 128     -------
 129     freq_axis,trans : numpy.ndarray
 130         Arrays ready for plotting.
 131
 132     Notes
 133     -----
 134
 135     If the units on your data are Volts,
 136     and your sampling frequency is in Hz,
 137     then `freq_axis` will be in Hz,
 138     and `trans` will be in Volts.
 139     """
 140     nsamps = floor_pow_of_two(len(data))
 141     # Which should satisfy the discrete form of Parseval's theorem
 142     #   n-1               n-1
 143     #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
 144     #   m=0               k=0
 145     # However, we want our FFT to satisfy the continuous Parseval eqn
 146     #   int_{-infty}^{infty} |x(t)|^2 dt = int_{-infty}^{infty} |X(f)|^2 df
 147     # which has the discrete form
 148     #   n-1              n-1
 149     #   SUM |x_m|^2 dt = SUM |X'_k|^2 df
 150     #   m=0              k=0
 151     # with X'_k = AX, this gives us
 152     #   n-1                     n-1
 153     #   SUM |x_m|^2 = A^2 df/dt SUM |X'_k|^2
 154     #   m=0                     k=0
 155     # so we see
 156     #   A^2 df/dt = 1/n
 157     #   A^2 = 1/n dt/df
 158     # From Numerical Recipes (http://www.fizyka.umk.pl/nrbook/bookcpdf.html),
 159     # Section 12.1, we see that for a sampling rate dt, the maximum frequency
 160     # f_c in the transformed data is the Nyquist frequency (12.1.2)
 161     #   f_c = 1/2dt
 162     # and the points are spaced out by (12.1.5)
 163     #   df = 1/ndt
 164     # so
 165     #   dt = 1/ndf
 166     #   dt/df = 1/ndf^2
 167     #   A^2 = 1/n^2df^2
 168     #   A = 1/ndf = ndt/n = dt
 169     # so we can convert the Numpy transformed data to match our unitary
 170     # continuous transformed data with (also NR 12.1.8)
 171     #   X'_k = dtX = X / <sampling freq>
 172     trans = rfft(data[0:nsamps]) / float(freq)
 173     freq_axis = linspace(0, freq/2, nsamps/2+1)
 174     return (freq_axis, trans)
 175
 176 def power_spectrum(data, freq=1.0):
 177     """Compute the power spectrum of the time series `data`.
 178
 179     Parameters
 180     ----------
 181     data : iterable
 182         Real (not complex) data taken with a sampling frequency `freq`.
 183     freq : real
 184         Sampling frequency.
 185
 186     Returns
 187     -------
 188     freq_axis,power : numpy.ndarray
 189         Arrays ready for plotting.
 190
 191     Notes
 192     -----
 193     If the number of samples in `data` is not an integer power of two,
 194     the FFT ignores some of the later points.
 195
 196     See Also
 197     --------
 198     unitary_power_spectrum,avg_power_spectrum
 199     """
 200     nsamps = floor_pow_of_two(len(data))
 201
 202     freq_axis = linspace(0, freq/2, nsamps/2+1)
 203     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
 204     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
 205     # See Numerical Recipies for a details.
 206     trans = rfft(data[0:nsamps])
 207     power = trans * trans.conj() # We want the square of the amplitude.
 208     return (freq_axis, power)
 209
 210 def unitary_power_spectrum(data, freq=1.0):
 211     """Compute the unitary power spectrum of the time series `data`.
 212
 213     See Also
 214     --------
 215     power_spectrum,unitary_avg_power_spectrum
 216     """
 217     freq_axis,power = power_spectrum(data, freq)
 218     # One sided power spectral density, so 2|H(f)|**2 (see NR 2nd edition 12.0.14, p498)
 219     #
 220     # numpy normalizes with 1/N on the inverse transform ifft,
 221     # so we should normalize the freq-space representation with 1/sqrt(N).
 222     # But we're using the rfft, where N points are like N/2 complex points, so 1/sqrt(N/2)
 223     # So the power gets normalized by that twice and we have 2/N
 224     #
 225     # On top of this, the FFT assumes a sampling freq of 1 per second,
 226     # and we want to preserve area under our curves.
 227     # If our total time T = len(data)/freq is smaller than 1,
 228     # our df_real = freq/len(data) is bigger that the FFT expects (dt_fft = 1/len(data)),
 229     # and we need to scale the powers down to conserve area.
 230     # df_fft * F_fft(f) = df_real *F_real(f)
 231     # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq
 232     # So the power gets normalized by *that* twice and we have 2/N * freq**2
 233
 234     # power per unit time
 235     # measure x(t) for time T
 236     # X(f)   = int_0^T x(t) exp(-2 pi ift) dt
 237     # PSD(f) = 2 |X(f)|**2 / T
 238
 239     # total_time = len(data)/float(freq)
 240     # power *= 2.0 / float(freq)**2   /   total_time
 241     # power *= 2.0 / freq**2   *   freq / len(data)
 242     power *= 2.0 / (freq * float(len(data)))
 243
 244     return (freq_axis, power)
 245
 246 def window_hann(length):
 247     r"""Returns a Hann window array with length entries
 248
 249     Notes
 250     -----
 251     The Hann window with length :math:`L` is defined as
 252
 253     .. math:: w_i = \frac{1}{2} (1-\cos(2\pi i/L))
 254     """
 255     win = zeros((length,), dtype=float)
 256     for i in range(length):
 257         win[i] = 0.5*(1.0-cos(2.0*pi*float(i)/(length)))
 258     # avg value of cos over a period is 0
 259     # so average height of Hann window is 0.5
 260     return win
 261
 262 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
 263                        overlap=True, window=window_hann):
 264     """Compute the avgerage power spectrum of `data`.
 265
 266     Parameters
 267     ----------
 268     data : iterable
 269         Real (not complex) data taken with a sampling frequency `freq`.
 270     freq : real
 271         Sampling frequency.
 272     chunk_size : int
 273         Number of samples per chunk.  Use a power of two.
 274     overlap: {True,False}
 275         If `True`, each chunk overlaps the previous chunk by half its
 276         length.  Otherwise, the chunks are end-to-end, and not
 277         overlapping.
 278     window: iterable
 279         Weights used to "smooth" the chunks, there is a whole science
 280         behind windowing, but if you're not trying to squeeze every
 281         drop of information out of your data, you'll be OK with the
 282         default Hann window.
 283
 284     Returns
 285     -------
 286     freq_axis,power : numpy.ndarray
 287         Arrays ready for plotting.
 288
 289     Notes
 290     -----
 291     The average power spectrum is computed by breaking `data` into
 292     chunks of length `chunk_size`.  These chunks are transformed
 293     individually into frequency space and then averaged together.
 294
 295     See Numerical Recipes 2 section 13.4 for a good introduction to
 296     the theory.
 297
 298     If the number of samples in `data` is not a multiple of
 299     `chunk_size`, we ignore the extra points.
 300     """
 301     assert chunk_size == floor_pow_of_two(chunk_size), \
 302         "chunk_size %d should be a power of 2" % chunk_size
 303
 304     nchunks = len(data)/chunk_size # integer division = implicit floor
 305     if overlap:
 306         chunk_step = chunk_size/2
 307     else:
 308         chunk_step = chunk_size
 309
 310     win = window(chunk_size) # generate a window of the appropriate size
 311     freq_axis = linspace(0, freq/2, chunk_size/2+1)
 312     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
 313     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
 314     # See Numerical Recipies for a details.
 315     power = zeros((chunk_size/2+1,), dtype=float)
 316     for i in range(nchunks):
 317         starti = i*chunk_step
 318         stopi = starti+chunk_size
 319         fft_chunk = rfft(data[starti:stopi]*win)
 320         p_chunk = fft_chunk * fft_chunk.conj()
 321         power += p_chunk.astype(float)
 322     power /= float(nchunks)
 323     return (freq_axis, power)
 324
 325 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
 326                                overlap=True, window=window_hann):
 327     """Compute the unitary avgerage power spectrum of `data`.
 328
 329     See Also
 330     --------
 331     avg_power_spectrum,unitary_power_spectrum
 332     """
 333     freq_axis,power = avg_power_spectrum(data, freq, chunk_size,
 334                                          overlap, window)
 335     #        2.0 / (freq * chunk_size)          |rfft()|**2 --> unitary_power_spectrum
 336     power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum()
 337     #                                       * 8/3  to remove power from windowing
 338     #  <[x(t)*w(t)]**2> = <x(t)**2 * w(t)**2> ~= <x(t)**2> * <w(t)**2>
 339     # where the ~= is because the frequency of x(t) >> the frequency of w(t).
 340     # So our calulated power has and extra <w(t)**2> in it.
 341     # For the Hann window, <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
 342     # For low frequency components, where the frequency of x(t) is ~= the frequency of w(t),
 343     # The normalization is not perfect. ??
 344     # The normalization approaches perfection as chunk_size -> infinity.
 345     return (freq_axis, power)
 346
 347
 348
 349 class TestRFFT (unittest.TestCase):
 350     r"""Ensure Numpy's FFT algorithm acts as expected.
 351
 352     Notes
 353     -----
 354
 355     The expected return values are [#numpybook]_:
 356
 357     .. math:: X_k = \sum_{m=0}^{n-1} x_m \exp^{-j 2\pi k_m/n}
 358
 359     .. [#numpybook] http://www.tramy.us/numpybook.pdf
 360     """
 361     def run_rfft(self, xs, Xs):
 362         j = complex(0,1)
 363         n = len(xs)
 364         Xa = []
 365         for k in range(n):
 366             Xa.append(sum([x*exp(-j*2*pi*k*m/n) for x,m in zip(xs,range(n))]))
 367             if k < len(Xs):
 368                 assert (Xs[k]-Xa[k])/abs(Xa[k]) < 1e-6, \
 369                     "rfft mismatch on element %d: %g != %g, relative error %g" \
 370                     % (k, Xs[k], Xa[k], (Xs[k]-Xa[k])/abs(Xa[k]))
 371         # Which should satisfy the discrete form of Parseval's theorem
 372         #   n-1               n-1
 373         #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
 374         #   m=0               k=0
 375         timeSum = sum([abs(x)**2 for x in xs])
 376         freqSum = sum([abs(X)**2 for X in Xa])
 377         assert abs(freqSum/float(n) - timeSum)/timeSum < 1e-6, \
 378             "Mismatch on Parseval's, %g != 1/%d * %g" % (timeSum, n, freqSum)
 379
 380     def test_rfft(self):
 381         xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
 382         self.run_rfft(xs, rfft(xs))
 383
 384 class TestUnitaryRFFT (unittest.TestCase):
 385     """Verify `unitary_rfft`.
 386     """
 387     def run_unitary_rfft_parsevals(self, xs, freq, freqs, Xs):
 388         """Check the discretized integral form of Parseval's theorem
 389
 390         Notes
 391         -----
 392
 393         Which is:
 394
 395         .. math:: \sum_{m=0}^{n-1} |x_m|^2 dt = \sum_{k=0}^{n-1} |X_k|^2 df
 396         """
 397         dt = 1.0/freq
 398         df = freqs[1]-freqs[0]
 399         assert (df - 1/(len(xs)*dt))/df < 1e-6, \
 400             "Mismatch in spacing, %g != 1/(%d*%g)" % (df, len(xs), dt)
 401         Xa = list(Xs)
 402         for k in range(len(Xs)-1,1,-1):
 403             Xa.append(Xa[k])
 404         assert len(xs) == len(Xa), "Length mismatch %d != %d" % (len(xs), len(Xa))
 405         lhs = sum([abs(x)**2 for x in xs]) * dt
 406         rhs = sum([abs(X)**2 for X in Xa]) * df
 407         assert abs(lhs - rhs)/lhs < 1e-4, "Mismatch on Parseval's, %g != %g" \
 408             % (lhs, rhs)
 409
 410     def test_unitary_rfft_parsevals(self):
 411         "Test unitary rfft on Parseval's theorem"
 412         xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
 413         dt = pi
 414         freqs,Xs = unitary_rfft(xs, 1.0/dt)
 415         self.run_unitary_rfft_parsevals(xs, 1.0/dt, freqs, Xs)
 416
 417     def rect(self, t):
 418         r"""Rectangle function.
 419
 420         Notes
 421         -----
 422
 423         .. math::
 424
 425             \rect(t) = \begin{cases}
 426                1& \text{if $|t| < 0.5$}, \\
 427                0& \text{if $|t| \ge 0.5$}.
 428                        \end{cases}
 429         """
 430         if abs(t) < 0.5:
 431             return 1
 432         else:
 433             return 0
 434
 435     def run_unitary_rfft_rect(self, a=1.0, time_shift=5.0, samp_freq=25.6,
 436                               samples=256):
 437         r"""Test `unitary_rttf` on known function `rect(at)`.
 438
 439         Notes
 440         -----
 441
 442         Analytic result:
 443
 444         .. math:: \rfft(\rect(at)) = 1/|a|\cdot\sinc(f/a)
 445         """
 446         samp_freq = float(samp_freq)
 447         a = float(a)
 448
 449         x = zeros((samples,), dtype=float)
 450         dt = 1.0/samp_freq
 451         for i in range(samples):
 452             t = i*dt
 453             x[i] = self.rect(a*(t-time_shift))
 454         freq_axis, X = unitary_rfft(x, samp_freq)
 455         #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
 456
 457         # remove the phase due to our time shift
 458         j = complex(0.0,1.0) # sqrt(-1)
 459         for i in range(len(freq_axis)):
 460             f = freq_axis[i]
 461             inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
 462             X[i] *= inverse_phase_shift
 463
 464         expected = zeros((len(freq_axis),), dtype=float)
 465         # normalized sinc(x) = sin(pi x)/(pi x)
 466         # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
 467         assert sinc(0.5) == 2.0/pi, "abnormal sinc()"
 468         for i in range(len(freq_axis)):
 469             f = freq_axis[i]
 470             expected[i] = 1.0/abs(a) * sinc(f/a)
 471
 472         if TEST_PLOTS:
 473             pylab.figure()
 474             pylab.subplot(211)
 475             pylab.plot(arange(0, dt*samples, dt), x)
 476             pylab.title('time series')
 477             pylab.subplot(212)
 478             pylab.plot(freq_axis, X.real, 'r.')
 479             pylab.plot(freq_axis, X.imag, 'g.')
 480             pylab.plot(freq_axis, expected, 'b-')
 481             pylab.title('freq series')
 482
 483     def test_unitary_rfft_rect(self):
 484         "Test unitary FFTs on variously shaped rectangular functions."
 485         self.run_unitary_rfft_rect(a=0.5)
 486         self.run_unitary_rfft_rect(a=2.0)
 487         self.run_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
 488         self.run_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
 489
 490     def gaussian(self, a, t):
 491         r"""Gaussian function.
 492
 493         Notes
 494         -----
 495
 496         .. math:: \gaussian(a,t) = \exp^{-at^2}
 497         """
 498         return exp(-a * t**2)
 499
 500     def run_unitary_rfft_gaussian(self, a=1.0, time_shift=5.0, samp_freq=25.6,
 501                                   samples=256):
 502         r"""Test `unitary_rttf` on known function `gaussian(a,t)`.
 503
 504         Notes
 505         -----
 506
 507         Analytic result:
 508
 509         .. math::
 510
 511             \rfft(\gaussian(a,t)) = \sqrt{\pi/a} \cdot \gaussian(1/a,\pi f)
 512         """
 513         samp_freq = float(samp_freq)
 514         a = float(a)
 515
 516         x = zeros((samples,), dtype=float)
 517         dt = 1.0/samp_freq
 518         for i in range(samples):
 519             t = i*dt
 520             x[i] = self.gaussian(a, (t-time_shift))
 521         freq_axis, X = unitary_rfft(x, samp_freq)
 522         #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
 523
 524         # remove the phase due to our time shift
 525         j = complex(0.0,1.0) # sqrt(-1)
 526         for i in range(len(freq_axis)):
 527             f = freq_axis[i]
 528             inverse_phase_shift = exp(j*2.0*pi*time_shift*f)
 529             X[i] *= inverse_phase_shift
 530
 531         expected = zeros((len(freq_axis),), dtype=float)
 532         for i in range(len(freq_axis)):
 533             f = freq_axis[i]
 534             expected[i] = sqrt(pi/a) * self.gaussian(1.0/a, pi*f) # see Wikipedia, or do the integral yourself.
 535
 536         if TEST_PLOTS:
 537             pylab.figure()
 538             pylab.subplot(211)
 539             pylab.plot(arange(0, dt*samples, dt), x)
 540             pylab.title('time series')
 541             pylab.subplot(212)
 542             pylab.plot(freq_axis, X.real, 'r.')
 543             pylab.plot(freq_axis, X.imag, 'g.')
 544             pylab.plot(freq_axis, expected, 'b-')
 545             pylab.title('freq series')
 546
 547     def test_unitary_rfft_gaussian(self):
 548         "Test unitary FFTs on variously shaped gaussian functions."
 549         self.run_unitary_rfft_gaussian(a=0.5)
 550         self.run_unitary_rfft_gaussian(a=2.0)
 551         self.run_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
 552         self.run_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
 553
 554 class TestUnitaryPowerSpectrum (unittest.TestCase):
 555     def run_unitary_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
 556                                        samples=1024):
 557         x = zeros((samples,), dtype=float)
 558         samp_freq = float(samp_freq)
 559         for i in range(samples):
 560             x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
 561         freq_axis, power = unitary_power_spectrum(x, samp_freq)
 562         imax = argmax(power)
 563
 564         expected = zeros((len(freq_axis),), dtype=float)
 565         df = samp_freq/float(samples) # df = 1/T, where T = total_time
 566         i = int(sin_freq/df)
 567         # average power per unit time is
 568         #  P = <x(t)**2>
 569         # average value of sin(t)**2 = 0.5    (b/c sin**2+cos**2 == 1)
 570         # so average value of (int sin(t)**2 dt) per unit time is 0.5
 571         #  P = 0.5
 572         # we spread that power over a frequency bin of width df, sp
 573         #  P(f0) = 0.5/df
 574         # where f0 is the sin's frequency
 575         #
 576         # or:
 577         # FFT of sin(2*pi*t*f0) gives
 578         #  X(f) = 0.5 i (delta(f-f0) - delta(f-f0)),
 579         # (area under x(t) = 0, area under X(f) = 0)
 580         # so one sided power spectral density (PSD) per unit time is
 581         #  P(f) = 2 |X(f)|**2 / T
 582         #       = 2 * |0.5 delta(f-f0)|**2 / T
 583         #       = 0.5 * |delta(f-f0)|**2 / T
 584         # but we're discrete and want the integral of the 'delta' to be 1,
 585         # so 'delta'*df = 1  --> 'delta' = 1/df, and
 586         #  P(f) = 0.5 / (df**2 * T)
 587         #       = 0.5 / df                (T = 1/df)
 588         expected[i] = 0.5 / df
 589
 590         print "The power should be a peak at %g Hz of %g (%g, %g)" % \
 591             (sin_freq, expected[i], freq_axis[imax], power[imax])
 592         Pexp = 0
 593         P    = 0
 594         for i in range(len(freq_axis)):
 595             Pexp += expected[i] *df
 596             P    += power[i] * df
 597         print " The total power should be %g (%g)" % (Pexp, P)
 598
 599         if TEST_PLOTS:
 600             pylab.figure()
 601             pylab.subplot(211)
 602             pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
 603             pylab.title('time series')
 604             pylab.subplot(212)
 605             pylab.plot(freq_axis, power, 'r.')
 606             pylab.plot(freq_axis, expected, 'b-')
 607             pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
 608
 609     def test_unitary_power_spectrum_sin(self):
 610         "Test unitary power spectrums on variously shaped sin functions"
 611         self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
 612         self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
 613         self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
 614         self.run_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024)
 615         self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
 616         # finally, with some irrational numbers, to check that I'm not getting lucky
 617         self.run_unitary_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)
 618         # test with non-integer number of periods
 619         self.run_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256)
 620
 621     def run_unitary_power_spectrum_delta(self, amp=1, samp_freq=1,
 622                                          samples=256):
 623         """TODO
 624         """
 625         x = zeros((samples,), dtype=float)
 626         samp_freq = float(samp_freq)
 627         x[0] = amp
 628         freq_axis, power = unitary_power_spectrum(x, samp_freq)
 629
 630         # power = <x(t)**2> = (amp)**2 * dt/T
 631         # we spread that power over the entire freq_axis [0,fN], so
 632         #  P(f)  = (amp)**2 dt / (T fN)
 633         # where
 634         #  dt = 1/samp_freq        (sample period)
 635         #  T  = samples/samp_freq  (total time of data aquisition)
 636         #  fN = 0.5 samp_freq      (Nyquist frequency)
 637         # so
 638         #  P(f) = amp**2 / (samp_freq * samples/samp_freq * 0.5 samp_freq)
 639         #       = 2 amp**2 / (samp_freq*samples)
 640         expected_amp = 2.0 * amp**2 / (samp_freq * samples)
 641         expected = ones((len(freq_axis),), dtype=float) * expected_amp
 642
 643         print "The power should be flat at y = %g (%g)" % (expected_amp, power[0])
 644
 645         if TEST_PLOTS:
 646             pylab.figure()
 647             pylab.subplot(211)
 648             pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
 649             pylab.title('time series')
 650             pylab.subplot(212)
 651             pylab.plot(freq_axis, power, 'r.')
 652             pylab.plot(freq_axis, expected, 'b-')
 653             pylab.title('%g samples of delta amp %g' % (samples, amp))
 654
 655     def _test_unitary_power_spectrum_delta(self):
 656         "Test unitary power spectrums on various delta functions"
 657         _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024)
 658         _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048)
 659         _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)# expected = 2*computed
 660         _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)# expected = 0.5*computed
 661         _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024)
 662         _test_unitary_power_spectrum_delta(amp=pi, samp_freq=exp(1), samples=1024)
 663
 664     def gaussian(self, area, mean, std, t):
 665         "Integral over all time = area (i.e. normalized for area=1)"
 666         return area/(std*sqrt(2.0*pi)) * exp(-0.5*((t-mean)/std)**2)
 667
 668     def run_unitary_power_spectrum_gaussian(self, area=2.5, mean=5, std=1,
 669                                             samp_freq=10.24 ,samples=512):
 670         """TODO.
 671         """
 672         x = zeros((samples,), dtype=float)
 673         mean = float(mean)
 674         for i in range(samples):
 675             t = i/float(samp_freq)
 676             x[i] = self.gaussian(area, mean, std, t)
 677         freq_axis, power = unitary_power_spectrum(x, samp_freq)
 678
 679         # generate the predicted curve
 680         # by comparing our self.gaussian() form to _gaussian(),
 681         # we see that the Fourier transform of x(t) has parameters:
 682         #  std'  = 1/(2 pi std)    (references declaring std' = 1/std are converting to angular frequency, not frequency like we are)
 683         #  area' = area/[std sqrt(2*pi)]   (plugging into FT of _gaussian() above)
 684         #  mean' = 0               (changing the mean in the time-domain just changes the phase in the freq-domain)
 685         # So our power spectral density per unit time is given by
 686         #  P(f) = 2 |X(f)|**2 / T
 687         # Where
 688         #  T  = samples/samp_freq  (total time of data aquisition)
 689         mean = 0.0
 690         area = area /(std*sqrt(2.0*pi))
 691         std = 1.0/(2.0*pi*std)
 692         expected = zeros((len(freq_axis),), dtype=float)
 693         df = float(samp_freq)/samples # 1/total_time ( = freq_axis[1]-freq_axis[0] = freq_axis[1])
 694         for i in range(len(freq_axis)):
 695             f = i*df
 696             gaus = self.gaussian(area, mean, std, f)
 697             expected[i] = 2.0 * gaus**2 * samp_freq/samples
 698         print "The power should be a half-gaussian, ",
 699         print "with a peak at 0 Hz with amplitude %g (%g)" % (expected[0], power[0])
 700
 701         if TEST_PLOTS:
 702             pylab.figure()
 703             pylab.subplot(211)
 704             pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
 705             pylab.title('time series')
 706             pylab.subplot(212)
 707             pylab.plot(freq_axis, power, 'r.')
 708             pylab.plot(freq_axis, expected, 'b-')
 709             pylab.title('freq series')
 710
 711     def test_unitary_power_spectrum_gaussian(self):
 712         "Test unitary power spectrums on various gaussian functions"
 713         for area in [1,pi]:
 714             for std in [1,sqrt(2)]:
 715                 for samp_freq in [10.0, exp(1)]:
 716                     for samples in [1024,2048]:
 717                         self.run_unitary_power_spectrum_gaussian(
 718                             area=area, std=std, samp_freq=samp_freq,
 719                             samples=samples)
 720
 721 class TestUnitaryAvgPowerSpectrum (unittest.TestCase):
 722     def run_unitary_avg_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
 723                                            samples=1024, chunk_size=512,
 724                                            overlap=True, window=window_hann):
 725         """TODO
 726         """
 727         x = zeros((samples,), dtype=float)
 728         samp_freq = float(samp_freq)
 729         for i in range(samples):
 730             x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
 731         freq_axis, power = unitary_avg_power_spectrum(x, samp_freq, chunk_size,
 732                                                       overlap, window)
 733         imax = argmax(power)
 734
 735         expected = zeros((len(freq_axis),), dtype=float)
 736         df = samp_freq/float(chunk_size) # df = 1/T, where T = total_time
 737         i = int(sin_freq/df)
 738         expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
 739
 740         print "The power should be a peak at %g Hz of %g (%g, %g)" % \
 741             (sin_freq, expected[i], freq_axis[imax], power[imax])
 742         Pexp = 0
 743         P    = 0
 744         for i in range(len(freq_axis)):
 745             Pexp += expected[i] * df
 746             P    += power[i] * df
 747         print " The total power should be %g (%g)" % (Pexp, P)
 748
 749         if TEST_PLOTS:
 750             pylab.figure()
 751             pylab.subplot(211)
 752             pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
 753             pylab.title('time series')
 754             pylab.subplot(212)
 755             pylab.plot(freq_axis, power, 'r.')
 756             pylab.plot(freq_axis, expected, 'b-')
 757             pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
 758
 759     def test_unitary_avg_power_spectrum_sin(self):
 760         "Test unitary avg power spectrums on variously shaped sin functions."
 761         self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
 762         self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
 763         self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
 764         self.run_unitary_avg_power_spectrum_sin(sin_freq=17, samp_freq=512, samples=1024)
 765         self.run_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
 766         # test long wavelenth sin, so be closer to window frequency
 767         self.run_unitary_avg_power_spectrum_sin(sin_freq=1, samp_freq=1024, samples=2048)
 768         # finally, with some irrational numbers, to check that I'm not getting lucky
 769         self.run_unitary_avg_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024)