Merge pull request #14 from int-brain-lab/peaktrough_waveforms

wav peak / trough inversion in near cases

src/neurodsp/waveforms.py

@@ -119,25 +119,30 @@ def find_peak(arr_in):
     return df
 
 
-def find_tip_trough(arr_peak, df):
-    '''
-    :param arr_in: inverted
-    :param df:
-    :return:
-    '''
-    # 2. Find trough and tip (at peak waveform)
+def find_trough(arr_peak, df):
+    # Find tip (at peak waveform)
+
     # Create masks pre/post
     arr_pre, arr_post = arr_pre_post(arr_peak, df['peak_time_idx'].to_numpy())
 
     # Find trough
     # indx_trough = np.nanargmin(arr_post * np.sign(val_peak)[:, np.newaxis], axis=1)
     indx_trough = np.nanargmax(arr_post, axis=1)
     val_trough = arr_peak[np.arange(0, arr_peak.shape[0], 1), indx_trough] * df['invert_sign_peak'].to_numpy()
-    del arr_post
-    # TODO spin function above into a function
-    #  if ratio or diff peak/trough smaller than x AND time diff is smaller than xx :
-    #  Assign trough as peak
-    #  Call this function again to compute trough with new peak assigned
+
+    # Put values into df
+    df['trough_time_idx'] = indx_trough
+    df['trough_val'] = val_trough
+
+    return df
+
+
+def find_tip(arr_peak, df):
+    # Find tip (at peak waveform)
+
+    # Create masks pre/post
+    arr_pre, arr_post = arr_pre_post(arr_peak, df['peak_time_idx'].to_numpy())
+
     # Find tip
     '''
     # 02-06-2023 ; Decided not to use the inflection point but rather maximum
@@ -156,18 +161,60 @@ def find_tip_trough(arr_peak, df):
     # Maximum
     indx_tip = np.nanargmax(arr_pre, axis=1)
     val_tip = arr_peak[np.arange(0, arr_peak.shape[0], 1), indx_tip] * df['invert_sign_peak'].to_numpy()
-    del arr_pre
 
     # Put values into df
-    df['trough_time_idx'] = indx_trough
-    df['trough_val'] = val_trough
-
     df['tip_time_idx'] = indx_tip
     df['tip_val'] = val_tip
 
     return df
 
 
+def find_tip_trough(arr_peak, arr_peak_real, df):
+    '''
+    :param arr_in: inverted
+    :param df:
+    :return:
+    '''
+    # 2. Find trough and tip (at peak waveform)
+
+    # Find trough
+    df = find_trough(arr_peak, df)
+    df = peak_to_trough_ratio(df)
+    # If ratio of peak/trough is near 1, and peak is positive :
+    # Assign trough as peak on same waveform channel
+    # Call the function again to compute trough etc. with new peak assigned
+
+    # Find df rows to be changed
+    df_index = df.index[(df['peak_val'] > 0) & (df['peak_to_trough_ratio'] <= 1.2)]
+    df_rows = df.iloc[df_index]
+    if len(df_index) > 0:
+        # New peak - Swap peak for trough values
+        df_rows = df_rows.drop(['peak_val', 'peak_time_idx'], axis=1)
+        df_rows['peak_val'] = df_rows['trough_val']
+        df_rows['peak_time_idx'] = df_rows['trough_time_idx']
+
+        # df_trials.loc[iss, f] = predicted[f].values
+
+        # Drop trough columns
+        df_rows = df_rows.drop(['trough_time_idx', 'trough_val'], axis=1)
+        # Create mini arr_peak for those rows uniquely (take the real waveforms value in, not inverted ones)
+        arr_peak_rows = arr_peak_real[df_index, :]
+        # Place into "inverted" array peak for return
+        arr_peak[df_index, :] = arr_peak_rows
+        # Get new sign for the peak
+        arr_peak_rows, df_rows = invert_peak_waveform(arr_peak_rows, df_rows)
+        # New trough
+        df_rows = find_trough(arr_peak_rows, df_rows)
+        # New peak-trough ratio
+        df_rows = peak_to_trough_ratio(df_rows)
+        # Assign back into the dataframe
+        df.loc[df_index] = df_rows
+    # Find tip
+    df = find_tip(arr_peak, df)
+
+    return df, arr_peak
+
+
 def plot_peaktiptrough(df, arr, ax, nth_wav=0, plot_grey=True):
     if plot_grey:
         ax.plot(arr[nth_wav], c='gray', alpha=0.5)
@@ -239,17 +286,29 @@ def half_peak_duration(df, fs=30000):
     return df
 
 
-def peak_to_trough(df, fs=30000):
+def peak_to_trough_duration(df, fs=30000):
     '''
-    Compute the duration (second) and ratio of the peak-to-trough
+    Compute the duration (second) of the peak-to-trough
     :param df: dataframe of waveforms features
     :param fs: sampling rate (Hz)
-    :return:
+    :return: df
     '''
     # Duration
     df['peak_to_trough_duration'] = (df['trough_time_idx'] - df['peak_time_idx']) / fs
+    return df
+
+
+def peak_to_trough_ratio(df):
+    '''
+    Compute the ratio of the peak-to-trough
+    :param df: dataframe of waveforms features
+    :param fs: sampling rate (Hz)
+    :return:
+    '''
     # Ratio
-    df['peak_to_trough_ratio'] = np.log(np.abs(df['peak_val'] / df['trough_val']))
+    df['peak_to_trough_ratio'] = np.abs(df['peak_val'] / df['trough_val'])  # Division by 0 returns NaN
+    # Ratio log-scale
+    df['peak_to_trough_ratio_log'] = np.log(df['peak_to_trough_ratio'])
     return df
 
 
@@ -418,11 +477,13 @@ def compute_spike_features(arr_in, fs=30000, recovery_duration_ms=0.16, return_p
     """
     df = find_peak(arr_in)
     # Per waveform, keep only trace that contains the peak
-    arr_peak = get_array_peak(arr_in, df)
+    arr_peak_real = get_array_peak(arr_in, df)
     # Invert positive spikes
-    arr_peak, df = invert_peak_waveform(arr_peak, df)
-    # Tip-trough
-    df = find_tip_trough(arr_peak, df)
+    arr_peak, df = invert_peak_waveform(arr_peak_real.copy(), df)  # Copy otherwise overwrite the variable in memory
+    # Tip-trough (this also computes the peak_to_trough_ratio)
+    df, arr_peak = find_tip_trough(arr_peak, arr_peak_real, df)
+    # Peak to trough duration
+    df = peak_to_trough_duration(df, fs=30000)
     # Half peak points
     df = half_peak_point(arr_peak, df)
     # Half peak duration
@@ -432,7 +493,8 @@ def compute_spike_features(arr_in, fs=30000, recovery_duration_ms=0.16, return_p
     # Slopes
     df = polarisation_slopes(df, fs=fs)
     df = recovery_slope(df, fs=fs)
+
     if return_peak_channel:
-        return df, arr_peak
+        return df, arr_peak_real
     else:
         return df

src/tests/unit/cpu/fixtures/waveform_sample/test_df.csv

@@ -1,5 +1,7 @@
-,peak_trace_idx,peak_time_idx,peak_val,invert_sign_peak,trough_time_idx,trough_val,tip_time_idx,tip_val,half_peak_post_time_idx,half_peak_pre_time_idx,half_peak_post_val,half_peak_pre_val,half_peak_duration,recovery_time_idx,recovery_val,depolarisation_slope,repolarisation_slope,recovery_slope
-0,20,44,-4.883662,1.0,54,1.2735074,24,0.77733666,47,40,-1.9125216,-1.7436237,0.00023333333333333333,59,0.98473907,-8491.498231887817,18471.508026123047,-1732.609748840332
-1,20,45,5.153482,-1.0,69,-1.0805972,41,-2.9835296,49,43,2.1435728,2.3065486,0.0002,74,-0.9415057,61027.586460113525,-7792.599201202393,834.5489501953125
-2,20,44,-4.48732,1.0,54,1.5353007,39,0.71306074,48,42,-1.6273819,-2.058325,0.0002,59,0.7605876,-31202.284812927246,18067.862033843994,-4648.2789516448975
-3,20,40,1.8537197,-1.0,61,-0.49219757,23,-0.66024184,46,38,0.7720857,0.6908671,0.0002666666666666667,66,-0.35654825,4436.402741600485,-3351.310321262905,813.8959407806396
+,peak_trace_idx,peak_time_idx,peak_val,invert_sign_peak,trough_time_idx,trough_val,peak_to_trough_ratio,peak_to_trough_ratio_log,tip_time_idx,tip_val,peak_to_trough_duration,half_peak_post_time_idx,half_peak_pre_time_idx,half_peak_post_val,half_peak_pre_val,half_peak_duration,recovery_time_idx,recovery_val,depolarisation_slope,repolarisation_slope,recovery_slope
+0,20,44,-4.883662,1.0,54,1.2735074,3.8348126,1.3441206,24,0.77733666,0.0003333333333333333,47,40,-1.9125216,-1.7436237,0.00023333333333333333,59,0.98473907,-8491.498231887817,18471.508026123047,-1732.609748840332
+1,20,45,5.153482,-1.0,69,-1.0805972,4.7691054,1.5621588,41,-2.9835296,0.0008,49,43,2.1435728,2.3065486,0.0002,74,-0.9415057,61027.586460113525,-7792.599201202393,834.5489501953125
+2,20,44,-4.48732,1.0,54,1.5353007,2.9227629,1.0725293,39,0.71306074,0.0003333333333333333,48,42,-1.6273819,-2.058325,0.0002,59,0.7605876,-31202.284812927246,18067.862033843994,-4648.2789516448975
+3,20,40,1.8537197,-1.0,61,-0.49219757,3.7662106,1.3260694,23,-0.66024184,0.0007,46,38,0.7720857,0.6908671,0.0002666666666666667,66,-0.35654825,4436.402741600485,-3351.310321262905,813.8959407806396
+4,19,48,-2.7637115,1.0,65,1.1901706,2.3221135,0.8424778,38,3.0702133,0.0005666666666666667,54,45,-1.362757,-1.3184559,0.0003,70,0.39087203,-17501.774311065674,6977.4392071892225,-4795.791864395142
+5,19,45,-2.240881,1.0,63,0.5131445,4.366959,1.474067,39,2.18663,0.0006,50,43,-0.93143564,-0.93315697,0.00023333333333333333,68,0.26503378,-22137.556076049805,4590.042432149252,-1488.6642694473267

src/tests/unit/cpu/fixtures/waveform_sample/test_df_wavinfo.csv

@@ -3,3 +3,5 @@
 1,T00500,da8dfec1-d265-44e8-84ce-6ae9c109b8bd,11
 0,T00500,1a276285-8b0e-4cc9-9f0a-a3a002978724,1
 1,T00500,1a276285-8b0e-4cc9-9f0a-a3a002978724,2
+0,T00500,5d570bf6-a4c6-4bf1-a14b-2c878c84ef0e,10
+1,T00500,5d570bf6-a4c6-4bf1-a14b-2c878c84ef0e,11