working xdf import with markers

Author: patricefortin

hypyp/xdf/xdf_import.py

@@ -1,9 +1,10 @@
-from typing import List
+from typing import List, Dict
 import warnings
 
 import pyxdf
+import numpy as np
 
-from .xdf_stream import XDFStream
+from .xdf_stream import XDFStream, Markers
 
 class XDFImport():
     """
@@ -28,7 +29,7 @@ def __init__(self,
         select_type: str = None,
         select_matches: list = None,
         mne_type_map: dict = None,
-        scale: float | str = None,
+        scale: float | str | None = None,
         verbose: bool = False,
         convert_to_mne: bool = True,
         ):
@@ -73,22 +74,21 @@ def raws(self):
         return [stream.raw for stream in self.selected_data_streams]
 
     @property
-    def annotations_dict(self):
+    def markers_dict(self) -> Dict[str, Markers]:
         ret = dict()
         for stream in self.selected_markers_streams:
-            ret[stream.name] = stream.annotations
+            ret[stream.name] = stream.markers
         return ret
 
     @property
-    def annotations(self):
-        return [stream.annotations for stream in self.selected_markers_streams]
-    
-    @property
-    def annotations_flat(self):
-        ret = []
-        for x in self.annotations:
-            ret += x    
-        return ret
+    def markers(self) -> Markers:
+        markers = None
+        for stream in self.selected_markers_streams:
+            if markers is None:
+                markers = stream.markers
+            else:
+                markers += stream.markers
+        return markers
 
     @property
     def selected_streams(self):
@@ -200,6 +200,11 @@ def convert_streams_to_mne(self):
                 warnings.warn("Multiple streams have the same name. Adding original stream_id as suffixes to the generated raws") 
 
             stream.convert_to_mne(self.scale, append_stream_id=has_duplicate_names)
+        
+        # add annotations from markers streams to data streams
+        for data_stream in self.selected_data_streams:
+            if self.markers is not None:
+                data_stream.raw.set_annotations(self.markers.as_mne_annotations_for_stream(data_stream.stream))
 
         if self.verbose:
             print("Convertion done.")

hypyp/xdf/xdf_stream.py

@@ -1,16 +1,49 @@
 from typing import List
 import os
+import warnings
+from dataclasses import dataclass
 
 import mne
 import numpy as np
 from pyxdf.pyxdf import StreamData
 
+@dataclass
+class Markers():
+    timestamps: np.array
+    durations: np.array
+    descriptions: List[str]
+
+    def as_mne_annotations(self, reference_time: float = 0.0):
+        onset_times = np.array(self.timestamps) - reference_time
+        empty = mne.Annotations(onset=[], duration=[], description=[])
+        if len(onset_times) == 0:
+            return empty
+
+        i = 0
+        while onset_times[i] < 0:
+            i += 1
+            if i == len(onset_times):
+                return empty
+                
+        return mne.Annotations(onset=onset_times[i:], duration=self.durations[i:], description=self.descriptions[i:])
+        
+    def as_mne_annotations_for_stream(self, stream: StreamData):
+        reference_time = stream['time_stamps'][0]
+        return self.as_mne_annotations(reference_time)
+    
+    def __add__(self, other):
+        if not isinstance(other, Markers):
+            return NotImplemented
+        timestamps = np.concatenate([self.timestamps, self.timestamps])
+        durations = np.concatenate([self.durations, self.durations])
+        descriptions = self.descriptions + other.descriptions
+        return Markers(timestamps, durations, descriptions)
+
 class XDFStream():
     stream: StreamData
     metadata_desc: dict
     raw: mne.io.RawArray | None
-    annotations: mne.Annotations | None
-    time_offset: float
+    markers: mne.Annotations | None
 
     @staticmethod
     def stream_type_to_mne_ch_type(stream_type: str, type_map: dict = None):
@@ -64,12 +97,12 @@ def __init__(self, stream, mne_type_map: dict = None):
         self.mne_type_map = mne_type_map
         self.metadata_desc = {}
         self.raw = None
-        self.annotations = None
+
         if self.stream["info"]["desc"] and self.stream["info"]["desc"][0]:
             self.metadata_desc = self.stream["info"]["desc"][0]
 
-        self.time_offset = float(self.stream['info']['created_at'][0])
-        
+        if self.is_markers_compatible:
+            self.init_markers()
 
     @property
     def name(self):
@@ -87,6 +120,10 @@ def id(self):
     def srate(self):
         return float(self.stream["info"]["nominal_srate"][0])
 
+    @property
+    def is_markers_compatible(self):
+        return self.srate == 0.0
+        
     @property
     def is_mne_compatible(self):
         return self.is_mne_raw_compatible or self.is_mne_annotations_compatible
@@ -99,7 +136,7 @@ def is_mne_raw_compatible(self):
 
     @property
     def is_mne_annotations_compatible(self):
-        return self.srate == 0.0
+        return self.is_markers_compatible
 
     @property
     def time_series(self):
@@ -143,27 +180,16 @@ def ch_names(self):
     def ch_types(self):
         return XDFStream.get_mne_ch_types(self.type, self.ch_names, self.mne_type_map)
 
-    def create_mne_annotations(self):
-        if self.type == 'Markers':
-            timestamps = self.stream['time_stamps']  # Time stamps of the markers
-            descriptions = self.stream['time_series']  # Marker descriptions (the event labels)
-
-            # Ensure timestamps are in seconds (if they aren't already)
-            timestamps = np.array(timestamps) - self.time_offset
-            onset_times = timestamps
-            duration = np.zeros_like(onset_times)
+    def init_markers(self):
+        assert self.is_markers_compatible
 
-            # Create the description for each marker (can be customized, here we use the marker text)
-            description = [str(desc[0]) for desc in descriptions]  # Flatten description if necessary
+        timestamps = np.array(self.stream['time_stamps'])
+        descriptions = np.array(self.stream['time_series'])
+        durations = np.zeros_like(timestamps)
 
-            # Create the annotations object
-            self.annotations = mne.Annotations(onset=onset_times, duration=duration, description=description)
-            # TODO debug time stamps
-            # print("AAAAAAAAAAAAAAAAAAAAAAAAAaaa")
-            # print(onset_times)
+        descriptions = [str(desc[0]) for desc in descriptions]  # Flatten description
 
-        else:
-            self.markers = None
+        self.markers = Markers(timestamps, durations, descriptions)
 
     def get_unique_name(self, append_stream_id: bool = False):
         unique_stream_name = self.name
@@ -172,15 +198,17 @@ def get_unique_name(self, append_stream_id: bool = False):
         return unique_stream_name
 
 
-    def convert_to_mne(self, scale: float | str | None, append_stream_id: bool = False, verbose: bool = False):
+    def convert_to_mne(
+        self,
+        scale: float | str | None,
+        append_stream_id: bool = False,
+        verbose: bool = False
+    ):
         self.unique_name = self.get_unique_name(append_stream_id=append_stream_id)
         if self.is_mne_raw_compatible:
             mne_info = self.create_mne_info(verbose=verbose)
             self.raw = self.create_mne_raw(mne_info, scale, verbose=verbose)
 
-        if self.is_mne_annotations_compatible:
-            self.create_mne_annotations()
-        
     def create_mne_info(self, verbose: bool = False):
         """
         Create a mne.info object from the XDF's EEG stream metadata.

tests/test_xdf.py

@@ -189,21 +189,41 @@ def test_subject_unique_name():
 
 def test_markers():
     xdf = XDFImport(file_path_with_markers)
-    assert xdf.selected_markers_streams[0].time_offset >= 0
-    keys = list(xdf.annotations_dict)
+    keys = list(xdf.markers_dict)
     assert len(keys) == 1
     assert keys[0] == XDFImport(file_path_with_markers, select_type='Markers').selected_stream_names[0]
-    assert len(list(xdf.annotations_dict.values())[0]) > 0
-    n_annotations_0 = len(list(xdf.annotations_dict.values())[0])
-    assert len(xdf.annotations[0]) == n_annotations_0
-    assert len(xdf.annotations_flat) == n_annotations_0
+    markers = list(xdf.markers_dict.values())[0]
+    annotations = markers.as_mne_annotations()
+    assert len(annotations) > 0
 
-    xdf.raws[0].set_annotations(xdf.annotations[0])
+def test_merge_markers():
+    xdf = XDFImport(file_path_with_markers)
+    annotations = xdf.markers.as_mne_annotations()
+    markers_dup = xdf.markers + xdf.markers
+    assert len(markers_dup.as_mne_annotations()) == 2 * len(annotations)
+
+def test_markers_timing():
+    xdf = XDFImport(file_path_with_markers)
+    markers = xdf.markers
+    offset = markers.timestamps[0]
+    annotations = markers.as_mne_annotations(offset)
+    assert annotations[0]['onset'] == 0.0
+
+    annotations_for_stream = markers.as_mne_annotations_for_stream(xdf.selected_data_streams[0].stream)
+    assert annotations_for_stream[0]['onset'] > 0
+    assert annotations_for_stream[0]['onset'] < 10
+    xdf.raws[0].set_annotations(annotations_for_stream)
 
-def test_time_alignment():
+def test_markers_negative():
     xdf = XDFImport(file_path_with_markers)
-    assert xdf.selected_markers_streams[0].time_offset >= 0
-    assert 'TODO: Need to check that the time of annotations is correct when adding them to the RawArray' == True
+    annotations_all = xdf.markers.as_mne_annotations()
+    first_time = xdf.markers.timestamps[0]
+    # let's say a data stream started AFTER the first marker, we want to exclude the first markers
+    annotations_partial = xdf.markers.as_mne_annotations(first_time + 2)
+    assert len(annotations_partial) < len(annotations_all)
+
+    # have a big time offset to make sure we receive no annotations
+    assert len(xdf.markers.as_mne_annotations(first_time + 2000)) == 0
 
 def test_fif_files():
     xdf = XDFImport(file_path, select_type='EEG')
@@ -218,8 +238,3 @@ def test_montage():
     xdf = XDFImport(file_path, select_type='EEG')
     xdf.rename_channels(['Fp1', 'Fp2', 'F3', 'F4', 'C3', 'C4', 'P3', 'P4', 'O1', 'O2', 'F7', 'F8', 'T7', 'T8', 'P7', 'P8', 'Fz', 'Cz', 'Pz', 'POz', 'FC1', 'FC2', 'CP1', 'CP2', 'FC5', 'FC6', 'CP5', 'CP6', 'FT9', 'FT10', 'TP9', 'TP10'])
     xdf.set_montage('standard_1020')
-
-
-
-##
-### test have markers in raw