From ce1c0fa31951579b944d006ebe52a32b0ec90b03 Mon Sep 17 00:00:00 2001
From: weygoldt <88969563+weygoldt@users.noreply.github.com>
Date: Thu, 12 Jan 2023 08:50:28 +0100
Subject: [PATCH] added more docstrings and moved chirpsim

---
 code/analysis.py            |  41 +++++-----
 code/chirpdetection.py      |  26 +++---
 code/modules/filters.py     | 153 +++++++++++++++++++-----------------
 code/modules/simulations.py |  61 ++++++++++++++
 4 files changed, 176 insertions(+), 105 deletions(-)
 create mode 100644 code/modules/simulations.py

diff --git a/code/analysis.py b/code/analysis.py
index cc6d5c8..787a53e 100644
--- a/code/analysis.py
+++ b/code/analysis.py
@@ -30,44 +30,45 @@ def main(folder):
             spec_power), vmin=-100, vmax=-50)
 
         for track_id in np.unique(ident):
-            # window_index for time array in time window 
+            # window_index for time array in time window
             window_index = np.arange(len(idx))[(ident == track_id) &
-                                        (time[idx] >= t0) &
-                                        (time[idx] <= (t0+dt))]
+                                               (time[idx] >= t0) &
+                                               (time[idx] <= (t0+dt))]
             freq_temp = freq[window_index]
-            time_temp = time[idx[window_index]] 
+            time_temp = time[idx[window_index]]
             #mean_freq = np.mean(freq_temp)
             #fdata = bandpass_filter(data_oi[:, track_id], data.samplerate, mean_freq-5, mean_freq+200)
             ax.plot(time_temp - t0, freq_temp)
 
     ax.set_ylim(500, 1000)
     plt.show()
-    # filter plot 
+    # filter plot
     id = 10.
     i = 10
     window_index = np.arange(len(idx))[(ident == id) &
-                                        (time[idx] >= t0) &
-                                        (time[idx] <= (t0+dt))] 
+                                       (time[idx] >= t0) &
+                                       (time[idx] <= (t0+dt))]
     freq_temp = freq[window_index]
-    time_temp = time[idx[window_index]] 
+    time_temp = time[idx[window_index]]
     mean_freq = np.mean(freq_temp)
-    fdata = bandpass_filter(data_oi[:, i], rate=data.samplerate, lowf=mean_freq-5, highf=mean_freq+200)
+    fdata = bandpass_filter(
+        data_oi[:, i], rate=data.samplerate, lowf=mean_freq-5, highf=mean_freq+200)
     fig, ax = plt.subplots()
     ax.plot(np.arange(len(fdata))/data.samplerate, fdata, marker='*')
-    #plt.show()
-    #freqency analyis of filtered data
-    
+    # plt.show()
+    # freqency analyis of filtered data
+
     time_fdata = np.arange(len(fdata))/data.samplerate
     roll_fdata = np.roll(fdata, shift=1)
-    period_index = np.arange(len(fdata))[(roll_fdata < 0) & (fdata>=0)]
+    period_index = np.arange(len(fdata))[(roll_fdata < 0) & (fdata >= 0)]
 
     plt.plot(time_fdata, fdata)
     plt.scatter(time_fdata[period_index], fdata[period_index], c='r')
     plt.scatter(time_fdata[period_index-1], fdata[period_index-1], c='r')
-    
+
     upper_bound = np.abs(fdata[period_index])
     lower_bound = np.abs(fdata[period_index-1])
-    
+
     upper_times = np.abs(time_fdata[period_index])
     lower_times = np.abs(time_fdata[period_index-1])
 
@@ -78,17 +79,14 @@ def main(folder):
     true_zero = lower_times + time_delta*lower_ratio
 
     plt.scatter(true_zero, np.zeros(len(true_zero)))
-    
+
     # calculate the frequency
     inst_freq = 1 / np.diff(true_zero)
     filtered_inst_freq = gaussian_filter1d(inst_freq, 0.005)
-    fig, ax  = plt.subplots()
+    fig, ax = plt.subplots()
     ax.plot(filtered_inst_freq, marker='.')
-    # in 5 sekunden welcher fisch auf einer elektrode am 
-
+    # in 5 sekunden welcher fisch auf einer elektrode am
 
-
-    
     embed()
     exit()
 
@@ -98,7 +96,6 @@ def main(folder):
     # fig, ax = plt.subplots(figsize=(20/2.54, 12/2.54))
     # ax.plot(np.arange(len(data_oi[:, i])), data_oi[:, i])
 
-
     pass
 
 
diff --git a/code/chirpdetection.py b/code/chirpdetection.py
index 3735973..889004f 100644
--- a/code/chirpdetection.py
+++ b/code/chirpdetection.py
@@ -121,7 +121,7 @@ def double_bandpass(
 ) -> tuple[np.ndarray, np.ndarray]:
     """
     Apply a bandpass filter to the baseline of a signal and a second bandpass
-    filter above or below the baseline.
+    filter above or below the baseline, as specified by the search frequency.
 
     Parameters
     ----------
@@ -171,8 +171,13 @@ def main(datapath: str) -> None:
     ident = np.load(datapath + "ident_v.npy", allow_pickle=True)
 
     # set time window # <------------------------ Iterate through windows here
-    window_duration = 60 * data.samplerate
-    window_overlap = 0.3
+    window_duration = 60 * 5    # 5 minutes window
+    window_overlap = 30         # 30 seconds overlap
+    window_starts = np.arange(
+        time[0], time[-1], window_duration)
+
+    embed()
+    exit()
 
     t0 = 3 * 60 * 60 + 6 * 60 + 43.5
     dt = 60
@@ -187,6 +192,14 @@ def main(datapath: str) -> None:
 
         # <------------------------------------------ Find best electrodes here
         # <------------------------------------------ Iterate through electrodes
+        # get indices for time array in time window
+        window_index = np.arange(len(idx))[
+            (ident == track_id) & (time[idx] >= t0) & (time[idx] <= (t0 + dt))
+        ]
+
+        # filter baseline and above
+        freq_temp = freq[window_index]
+        time_temp = time[idx[window_index]]
 
         electrode = 10
 
@@ -224,14 +237,7 @@ def main(datapath: str) -> None:
         # frequency where second filter filters
         search_freq = -50
 
-        # get indices for time array in time window
-        window_index = np.arange(len(idx))[
-            (ident == track_id) & (time[idx] >= t0) & (time[idx] <= (t0 + dt))
-        ]
-
         # filter baseline and above
-        freq_temp = freq[window_index]
-        time_temp = time[idx[window_index]]
         baseline, search = double_bandpass(
             data_oi[:, electrode], data.samplerate, freq_temp, search_freq
         )
diff --git a/code/modules/filters.py b/code/modules/filters.py
index bf311d7..499a3dc 100644
--- a/code/modules/filters.py
+++ b/code/modules/filters.py
@@ -2,96 +2,103 @@ from scipy.signal import butter, sosfiltfilt
 import numpy as np
 
 
-def bandpass_filter(data, rate, lowf=100, highf=1100):
+def bandpass_filter(
+        data: np.ndarray,
+        rate: float,
+        lowf: float,
+        highf: float,
+) -> np.ndarray:
+    """Bandpass filter a signal.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        The data to be filtered
+    rate : float
+        The sampling rate
+    lowf : float
+        The low frequency cutoff
+    highf : float
+        The high frequency cutoff
+
+    Returns
+    -------
+    np.ndarray : The filtered data
+    """
     sos = butter(2, (lowf, highf), "bandpass", fs=rate, output="sos")
     fdata = sosfiltfilt(sos, data)
     return fdata
 
 
 def highpass_filter(
-    data,
-    rate,
-    cutoff=100,
-):
-
+    data: np.ndarray,
+    rate: float,
+    cutoff: float,
+) -> np.ndarray:
+    """Highpass filter a signal.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        The data to be filtered
+    rate : float
+        The sampling rate
+    cutoff : float
+        The cutoff frequency
+
+    Returns
+    -------
+    np.ndarray
+        The filtered data
+    """
     sos = butter(2, cutoff, "highpass", fs=rate, output="sos")
     fdata = sosfiltfilt(sos, data)
     return fdata
 
 
 def lowpass_filter(
-    data,
-    rate,
-    cutoff=100,
-):
-
+    data: np.ndarray,
+    rate: float,
+    cutoff: float
+) -> np.ndarray:
+    """Lowpass filter a signal.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        The data to be filtered
+    rate : float
+        The sampling rate
+    cutoff : float
+        The cutoff frequency
+
+    Returns
+    -------
+    np.ndarray
+        The filtered data
+    """
     sos = butter(2, cutoff, "lowpass", fs=rate, output="sos")
     fdata = sosfiltfilt(sos, data)
     return fdata
 
 
-def envelope(data, rate, freq=100):
-    sos = butter(2, freq, "lowpass", fs=rate, output="sos")
-    envelope = np.sqrt(2) * sosfiltfilt(sos, np.abs(data))
-    return envelope
+def envelope(data: np.ndarray, rate: float, freq: float) -> np.ndarray:
+    """Calculate the envelope of a signal using a lowpass filter.
 
+    Parameters
+    ----------
+    data : np.ndarray
+        The signal to calculate the envelope of
+    rate : float
+        The sampling rate of the signal
+    freq : float
+        The cutoff frequency of the lowpass filter
 
-def create_chirp(
-    eodf=500,
-    chirpsize=100,
-    chirpduration=0.015,
-    ampl_reduction=0.05,
-    chirptimes=[0.05, 0.2],
-    kurtosis=1.0,
-    duration=1.0,
-    dt=0.00001,
-):
-    """create a fake fish eod that contains chirps at the given times. EOF is a simple sinewave. Chirps are modeled with Gaussian profiles in amplitude reduction and frequency ecxcursion.
-
-    Args:
-        eodf (int, optional): The chriping fish's EOD frequency. Defaults to 500 Hz.
-        chirpsize (int, optional): the size of the chrip's frequency excursion. Defaults to 100 Hz.
-        chirpwidth (float, optional): the duration of the chirp. Defaults to 0.015 s.
-        ampl_reduction (float, optional): Amount of amplitude reduction during the chrips. Defaults to 0.05, i.e. 5\%
-        chirptimes (list, optional): Times of chirp centers. Defaults to [0.05, 0.2].
-        kurtosis (float, optional): The kurtosis of the Gaussian profiles. Defaults to 1.0
-        dt (float, optional): the stepsize of the simulation. Defaults to 0.00001 s.
-
-    Returns:
-        np.ndarray: the time
-        np.ndarray: the eod
-        np.ndarray: the amplitude profile
-        np.adarray: tha frequency profile
+    Returns
+    -------
+    np.ndarray
+        The envelope of the signal
     """
-    p = 0.0
-
-    time = np.arange(0.0, duration, dt)
-    signal = np.zeros_like(time)
-    ampl = np.ones_like(time)
-    freq = np.ones_like(time)
-
-    ck = 0
-    csig = 0.5 * chirpduration / np.power(2.0 * np.log(10.0), 0.5 / kurtosis)
-    #csig = csig*-1
-    for k, t in enumerate(time):
-        a = 1.0
-        f = eodf
-
-        if ck < len(chirptimes):
-            if np.abs(t - chirptimes[ck]) < 2.0 * chirpduration:
-                x = t - chirptimes[ck]
-                gg = np.exp(-0.5 * np.power((x / csig) ** 2, kurtosis))
-                cc = chirpsize * gg
-
-                # g = np.exp( -0.5 * (x/csig)**2 )
-                f = chirpsize * gg + eodf
-                a *= 1.0 - ampl_reduction * gg
-            elif t > chirptimes[ck] + 2.0 * chirpduration:
-                ck += 1
-        freq[k] = f
-        ampl[k] = a
-        p += f * dt
-        signal[k] = -1 * a * np.sin(2 * np.pi * p)
-
-    return time, signal, ampl, freq
-
+    sos = butter(2, freq, "lowpass", fs=rate, output="sos")
+    envelope = np.sqrt(2) * sosfiltfilt(sos, np.abs(data))
+    return envelope
diff --git a/code/modules/simulations.py b/code/modules/simulations.py
new file mode 100644
index 0000000..473bac8
--- /dev/null
+++ b/code/modules/simulations.py
@@ -0,0 +1,61 @@
+import numpy as np
+
+
+def create_chirp(
+    eodf=500,
+    chirpsize=100,
+    chirpduration=0.015,
+    ampl_reduction=0.05,
+    chirptimes=[0.05, 0.2],
+    kurtosis=1.0,
+    duration=1.0,
+    dt=0.00001,
+):
+    """create a fake fish eod that contains chirps at the given times. EOF is a simple sinewave. Chirps are modeled with Gaussian profiles in amplitude reduction and frequency ecxcursion.
+
+    Args:
+        eodf (int, optional): The chriping fish's EOD frequency. Defaults to 500 Hz.
+        chirpsize (int, optional): the size of the chrip's frequency excursion. Defaults to 100 Hz.
+        chirpwidth (float, optional): the duration of the chirp. Defaults to 0.015 s.
+        ampl_reduction (float, optional): Amount of amplitude reduction during the chrips. Defaults to 0.05, i.e. 5\%
+        chirptimes (list, optional): Times of chirp centers. Defaults to [0.05, 0.2].
+        kurtosis (float, optional): The kurtosis of the Gaussian profiles. Defaults to 1.0
+        dt (float, optional): the stepsize of the simulation. Defaults to 0.00001 s.
+
+    Returns:
+        np.ndarray: the time
+        np.ndarray: the eod
+        np.ndarray: the amplitude profile
+        np.adarray: tha frequency profile
+    """
+    p = 0.0
+
+    time = np.arange(0.0, duration, dt)
+    signal = np.zeros_like(time)
+    ampl = np.ones_like(time)
+    freq = np.ones_like(time)
+
+    ck = 0
+    csig = 0.5 * chirpduration / np.power(2.0 * np.log(10.0), 0.5 / kurtosis)
+    #csig = csig*-1
+    for k, t in enumerate(time):
+        a = 1.0
+        f = eodf
+
+        if ck < len(chirptimes):
+            if np.abs(t - chirptimes[ck]) < 2.0 * chirpduration:
+                x = t - chirptimes[ck]
+                gg = np.exp(-0.5 * np.power((x / csig) ** 2, kurtosis))
+                cc = chirpsize * gg
+
+                # g = np.exp( -0.5 * (x/csig)**2 )
+                f = chirpsize * gg + eodf
+                a *= 1.0 - ampl_reduction * gg
+            elif t > chirptimes[ck] + 2.0 * chirpduration:
+                ck += 1
+        freq[k] = f
+        ampl[k] = a
+        p += f * dt
+        signal[k] = -1 * a * np.sin(2 * np.pi * p)
+
+    return time, signal, ampl, freq