FFT_tools: Catch RuntimeErrors when importing matplotlib
[FFT-tools.git] / FFT_tools.py
index 933e064b884e9ea965218c8519427aac0dba44ae..e6d058cc178416892a8425bcc05cc67b36d7aeb7 100644 (file)
 # 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.
+"""Wrap Numpy's :py:mod:`~numpy.fft` module to reduce clutter.
 
 Provides a unitary discrete FFT and a windowed version based on
-:func:`numpy.fft.rfft`.
+:py:func:`numpy.fft.rfft`.
 
 Main entry functions:
 
-* :func:`unitary_rfft`
-* :func:`power_spectrum`
-* :func:`unitary_power_spectrum`
-* :func:`avg_power_spectrum`
-* :func:`unitary_avg_power_spectrum`
+* :py:func:`unitary_rfft`
+* :py:func:`power_spectrum`
+* :py:func:`unitary_power_spectrum`
+* :py:func:`avg_power_spectrum`
+* :py:func:`unitary_avg_power_spectrum`
 """
 
+import logging as _logging
 import unittest as _unittest
 
 import numpy as _numpy
+try:
+    import matplotlib.pyplot as _pyplot
+except (ImportError, RuntimeError) as e:
+    _pyplot = None
+    _pyplot_import_error = e
 
 
-__version__ = '0.4'
+__version__ = '0.5'
+
+
+LOG = _logging.getLogger('FFT-tools')
+LOG.addHandler(_logging.StreamHandler())
+LOG.setLevel(_logging.ERROR)
+
 
 # Display time- and freq-space plots of the test transforms if True
 TEST_PLOTS = False
@@ -348,7 +360,8 @@ class TestRFFT (_unittest.TestCase):
 
     .. math:: X_k = \sum_{m=0}^{n-1} x_m \exp^{-2\pi imk/n}
 
-    .. [#dft] See the *Background information* section of :mod:`numpy.fft`.
+    .. [#dft] See the *Background information* section of
+       :py:mod:`numpy.fft`.
     """
     def run_rfft(self, xs, Xs):
         i = _numpy.complex(0, 1)
@@ -358,32 +371,32 @@ class TestRFFT (_unittest.TestCase):
             Xa.append(sum([x * _numpy.exp(-2 * _numpy.pi * i * m * k / n)
                            for x,m in zip(xs, range(n))]))
             if k < len(Xs):
-                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])))
+                self.assertAlmostEqual(
+                    (Xs[k] - Xa[k]) / _numpy.abs(Xa[k]), 0, 6,
+                    ('rfft mismatch on element {}: {} != {}, '
+                     'relative error {}').format(
+                        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:
-            raise ValueError(
-                "Mismatch on Parseval's, {} != 1/{} * {}".format(
-                    timeSum, n, freqSum))
+        self.assertAlmostEqual(
+            _numpy.abs(freqSum / _numpy.float(n) - timeSum) / timeSum, 0, 6,
+            "Mismatch on Parseval's, {} != 1/{} * {}".format(
+                timeSum, n, freqSum))
 
     def test_rfft(self):
+        "Test NumPy's builtin :py:func:`numpy.fft.rfft`"
         xs = [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1]
         self.run_rfft(xs, _numpy.fft.rfft(xs))
 
 
 class TestUnitaryRFFT (_unittest.TestCase):
-    """Verify `unitary_rfft`.
+    """Verify :py:func:`unitary_rfft`.
     """
-    def run_unitary_rfft_parsevals(self, xs, freq, freqs, Xs):
+    def run_parsevals(self, xs, freq, freqs, Xs):
         """Check the discretized integral form of Parseval's theorem
 
         Notes
@@ -394,27 +407,25 @@ class TestUnitaryRFFT (_unittest.TestCase):
         """
         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))
+        self.assertAlmostEqual(
+            (df - 1 / (len(xs) * dt)) / df, 0, 6,
+            '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])
-        if len(xs) != len(Xa):
-            raise ValueError(
-                'Length mismatch {} != {}'.format(len(xs), len(Xa)))
+        self.assertEqual(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:
-            raise ValueError(
-                "Mismatch on Parseval's, {} != {}".format(lhs, rhs))
+        self.assertAlmostEqual(
+            _numpy.abs(lhs - rhs) / lhs, 0, 3,
+            "Mismatch on Parseval's, {} != {}".format(lhs, rhs))
 
-    def test_unitary_rfft_parsevals(self):
+    def test_parsevals(self):
         "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)
-        self.run_unitary_rfft_parsevals(xs, 1.0 / dt, freqs, Xs)
+        self.run_parsevals(xs, 1.0 / dt, freqs, Xs)
 
     def rect(self, t):
         r"""Rectangle function.
@@ -427,16 +438,15 @@ class TestUnitaryRFFT (_unittest.TestCase):
             \rect(t) = \begin{cases}
                1& \text{if $|t| < 0.5$}, \\
                0& \text{if $|t| \ge 0.5$}.
-                       \end{cases}
+                             \end{cases}
         """
         if _numpy.abs(t) < 0.5:
             return 1
         else:
             return 0
 
-    def run_unitary_rfft_rect(self, a=1.0, time_shift=5.0, samp_freq=25.6,
-                              samples=256):
-        r"""Test `unitary_rttf` on known function `rect(at)`.
+    def run_rect(self, a=1.0, time_shift=5.0, samp_freq=25.6, samples=256):
+        r"""Test :py:func:`unitary_rfft` on known function :py:meth:`rect`.
 
         Notes
         -----
@@ -465,13 +475,14 @@ class TestUnitaryRFFT (_unittest.TestCase):
         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:
-            raise ValueError('abnormal sinc()')
+        self.assertEqual(_numpy.sinc(0.5), 2.0 / _numpy.pi)
         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 _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(_numpy.arange(0, dt * samples, dt), x)
@@ -482,12 +493,12 @@ class TestUnitaryRFFT (_unittest.TestCase):
             freq_axes.plot(freq_axis, expected, 'b-')
             freq_axes.set_title('freq series')
 
-    def test_unitary_rfft_rect(self):
+    def test_rect(self):
         "Test unitary FFTs on variously shaped rectangular functions."
-        self.run_unitary_rfft_rect(a=0.5)
-        self.run_unitary_rfft_rect(a=2.0)
-        self.run_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512)
-        self.run_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024)
+        self.run_rect(a=0.5)
+        self.run_rect(a=2.0)
+        self.run_rect(a=0.7, samp_freq=50, samples=512)
+        self.run_rect(a=3.0, samp_freq=60, samples=1024)
 
     def gaussian(self, a, t):
         r"""Gaussian function.
@@ -499,9 +510,8 @@ class TestUnitaryRFFT (_unittest.TestCase):
         """
         return _numpy.exp(-a * t**2)
 
-    def run_unitary_rfft_gaussian(self, a=1.0, time_shift=5.0, samp_freq=25.6,
-                                  samples=256):
-        r"""Test `unitary_rttf` on known function `gaussian(a,t)`.
+    def run_gaussian(self, a=1.0, time_shift=5.0, samp_freq=25.6, samples=256):
+        r"""Test :py:func:`unitary_rttf` on known function :py:meth:`gaussian`.
 
         Notes
         -----
@@ -509,7 +519,8 @@ class TestUnitaryRFFT (_unittest.TestCase):
 
         .. math::
 
-            \rfft(\gaussian(a,t)) = \sqrt{\pi/a} \cdot \gaussian(1/a,\pi f)
+            \rfft(\gaussian(a,t))
+              = \sqrt{\pi/a} \cdot \gaussian(1/a,\pi f)
         """
         samp_freq = _numpy.float(samp_freq)
         a = _numpy.float(a)
@@ -537,6 +548,8 @@ class TestUnitaryRFFT (_unittest.TestCase):
                 1.0 / a, _numpy.pi * f)
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(_numpy.arange(0, dt * samples, dt), x)
@@ -547,17 +560,16 @@ class TestUnitaryRFFT (_unittest.TestCase):
             freq_axes.plot(freq_axis, expected, 'b-')
             freq_axes.set_title('freq series')
 
-    def test_unitary_rfft_gaussian(self):
+    def test_gaussian(self):
         "Test unitary FFTs on variously shaped gaussian functions."
-        self.run_unitary_rfft_gaussian(a=0.5)
-        self.run_unitary_rfft_gaussian(a=2.0)
-        self.run_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512)
-        self.run_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024)
+        self.run_gaussian(a=0.5)
+        self.run_gaussian(a=2.0)
+        self.run_gaussian(a=0.7, samp_freq=50, samples=512)
+        self.run_gaussian(a=3.0, samp_freq=60, samples=1024)
 
 
 class TestUnitaryPowerSpectrum (_unittest.TestCase):
-    def run_unitary_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
-                                       samples=1024):
+    def run_sin(self, 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):
@@ -592,15 +604,19 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
         #       = 0.5 / df                (T = 1/df)
         expected[i] = 0.5 / df
 
-        print('The power should be a peak at {} Hz of {} ({}, {})'.format(
+        LOG.debug('The power should be a peak at {} Hz of {} ({}, {})'.format(
                 sin_freq, expected[i], freq_axis[imax], power[imax]))
         Pexp = P = 0
         for i in range(len(freq_axis)):
             Pexp += expected[i] * df
             P += power[i] * df
-        print('The total power should be {} ({})'.format(Pexp, P))
+        self.assertAlmostEqual(
+            _numpy.abs((P - Pexp) / Pexp), 0, 1,
+            'The total power should be {} ({})'.format(Pexp, P))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(
@@ -612,29 +628,21 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
             freq_axes.set_title(
                 '{} samples of sin at {} Hz'.format(samples, sin_freq))
 
-
-    def test_unitary_power_spectrum_sin(self):
+    def test_sin(self):
         "Test unitary power spectrums on variously shaped sin functions"
-        self.run_unitary_power_spectrum_sin(
-            sin_freq=5, samp_freq=512, samples=1024)
-        self.run_unitary_power_spectrum_sin(
-            sin_freq=5, samp_freq=512, samples=2048)
-        self.run_unitary_power_spectrum_sin(
-            sin_freq=5, samp_freq=512, samples=4098)
-        self.run_unitary_power_spectrum_sin(
-            sin_freq=7, samp_freq=512, samples=1024)
-        self.run_unitary_power_spectrum_sin(
-            sin_freq=5, samp_freq=1024, samples=2048)
+        self.run_sin(sin_freq=5, samp_freq=512, samples=1024)
+        self.run_sin(sin_freq=5, samp_freq=512, samples=2048)
+        self.run_sin(sin_freq=5, samp_freq=512, samples=4098)
+        self.run_sin(sin_freq=7, samp_freq=512, samples=1024)
+        self.run_sin(sin_freq=5, samp_freq=1024, samples=2048)
         # finally, with some irrational numbers, to check that I'm not
         # getting lucky
-        self.run_unitary_power_spectrum_sin(
+        self.run_sin(
             sin_freq=_numpy.pi, samp_freq=100 * _numpy.exp(1), samples=1024)
         # test with non-integer number of periods
-        self.run_unitary_power_spectrum_sin(
-            sin_freq=5, samp_freq=512, samples=256)
+        self.run_sin(sin_freq=5, samp_freq=512, samples=256)
 
-    def run_unitary_power_spectrum_delta(self, amp=1, samp_freq=1,
-                                         samples=256):
+    def run_delta(self, amp=1, samp_freq=1, samples=256):
         """TODO
         """
         x = _numpy.zeros((samples,), dtype=_numpy.float)
@@ -656,10 +664,14 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
         expected = _numpy.ones(
             (len(freq_axis),), dtype=_numpy.float) * expected_amp
 
-        print('The power should be flat at y = {} ({})'.format(
-            expected_amp, power[0]))
+        self.assertAlmostEqual(
+            expected_amp, power[0], 4,
+            'The power should be flat at y = {} ({})'.format(
+                expected_amp, power[0]))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(
@@ -670,30 +682,24 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
             freq_axes.plot(freq_axis, expected, 'b-')
             freq_axes.set_title('{} samples of delta amp {}'.format(samples, amp))
 
-    def test_unitary_power_spectrum_delta(self):
+    def test_delta(self):
         "Test unitary power spectrums on various delta functions"
-        self.run_unitary_power_spectrum_delta(
-            amp=1, samp_freq=1.0, samples=1024)
-        self.run_unitary_power_spectrum_delta(
-            amp=1, samp_freq=1.0, samples=2048)
+        self.run_delta(amp=1, samp_freq=1.0, samples=1024)
+        self.run_delta(amp=1, samp_freq=1.0, samples=2048)
         # expected = 2*computed
-        self.run_unitary_power_spectrum_delta(
-            amp=1, samp_freq=0.5, samples=2048)
+        self.run_delta(amp=1, samp_freq=0.5, samples=2048)
         # expected = 0.5*computed
-        self.run_unitary_power_spectrum_delta(
-            amp=1, samp_freq=2.0, samples=2048)
-        self.run_unitary_power_spectrum_delta(
-            amp=3, samp_freq=1.0, samples=1024)
-        self.run_unitary_power_spectrum_delta(
-            amp=_numpy.pi, samp_freq=_numpy.exp(1), samples=1024)
+        self.run_delta(amp=1, samp_freq=2.0, samples=2048)
+        self.run_delta(amp=3, samp_freq=1.0, samples=1024)
+        self.run_delta(amp=_numpy.pi, samp_freq=_numpy.exp(1), samples=1024)
 
     def gaussian(self, 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 run_unitary_power_spectrum_gaussian(self, area=2.5, mean=5, std=1,
-                                            samp_freq=10.24 ,samples=512):
+    def run_gaussian(self, area=2.5, mean=5, std=1, samp_freq=10.24,
+                     samples=512):
         """TODO.
         """
         x = _numpy.zeros((samples,), dtype=_numpy.float)
@@ -728,11 +734,15 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
             f = i * df
             gaus = self.gaussian(area, mean, std, f)
             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(
+        self.assertAlmostEqual(
+            expected[0], power[0], 3,
+            ('The power should be a half-gaussian, '
+             'with a peak at 0 Hz with amplitude {} ({})').format(
                 expected[0], power[0]))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(
@@ -744,21 +754,20 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
             freq_axes.plot(freq_axis, expected, 'b-')
             freq_axes.set_title('freq series')
 
-    def test_unitary_power_spectrum_gaussian(self):
+    def test_gaussian(self):
         "Test unitary power spectrums on various gaussian functions"
         for area in [1, _numpy.pi]:
             for std in [1, _numpy.sqrt(2)]:
                 for samp_freq in [10.0, _numpy.exp(1)]:
                     for samples in [1024, 2048]:
-                        self.run_unitary_power_spectrum_gaussian(
+                        self.run_gaussian(
                             area=area, std=std, samp_freq=samp_freq,
                             samples=samples)
 
 
 class TestUnitaryAvgPowerSpectrum (_unittest.TestCase):
-    def run_unitary_avg_power_spectrum_sin(self, sin_freq=10, samp_freq=512,
-                                           samples=1024, chunk_size=512,
-                                           overlap=True, window=window_hann):
+    def run_sin(self, sin_freq=10, samp_freq=512, samples=1024, chunk_size=512,
+                overlap=True, window=window_hann, places=3):
         """TODO
         """
         x = _numpy.zeros((samples,), dtype=_numpy.float)
@@ -776,15 +785,19 @@ class TestUnitaryAvgPowerSpectrum (_unittest.TestCase):
         # see TestUnitaryPowerSpectrum.run_unitary_power_spectrum_sin()
         expected[i] = 0.5 / df
 
-        print('The power should peak at {} Hz of {} ({}, {})'.format(
-            sin_freq, expected[i], freq_axis[imax], power[imax]))
+        LOG.debug('The power should peak at {} Hz of {} ({}, {})'.format(
+                sin_freq, expected[i], freq_axis[imax], power[imax]))
         Pexp = P = 0
         for i in range(len(freq_axis)):
             Pexp += expected[i] * df
             P += power[i] * df
-        print('The total power should be {} ({})'.format(Pexp, P))
+        self.assertAlmostEqual(
+            Pexp, P, places,
+            'The total power should be {} ({})'.format(Pexp, P))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(
@@ -797,22 +810,16 @@ class TestUnitaryAvgPowerSpectrum (_unittest.TestCase):
             freq_axes.set_title(
                 '{} samples of sin at {} Hz'.format(samples, sin_freq))
 
-    def test_unitary_avg_power_spectrum_sin(self):
+    def test_sin(self):
         "Test unitary avg power spectrums on variously shaped sin functions."
-        self.run_unitary_avg_power_spectrum_sin(
-            sin_freq=5, samp_freq=512, samples=1024)
-        self.run_unitary_avg_power_spectrum_sin(
-            sin_freq=5, samp_freq=512, samples=2048)
-        self.run_unitary_avg_power_spectrum_sin(
-            sin_freq=5, samp_freq=512, samples=4098)
-        self.run_unitary_avg_power_spectrum_sin(
-            sin_freq=17, samp_freq=512, samples=1024)
-        self.run_unitary_avg_power_spectrum_sin(
-            sin_freq=5, samp_freq=1024, samples=2048)
+        self.run_sin(sin_freq=5, samp_freq=512, samples=1024)
+        self.run_sin(sin_freq=5, samp_freq=512, samples=2048)
+        self.run_sin(sin_freq=5, samp_freq=512, samples=4098)
+        self.run_sin(sin_freq=17, samp_freq=512, samples=1024)
+        self.run_sin(sin_freq=5, samp_freq=1024, samples=2048)
         # test long wavelenth sin, so be closer to window frequency
-        self.run_unitary_avg_power_spectrum_sin(
-            sin_freq=1, samp_freq=1024, samples=2048)
+        self.run_sin(sin_freq=1, samp_freq=1024, samples=2048, places=0)
         # finally, with some irrational numbers, to check that I'm not
         # getting lucky
-        self.run_unitary_avg_power_spectrum_sin(
+        self.run_sin(
             sin_freq=_numpy.pi, samp_freq=100 * _numpy.exp(1), samples=1024)