From 93026e27e92a39faeb05d758d8ea6761e7606af0 Mon Sep 17 00:00:00 2001
From: Per A Brodtkorb <per.andreas.brodtkorb@gmail.com>
Date: Sat, 11 Feb 2017 10:54:01 +0100
Subject: [PATCH] Updated doctests

---
 wafo/demo_sg.py                         | 52 ++++++++++++++-----------
 wafo/demos.py                           |  8 +---
 wafo/wave_theory/dispersion_relation.py |  9 +----
 3 files changed, 34 insertions(+), 35 deletions(-)

diff --git a/wafo/demo_sg.py b/wafo/demo_sg.py
index 178a5b5..8e6d3d8 100644
--- a/wafo/demo_sg.py
+++ b/wafo/demo_sg.py
@@ -1,46 +1,54 @@
 import matplotlib.pyplot as plt
-from pylab import subplot, plot, title, figure
-from numpy import random, arange, sin
-from .sg_filter import SavitzkyGolay, smoothn  # calc_coeff, smooth
+import numpy as np
+from wafo.sg_filter import SavitzkyGolay, smoothn  # calc_coeff, smooth
 
 
 def example_reconstruct_noisy_chirp():
-    figure(figsize=(7, 12))
+    """
+    Example
+    -------
+    >>> example_reconstruct_noisy_chirp()
+    """
+    plt.figure(figsize=(7, 12))
 
     # generate chirp signal
-    tvec = arange(0, 6.28, .02)
-    true_signal = sin(tvec * (2.0 + tvec))
+    tvec = np.arange(0, 6.28, .02)
+    true_signal = np.sin(tvec * (2.0 + tvec))
+    true_d_signal = (2+tvec) * np.cos(tvec * (2.0 + tvec))
 
     # add noise to signal
-    noise = random.normal(size=true_signal.shape)
+    noise = np.random.normal(size=true_signal.shape)
     signal = true_signal + .15 * noise
 
     # plot signal
-    subplot(311)
-    plot(signal)
-    title('signal')
+    plt.subplot(311)
+    plt.plot(signal)
+    plt.title('signal')
 
     # smooth and plot signal
-    subplot(312)
+    plt.subplot(312)
     savgol = SavitzkyGolay(n=8, degree=4)
     s_signal = savgol.smooth(signal)
     s2 = smoothn(signal, robust=True)
-    plot(s_signal)
-    plot(s2)
-    plot(true_signal, 'r--')
-    title('smoothed signal')
+    plt.plot(s_signal)
+    plt.plot(s2)
+    plt.plot(true_signal, 'r--')
+    plt.title('smoothed signal')
 
     # smooth derivative of signal and plot it
-    subplot(313)
+    plt.subplot(313)
     savgol1 = SavitzkyGolay(n=8, degree=1, diff_order=1)
 
-    d_signal = savgol1.smooth(signal)
+    dt = tvec[1]-tvec[0]
+    d_signal = savgol1.smooth(signal) / dt
 
-    plot(d_signal)
-    title('smoothed derivative of signal')
-
-    plt.show('hold')  # show plot
+    plt.plot(d_signal)
+    plt.plot(true_d_signal, 'r--')
+    plt.title('smoothed derivative of signal')
 
 
 if __name__ == '__main__':
-    example_reconstruct_noisy_chirp()
+    from wafo.testing import test_docstrings
+    test_docstrings(__file__)
+    # example_reconstruct_noisy_chirp()
+    # plt.show('hold')  # show plot
diff --git a/wafo/demos.py b/wafo/demos.py
index 2d73dd1..8f73f18 100644
--- a/wafo/demos.py
+++ b/wafo/demos.py
@@ -132,10 +132,6 @@ def humps(x=None):
         2 * y - 5.2
 
 
-def test_docstrings():
-    import doctest
-    print('Testing docstrings in %s' % __file__)
-    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
-
 if __name__ == '__main__':
-    test_docstrings()
+    from wafo.testing import test_docstrings
+    test_docstrings(__file__)
diff --git a/wafo/wave_theory/dispersion_relation.py b/wafo/wave_theory/dispersion_relation.py
index a24a661..86b5d64 100644
--- a/wafo/wave_theory/dispersion_relation.py
+++ b/wafo/wave_theory/dispersion_relation.py
@@ -204,11 +204,6 @@ def w2k(w, theta=0.0, h=inf, g=9.81, count_limit=100):
     return k1, k2
 
 
-def test_docstrings():
-    import doctest
-    print('Testing docstrings in %s' % __file__)
-    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)
-
-
 if __name__ == '__main__':
-    test_docstrings()
+    from wafo.testing import test_docstrings
+    test_docstrings(__file__)