From: W. Trevor King Date: Sun, 18 Nov 2012 21:40:44 +0000 (-0500) Subject: FFT_tools: replace " :" with ":" for PEP 8 compliance X-Git-Tag: 0.4~13 X-Git-Url: http://git.tremily.us/?a=commitdiff_plain;h=ba286085ddb39a10318ff4d6bca54fd122a11d9d;p=FFT-tools.git FFT_tools: replace " :" with ":" for PEP 8 compliance --- diff --git a/FFT_tools.py b/FFT_tools.py index 1c3455f..8e9bbd7 100644 --- a/FFT_tools.py +++ b/FFT_tools.py @@ -43,28 +43,28 @@ __version__ = '0.3' TEST_PLOTS = False -def floor_pow_of_two(num) : +def floor_pow_of_two(num): "Round num down to the closest exact a power of two." lnum = _numpy.log2(num) - if int(lnum) != lnum : + if int(lnum) != lnum: num = 2**_numpy.floor(lnum) return num -def round_pow_of_two(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) - if int(lnum) != lnum : + if int(lnum) != lnum: num = 2**_numpy.round(lnum) return num -def ceil_pow_of_two(num) : +def ceil_pow_of_two(num): "Round num up to the closest exact a power of two." lnum = _numpy.log2(num) - if int(lnum) != lnum : + if int(lnum) != lnum: num = 2**_numpy.ceil(lnum) return num -def _test_rfft(xs, Xs) : +def _test_rfft(xs, Xs): # Numpy's FFT algoritm returns # n-1 # X[k] = SUM x[m] exp (-j 2pi km /n) @@ -73,7 +73,7 @@ def _test_rfft(xs, Xs) : j = _numpy.complex(0,1) n = len(xs) Xa = [] - for k in range(n) : + 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))])) if k < len(Xs): @@ -93,12 +93,12 @@ def _test_rfft(xs, Xs) : "Mismatch on Parseval's, {} != 1/{} * {}".format( timeSum, n, freqSum)) -def _test_rfft_suite() : +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] _test_rfft(xs, _numpy.fft.rfft(xs)) -def unitary_rfft(data, freq=1.0) : +def unitary_rfft(data, freq=1.0): """Compute the Fourier transform of real data. Unitary (preserves power [Parseval's theorem]). @@ -155,7 +155,7 @@ def _test_unitary_rfft_parsevals(xs, freq, freqs, Xs): 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))) @@ -171,20 +171,20 @@ def _test_unitary_rfft_parsevals_suite(): 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 : +def _rect(t): + if _numpy.abs(t) < 0.5: return 1 - else : + else: return 0 -def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) : +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)" samp_freq = _numpy.float(samp_freq) a = _numpy.float(a) x = _numpy.zeros((samples,), dtype=_numpy.float) dt = 1.0/samp_freq - for i in range(samples) : + for i in range(samples): t = i*dt x[i] = _rect(a*(t-time_shift)) freq_axis, X = unitary_rfft(x, samp_freq) @@ -192,7 +192,7 @@ def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) # remove the phase due to our time shift j = _numpy.complex(0.0,1.0) # sqrt(-1) - for i in range(len(freq_axis)) : + for i in range(len(freq_axis)): f = freq_axis[i] inverse_phase_shift = _numpy.exp(j*2.0*_numpy.pi*time_shift*f) X[i] *= inverse_phase_shift @@ -202,11 +202,11 @@ def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi if _numpy.sinc(0.5) != 2.0/_numpy.pi: raise ValueError('abnormal sinc()') - for i in range(len(freq_axis)) : + for i in range(len(freq_axis)): f = freq_axis[i] expected[i] = 1.0/_numpy.abs(a) * _numpy.sinc(f/a) - if TEST_PLOTS : + if TEST_PLOTS: figure = _pyplot.figure() time_axes = figure.add_subplot(2, 1, 1) time_axes.plot(_numpy.arange(0, dt*samples, dt), x) @@ -217,24 +217,24 @@ def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) freq_axes.plot(freq_axis, expected, 'b-') freq_axes.set_title('freq series') -def _test_unitary_rfft_rect_suite() : +def _test_unitary_rfft_rect_suite(): 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) _test_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024) -def _gaussian(a, t) : +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) : +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)" samp_freq = _numpy.float(samp_freq) a = _numpy.float(a) x = _numpy.zeros((samples,), dtype=_numpy.float) dt = 1.0/samp_freq - for i in range(samples) : + for i in range(samples): t = i*dt x[i] = _gaussian(a, (t-time_shift)) freq_axis, X = unitary_rfft(x, samp_freq) @@ -242,19 +242,19 @@ def _test_unitary_rfft_gaussian(a=1.0, time_shift=5.0, samp_freq=25.6, samples=2 # remove the phase due to our time shift j = _numpy.complex(0.0,1.0) # sqrt(-1) - for i in range(len(freq_axis)) : + for i in range(len(freq_axis)): f = freq_axis[i] 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)) : + 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) - if TEST_PLOTS : + if TEST_PLOTS: figure = _pyplot.figure() time_axes = figure.add_subplot(2, 1, 1) time_axes.plot(_numpy.arange(0, dt*samples, dt), x) @@ -265,7 +265,7 @@ def _test_unitary_rfft_gaussian(a=1.0, time_shift=5.0, samp_freq=25.6, samples=2 freq_axes.plot(freq_axis, expected, 'b-') freq_axes.set_title('freq series') -def _test_unitary_rfft_gaussian_suite() : +def _test_unitary_rfft_gaussian_suite(): print("Test unitary FFTs on variously shaped gaussian functions") _test_unitary_rfft_gaussian(a=0.5) _test_unitary_rfft_gaussian(a=2.0) @@ -274,7 +274,7 @@ def _test_unitary_rfft_gaussian_suite() : -def power_spectrum(data, freq=1.0) : +def power_spectrum(data, freq=1.0): """Compute the power spectrum of DATA taken with a sampling frequency FREQ. DATA must be real (not complex). @@ -292,7 +292,7 @@ def power_spectrum(data, freq=1.0) : power = (trans * trans.conj()).real # We want the square of the amplitude. return (freq_axis, power) -def unitary_power_spectrum(data, freq=1.0) : +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) # @@ -322,10 +322,10 @@ 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) : +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) : + for i in range(samples): 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) @@ -342,7 +342,7 @@ def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) : # P(f0) = 0.5/df # where f0 is the sin's frequency # - # or : + # or: # FFT of sin(2*pi*t*f0) gives # X(f) = 0.5 i (delta(f-f0) - delta(f-f0)), # (area under x(t) = 0, area under X(f) = 0) @@ -360,12 +360,12 @@ def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) : sin_freq, expected[i], freq_axis[imax], power[imax])) Pexp = 0 P = 0 - for i in range(len(freq_axis)) : + for i in range(len(freq_axis)): Pexp += expected[i] *df P += power[i] * df print('The total power should be {} ({})'.format(Pexp, P)) - if TEST_PLOTS : + if TEST_PLOTS: figure = _pyplot.figure() time_axes = figure.add_subplot(2, 1, 1) time_axes.plot( @@ -377,7 +377,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() : +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) _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048) @@ -390,7 +390,7 @@ def _test_unitary_power_spectrum_sin_suite() : # 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) : +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) x[0] = amp @@ -413,7 +413,7 @@ def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256) : print('The power should be flat at y = {} ({})'.format( expected_amp, power[0])) - if TEST_PLOTS : + if TEST_PLOTS: figure = _pyplot.figure() time_axes = figure.add_subplot(2, 1, 1) time_axes.plot( @@ -424,7 +424,7 @@ def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256) : 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() : +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) @@ -434,15 +434,15 @@ def _test_unitary_power_spectrum_delta_suite() : _test_unitary_power_spectrum_delta( amp=_numpy.pi, samp_freq=_numpy.exp(1), samples=1024) -def _gaussian2(area, mean, std, t) : +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) -def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10.24 ,samples=512) : #1024 +def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10.24 ,samples=512): #1024 x = _numpy.zeros((samples,), dtype=_numpy.float) mean = _numpy.float(mean) - for i in range(samples) : + for i in range(samples): t = i/_numpy.float(samp_freq) x[i] = _gaussian2(area, mean, std, t) freq_axis, power = unitary_power_spectrum(x, samp_freq) @@ -463,7 +463,7 @@ def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10. 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 - for i in range(len(freq_axis)) : + for i in range(len(freq_axis)): f = i*df gaus = _gaussian2(area, mean, std, f) expected[i] = 2.0 * gaus**2 * samp_freq/samples @@ -471,7 +471,7 @@ def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10. 'with a peak at 0 Hz with amplitude {} ({})').format( expected[0], power[0])) - if TEST_PLOTS : + if TEST_PLOTS: figure = _pyplot.figure() time_axes = figure.add_subplot(2, 1, 1) time_axes.plot( @@ -482,7 +482,7 @@ def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10. freq_axes.plot(freq_axis, expected, 'b-') freq_axes.set_title('freq series') -def _test_unitary_power_spectrum_gaussian_suite() : +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) @@ -493,17 +493,17 @@ def _test_unitary_power_spectrum_gaussian_suite() : area=_numpy.pi, std=_numpy.sqrt(2), samp_freq=_numpy.exp(1), samples=1024) -def window_hann(length) : +def window_hann(length): "Returns a Hann window array with length entries" win = _numpy.zeros((length,), dtype=_numpy.float) - for i in range(length) : + for i in range(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) : + overlap=True, window=window_hann): """Compute the avgerage power spectrum of DATA. Taken with a sampling frequency FREQ. @@ -524,9 +524,9 @@ def avg_power_spectrum(data, freq=1.0, chunk_size=2048, 'chunk_size {} should be a power of 2'.format(chunk_size)) nchunks = len(data)/chunk_size # integer division = implicit floor - if overlap : + if overlap: chunk_step = chunk_size/2 - else : + else: chunk_step = chunk_size win = window(chunk_size) # generate a window of the appropriate size @@ -535,7 +535,7 @@ def avg_power_spectrum(data, freq=1.0, chunk_size=2048, # >>> 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) - for i in range(nchunks) : + for i in range(nchunks): starti = i*chunk_step stopi = starti+chunk_size fft_chunk = _numpy.fft.rfft(data[starti:stopi]*win) @@ -545,7 +545,7 @@ def avg_power_spectrum(data, freq=1.0, chunk_size=2048, return (freq_axis, power) def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048, - overlap=True, window=window_hann) : + overlap=True, window=window_hann): """Compute the average power spectrum, preserving normalization """ freq_axis,power = avg_power_spectrum( @@ -565,10 +565,10 @@ def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048, def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024, chunk_size=512, overlap=True, - window=window_hann) : + window=window_hann): x = _numpy.zeros((samples,), dtype=_numpy.float) samp_freq = _numpy.float(samp_freq) - for i in range(samples) : + for i in range(samples): 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) @@ -583,12 +583,12 @@ def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=102 sin_freq, expected[i], freq_axis[imax], power[imax])) Pexp = 0 P = 0 - for i in range(len(freq_axis)) : + for i in range(len(freq_axis)): Pexp += expected[i] * df P += power[i] * df print('The total power should be {} ({})'.format(Pexp, P)) - if TEST_PLOTS : + if TEST_PLOTS: figure = _pyplot.figure() time_axes = figure.add_subplot(2, 1, 1) time_axes.plot( @@ -600,7 +600,7 @@ def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=102 freq_axes.set_title( '{} samples of sin at {} Hz'.format(samples, sin_freq)) -def _test_unitary_avg_power_spectrum_sin_suite() : +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) @@ -614,7 +614,7 @@ def _test_unitary_avg_power_spectrum_sin_suite() : sin_freq=_numpy.pi, samp_freq=100*_numpy.exp(1), samples=1024) -def test() : +def test(): _test_rfft_suite() _test_unitary_rfft_parsevals_suite() _test_unitary_rfft_rect_suite()