Done with sleep staging but needs checking

portiloop/src/demo/offline.py

@@ -2,7 +2,7 @@ import numpy as np
 from portiloop.src.detection import SleepSpindleRealTimeDetector
 from portiloop.src.stimulation import UpStateDelayer
 from portiloop.src.processing import FilterPipeline
-from portiloop.src.demo.utils import compute_output_table, xdf2array, offline_detect, offline_filter, OfflineSleepSpindleRealTimeStimulator
+from portiloop.src.demo.utils import compute_output_table, sleep_stage, xdf2array, offline_detect, offline_filter, OfflineSleepSpindleRealTimeStimulator
 import gradio as gr
 
 
@@ -29,6 +29,7 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq, stim
     # Read the xdf file to a numpy array
     print("Loading xdf file...")
     data_whole, columns = xdf2array(xdf_file.name, int(channel_num))
+
     # Do the offline filtering of the data
     if offline_filtering:
         print("Filtering offline...")
@@ -38,13 +39,18 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq, stim
         data_whole = np.concatenate((data_whole, offline_filtered_data), axis=1)
         columns.append("offline_filtered_signal")
 
+    # Do the sleep staging approximation
+    if wamsley or lacourse:
+        print("Sleep staging...")
+        mask = sleep_stage(data_whole[:, columns.index("offline_filtered_signal")], threshold=150, group_size=100)
+
     #  Do Wamsley's method
     if wamsley:
         print("Running Wamsley detection...")
         wamsley_data = offline_detect("Wamsley", \
             data_whole[:, columns.index("offline_filtered_signal")],\
                 data_whole[:, columns.index("time_stamps")],\
-                    freq)
+                    freq, mask)
         wamsley_data = np.expand_dims(wamsley_data, axis=1)
         data_whole = np.concatenate((data_whole, wamsley_data), axis=1)
         columns.append("wamsley_spindles")
@@ -55,7 +61,7 @@ def run_offline(xdf_file, detect_filter_opts, threshold, channel_num, freq, stim
         lacourse_data = offline_detect("Lacourse", \
             data_whole[:, columns.index("offline_filtered_signal")],\
                 data_whole[:, columns.index("time_stamps")],\
-                    freq)
+                    freq, mask)
         lacourse_data = np.expand_dims(lacourse_data, axis=1)
         data_whole = np.concatenate((data_whole, lacourse_data), axis=1)
         columns.append("lacourse_spindles")

portiloop/src/demo/test_offline.py

@@ -4,6 +4,8 @@ from portiloop.src.demo.offline import run_offline
 from pathlib import Path
 import matplotlib.pyplot as plt
 
+from portiloop.src.demo.utils import sleep_stage, xdf2array
+
 class TestOffline(unittest.TestCase):
 
     def setUp(self):
@@ -21,7 +23,7 @@ class TestOffline(unittest.TestCase):
         all_options_iterator = itertools.product(*map(combinatorial_config.get, keys))
         all_options_dicts = [dict(zip(keys, values)) for values in all_options_iterator]
         self.filtered_options = [value for value in all_options_dicts if (value['online_detection'] and value['online_filtering']) or not value['online_detection']]
-        self.xdf_file = Path(__file__).parents[3] / "test_xdf" / "test_file.xdf"
+        self.xdf_file = Path(__file__).parents[3] / "test_file.xdf"
 
 
     def test_all_options(self):
@@ -32,7 +34,7 @@ class TestOffline(unittest.TestCase):
     def test_single_option(self):
 
         # Test options correspond to an index in the possible checkbox group options
-        test_options = [3, 4]
+        test_options = [0, 1, 2]
 
         res = list(run_offline(
                 self.xdf_file,
@@ -43,6 +45,7 @@ class TestOffline(unittest.TestCase):
                 stimulation_phase="Peak",
                 buffer_time=0.3))
         print(res)
+        pass
 
     def tearDown(self):
         pass

portiloop/src/demo/utils.py

@@ -13,6 +13,32 @@ STREAM_NAMES = {
 }
 
 
+def sleep_stage(data, threshold=150, group_size=2):
+    """Sleep stage approximation using a threshold and a group size.
+        Returns a numpy array containing all indices in the input data which CAN be used for offline detection. 
+        These indices can then be used to reconstruct the signal from the original data.
+    """
+    # Find all indexes where the signal is above or below the threshold
+    above = np.where(data > threshold)
+    below = np.where(data < -threshold)
+    indices = np.concatenate((above, below), axis=1)[0]
+
+    indices = np.sort(indices)
+    # Get all the indices where the difference between two consecutive indices is larger than 100
+    groups = np.where(np.diff(indices) <= group_size)[0] + 1
+    # Get the important indices
+    important_indices = indices[groups]
+    # Get all the indices between the important indices
+    group_filler = [np.arange(indices[groups[n] - 1] + 1, index) for n, index in enumerate(important_indices)]
+    # Create flat array from fillers
+    group_filler = np.concatenate(group_filler)
+    # Append all group fillers to the indices
+    masked_indices = np.sort(np.concatenate((indices, group_filler)))
+    unmasked_indices = np.setdiff1d(np.arange(len(data)), masked_indices)
+
+    return unmasked_indices
+
+
 class OfflineSleepSpindleRealTimeStimulator(Stimulator):
     def __init__(self):
         self.last_detected_ts = time.time()
@@ -87,15 +113,19 @@ def xdf2array(xdf_path, channel):
     return np.array(csv_list), columns
     
 
-def offline_detect(method, data, timesteps, freq):
+def offline_detect(method, data, timesteps, freq, mask):
+    # Extract only the interesting elements from the mask
+    data_masked = data[mask]
+
     # Get the spindle data from the offline methods
     time = np.arange(0, len(data)) / freq
+    time_masked = time[mask] 
     if method == "Lacourse":
         detector = DetectSpindle(method='Lacourse2018')
-        spindles, _, _ = detect_Lacourse2018(data, freq, time, detector)
+        spindles, _, _ = detect_Lacourse2018(data_masked, freq, time_masked, detector)
     elif method == "Wamsley":
         detector = DetectSpindle(method='Wamsley2012')
-        spindles, _, _ = detect_Wamsley2012(data, freq, time, detector)
+        spindles, _, _ = detect_Wamsley2012(data_masked, freq, time_masked, detector)
     else:
         raise ValueError("Invalid method")
 
@@ -155,3 +185,4 @@ def compute_output_table(online_stimulation, lacourse_spindles, wamsley_spindles
     if wamsley_spindles is not None:
         table += f"| Wamsley | {wamsley_spindles_count} | {both_online_wamsley} |\n"
     return table
+