electrode loop and adjusted plot

code/chirpdetection.py

@@ -217,21 +217,10 @@ def main(datapath: str) -> None:
         # start_index = t0 * data.samplerate
         # stop_index = (t0 + dt) * data.samplerate
 
-        fig, axs = plt.subplots(
-            7,
-            2,
-            figsize=(20 / 2.54, 12 / 2.54),
-            constrained_layout=True,
-            sharex=True,
-            sharey='row',
-        )
 
         # iterate through all fish
         for i, track_id in enumerate(np.unique(ident[~np.isnan(ident)])[:2]):
 
-            # load region of interest of raw data file
-            data_oi = data[start_index:stop_index, :]
-            time_oi = raw_time[start_index:stop_index]
 
             # get indices for time array in time window
             window_index = np.arange(len(idx))[
@@ -250,11 +239,25 @@ def main(datapath: str) -> None:
             if len(freq_temp) < expected_duration * 0.9:
                 continue
 
+            # Create plot (three electrodes per fish)
+            fig, axs = plt.subplots(
+                7,
+                3,
+                figsize=(20 / 2.54, 12 / 2.54),
+                constrained_layout=True,
+                sharex=True,
+                sharey='row',
+            )
+
             # get best electrode
-            electrode = np.argsort(np.nanmean(powers_temp, axis=0))[-1]
+            best_electrodes = np.argsort(np.nanmean(powers_temp, axis=0))[-3:]
 
             # <------------------------------------------ Iterate through electrodes
+            for e, electrode in enumerate(best_electrodes):
 
+                # load region of interest of raw data file
+                data_oi = data[start_index:stop_index, :]
+                time_oi = raw_time[start_index:stop_index]
                 # plot wavetracker tracks to spectrogram
                 # for track_id in np.unique(ident):  # <---------- Find freq gaps later
                 # here
@@ -371,72 +374,72 @@ def main(datapath: str) -> None:
                 # PLOT ------------------------------------------------------------
 
                 # plot spectrogram
-            plot_spectrogram(axs[0, i], data_oi[:, electrode], data.samplerate, t0)
+                plot_spectrogram(axs[0, e], data_oi[:, electrode], data.samplerate, t0)
 
                 # plot baseline instantaneos frequency
-            axs[1, i].plot(baseline_freq_time, baseline_freq -
+                axs[1, e].plot(baseline_freq_time, baseline_freq -
                             np.median(baseline_freq), marker=".")
 
                 # plot waveform of filtered signal
-            axs[2, i].plot(time_oi, baseline, c="k")
+                axs[2, e].plot(time_oi, baseline, c="k")
 
                 # plot narrow filtered baseline
-            axs[2, i].plot(
+                axs[2, e].plot(
                     time_oi,
                     baseline_envelope,
                     c="orange",
                 )
 
                 # plot broad filtered baseline
-            axs[2, i].plot(
+                axs[2, e].plot(
                     time_oi,
                     broad_baseline,
                     c="green",
                 )
 
                 # plot waveform of filtered search signal
-            axs[3, i].plot(time_oi, search)
+                axs[3, e].plot(time_oi, search)
 
                 # plot envelope of search signal
-            axs[3, i].plot(
+                axs[3, e].plot(
                     time_oi,
                     search_envelope,
                     c="orange",
                 )
 
                 # plot filtered and rectified envelope
-            axs[4, i].plot(time_oi, baseline_envelope)
-            axs[4, i].scatter(
+                axs[4, e].plot(time_oi, baseline_envelope)
+                axs[4, e].scatter(
                     (time_oi)[baseline_peaks],
                     baseline_envelope[baseline_peaks],
                     c="red",
                 )
 
                 # plot envelope of search signal
-            axs[5, i].plot(time_oi, search_envelope)
-            axs[5, i].scatter(
+                axs[5, e].plot(time_oi, search_envelope)
+                axs[5, e].scatter(
                     (time_oi)[search_peaks],
                     search_envelope[search_peaks],
                     c="red",
                 )
 
                 # plot filtered instantaneous frequency
-            axs[6, i].plot(baseline_freq_time, np.abs(inst_freq_filtered))
-            axs[6, i].scatter(
+                axs[6, e].plot(baseline_freq_time, np.abs(inst_freq_filtered))
+                axs[6, e].scatter(
                     baseline_freq_time[inst_freq_peaks],
                     np.abs(inst_freq_filtered)[inst_freq_peaks],
                     c="red",
                 )
 
-            axs[6, i].set_xlabel("Time [s]")
-            axs[0, i].set_title("Spectrogram")
-            axs[1, i].set_title("Fitered baseline instanenous frequency")
-            axs[2, i].set_title("Fitered baseline")
-            axs[3, i].set_title("Fitered above")
-            axs[4, i].set_title("Filtered envelope of baseline envelope")
-            axs[5, i].set_title("Search envelope")
-            axs[6, i].set_title("Filtered absolute instantaneous frequency")
-
+                axs[6, e].set_xlabel("Time [s]")
+                axs[0, e].set_title("Spectrogram")
+                axs[1, e].set_title("Fitered baseline instanenous frequency")
+                axs[2, e].set_title("Fitered baseline")
+                axs[3, e].set_title("Fitered above")
+                axs[4, e].set_title("Filtered envelope of baseline envelope")
+                axs[5, e].set_title("Search envelope")
+                axs[6, e].set_title("Filtered absolute instantaneous frequency")
+                fig.suptitle('Fish ID %i' %track_id)
             plt.show()