FFT-tools: modernize string formatting and exception handling
authorW. Trevor King <wking@tremily.us>
Sun, 18 Nov 2012 20:48:52 +0000 (15:48 -0500)
committerW. Trevor King <wking@tremily.us>
Sun, 18 Nov 2012 21:49:21 +0000 (16:49 -0500)
Changes:
* "print ..." -> "print(...)" (required for Python 3+)
* "'%s' % (...)" -> "'{}'.format(...)" for Python 2.7+
* prefer "raise ValueError(...)" to "assert ..."
* a few minor docstring cleanups for PEP 257 compliance

FFT_tools.py

index 4cb0432de73972fc48777c5468b8910c0d1d6c2d..75a281903c06b84f63e4cc09ef4ffd67372a7485 100644 (file)
@@ -70,28 +70,31 @@ def _test_rfft(xs, Xs) :
     for k in range(n) :
         Xa.append(sum([x*exp(-j*2*pi*k*m/n) for x,m in zip(xs,range(n))]))
         if k < len(Xs):
-            assert (Xs[k]-Xa[k])/abs(Xa[k]) < 1e-6, \
-                "rfft mismatch on element %d: %g != %g, relative error %g" \
-                % (k, Xs[k], Xa[k], (Xs[k]-Xa[k])/abs(Xa[k]))
+            if (Xs[k]-Xa[k])/abs(Xa[k]) >= 1e-6:
+                raise ValueError(
+                    ('rfft mismatch on element {}: {} != {}, relative error {}'
+                     ).format(k, Xs[k], Xa[k], (Xs[k]-Xa[k])/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. 
+    #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
     #   m=0               k=0
     timeSum = sum([abs(x)**2 for x in xs])
     freqSum = sum([abs(X)**2 for X in Xa])
-    assert abs(freqSum/float(n) - timeSum)/timeSum < 1e-6, \
-        "Mismatch on Parseval's, %g != 1/%d * %g" % (timeSum, n, freqSum)
+    if abs(freqSum/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"
+    print('Test numpy rfft definition')
     xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1]
     _test_rfft(xs, rfft(xs))
 
 def unitary_rfft(data, freq=1.0) :
-    """
-    Compute the Fourier transform of real data.
+    """Compute the Fourier transform of real data.
+
     Unitary (preserves power [Parseval's theorem]).
-   
+
     If the units on your data are Volts,
     and your sampling frequency is in Hz,
     then freq_axis will be in Hz,
@@ -100,7 +103,7 @@ def unitary_rfft(data, freq=1.0) :
     nsamps = floor_pow_of_two(len(data))
     # Which should satisfy the discrete form of Parseval's theorem
     #   n-1               n-1
-    #   SUM |x_m|^2 = 1/n SUM |X_k|^2. 
+    #   SUM |x_m|^2 = 1/n SUM |X_k|^2.
     #   m=0               k=0
     # However, we want our FFT to satisfy the continuous Parseval eqn
     #   int_{-infty}^{infty} |x(t)|^2 dt = int_{-infty}^{infty} |X(f)|^2 df
@@ -140,19 +143,21 @@ def _test_unitary_rfft_parsevals(xs, freq, freqs, Xs):
     #   m=0              k=0
     dt = 1.0/freq
     df = freqs[1]-freqs[0]
-    assert (df - 1/(len(xs)*dt))/df < 1e-6, \
-        "Mismatch in spacing, %g != 1/(%d*%g)" % (df, len(xs), dt)
+    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) :
         Xa.append(Xa[k])
-    assert len(xs) == len(Xa), "Length mismatch %d != %d" % (len(xs), len(Xa))
+    if len(xs) != len(Xa):
+        raise ValueError('Length mismatch {} != {}'.format(len(xs), len(Xa)))
     lhs = sum([abs(x)**2 for x in xs]) * dt
     rhs = sum([abs(X)**2 for X in Xa]) * df
-    assert abs(lhs - rhs)/lhs < 1e-4, "Mismatch on Parseval's, %g != %g" \
-        % (lhs, rhs)
+    if 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"
+    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]
     dt = pi
     freqs,Xs = unitary_rfft(xs, 1.0/dt)
@@ -187,7 +192,8 @@ def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256)
     expected = zeros((len(freq_axis),), dtype=float)
     # normalized sinc(x) = sin(pi x)/(pi x)
     # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi
-    assert sinc(0.5) == 2.0/pi, "abnormal sinc()"
+    if sinc(0.5) != 2.0/pi:
+        raise ValueError('abnormal sinc()')
     for i in range(len(freq_axis)) :
         f = freq_axis[i]
         expected[i] = 1.0/abs(a) * sinc(f/a)
@@ -204,7 +210,7 @@ def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256)
         pylab.title('freq series')
 
 def _test_unitary_rfft_rect_suite() :
-    print "Test unitary FFTs on variously shaped rectangular functions"
+    print('Test unitary FFTs on variously shaped rectangular functions')
     _test_unitary_rfft_rect(a=0.5)
     _test_unitary_rfft_rect(a=2.0)
     _test_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
@@ -250,7 +256,7 @@ def _test_unitary_rfft_gaussian(a=1.0, time_shift=5.0, samp_freq=25.6, samples=2
         pylab.title('freq series')
 
 def _test_unitary_rfft_gaussian_suite() :
-    print "Test unitary FFTs on variously shaped gaussian functions"
+    print("Test unitary FFTs on variously shaped gaussian functions")
     _test_unitary_rfft_gaussian(a=0.5)
     _test_unitary_rfft_gaussian(a=2.0)
     _test_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
@@ -259,15 +265,15 @@ def _test_unitary_rfft_gaussian_suite() :
 
 
 def power_spectrum(data, freq=1.0) :
-    """
-    Compute the power spectrum of DATA taken with a sampling frequency FREQ.
+    """Compute the power spectrum of DATA taken with a sampling frequency FREQ.
+
     DATA must be real (not complex).
     Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
     If the number of samples in data is not an integer power of two,
     the FFT ignores some of the later points.
     """
     nsamps = floor_pow_of_two(len(data))
-    
+
     freq_axis = 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.
@@ -288,7 +294,7 @@ def unitary_power_spectrum(data, freq=1.0) :
     # 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
@@ -317,7 +323,7 @@ def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) :
     expected = zeros((len(freq_axis),), dtype=float)
     df = samp_freq/float(samples) # df = 1/T, where T = total_time
     i = int(sin_freq/df)
-    # average power per unit time is 
+    # 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)
     # so average value of (int sin(t)**2 dt) per unit time is 0.5
@@ -334,20 +340,20 @@ def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) :
     #  P(f) = 2 |X(f)|**2 / T
     #       = 2 * |0.5 delta(f-f0)|**2 / T
     #       = 0.5 * |delta(f-f0)|**2 / T
-    # but we're discrete and want the integral of the 'delta' to be 1, 
+    # but we're discrete and want the integral of the 'delta' to be 1,
     # so 'delta'*df = 1  --> 'delta' = 1/df, and
     #  P(f) = 0.5 / (df**2 * T)
     #       = 0.5 / df                (T = 1/df)
     expected[i] = 0.5 / df
 
-    print "The power should be a peak at %g Hz of %g (%g, %g)" % \
-        (sin_freq, expected[i], freq_axis[imax], power[imax])
+    print('The power should be a peak at {} Hz of {} ({}, {})'.format(
+            sin_freq, expected[i], freq_axis[imax], power[imax]))
     Pexp = 0
     P    = 0
     for i in range(len(freq_axis)) :
         Pexp += expected[i] *df
         P    += power[i] * df
-    print " The total power should be %g (%g)" % (Pexp, P)
+    print('The total power should be {} ({})'.format(Pexp, P))
 
     if TEST_PLOTS :
         pylab.figure()
@@ -357,10 +363,10 @@ def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) :
         pylab.subplot(212)
         pylab.plot(freq_axis, power, 'r.')
         pylab.plot(freq_axis, expected, 'b-')
-        pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
+        pylab.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"
+    print('Test unitary power spectrums on variously shaped sin functions')
     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024)
     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048)
     _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098)
@@ -390,8 +396,9 @@ def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256) :
     expected_amp = 2.0 * amp**2 / (samp_freq * samples)
     expected = ones((len(freq_axis),), dtype=float) * expected_amp
 
-    print "The power should be flat at y = %g (%g)" % (expected_amp, power[0])
-    
+    print('The power should be flat at y = {} ({})'.format(
+        expected_amp, power[0]))
+
     if TEST_PLOTS :
         pylab.figure()
         pylab.subplot(211)
@@ -400,10 +407,10 @@ def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256) :
         pylab.subplot(212)
         pylab.plot(freq_axis, power, 'r.')
         pylab.plot(freq_axis, expected, 'b-')
-        pylab.title('%g samples of delta amp %g' % (samples, amp))
+        pylab.title('{} samples of delta amp {}'.format(samples, amp))
 
 def _test_unitary_power_spectrum_delta_suite() :
-    print "Test unitary power spectrums on various delta functions"
+    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
@@ -414,7 +421,7 @@ def _test_unitary_power_spectrum_delta_suite() :
 def _gaussian2(area, mean, std, t) :
     "Integral over all time = area (i.e. normalized for area=1)"
     return area/(std*sqrt(2.0*pi)) * 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) : #1024
     x = zeros((samples,), dtype=float)
     mean = float(mean)
@@ -442,8 +449,9 @@ def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10.
         f = i*df
         gaus = _gaussian2(area, mean, std, f)
         expected[i] = 2.0 * gaus**2 * samp_freq/samples
-    print "The power should be a half-gaussian, ",
-    print "with a peak at 0 Hz with amplitude %g (%g)" % (expected[0], power[0])
+    print(('The power should be a half-gaussian, '
+           'with a peak at 0 Hz with amplitude {} ({})').format(
+            expected[0], power[0]))
 
     if TEST_PLOTS :
         pylab.figure()
@@ -456,7 +464,7 @@ def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10.
         pylab.title('freq series')
 
 def _test_unitary_power_spectrum_gaussian_suite() :
-    print "Test unitary power spectrums on various gaussian functions"
+    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)
@@ -475,28 +483,31 @@ def window_hann(length) :
 
 def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
                        overlap=True, window=window_hann) :
-    """
-    Compute the avg power spectrum of DATA taken with a sampling frequency FREQ.
+    """Compute the avgerage power spectrum of DATA.
+
+    Taken with a sampling frequency FREQ.
+
     DATA must be real (not complex) by breaking DATA into chunks.
     The chunks may or may not be overlapping (by setting OVERLAP).
     The chunks are windowed by dotting with WINDOW(CHUNK_SIZE), FFTed,
     and the resulting spectra are averaged together.
     See NR 13.4 for rational.
-    
+
     Returns a tuple of two arrays, (freq_axis, power), suitable for plotting.
     CHUNK_SIZE should really be a power of 2.
     If the number of samples in DATA is not an integer power of CHUNK_SIZE,
     the FFT ignores some of the later points.
     """
-    assert chunk_size == floor_pow_of_two(chunk_size), \
-        "chunk_size %d should be a power of 2" % chunk_size
+    if chunk_size != floor_pow_of_two(chunk_size):
+        raise ValueError(
+            'chunk_size {} should be a power of 2'.format(chunk_size))
 
     nchunks = len(data)/chunk_size # integer division = implicit floor
     if overlap :
         chunk_step = chunk_size/2
     else :
         chunk_step = chunk_size
-    
+
     win = window(chunk_size) # generate a window of the appropriate size
     freq_axis = linspace(0, freq/2, chunk_size/2+1)
     # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real.
@@ -514,13 +525,12 @@ def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
 
 def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048,
                                overlap=True, window=window_hann) :
+    """Compute the average power spectrum, preserving normalization
     """
-    compute the average power spectrum, preserving normalization
-    """
-    freq_axis,power = avg_power_spectrum(data, freq, chunk_size,
-                                         overlap, window)
+    freq_axis,power = avg_power_spectrum(
+        data, freq, chunk_size, overlap, window)
     #        2.0 / (freq * chunk_size)          |rfft()|**2 --> unitary_power_spectrum
-    power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum()            
+    power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum()
     #                                       * 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).
@@ -538,8 +548,8 @@ def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=102
     samp_freq = float(samp_freq)
     for i in range(samples) :
         x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq)
-    freq_axis, power = unitary_avg_power_spectrum(x, samp_freq, chunk_size,
-                                                  overlap, window)
+    freq_axis, power = unitary_avg_power_spectrum(
+        x, samp_freq, chunk_size, overlap, window)
     imax = argmax(power)
 
     expected = zeros((len(freq_axis),), dtype=float)
@@ -547,14 +557,14 @@ def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=102
     i = int(sin_freq/df)
     expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin()
 
-    print "The power should be a peak at %g Hz of %g (%g, %g)" % \
-        (sin_freq, expected[i], freq_axis[imax], power[imax])
+    print('The power should peak at {} Hz of {} ({}, {})'.format(
+        sin_freq, expected[i], freq_axis[imax], power[imax]))
     Pexp = 0
     P    = 0
     for i in range(len(freq_axis)) :
         Pexp += expected[i] * df
         P    += power[i] * df
-    print " The total power should be %g (%g)" % (Pexp, P)
+    print('The total power should be {} ({})'.format(Pexp, P))
 
     if TEST_PLOTS :
         pylab.figure()
@@ -564,10 +574,10 @@ def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=102
         pylab.subplot(212)
         pylab.plot(freq_axis, power, 'r.')
         pylab.plot(freq_axis, expected, 'b-')
-        pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq))
+        pylab.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"
+    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)
@@ -589,11 +599,11 @@ def test() :
     _test_unitary_power_spectrum_gaussian_suite()
     _test_unitary_avg_power_spectrum_sin_suite()
 
-if __name__ == "__main__" :
+if __name__ == '__main__':
     from optparse import OptionParser
 
     p = OptionParser('%prog [options]', epilog='Run FFT-tools unit tests.')
-    p.add_option('-p', '--plot', dest='plot', action='store_true', 
+    p.add_option('-p', '--plot', dest='plot', action='store_true',
                  help='Display time- and freq-space plots of test transforms.')
 
     options,args = p.parse_args()