From: W. Trevor King Date: Sun, 18 Nov 2012 22:48:17 +0000 (-0500) Subject: FFT_tools: add spaces around operators (PEP8) except ** X-Git-Tag: 0.4~3 X-Git-Url: http://git.tremily.us/?a=commitdiff_plain;h=ca676d255ab6b52856bfbfa4ccff28283a3e8f05;p=FFT-tools.git FFT_tools: add spaces around operators (PEP8) except ** --- diff --git a/FFT_tools.py b/FFT_tools.py index 7df0dbe..7eb06a7 100644 --- a/FFT_tools.py +++ b/FFT_tools.py @@ -77,21 +77,21 @@ def _test_rfft(xs, Xs): n = len(xs) Xa = [] for k in range(n): - Xa.append(sum([x*_numpy.exp(-j*2*_numpy.pi*k*m/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)) @@ -146,7 +146,7 @@ def unitary_rfft(data, freq=1.0): # continuous transformed data with (also NR 12.1.8) # X'_k = dtX = X / 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) @@ -155,19 +155,19 @@ def _test_unitary_rfft_parsevals(xs, freq, freqs, Xs): # 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)) @@ -175,8 +175,8 @@ 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] 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): @@ -193,10 +193,10 @@ 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) @@ -204,22 +204,22 @@ def _test_unitary_rfft_rect( 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.') @@ -247,10 +247,10 @@ 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) @@ -258,20 +258,20 @@ def _test_unitary_rfft_gaussian( 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.') @@ -298,7 +298,7 @@ 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. @@ -345,13 +345,13 @@ 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 = # average value of sin(t)**2 = 0.5 (b/c sin**2+cos**2 == 1) @@ -377,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.') @@ -406,7 +405,7 @@ def _test_unitary_power_spectrum_sin_suite(): _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048) # 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) @@ -438,7 +437,7 @@ 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.') @@ -461,8 +460,8 @@ 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*_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( @@ -470,7 +469,7 @@ def _test_unitary_power_spectrum_gaussian( 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) @@ -488,15 +487,15 @@ def _test_unitary_power_spectrum_gaussian( # 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])) @@ -505,7 +504,7 @@ 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.') @@ -534,7 +533,8 @@ 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 @@ -561,22 +561,22 @@ 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) + 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) @@ -591,7 +591,7 @@ def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048, data, freq, chunk_size, overlap, window) # 2.0 / (freq * chunk_size) |rfft()|**2 --> unitary_power_spectrum # see unitary_power_spectrum() - power *= 2.0 / (freq*_numpy.float(chunk_size)) * 8/3 + power *= 2.0 / (freq * _numpy.float(chunk_size)) * 8.0 / 3 # * 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). @@ -611,30 +611,29 @@ def _test_unitary_avg_power_spectrum_sin( 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) + 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.') @@ -661,7 +660,9 @@ def _test_unitary_avg_power_spectrum_sin_suite(): # 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=_numpy.pi, + samp_freq=100 * _numpy.exp(1), + samples=1024) def test():