[eodanalysis] added wave_spectrum_plot()

thunderfish/eodanalysis.py

@@ -8,7 +8,8 @@
 
 ## Visualization
 - `eod_waveform_plot()`: plot and annotate the averaged EOD-waveform with standard deviation.
-- `pulse_spectrum_plot()`: plot and annotate spectrum of single pulse.
+- `wave_spectrum_plot()`: plot and annotate spectrum of wave-type EODs.
+- `pulse_spectrum_plot()`: plot and annotate spectrum of single pulse-type EOD.
 
 ## Fit functions
 - `fourier_series()`: Fourier series of sine waves with amplitudes and phases.
@@ -645,15 +646,67 @@ def eod_waveform_plot(eod_waveform, peaks, ax, unit=None, tau=None,
                         zorder=10)
     ax.set_xlim(time[0], time[-1])
     ax.set_xlabel('Time [msec]')
-    if unit is not None and len(unit)>0:
+    if unit:
         ax.set_ylabel('Amplitude [%s]' % unit)
     else:
         ax.set_ylabel('Amplitude')
 
 
+def wave_spectrum_plot(spec, props, axa, axp, unit=None, color='b', lw=2, markersize=12):
+    """Plot and annotate spectrum of wave-type EOD.
+
+    Parameters
+    ----------
+    spec: 2-D array
+        The amplitude spectrum of a single pulse as returned by `analyze_wave()`.
+        First column is the index of the harmonics, second column its frequency,
+        third column its amplitude, fourth column its amplitude relative to the fundamental,
+        and fifth column the phase shift relative to the fundamental.
+    props: dict
+        A dictionary with properties of the analyzed EOD waveform as
+        returned by `analyze_wave()`.
+    axa:
+        Axis for amplitude plot.
+    axa:
+        Axis for phase plot.
+    unit: string
+        Optional unit of the data used for y-label.
+    color:
+        Color for line and points of spectrum.
+    lw: float
+        Linewidth for spectrum.
+    markersize: float
+        Size of points on spectrum.
+    """
+    n = 9
+    # amplitudes:
+    markers, stemlines, baseline = axa.stem(spec[:n,0], spec[:n,2])
+    plt.setp(markers, color=color, markersize=markersize, clip_on=False)
+    plt.setp(stemlines, color=color, lw=lw)
+    axa.set_xlim(-0.5, n-0.5)
+    axa.set_xticks(np.arange(0, n, 1))
+    axa.set_xticklabels([])
+    if unit:
+        axa.set_ylabel('Amplitude [%s]' % unit)
+    else:
+        axa.set_ylabel('Amplitude')
+    # phases:
+    phases = spec[:,4]
+    phases[phases<0.0] = phases[phases<0.0] + 2.0*np.pi
+    markers, stemlines, baseline = axp.stem(spec[:n,0], phases[:n])
+    plt.setp(markers, color=color, markersize=markersize, clip_on=False)
+    plt.setp(stemlines, color=color, lw=lw)
+    axp.set_xlim(-0.5, n-0.5)
+    axp.set_xticks(np.arange(0, n, 1))
+    axp.set_ylim(0, 2.0*np.pi)
+    axp.set_yticks([0, np.pi, 2.0*np.pi])
+    axp.set_yticklabels([u'0', u'\u03c0', u'2\u03c0'])
+    axp.set_xlabel('Harmonics')
+    axp.set_ylabel('Phase')
+
 
 def pulse_spectrum_plot(power, props, ax, color='b', lw=3, markersize=80):
-    """Plot and annotate spectrum of single pulse.
+    """Plot and annotate spectrum of single pulse-type EOD.
 
     Parameters
     ----------

thunderfish/thunderfish.py

@@ -22,7 +22,7 @@
 from .harmonicgroups import harmonic_groups, harmonic_groups_args, psd_peak_detection_args, fundamental_freqs, fundamental_freqs_and_power, colors_markers, plot_harmonic_groups
 from .consistentfishes import consistent_fishes
 from .eodanalysis import eod_waveform, analyze_wave, analyze_pulse
-from .eodanalysis import eod_waveform_plot, pulse_spectrum_plot
+from .eodanalysis import eod_waveform_plot, pulse_spectrum_plot, wave_spectrum_plot
 from .eodanalysis import add_eod_analysis_config, eod_waveform_args
 from .eodanalysis import analyze_wave_args, analyze_pulse_args
 from .tabledata import TableData, add_write_table_config, write_table_args
@@ -100,7 +100,9 @@ def keypress(event):
     ax2 = fig.add_axes([0.075, 0.06, 0.9, 0.09]) # whole trace
     ax3 = fig.add_axes([0.075, 0.6, 0.7, 0.3])   # psd
     ax4 = fig.add_axes([0.075, 0.2, 0.4, 0.3])   # mean eod
-    ax5 = fig.add_axes([0.575, 0.2, 0.4, 0.3])   # pusle spectrum
+    ax5 = fig.add_axes([0.575, 0.2, 0.4, 0.3])   # pulse spectrum
+    ax6 = fig.add_axes([0.575, 0.36, 0.4, 0.14]) # amplitude spectrum
+    ax7 = fig.add_axes([0.575, 0.2, 0.4, 0.14])  # phase spectrum
 
     if show_plot:
         fig.canvas.mpl_connect('key_press_event', keypress)
@@ -183,10 +185,20 @@ def keypress(event):
 
     ################
 
-    if not usedax5 and len(eod_props) > 0 and eod_props[0]['type'] == 'pulse':
+    ax5.set_visible(True)
+    ax6.set_visible(False)
+    ax7.set_visible(False)
+    if not usedax5 and len(eod_props) > 0:
         usedax5 = True
-        pulse_spectrum_plot(spec_data[0], eod_props[0], ax5)
-        ax5.set_title('Single pulse spectrum', fontsize=14, y=1.05)
+        if  eod_props[0]['type'] == 'pulse':
+            pulse_spectrum_plot(spec_data[0], eod_props[0], ax5)
+            ax5.set_title('Single pulse spectrum', fontsize=14, y=1.05)
+        else:
+            ax5.set_visible(False)
+            ax6.set_visible(True)
+            ax7.set_visible(True)
+            wave_spectrum_plot(spec_data[0], eod_props[0], ax6, ax7, unit)
+            ax6.set_title('Amplitude and phase spectrum', fontsize=14, y=1.05)
 
     ## # plot inter EOD interval histogram
     ## if len(inter_eod_intervals)>2:
@@ -219,7 +231,7 @@ def keypress(event):
     ##                 loc='upper right', frameon=False)
 
     # cosmetics
-    for ax in [ax2, ax3, ax4, ax5]:
+    for ax in [ax2, ax3, ax4, ax5, ax6, ax7]:
         ax.spines['top'].set_visible(False)
         ax.spines['right'].set_visible(False)
         ax.get_xaxis().tick_bottom()