FFT_tools.py

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