FFT_tools: add spaces after commas (PEP8) except tuple assignment
[FFT-tools.git] / FFT_tools.py
index 352083a0925beac8d3b7a13a000f831b73ccfa82..13fd5a0c9e035936aab5b035750c14690712bf66 100644 (file)
@@ -50,6 +50,7 @@ def floor_pow_of_two(num):
         num = 2**_numpy.floor(lnum)
     return num
 
+
 def round_pow_of_two(num):
     "Round num to the closest exact a power of two on a log scale."
     lnum = _numpy.log2(num)
@@ -57,6 +58,7 @@ def round_pow_of_two(num):
         num = 2**_numpy.round(lnum)
     return num
 
+
 def ceil_pow_of_two(num):
     "Round num up to the closest exact a power of two."
     lnum = _numpy.log2(num)
@@ -64,40 +66,43 @@ def ceil_pow_of_two(num):
         num = 2**_numpy.ceil(lnum)
     return num
 
+
 def _test_rfft(xs, Xs):
     # Numpy's FFT algoritm returns
     #          n-1
     #   X[k] = SUM x[m] exp (-j 2pi km /n)
     #          m=0
     # (see http://www.tramy.us/numpybook.pdf)
-    j = _numpy.complex(0,1)
+    j = _numpy.complex(0, 1)
     n = len(xs)
     Xa = []
     for k in range(n):
-        Xa.append(sum([x*_numpy.exp(-j*2*_numpy.pi*k*m/n)
-                       for x,m in zip(xs,range(n))]))
+        Xa.append(sum([x * _numpy.exp(-j * 2 * _numpy.pi * k * m / n)
+                       for x,m in zip(xs, range(n))]))
         if k < len(Xs):
-            if (Xs[k]-Xa[k])/_numpy.abs(Xa[k]) >= 1e-6:
+            if (Xs[k] - Xa[k]) / _numpy.abs(Xa[k]) >= 1e-6:
                 raise ValueError(
                     ('rfft mismatch on element {}: {} != {}, relative error {}'
                      ).format(
-                        k, Xs[k], Xa[k], (Xs[k]-Xa[k])/_numpy.abs(Xa[k])))
+                        k, Xs[k], Xa[k], (Xs[k] - Xa[k]) / _numpy.abs(Xa[k])))
     # Which should satisfy the discrete form of Parseval's theorem
     #   n-1               n-1
     #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
     #   m=0               k=0
     timeSum = sum([_numpy.abs(x)**2 for x in xs])
     freqSum = sum([_numpy.abs(X)**2 for X in Xa])
-    if _numpy.abs(freqSum/_numpy.float(n) - timeSum)/timeSum >= 1e-6:
+    if _numpy.abs(freqSum / _numpy.float(n) - timeSum) / timeSum >= 1e-6:
         raise ValueError(
             "Mismatch on Parseval's, {} != 1/{} * {}".format(
                 timeSum, n, freqSum))
 
+
 def _test_rfft_suite():
     print('Test numpy rfft definition')
-    xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
+    xs = [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1]
     _test_rfft(xs, _numpy.fft.rfft(xs))
 
+
 def unitary_rfft(data, freq=1.0):
     """Compute the Fourier transform of real data.
 
@@ -141,35 +146,38 @@ def unitary_rfft(data, freq=1.0):
     # continuous transformed data with (also NR 12.1.8)
     #   X'_k = dtX = X / <sampling freq>
     trans = _numpy.fft.rfft(data[0:nsamps]) / _numpy.float(freq)
-    freq_axis = _numpy.linspace(0, freq/2, nsamps/2+1)
+    freq_axis = _numpy.linspace(0, freq / 2, nsamps / 2 + 1)
     return (freq_axis, trans)
 
+
 def _test_unitary_rfft_parsevals(xs, freq, freqs, Xs):
     # Which should satisfy the discretized integral form of Parseval's theorem
     #   n-1              n-1
     #   SUM |x_m|^2 dt = SUM |X_k|^2 df
     #   m=0              k=0
-    dt = 1.0/freq
-    df = freqs[1]-freqs[0]
-    if df - 1/(len(xs)*dt)/df >= 1e-6:
+    dt = 1.0 / freq
+    df = freqs[1] - freqs[0]
+    if df - 1 / (len(xs) * dt * df) >= 1e-6:
         raise ValueError(
             'Mismatch in spacing, {} != 1/({}*{})'.format(df, len(xs), dt))
     Xa = list(Xs)
-    for k in range(len(Xs)-1,1,-1):
+    for k in range(len(Xs) - 1, 1, -1):
         Xa.append(Xa[k])
     if len(xs) != len(Xa):
         raise ValueError('Length mismatch {} != {}'.format(len(xs), len(Xa)))
     lhs = sum([_numpy.abs(x)**2 for x in xs]) * dt
     rhs = sum([_numpy.abs(X)**2 for X in Xa]) * df
-    if _numpy.abs(lhs - rhs)/lhs >= 1e-4:
+    if _numpy.abs(lhs - rhs) / lhs >= 1e-4:
         raise ValueError("Mismatch on Parseval's, {} != {}".format(lhs, rhs))
 
+
 def _test_unitary_rfft_parsevals_suite():
     print("Test unitary rfft on Parseval's theorem")
-    xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
+    xs = [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1]
     dt = _numpy.pi
-    freqs,Xs = unitary_rfft(xs, 1.0/dt)
-    _test_unitary_rfft_parsevals(xs, 1.0/dt, freqs, Xs)
+    freqs,Xs = unitary_rfft(xs, 1.0 / dt)
+    _test_unitary_rfft_parsevals(xs, 1.0 / dt, freqs, Xs)
+
 
 def _rect(t):
     if _numpy.abs(t) < 0.5:
@@ -177,6 +185,7 @@ def _rect(t):
     else:
         return 0
 
+
 def _test_unitary_rfft_rect(
     a=1.0, time_shift=5.0, samp_freq=25.6, samples=256):
     "Show fft(rect(at)) = 1/abs(a) * _numpy.sinc(f/a)"
@@ -184,33 +193,33 @@ def _test_unitary_rfft_rect(
     a = _numpy.float(a)
 
     x = _numpy.zeros((samples,), dtype=_numpy.float)
-    dt = 1.0/samp_freq
+    dt = 1.0 / samp_freq
     for i in range(samples):
-        t = i*dt
-        x[i] = _rect(a*(t-time_shift))
+        t = i * dt
+        x[i] = _rect(a * (t - time_shift))
     freq_axis, X = unitary_rfft(x, samp_freq)
     #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
 
     # remove the phase due to our time shift
-    j = _numpy.complex(0.0,1.0) # sqrt(-1)
+    j = _numpy.complex(0.0, 1.0)  # sqrt(-1)
     for i in range(len(freq_axis)):
         f = freq_axis[i]
-        inverse_phase_shift = _numpy.exp(j*2.0*_numpy.pi*time_shift*f)
+        inverse_phase_shift = _numpy.exp(j * 2.0 * _numpy.pi * time_shift * f)
         X[i] *= inverse_phase_shift
 
     expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
     # normalized sinc(x) = sin(pi x)/(pi x)
     # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
-    if _numpy.sinc(0.5) != 2.0/_numpy.pi:
+    if _numpy.sinc(0.5) != 2.0 / _numpy.pi:
         raise ValueError('abnormal sinc()')
     for i in range(len(freq_axis)):
         f = freq_axis[i]
-        expected[i] = 1.0/_numpy.abs(a) * _numpy.sinc(f/a)
+        expected[i] = 1.0 / _numpy.abs(a) * _numpy.sinc(f / a)
 
     if TEST_PLOTS:
         figure = _pyplot.figure()
         time_axes = figure.add_subplot(2, 1, 1)
-        time_axes.plot(_numpy.arange(0, dt*samples, dt), x)
+        time_axes.plot(_numpy.arange(0, dt * samples, dt), x)
         time_axes.set_title('time series')
         freq_axes = figure.add_subplot(2, 1, 2)
         freq_axes.plot(freq_axis, X.real, 'r.')
@@ -218,6 +227,7 @@ def _test_unitary_rfft_rect(
         freq_axes.plot(freq_axis, expected, 'b-')
         freq_axes.set_title('freq series')
 
+
 def _test_unitary_rfft_rect_suite():
     print('Test unitary FFTs on variously shaped rectangular functions')
     _test_unitary_rfft_rect(a=0.5)
@@ -225,9 +235,11 @@ def _test_unitary_rfft_rect_suite():
     _test_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
     _test_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
 
+
 def _gaussian(a, t):
     return _numpy.exp(-a * t**2)
 
+
 def _test_unitary_rfft_gaussian(
     a=1.0, time_shift=5.0, samp_freq=25.6, samples=256):
     "Show fft(rect(at)) = 1/abs(a) * sinc(f/a)"
@@ -235,31 +247,31 @@ def _test_unitary_rfft_gaussian(
     a = _numpy.float(a)
 
     x = _numpy.zeros((samples,), dtype=_numpy.float)
-    dt = 1.0/samp_freq
+    dt = 1.0 / samp_freq
     for i in range(samples):
-        t = i*dt
-        x[i] = _gaussian(a, (t-time_shift))
+        t = i * dt
+        x[i] = _gaussian(a, (t - time_shift))
     freq_axis, X = unitary_rfft(x, samp_freq)
     #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X)
 
     # remove the phase due to our time shift
-    j = _numpy.complex(0.0,1.0) # sqrt(-1)
+    j = _numpy.complex(0.0, 1.0)  # sqrt(-1)
     for i in range(len(freq_axis)):
         f = freq_axis[i]
-        inverse_phase_shift = _numpy.exp(j*2.0*_numpy.pi*time_shift*f)
+        inverse_phase_shift = _numpy.exp(j * 2.0 * _numpy.pi * time_shift * f)
         X[i] *= inverse_phase_shift
 
     expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
     for i in range(len(freq_axis)):
         f = freq_axis[i]
         # see Wikipedia, or do the integral yourself.
-        expected[i] = _numpy.sqrt(_numpy.pi/a) * _gaussian(
-            1.0/a, _numpy.pi*f)
+        expected[i] = _numpy.sqrt(_numpy.pi / a) * _gaussian(
+            1.0 / a, _numpy.pi * f)
 
     if TEST_PLOTS:
         figure = _pyplot.figure()
         time_axes = figure.add_subplot(2, 1, 1)
-        time_axes.plot(_numpy.arange(0, dt*samples, dt), x)
+        time_axes.plot(_numpy.arange(0, dt * samples, dt), x)
         time_axes.set_title('time series')
         freq_axes = figure.add_subplot(2, 1, 2)
         freq_axes.plot(freq_axis, X.real, 'r.')
@@ -267,6 +279,7 @@ def _test_unitary_rfft_gaussian(
         freq_axes.plot(freq_axis, expected, 'b-')
         freq_axes.set_title('freq series')
 
+
 def _test_unitary_rfft_gaussian_suite():
     print("Test unitary FFTs on variously shaped gaussian functions")
     _test_unitary_rfft_gaussian(a=0.5)
@@ -275,7 +288,6 @@ def _test_unitary_rfft_gaussian_suite():
     _test_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
 
 
-
 def power_spectrum(data, freq=1.0):
     """Compute the power spectrum of DATA taken with a sampling frequency FREQ.
 
@@ -286,27 +298,31 @@ def power_spectrum(data, freq=1.0):
     """
     nsamps = floor_pow_of_two(len(data))
 
-    freq_axis = _numpy.linspace(0, freq/2, nsamps/2+1)
+    freq_axis = _numpy.linspace(0, freq / 2, nsamps / 2 + 1)
     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
     # See Numerical Recipies for a details.
     trans = _numpy.fft.rfft(data[0:nsamps])
-    power = (trans * trans.conj()).real # We want the square of the amplitude.
+    power = (trans * trans.conj()).real  # we want the square of the amplitude
     return (freq_axis, power)
 
+
 def unitary_power_spectrum(data, freq=1.0):
     freq_axis,power = power_spectrum(data, freq)
-    # One sided power spectral density, so 2|H(f)|**2 (see NR 2nd edition 12.0.14, p498)
+    # One sided power spectral density, so 2|H(f)|**2
+    # (see NR 2nd edition 12.0.14, p498)
     #
     # numpy normalizes with 1/N on the inverse transform ifft,
     # so we should normalize the freq-space representation with 1/sqrt(N).
-    # But we're using the rfft, where N points are like N/2 complex points, so 1/sqrt(N/2)
+    # But we're using the rfft, where N points are like N/2 complex points,
+    # so 1/sqrt(N/2)
     # So the power gets normalized by that twice and we have 2/N
     #
     # On top of this, the FFT assumes a sampling freq of 1 per second,
     # and we want to preserve area under our curves.
     # If our total time T = len(data)/freq is smaller than 1,
-    # our df_real = freq/len(data) is bigger that the FFT expects (dt_fft = 1/len(data)),
+    # our df_real = freq/len(data) is bigger that the FFT expects
+    # (dt_fft = 1/len(data)),
     # and we need to scale the powers down to conserve area.
     # df_fft * F_fft(f) = df_real *F_real(f)
     # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq
@@ -324,17 +340,18 @@ def unitary_power_spectrum(data, freq=1.0):
 
     return (freq_axis, power)
 
+
 def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024):
     x = _numpy.zeros((samples,), dtype=_numpy.float)
     samp_freq = _numpy.float(samp_freq)
     for i in range(samples):
-        x[i] = _numpy.sin(2.0 * _numpy.pi * (i/samp_freq) * sin_freq)
+        x[i] = _numpy.sin(2.0 * _numpy.pi * (i / samp_freq) * sin_freq)
     freq_axis, power = unitary_power_spectrum(x, samp_freq)
     imax = _numpy.argmax(power)
 
     expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
-    df = samp_freq/_numpy.float(samples) # df = 1/T, where T = total_time
-    i = int(sin_freq/df)
+    df = samp_freq / _numpy.float(samples)  # df = 1/T, where T = total_time
+    i = int(sin_freq / df)
     # average power per unit time is
     #  P = <x(t)**2>
     # average value of sin(t)**2 = 0.5    (b/c sin**2+cos**2 == 1)
@@ -360,18 +377,17 @@ def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024):
 
     print('The power should be a peak at {} Hz of {} ({}, {})'.format(
             sin_freq, expected[i], freq_axis[imax], power[imax]))
-    Pexp = 0
-    P    = 0
+    Pexp = P = 0
     for i in range(len(freq_axis)):
-        Pexp += expected[i] *df
-        P    += power[i] * df
+        Pexp += expected[i] * df
+        P += power[i] * df
     print('The total power should be {} ({})'.format(Pexp, P))
 
     if TEST_PLOTS:
         figure = _pyplot.figure()
         time_axes = figure.add_subplot(2, 1, 1)
         time_axes.plot(
-            _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
+            _numpy.arange(0, samples / samp_freq, 1.0 / samp_freq), x, 'b-')
         time_axes.set_title('time series')
         freq_axes = figure.add_subplot(2, 1, 2)
         freq_axes.plot(freq_axis, power, 'r.')
@@ -379,6 +395,7 @@ def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024):
         freq_axes.set_title(
             '{} samples of sin at {} Hz'.format(samples, sin_freq))
 
+
 def _test_unitary_power_spectrum_sin_suite():
     print('Test unitary power spectrums on variously shaped sin functions')
     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
@@ -386,12 +403,13 @@ def _test_unitary_power_spectrum_sin_suite():
     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
     _test_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024)
     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
-    # finally, with some irrational numbers, to check that I'm not getting lucky
+    # with some irrational numbers, to check that I'm not getting lucky
     _test_unitary_power_spectrum_sin(
-        sin_freq=_numpy.pi, samp_freq=100*_numpy.exp(1), samples=1024)
+        sin_freq=_numpy.pi, samp_freq=100 * _numpy.exp(1), samples=1024)
     # test with non-integer number of periods
     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256)
 
+
 def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256):
     x = _numpy.zeros((samples,), dtype=_numpy.float)
     samp_freq = _numpy.float(samp_freq)
@@ -419,57 +437,65 @@ def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256):
         figure = _pyplot.figure()
         time_axes = figure.add_subplot(2, 1, 1)
         time_axes.plot(
-            _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
+            _numpy.arange(0, samples / samp_freq, 1.0 / samp_freq), x, 'b-')
         time_axes.set_title('time series')
         freq_axes = figure.add_subplot(2, 1, 2)
         freq_axes.plot(freq_axis, power, 'r.')
         freq_axes.plot(freq_axis, expected, 'b-')
         freq_axes.set_title('{} samples of delta amp {}'.format(samples, amp))
 
+
 def _test_unitary_power_spectrum_delta_suite():
     print('Test unitary power spectrums on various delta functions')
     _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024)
     _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048)
-    _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)# expected = 2*computed
-    _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)# expected = 0.5*computed
+    # expected = 2*computed
+    _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)
+    # expected = 0.5*computed
+    _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)
     _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024)
     _test_unitary_power_spectrum_delta(
         amp=_numpy.pi, samp_freq=_numpy.exp(1), samples=1024)
 
+
 def _gaussian2(area, mean, std, t):
     "Integral over all time = area (i.e. normalized for area=1)"
-    return area/(std*_numpy.sqrt(2.0*_numpy.pi)) * _numpy.exp(
-        -0.5*((t-mean)/std)**2)
+    return area / (std * _numpy.sqrt(2.0 * _numpy.pi)) * _numpy.exp(
+        -0.5 * ((t-mean)/std)**2)
+
 
 def _test_unitary_power_spectrum_gaussian(
-    area=2.5, mean=5, std=1, samp_freq=10.24 ,samples=512):
+    area=2.5, mean=5, std=1, samp_freq=10.24samples=512):
     x = _numpy.zeros((samples,), dtype=_numpy.float)
     mean = _numpy.float(mean)
     for i in range(samples):
-        t = i/_numpy.float(samp_freq)
+        t = i / _numpy.float(samp_freq)
         x[i] = _gaussian2(area, mean, std, t)
     freq_axis, power = unitary_power_spectrum(x, samp_freq)
 
     # generate the predicted curve
     # by comparing our _gaussian2() form to _gaussian(),
     # we see that the Fourier transform of x(t) has parameters:
-    #  std'  = 1/(2 pi std)    (references declaring std' = 1/std are converting to angular frequency, not frequency like we are)
+    #  std'  = 1/(2 pi std)    (references declaring std' = 1/std are
+    #                           converting to angular frequency,
+    #                           not frequency like we are)
     #  area' = area/[std sqrt(2*pi)]   (plugging into FT of _gaussian() above)
-    #  mean' = 0               (changing the mean in the time-domain just changes the phase in the freq-domain)
+    #  mean' = 0               (changing the mean in the time-domain just
+    #                           changes the phase in the freq-domain)
     # So our power spectral density per unit time is given by
     #  P(f) = 2 |X(f)|**2 / T
     # Where
     #  T  = samples/samp_freq  (total time of data aquisition)
     mean = 0.0
-    area = area /(std*_numpy.sqrt(2.0*_numpy.pi))
-    std = 1.0/(2.0*_numpy.pi*std)
+    area = area / (std * _numpy.sqrt(2.0 * _numpy.pi))
+    std = 1.0 / (2.0 * _numpy.pi * std)
     expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
     # 1/total_time ( = freq_axis[1]-freq_axis[0] = freq_axis[1])
-    df = _numpy.float(samp_freq)/samples
+    df = _numpy.float(samp_freq) / samples
     for i in range(len(freq_axis)):
-        f = i*df
+        f = i * df
         gaus = _gaussian2(area, mean, std, f)
-        expected[i] = 2.0 * gaus**2 * samp_freq/samples
+        expected[i] = 2.0 * gaus**2 * samp_freq / samples
     print(('The power should be a half-gaussian, '
            'with a peak at 0 Hz with amplitude {} ({})').format(
             expected[0], power[0]))
@@ -478,33 +504,42 @@ def _test_unitary_power_spectrum_gaussian(
         figure = _pyplot.figure()
         time_axes = figure.add_subplot(2, 1, 1)
         time_axes.plot(
-            _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
+            _numpy.arange(0, samples / samp_freq, 1.0 / samp_freq), x, 'b-')
         time_axes.set_title('time series')
         freq_axes = figure.add_subplot(2, 1, 2)
         freq_axes.plot(freq_axis, power, 'r.')
         freq_axes.plot(freq_axis, expected, 'b-')
         freq_axes.set_title('freq series')
 
+
 def _test_unitary_power_spectrum_gaussian_suite():
     print('Test unitary power spectrums on various gaussian functions')
-    _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=1024)
-    _test_unitary_power_spectrum_gaussian(area=1, std=2, samp_freq=10.0, samples=1024)
-    _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=2048)
-    _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=20.0, samples=2048)
-    _test_unitary_power_spectrum_gaussian(area=3, std=1, samp_freq=10.0, samples=1024)
+    _test_unitary_power_spectrum_gaussian(
+        area=1, std=1, samp_freq=10.0, samples=1024)
+    _test_unitary_power_spectrum_gaussian(
+        area=1, std=2, samp_freq=10.0, samples=1024)
+    _test_unitary_power_spectrum_gaussian(
+        area=1, std=1, samp_freq=10.0, samples=2048)
+    _test_unitary_power_spectrum_gaussian(
+        area=1, std=1, samp_freq=20.0, samples=2048)
+    _test_unitary_power_spectrum_gaussian(
+        area=3, std=1, samp_freq=10.0, samples=1024)
     _test_unitary_power_spectrum_gaussian(
         area=_numpy.pi, std=_numpy.sqrt(2), samp_freq=_numpy.exp(1),
         samples=1024)
 
+
 def window_hann(length):
     "Returns a Hann window array with length entries"
     win = _numpy.zeros((length,), dtype=_numpy.float)
     for i in range(length):
-        win[i] = 0.5*(1.0-_numpy.cos(2.0*_numpy.pi*_numpy.float(i)/(length)))
+        win[i] = 0.5 * (
+            1.0 - _numpy.cos(2.0 * _numpy.pi * _numpy.float(i) / (length)))
     # avg value of cos over a period is 0
     # so average height of Hann window is 0.5
     return win
 
+
 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
                        overlap=True, window=window_hann):
     """Compute the avgerage power spectrum of DATA.
@@ -526,76 +561,79 @@ def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
         raise ValueError(
             'chunk_size {} should be a power of 2'.format(chunk_size))
 
-    nchunks = len(data)/chunk_size # integer division = implicit floor
+    nchunks = len(data) // chunk_size  # integer division = implicit floor
     if overlap:
-        chunk_step = chunk_size/2
+        chunk_step = chunk_size / 2
     else:
         chunk_step = chunk_size
 
-    win = window(chunk_size) # generate a window of the appropriate size
-    freq_axis = _numpy.linspace(0, freq/2, chunk_size/2+1)
+    win = window(chunk_size)  # generate a window of the appropriate size
+    freq_axis = _numpy.linspace(0, freq / 2, chunk_size / 2 + 1)
     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
     # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination.
     # See Numerical Recipies for a details.
-    power = _numpy.zeros((chunk_size/2+1,), dtype=_numpy.float)
+    power = _numpy.zeros((chunk_size / 2 + 1, ), dtype=_numpy.float)
     for i in range(nchunks):
-        starti = i*chunk_step
-        stopi = starti+chunk_size
-        fft_chunk = _numpy.fft.rfft(data[starti:stopi]*win)
+        starti = i * chunk_step
+        stopi = starti + chunk_size
+        fft_chunk = _numpy.fft.rfft(data[starti:stopi] * win)
         p_chunk = (fft_chunk * fft_chunk.conj()).real
         power += p_chunk.astype(_numpy.float)
     power /= _numpy.float(nchunks)
     return (freq_axis, power)
 
+
 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
                                overlap=True, window=window_hann):
     """Compute the average power spectrum, preserving normalization
     """
     freq_axis,power = avg_power_spectrum(
         data, freq, chunk_size, overlap, window)
-    #        2.0 / (freq * chunk_size)          |rfft()|**2 --> unitary_power_spectrum
+    #   2.0 / (freq * chunk_size)       |rfft()|**2 --> unitary_power_spectrum
     # see unitary_power_spectrum()
-    power *= 2.0 / (freq*_numpy.float(chunk_size)) * 8/3
-    #                                       * 8/3  to remove power from windowing
+    power *= 2.0 / (freq * _numpy.float(chunk_size)) * 8.0 / 3
+    #             * 8/3  to remove power from windowing
     #  <[x(t)*w(t)]**2> = <x(t)**2 * w(t)**2> ~= <x(t)**2> * <w(t)**2>
     # where the ~= is because the frequency of x(t) >> the frequency of w(t).
     # So our calulated power has and extra <w(t)**2> in it.
-    # For the Hann window, <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
-    # For low frequency components, where the frequency of x(t) is ~= the frequency of w(t),
-    # The normalization is not perfect. ??
+    # For the Hann window,
+    #   <w(t)**2> = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8
+    # For low frequency components,
+    # where the frequency of x(t) is ~= the frequency of w(t),
+    # the normalization is not perfect. ??
     # The normalization approaches perfection as chunk_size -> infinity.
     return (freq_axis, power)
 
+
 def _test_unitary_avg_power_spectrum_sin(
     sin_freq=10, samp_freq=512, samples=1024, chunk_size=512, overlap=True,
     window=window_hann):
     x = _numpy.zeros((samples,), dtype=_numpy.float)
     samp_freq = _numpy.float(samp_freq)
     for i in range(samples):
-        x[i] = _numpy.sin(2.0 * _numpy.pi * (i/samp_freq) * sin_freq)
+        x[i] = _numpy.sin(2.0 * _numpy.pi * (i / samp_freq) * sin_freq)
     freq_axis, power = unitary_avg_power_spectrum(
         x, samp_freq, chunk_size, overlap, window)
     imax = _numpy.argmax(power)
 
     expected = _numpy.zeros((len(freq_axis),), dtype=_numpy.float)
-    df = samp_freq/_numpy.float(chunk_size) # df = 1/T, where T = total_time
-    i = int(sin_freq/df)
-    expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
+    df = samp_freq / _numpy.float(chunk_size)  # df = 1/T, where T = total_time
+    i = int(sin_freq / df)
+    expected[i] = 0.5 / df  # see _test_unitary_power_spectrum_sin()
 
     print('The power should peak at {} Hz of {} ({}, {})'.format(
         sin_freq, expected[i], freq_axis[imax], power[imax]))
-    Pexp = 0
-    P    = 0
+    Pexp = P = 0
     for i in range(len(freq_axis)):
         Pexp += expected[i] * df
-        P    += power[i] * df
+        P += power[i] * df
     print('The total power should be {} ({})'.format(Pexp, P))
 
     if TEST_PLOTS:
         figure = _pyplot.figure()
         time_axes = figure.add_subplot(2, 1, 1)
         time_axes.plot(
-            _numpy.arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-')
+            _numpy.arange(0, samples / samp_freq, 1.0 / samp_freq), x, 'b-')
         time_axes.set_title('time series')
         freq_axes = figure.add_subplot(2, 1, 2)
         freq_axes.plot(freq_axis, power, 'r.')
@@ -603,18 +641,28 @@ def _test_unitary_avg_power_spectrum_sin(
         freq_axes.set_title(
             '{} samples of sin at {} Hz'.format(samples, sin_freq))
 
+
 def _test_unitary_avg_power_spectrum_sin_suite():
     print('Test unitary avg power spectrums on variously shaped sin functions')
-    _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
-    _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
-    _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
-    _test_unitary_avg_power_spectrum_sin(sin_freq=17, samp_freq=512, samples=1024)
-    _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048)
+    _test_unitary_avg_power_spectrum_sin(
+        sin_freq=5, samp_freq=512, samples=1024)
+    _test_unitary_avg_power_spectrum_sin(
+        sin_freq=5, samp_freq=512, samples=2048)
+    _test_unitary_avg_power_spectrum_sin(
+        sin_freq=5, samp_freq=512, samples=4098)
+    _test_unitary_avg_power_spectrum_sin(
+        sin_freq=17, samp_freq=512, samples=1024)
+    _test_unitary_avg_power_spectrum_sin(
+        sin_freq=5, samp_freq=1024, samples=2048)
     # test long wavelenth sin, so be closer to window frequency
-    _test_unitary_avg_power_spectrum_sin(sin_freq=1, samp_freq=1024, samples=2048)
-    # finally, with some irrational numbers, to check that I'm not getting lucky
     _test_unitary_avg_power_spectrum_sin(
-        sin_freq=_numpy.pi, samp_freq=100*_numpy.exp(1), samples=1024)
+        sin_freq=1, samp_freq=1024, samples=2048)
+    # finally, with some irrational numbers, to check that I'm not
+    # getting lucky
+    _test_unitary_avg_power_spectrum_sin(
+        sin_freq=_numpy.pi,
+        samp_freq=100 * _numpy.exp(1),
+        samples=1024)
 
 
 def test():
@@ -627,6 +675,7 @@ def test():
     _test_unitary_power_spectrum_gaussian_suite()
     _test_unitary_avg_power_spectrum_sin_suite()
 
+
 if __name__ == '__main__':
     from optparse import OptionParser