bci_essentials.data_tank.data_tank

  1import numpy as np
  2from ..utils.logger import Logger  # Logger wrapper
  3
  4# Instantiate a logger for the module at the default level of logging.INFO
  5# Logs to bci_essentials.__module__) where __module__ is the name of the module
  6logger = Logger(name=__name__)
  7
  8
  9# Will eventually move somewhere else
 10class DataTank:
 11    """
 12    DataTank class is for storing raw EEG, markers, epochs, and rtesting state data
 13
 14    TODO: Add your desired flavour of save output from here.
 15    To be added:
 16    - MNE Raw
 17    - MNE Epochs
 18    - BIDS
 19    - XDF
 20    """
 21
 22    def __init__(self):
 23        """
 24        Initialize the DataTank.
 25        """
 26
 27        # Initialize np arrays to store the data
 28        self.__raw_eeg = np.zeros((0, 0))
 29        self.__raw_eeg_timestamps = np.zeros((0))
 30        self.__raw_marker_strings = np.zeros((0), dtype=str)
 31        self.__raw_marker_timestamps = np.zeros((0))
 32        # self.event_marker_strings = np.zeros((0), dtype=str)
 33        # self.event_marker_timestamps = np.zeros((0))
 34
 35        # Keep track of the latest timestamp
 36        self.latest_eeg_timestamp = 0
 37
 38        # Keep track of how many epochs have been sent, so it is possible to only send new ones
 39        self.epochs_sent = 0
 40        self.epochs = np.zeros((0, 0))
 41        self.labels = np.array([], dtype=str)
 42
 43    def set_source_data(
 44        self, headset_string, fsample, n_channels, ch_types, ch_units, channel_labels
 45    ):
 46        """
 47        Set the source data for the DataTank so that this metadata can be saved with the data.
 48
 49        Parameters
 50        ----------
 51
 52        headset_string : str
 53            The name of the headset used to collect the data.
 54        fsample : float
 55            The sampling frequency of the data.
 56        n_channels : int
 57            The number of channels in the data.
 58        ch_types : list of str
 59            The type of each channel.
 60        ch_units : list of str
 61            The units of each channel.
 62        channel_labels : list of str
 63            The labels of each channel.
 64
 65        Returns
 66        -------
 67        `None`
 68        """
 69        self.headset_string = headset_string
 70        self.fsample = fsample
 71        self.n_channels = n_channels
 72        self.ch_types = ch_types
 73        self.ch_units = ch_units
 74        self.channel_labels = channel_labels
 75
 76    def add_raw_eeg(self, new_raw_eeg, new_eeg_timestamps):
 77        """
 78        Add raw EEG data to the data tank.
 79
 80        Parameters
 81        ----------
 82        new_raw_eeg : np.array
 83            The new raw EEG data to add.
 84
 85        new_eeg_timestamps : np.array
 86            The timestamps of the new raw EEG data.
 87
 88        Returns
 89        -------
 90        `None`
 91        """
 92
 93        # If this is the first chunk of EEG, initialize the arrays
 94        if self.__raw_eeg.size == 0:
 95            self.__raw_eeg = new_raw_eeg
 96            self.__raw_eeg_timestamps = new_eeg_timestamps
 97        else:
 98            self.__raw_eeg = np.concatenate((self.__raw_eeg, new_raw_eeg), axis=1)
 99            self.__raw_eeg_timestamps = np.concatenate(
100                (self.__raw_eeg_timestamps, new_eeg_timestamps)
101            )
102
103        self.latest_eeg_timestamp = new_eeg_timestamps[-1]
104
105    def add_raw_markers(self, new_marker_strings, new_marker_timestamps):
106        """
107        Add raw markers to the data tank.
108
109        Parameters
110        ----------
111        new_marker_strings : np.array
112            The new marker strings to add.
113        new_marker_timestamps : np.array
114            The timestamps of the new marker strings.
115
116        Returns
117        -------
118        `None`
119        """
120
121        if self.__raw_marker_strings.size == 0:
122            self.__raw_marker_strings = new_marker_strings
123            self.__raw_marker_timestamps = new_marker_timestamps
124        else:
125            self.__raw_marker_strings = np.concatenate(
126                (self.__raw_marker_strings, new_marker_strings)
127            )
128            self.__raw_marker_timestamps = np.concatenate(
129                (self.__raw_marker_timestamps, new_marker_timestamps)
130            )
131
132    def get_raw_eeg(self):
133        """
134        Get the raw EEG data from the DataTank.
135
136        Returns
137        -------
138        np.array
139            The raw EEG data.
140        np.array
141            The timestamps of the raw EEG data.
142        """
143        # Get the EEG data between the start and end times
144        return self.__raw_eeg, self.__raw_eeg_timestamps
145
146    def get_raw_markers(self):
147        """
148        Get the raw markers from the DataTank.
149
150        Returns
151        -------
152        np.array
153            The raw marker strings.
154        np.array
155            The timestamps of the raw marker strings.
156        """
157        return self.__raw_marker_strings, self.__raw_marker_timestamps
158
159    def add_epochs(self, X, y):
160        """
161        Add epochs to the data tank.
162
163        Parameters
164        ----------
165        X : np.array
166            The new epochs to add. Shape is (n_epochs, n_channels, n_samples).
167        y : np.array
168            The labels of the epochs. Shape is (n_epochs).
169
170        Returns
171        -------
172        `None`
173
174        """
175        # Add new epochs to the data tank
176        if self.epochs.size == 0:
177            self.epochs = np.array(X)
178            self.labels = np.array(y)
179
180        else:
181            # Check the size of the new data
182            if X.shape[1:] != self.epochs.shape[1:]:
183                logger.warning("Epochs are not the same size, skipping this data.")
184            else:
185                self.epochs = np.concatenate((self.epochs, np.array(X)))
186                self.labels = np.concatenate((self.labels, np.array(y)))
187
188    def get_epochs(self, latest=False):
189        """
190        Get the epochs from the data tank.
191
192        Parameters
193        ----------
194        latest : bool
195            If `True`, only return the new data since the last call to this function.
196
197        Returns
198        -------
199        np.array
200            The epochs. Shape is (n_epochs, n_channels, n_samples).
201        np.array
202            The labels of the epochs. Shape is (n_epochs).
203
204        """
205        if self.epochs.size == 0:
206            logger.warning("Data tank contains no epochs, returning None.")
207            return [], []
208
209        if latest:
210            # Return only the new data
211            first_unsent = self.epochs_sent
212            self.epochs_sent = len(self.epochs)
213
214            return self.epochs[first_unsent:], self.labels[first_unsent:]
215        else:
216            # Return all
217            return self.epochs, self.labels
218
219    def add_resting_state_data(self, resting_state_data):
220        """
221        Add resting state data to the data tank.
222
223        Parameters
224        ----------
225        resting_state_data : dict
226            Dictionary containing resting state data.
227
228        Returns
229        -------
230        `None`
231
232        """
233        # Get the resting state data
234        self.__resting_state_data = resting_state_data
235
236    def get_resting_state_data(self):
237        """
238        Get the resting state data.
239
240        Returns
241        -------
242        dict
243            Dictionary containing resting state data.
244
245        """
246        return self.__resting_state_data
247
248    def save_epochs_as_npz(self, file_name: str):
249        """
250        Saves EEG trials and labels as a numpy file.
251
252        Parameters
253        ----------
254        file_name : str
255            The name of the file to save the EEG trials and labels to.
256
257        Returns
258        -------
259        `None`
260
261        """
262        # Check if file ends with .npz, if not add it
263        if file_name[-4:] != ".npz":
264            file_name += ".npz"
265
266        # Get the raw EEG trials and labels
267        X = self.epochs
268        y = self.labels
269
270        # Save the raw EEG trials and labels as a numpy file
271        np.savez(file_name, X=X, y=y)