-#!/usr/bin/python
-
-"""
-Define some FFT wrappers to reduce clutter.
-Provides a unitary discrete FFT and a windowed version.
-Based on numpy.fft.rfft.
+# Copyright (C) 2008-2012 W. Trevor King
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+"""Wrap Numpy's fft module to reduce clutter.
+
+Provides a unitary discrete FFT and a windowed version based on
+numpy.fft.rfft.
Main entry functions:
unitary_rfft(data, freq=1.0)
from numpy.fft import rfft
-# print time- and freq- space plots of the test transforms if True
+__version__ = '0.3'
+
+# Display time- and freq-space plots of the test transforms if True
TEST_PLOTS = False
-#TEST_PLOTS = True
+
def floor_pow_of_two(num) :
"Round num down to the closest exact a power of two."
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,
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
# 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)
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)
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)
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)
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.
# See Numerical Recipies for a details.
trans = rfft(data[0:nsamps])
- power = trans * trans.conj() # 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) :
# 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
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
# 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()
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)
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)
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
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)
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()
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)
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.
starti = i*chunk_step
stopi = starti+chunk_size
fft_chunk = rfft(data[starti:stopi]*win)
- p_chunk = fft_chunk * fft_chunk.conj()
+ p_chunk = (fft_chunk * fft_chunk.conj()).real
power += p_chunk.astype(float)
power /= 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
"""
- 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).
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)
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()
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)
_test_unitary_power_spectrum_gaussian_suite()
_test_unitary_avg_power_spectrum_sin_suite()
-if __name__ == "__main__" :
- if TEST_PLOTS :
+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',
+ help='Display time- and freq-space plots of test transforms.')
+
+ options,args = p.parse_args()
+
+ if options.plot:
import pylab
+ TEST_PLOTS = True
test()
- if TEST_PLOTS :
+ if options.plot:
pylab.show()
projects / FFT-tools.git / blobdiff