From abbb548463b442f9861ce9d8acc8af1623d76a3f Mon Sep 17 00:00:00 2001
From: ackurth <anno.kurth@posteo.de>
Date: Sun, 17 Jan 2021 01:58:34 +0100
Subject: [PATCH] Add example for comparison for multitaper and welch

---
 elephant/spectral.py | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/elephant/spectral.py b/elephant/spectral.py
index 983fbc0de..a90ded061 100644
--- a/elephant/spectral.py
+++ b/elephant/spectral.py
@@ -294,7 +294,7 @@ def multitaper_psd(signal, fs=1, NW=4, num_tapers='auto'):
     spectrum_estimates = np.abs(np.fft.rfft(signal * slepain_fcts, axis=1))**2
 
     # average to obtain Multitaper PSD estimate
-    multitaper_psd = np.mean(spectrum_estimates, axis=0)
+    multitaper_psd = np.mean(spectrum_estimates, axis=0) / fs
 
     return freqs, multitaper_psd
 
@@ -598,14 +598,34 @@ def generate_data(length, coeffs, variance, fs, transients=10):
     # Choose parameters, coeffs as in nitime
     length = 2**10
     coeffs = np.array([2.7607, -3.8106, 2.6535, -0.9238])
-    variance = 1
+    variance = 1.
     fs = 10
 
-    times, time_series, freqs, psd_time_series = generate_data(length,
-                                                               coeffs,
-                                                               variance,
-                                                               fs)
+    times, time_series, freqs, psd_time_series = generate_data(
+        length=length,
+        coeffs=coeffs,
+        variance=variance,
+        fs=fs)
 
+    freqs_multi, psd_multitaper = multitaper_psd(signal=time_series,
+                                                 fs=fs,
+                                                 NW=4)
 
-    import IPython
-    IPython.embed()
+    freqs_welch, psd_welch = welch_psd(signal=time_series,
+                                       n_segments=8,
+                                       fs=fs,
+                                       len_segment=length)
+
+    from matplotlib import pyplot as plt
+
+    plt.figure(figsize=(10, 10))
+    plt.yscale('log')
+
+    plt.plot(freqs, psd_time_series, color='blue', label='Ground truth')
+    plt.plot(freqs, psd_multitaper, color='orange',
+             label='Multitaper estimate')
+    plt.plot(freqs, psd_welch, color='black', label='Welch estimate')
+
+    plt.legend()
+
+    plt.show()