From 14544fb931b9db8c9118ed8ee08a36e0c51d58da Mon Sep 17 00:00:00 2001 From: "W. Trevor King" Date: Sun, 18 Nov 2012 15:48:52 -0500 Subject: [PATCH] FFT-tools: modernize string formatting and exception handling 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 | 128 +++++++++++++++++++++++++++------------------------ 1 file changed, 69 insertions(+), 59 deletions(-) diff --git a/FFT_tools.py b/FFT_tools.py index 4cb0432..75a2819 100644 --- a/FFT_tools.py +++ b/FFT_tools.py @@ -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 = # 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> = ~= * # 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() -- 2.26.2