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,194 +239,208 @@ def main(datapath: str) -> None:
             if len(freq_temp) < expected_duration * 0.9:
                 continue
 
-            # get best electrode
+            # Create plot (three electrodes per fish)
-            electrode = np.argsort(np.nanmean(powers_temp, axis=0))[-1]
+            fig, axs = plt.subplots(
-
+                7,
-            # <------------------------------------------ Iterate through electrodes
+                3,
-
+                figsize=(20 / 2.54, 12 / 2.54),
-            # plot wavetracker tracks to spectrogram
+                constrained_layout=True,
-            # for track_id in np.unique(ident):  # <---------- Find freq gaps later
+                sharex=True,
-            # here
+                sharey='row',
-
-            #     # get indices for time array in time window
-            #     window_index = np.arange(len(idx))[
-            #         (ident == track_id) &
-            #         (time[idx] >= t0) &
-            #         (time[idx] <= (t0 + dt))
-            #     ]
-
-            #     freq_temp = freq[window_index]
-            #     time_temp = time[idx[window_index]]
-
-            #     axs[0].plot(time_temp-t0, freq_temp, lw=2)
-            #     axs[0].set_ylim(500, 1000)
-
-            # track_id = ids
-
-            # frequency where second filter filters
-            search_freq = 50
-
-            # filter baseline and above
-            baseline, search = double_bandpass(
-                data_oi[:, electrode], data.samplerate, freq_temp, search_freq
-            )
-
-            # compute instantaneous frequency on broad signal
-            broad_baseline = bandpass_filter(
-                data_oi[:, electrode],
-                data.samplerate,
-                lowf=np.mean(freq_temp)-5,
-                highf=np.mean(freq_temp)+100
-            )
-            # compute instantaneous frequency on narrow signal
-            baseline_freq_time, baseline_freq = instantaneos_frequency(
-                baseline, data.samplerate
             )
 
-            # compute envelopes
+            # get best electrode
-            cutoff = 25
+            best_electrodes = np.argsort(np.nanmean(powers_temp, axis=0))[-3:]
-            baseline_envelope = envelope(baseline, data.samplerate, cutoff)
-            search_envelope = envelope(search, data.samplerate, cutoff)
-
-            # highpass filter envelopes
-            cutoff = 5
-            baseline_envelope = highpass_filter(
-                baseline_envelope, data.samplerate, cutoff=cutoff
-            )
-            baseline_envelope = np.abs(baseline_envelope)
-            # search_envelope = highpass_filter(
-            #     search_envelope, data.samplerate, cutoff=cutoff)
-
-            # envelopes of filtered envelope of filtered baseline
-            # baseline_envelope = envelope(
-            #     np.abs(baseline_envelope), data.samplerate, cutoff
-            # )
-
-            # search_envelope = bandpass_filter(
-            # search_envelope, data.samplerate, lowf=lowf, highf=highf)
-
-            # bandpass filter the instantaneous
-            inst_freq_filtered = bandpass_filter(
-                baseline_freq, data.samplerate, lowf=15, highf=8000
-            )
-
-            # CUT OFF OVERLAP -------------------------------------------------
-
-            # cut off first and last 0.5 * overlap at start and end
-            valid = np.arange(
-                int(0.5 * window_overlap), len(baseline_envelope) -
-                int(0.5 * window_overlap)
-            )
-            baseline_envelope = baseline_envelope[valid]
-            search_envelope = search_envelope[valid]
-
-            # get inst freq valid snippet
-            valid_t0 = int(0.5 * window_overlap) / data.samplerate
-            valid_t1 = baseline_freq_time[-1] - \
-                (int(0.5 * window_overlap) / data.samplerate)
-
-            inst_freq_filtered = inst_freq_filtered[(baseline_freq_time >= valid_t0) & (
-                baseline_freq_time <= valid_t1)]
-
-            baseline_freq = baseline_freq[(baseline_freq_time >= valid_t0) & (
-                baseline_freq_time <= valid_t1)]
-
-            baseline_freq_time = baseline_freq_time[(baseline_freq_time >= valid_t0) & (
-                baseline_freq_time <= valid_t1)] + t0
-
-            # overwrite raw time to valid region
-            time_oi = time_oi[valid]
-            baseline = baseline[valid]
-            broad_baseline = broad_baseline[valid]
-            search = search[valid]
-
-            # PEAK DETECTION --------------------------------------------------
-
-            # detect peaks baseline_enelope
-            prominence = np.percentile(baseline_envelope, 90)
-            baseline_peaks, _ = find_peaks(
-                np.abs(baseline_envelope), prominence=prominence)
-
-            # detect peaks search_envelope
-            prominence = np.percentile(search_envelope, 75)
-            search_peaks, _ = find_peaks(
-                search_envelope, prominence=prominence)
-
-            # detect peaks inst_freq_filtered
-            prominence = 2
-            inst_freq_peaks, _ = find_peaks(
-                np.abs(inst_freq_filtered), prominence=prominence)
-
-            # PLOT ------------------------------------------------------------
-
-            # plot spectrogram
-            plot_spectrogram(axs[0, i], data_oi[:, electrode], data.samplerate, t0)
-
-            # plot baseline instantaneos frequency
-            axs[1, i].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")
-
-            # plot narrow filtered baseline
-            axs[2, i].plot(
-                time_oi,
-                baseline_envelope,
-                c="orange",
-            )
-
-            # plot broad filtered baseline
-            axs[2, i].plot(
-                time_oi,
-                broad_baseline,
-                c="green",
-            )
-
-            # plot waveform of filtered search signal
-            axs[3, i].plot(time_oi, search)
-
-            # plot envelope of search signal
-            axs[3, i].plot(
-                time_oi,
-                search_envelope,
-                c="orange",
-            )
-
-            # plot filtered and rectified envelope
-            axs[4, i].plot(time_oi, baseline_envelope)
-            axs[4, i].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(
-                (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(
-                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")
 
-        plt.show()
+            # <------------------------------------------ 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
+
+                #     # get indices for time array in time window
+                #     window_index = np.arange(len(idx))[
+                #         (ident == track_id) &
+                #         (time[idx] >= t0) &
+                #         (time[idx] <= (t0 + dt))
+                #     ]
+
+                #     freq_temp = freq[window_index]
+                #     time_temp = time[idx[window_index]]
+
+                #     axs[0].plot(time_temp-t0, freq_temp, lw=2)
+                #     axs[0].set_ylim(500, 1000)
+
+                # track_id = ids
+
+                # frequency where second filter filters
+                search_freq = 50
+
+                # filter baseline and above
+                baseline, search = double_bandpass(
+                    data_oi[:, electrode], data.samplerate, freq_temp, search_freq
+                )
+
+                # compute instantaneous frequency on broad signal
+                broad_baseline = bandpass_filter(
+                    data_oi[:, electrode],
+                    data.samplerate,
+                    lowf=np.mean(freq_temp)-5,
+                    highf=np.mean(freq_temp)+100
+                )
+                # compute instantaneous frequency on narrow signal
+                baseline_freq_time, baseline_freq = instantaneos_frequency(
+                    baseline, data.samplerate
+                )
+
+                # compute envelopes
+                cutoff = 25
+                baseline_envelope = envelope(baseline, data.samplerate, cutoff)
+                search_envelope = envelope(search, data.samplerate, cutoff)
+
+                # highpass filter envelopes
+                cutoff = 5
+                baseline_envelope = highpass_filter(
+                    baseline_envelope, data.samplerate, cutoff=cutoff
+                )
+                baseline_envelope = np.abs(baseline_envelope)
+                # search_envelope = highpass_filter(
+                #     search_envelope, data.samplerate, cutoff=cutoff)
+
+                # envelopes of filtered envelope of filtered baseline
+                # baseline_envelope = envelope(
+                #     np.abs(baseline_envelope), data.samplerate, cutoff
+                # )
+
+                # search_envelope = bandpass_filter(
+                # search_envelope, data.samplerate, lowf=lowf, highf=highf)
+
+                # bandpass filter the instantaneous
+                inst_freq_filtered = bandpass_filter(
+                    baseline_freq, data.samplerate, lowf=15, highf=8000
+                )
+
+                # CUT OFF OVERLAP -------------------------------------------------
+
+                # cut off first and last 0.5 * overlap at start and end
+                valid = np.arange(
+                    int(0.5 * window_overlap), len(baseline_envelope) -
+                    int(0.5 * window_overlap)
+                )
+                baseline_envelope = baseline_envelope[valid]
+                search_envelope = search_envelope[valid]
+
+                # get inst freq valid snippet
+                valid_t0 = int(0.5 * window_overlap) / data.samplerate
+                valid_t1 = baseline_freq_time[-1] - \
+                    (int(0.5 * window_overlap) / data.samplerate)
+
+                inst_freq_filtered = inst_freq_filtered[(baseline_freq_time >= valid_t0) & (
+                    baseline_freq_time <= valid_t1)]
+
+                baseline_freq = baseline_freq[(baseline_freq_time >= valid_t0) & (
+                    baseline_freq_time <= valid_t1)]
+
+                baseline_freq_time = baseline_freq_time[(baseline_freq_time >= valid_t0) & (
+                    baseline_freq_time <= valid_t1)] + t0
+
+                # overwrite raw time to valid region
+                time_oi = time_oi[valid]
+                baseline = baseline[valid]
+                broad_baseline = broad_baseline[valid]
+                search = search[valid]
+
+                # PEAK DETECTION --------------------------------------------------
+
+                # detect peaks baseline_enelope
+                prominence = np.percentile(baseline_envelope, 90)
+                baseline_peaks, _ = find_peaks(
+                    np.abs(baseline_envelope), prominence=prominence)
+
+                # detect peaks search_envelope
+                prominence = np.percentile(search_envelope, 75)
+                search_peaks, _ = find_peaks(
+                    search_envelope, prominence=prominence)
+
+                # detect peaks inst_freq_filtered
+                prominence = 2
+                inst_freq_peaks, _ = find_peaks(
+                    np.abs(inst_freq_filtered), prominence=prominence)
+
+                # PLOT ------------------------------------------------------------
+
+                # plot spectrogram
+                plot_spectrogram(axs[0, e], data_oi[:, electrode], data.samplerate, t0)
+
+                # plot baseline instantaneos frequency
+                axs[1, e].plot(baseline_freq_time, baseline_freq -
+                            np.median(baseline_freq), marker=".")
+
+                # plot waveform of filtered signal
+                axs[2, e].plot(time_oi, baseline, c="k")
+
+                # plot narrow filtered baseline
+                axs[2, e].plot(
+                    time_oi,
+                    baseline_envelope,
+                    c="orange",
+                )
+
+                # plot broad filtered baseline
+                axs[2, e].plot(
+                    time_oi,
+                    broad_baseline,
+                    c="green",
+                )
+
+                # plot waveform of filtered search signal
+                axs[3, e].plot(time_oi, search)
+
+                # plot envelope of search signal
+                axs[3, e].plot(
+                    time_oi,
+                    search_envelope,
+                    c="orange",
+                )
+
+                # plot filtered and rectified envelope
+                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, 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, 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, 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()
 
 
 if __name__ == "__main__":