Explicitly convert powers to reals (they should have no imaginary portion.)

[FFT-tools.git] / FFT_tools.py

diff --git a/FFT_tools.py b/FFT_tools.py
index 7174d1de5ac6d06ce48c1668dd26f8929ce42976..12580c288bc930c2d241d8a3c6309fa553dafc7f 100644 (file)
--- a/FFT_tools.py
+++ b/FFT_tools.py
@@ -1,9 +1,21 @@
-#!/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-2011  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)
@@ -19,9 +31,11 @@ from numpy import log2, floor, round, ceil, abs, pi, exp, cos, sin, sqrt, \
 from numpy.fft import rfft
 
 
-# print time- and freq- space plots of the test transforms if True
+__version__ = '0.2'
+
+# 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."
@@ -259,7 +273,7 @@ def power_spectrum(data, freq=1.0) :
     # >>> 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) :
@@ -493,7 +507,7 @@ def avg_power_spectrum(data, freq=1.0, chunk_size=2048,
         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)
@@ -576,8 +590,17 @@ def test() :
     _test_unitary_avg_power_spectrum_sin_suite()
 
 if __name__ == "__main__" :
-    if TEST_PLOTS :
+    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()