py-neuromodulation 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

py_neuromodulation/nm_linelength.py

@@ -1,33 +1,21 @@
-import numpy as np
-from typing import Iterable
-
-from py_neuromodulation import nm_features_abc
-
-
-class LineLength(nm_features_abc.Feature):
-    def __init__(
-        self, settings: dict, ch_names: Iterable[str], sfreq: float
-    ) -> None:
-        self.s = settings
-        self.ch_names = ch_names
-
-    @staticmethod
-    def get_line_length(x: np.ndarray) -> np.ndarray:
-        return np.mean(np.abs(np.diff(x)) / (x.shape[0] - 1))
-
-    @staticmethod
-    def test_settings(
-        settings: dict,
-        ch_names: Iterable[str],
-        sfreq: int | float,
-    ):
-        # no settings to be checked
-        pass
-
-    def calc_feature(self, data: np.ndarray, features_compute: dict) -> dict:
-        for ch_idx, ch_name in enumerate(self.ch_names):
-            features_compute[
-                "_".join([ch_name, "LineLength"])
-            ] = self.get_line_length(data[ch_idx, :])
-
-        return features_compute
+import numpy as np
+from collections.abc import Sequence
+
+from py_neuromodulation.nm_features import NMFeature
+
+
+class LineLength(NMFeature):
+    def __init__(self, settings: dict, ch_names: Sequence[str], sfreq: float) -> None:
+        self.ch_names = ch_names
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+
+        line_length = np.mean(
+            np.abs(np.diff(data, axis=-1)) / (data.shape[1] - 1), axis=-1
+        )
+
+        feature_results = {}
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            feature_results[f"{ch_name}_LineLength"] = line_length[ch_idx]
+
+        return feature_results

py_neuromodulation/nm_logger.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+from py_neuromodulation.nm_types import _PathLike
+import logging
+
+INFOFORMAT = "%(name)s:\t%(message)s"
+DEBUGFORMAT = "%(asctime)s:%(levelname)s:%(name)s:%(filename)s:%(funcName)s:%(lineno)d:\t%(message)s"
+
+LOG_LEVELS = {
+    "DEBUG": (logging.DEBUG, DEBUGFORMAT),
+    "INFO": (logging.INFO, INFOFORMAT),
+    "WARNING": (logging.WARN, DEBUGFORMAT),
+    "ERROR": (logging.ERROR, DEBUGFORMAT),
+}
+
+
+class NMLogger(logging.Logger):
+    """
+    Subclass of logging.Logger with some extra functionality
+    """
+
+    def __init__(self, name: str, level: str = "INFO") -> None:
+        super().__init__(name, LOG_LEVELS[level][0])
+
+        self.setLevel(level)
+
+        self._console_handler = logging.StreamHandler()
+        self._console_handler.setLevel(level)
+        self._console_handler.setFormatter(logging.Formatter(LOG_LEVELS[level][1]))
+
+        self.addHandler(self._console_handler)
+
+    def set_level(self, level: str):
+        """
+        Set console logging level
+        """
+        self.setLevel(level)
+        self._console_handler.setLevel(level)
+        self._console_handler.setFormatter(logging.Formatter(LOG_LEVELS[level][1]))
+
+    def log_to_file(self, path: _PathLike, mode: str = "w"):
+        """
+        Add file handlers to the logger
+
+        Parameters
+        ----------
+        path: directory where to save logfiles
+        mode : str, ('w', 'a')
+            w: overwrite files
+            a: append to files
+        """
+
+        path = Path(path)
+        path.mkdir(parents=True, exist_ok=True)
+
+        self.debug_file_handler = logging.FileHandler(path / "logfile_pynm_debug.log")
+        self.debug_file_handler.setLevel(logging.DEBUG)
+        self.debug_file_handler.setFormatter(logging.Formatter(DEBUGFORMAT))
+
+        self.info_file_handler = logging.FileHandler(
+            path / "logfile_pynm_info.log", mode=mode
+        )
+        self.info_file_handler.setLevel(logging.INFO)
+        self.info_file_handler.setFormatter(logging.Formatter(INFOFORMAT))
+
+        self.addHandler(self.info_file_handler)
+        self.addHandler(self.debug_file_handler)

py_neuromodulation/nm_mne_connectivity.py

@@ -1,112 +1,149 @@
-from typing import Iterable
-import numpy as np
-
-import mne
-from mne_connectivity import spectral_connectivity_epochs
-
-from py_neuromodulation import nm_features_abc
-
-
-class MNEConnectivity(nm_features_abc.Feature):
-    def __init__(
-        self,
-        settings: dict,
-        ch_names: Iterable[str],
-        sfreq: float,
-    ) -> None:
-        self.s = settings
-        self.ch_names = ch_names
-        self.mode = settings["mne_connectiviy"]["mode"]
-        self.method = settings["mne_connectiviy"]["method"]
-        self.sfreq = sfreq
-
-        self.fbands = list(self.s["frequency_ranges_hz"].keys())
-        self.fband_ranges = []
-
-    @staticmethod
-    def test_settings(
-        settings: dict,
-        ch_names: Iterable[str],
-        sfreq: int | float,
-    ):
-        # TODO: Double check passed parameters with mne_connectivity
-        pass
-
-    @staticmethod
-    def get_epoched_data(
-        raw: mne.io.RawArray, epoch_length: float = 1
-    ) -> np.array:
-        time_samples_s = raw.get_data().shape[1] / raw.info["sfreq"]
-        if epoch_length > time_samples_s:
-             raise ValueError(
-                f"the intended epoch length for mne connectivity: {epoch_length}s"
-                f" are longer than the passed data array {np.round(time_samples_s, 2)}s"
-             )
-        events = mne.make_fixed_length_events(
-            raw, duration=epoch_length, overlap=0
-        )
-        event_id = {"rest": 1}
-
-        epochs = mne.Epochs(
-            raw,
-            events=events,
-            event_id=event_id,
-            tmin=0,
-            tmax=epoch_length,
-            baseline=None,
-            reject_by_annotation=True,
-        )
-        if epochs.events.shape[0] < 2:
-            raise Exception(
-                f"A minimum of 2 epochs is required for mne_connectivity,"
-                f" got only {epochs.events.shape[0]}. Increase settings['segment_length_features_ms']"
-            )
-        return epochs
-
-    def estimate_connectivity(self, epochs: mne.Epochs):
-        # n_jobs is here kept to 1, since setup of the multiprocessing Pool 
-        # takes longer than most batch computing sizes
-
-        spec_out = spectral_connectivity_epochs(
-            data=epochs,
-            sfreq=self.sfreq,
-            n_jobs=1,
-            method=self.method,
-            mode=self.mode,
-            indices=(np.array([0, 0, 1, 1]), np.array([2, 3, 2, 3])),
-            faverage=False,
-            block_size=1000,
-        )
-        return spec_out
-
-    def calc_feature(self, data: np.array, features_compute: dict) -> dict:
-
-        raw = mne.io.RawArray(
-            data=data,
-            info=mne.create_info(ch_names=self.ch_names, sfreq=self.sfreq),
-        )
-        epochs = self.get_epoched_data(raw)
-        # there need to be minimum 2 of two epochs, otherwise mne_connectivity 
-        # is not correctly initialized
-
-        spec_out = self.estimate_connectivity(epochs)
-        if len(self.fband_ranges) == 0:
-            for fband in self.fbands:
-                self.fband_ranges.append(
-                    np.where(
-                        np.logical_and(
-                            np.array(spec_out.freqs)
-                            > self.s["frequency_ranges_hz"][fband][0],
-                            np.array(spec_out.freqs)
-                            < self.s["frequency_ranges_hz"][fband][1],
-                        )
-                    )[0]
-                )
-        dat_conn = spec_out.get_data()
-        for conn in np.arange(dat_conn.shape[0]):
-            for fband_idx, fband in enumerate(self.fbands):
-                features_compute[
-                    "_".join(["ch1", self.method, str(conn), fband])
-                ] = np.mean(dat_conn[conn, self.fband_ranges[fband_idx]])
-
-        return features_compute
+from collections.abc import Iterable
+import numpy as np
+import pandas as pd
+from typing import TYPE_CHECKING
+
+from py_neuromodulation.nm_features import NMFeature
+from py_neuromodulation.nm_types import NMBaseModel
+
+if TYPE_CHECKING:
+    from py_neuromodulation.nm_settings import NMSettings
+    from mne.io import RawArray
+    from mne import Epochs
+
+
+class MNEConnectivitySettings(NMBaseModel):
+    method: str = "plv"
+    mode: str = "multitaper"
+
+
+class MNEConnectivity(NMFeature):
+    def __init__(
+        self,
+        settings: "NMSettings",
+        ch_names: Iterable[str],
+        sfreq: float,
+    ) -> None:
+        from mne import create_info
+
+        self.settings = settings
+
+        self.ch_names = ch_names
+        self.sfreq = sfreq
+
+        # Params used by spectral_connectivity_epochs
+        self.mode = settings.mne_connectivity.mode
+        self.method = settings.mne_connectivity.method
+
+        self.fbands = settings.frequency_ranges_hz
+        self.fband_ranges: list = []
+        self.result_keys = []
+
+        self.raw_info = create_info(ch_names=self.ch_names, sfreq=self.sfreq)
+        self.raw_array: "RawArray"
+        self.epochs: "Epochs"
+        self.prev_batch_shape: tuple = (-1, -1)  # sentinel value
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        from mne.io import RawArray
+        from mne import Epochs
+        from mne_connectivity import spectral_connectivity_epochs
+
+        time_samples_s = data.shape[1] / self.sfreq
+        epoch_length: float = 1  # TODO: Make this a parameter?
+
+        if epoch_length > time_samples_s:
+            raise ValueError(
+                f"the intended epoch length for mne connectivity: {epoch_length}s"
+                f" are longer than the passed data array {np.round(time_samples_s, 2)}s"
+            )
+
+        # Only reinitialize the raw_array and epochs object if the data shape has changed
+        # That could mean that the channels have been re-selected, or we're in the last batch
+        # TODO: If sfreq or channels change, do we re-initialize the whole Stream object?
+        if data.shape != self.prev_batch_shape:
+            self.raw_array = RawArray(
+                data=data,
+                info=self.raw_info,
+                copy=None,  # type: ignore
+                verbose=False,
+            )
+
+            # self.events = make_fixed_length_events(self.raw_array, duration=epoch_length)
+            # Equivalent code for those parameters:
+            event_times = np.arange(
+                0, data.shape[-1], self.sfreq * epoch_length, dtype=int
+            )
+            events = np.column_stack(
+                (
+                    event_times,
+                    np.zeros_like(event_times, dtype=int),
+                    np.ones_like(event_times, dtype=int),
+                )
+            )
+
+            # there need to be minimum 2 of two epochs, otherwise mne_connectivity
+            # is not correctly initialized
+            if events.shape[0] < 2:
+                raise RuntimeError(
+                    f"A minimum of 2 epochs is required for mne_connectivity,"
+                    f" got only {events.shape[0]}. Increase settings['segment_length_features_ms']"
+                )
+
+            self.epochs = Epochs(
+                self.raw_array,
+                events=events,
+                event_id={"rest": 1},
+                tmin=0,
+                tmax=epoch_length,
+                baseline=None,
+                reject_by_annotation=True,
+                verbose=False,
+            )
+
+            # Trick the function "spectral_connectivity_epochs" into not calling "add_annotations_to_metadata"
+            # TODO: This is a hack, and maybe needs a fix in the mne_connectivity library
+            self.epochs._metadata = pd.DataFrame(index=np.arange(events.shape[0]))
+
+        else:
+            # As long as the initialization parameters, channels, sfreq and batch size are the same
+            # We can re-use the existing epochs object by updating the raw data
+            self.raw_array._data = data
+            self.epochs._raw = self.raw_array
+
+        # n_jobs is here kept to 1, since setup of the multiprocessing Pool
+        # takes longer than most batch computing sizes
+        spec_out = spectral_connectivity_epochs(
+            data=self.epochs,
+            sfreq=self.sfreq,
+            method=self.method,
+            mode=self.mode,
+            indices=(np.array([0, 0, 1, 1]), np.array([2, 3, 2, 3])),
+            verbose=False,
+        )
+        dat_conn: np.ndarray = spec_out.get_data()
+
+        # Get frequency band ranges only for the first batch, it's already the same
+        if len(self.fband_ranges) == 0:
+            for fband_range in self.fbands.values():
+                self.fband_ranges.append(
+                    np.where(
+                        (np.array(spec_out.freqs) > fband_range[0])
+                        & (np.array(spec_out.freqs) < fband_range[1])
+                    )[0]
+                )
+
+        # TODO: If I compute the mean for the entire fband, results are almost the same before
+        # normalization (0.9999999... vs 1.0), but some change wildly after normalization (-3 vs 0)
+        # Investigate why, is this a bug in normalization?
+        feature_results = {}
+        for conn in np.arange(dat_conn.shape[0]):
+            for fband_idx, fband in enumerate(self.fbands):
+                feature_results["_".join(["ch1", self.method, str(conn), fband])] = (
+                    np.mean(dat_conn[conn, self.fband_ranges[fband_idx]])
+                )
+
+        # Store current experiment parameters to check if re-initialization is needed
+        self.prev_batch_shape = data.shape
+
+        return feature_results

py_neuromodulation/nm_mnelsl_generator.py

@@ -0,0 +1,90 @@
+import numpy as np
+import mne
+from pathlib import Path
+
+from py_neuromodulation import logger, nm_types, nm_IO
+
+class LSLOfflinePlayer:
+
+    def __init__(
+        self,
+        stream_name: str | None = "lsl_offline_player",
+        f_name: str | nm_types.PathLike = None,
+        raw: mne.io.Raw | None = None,
+        sfreq: int | float | None = None,
+        data: np.ndarray | None = None,
+        ch_type: str | None = "dbs",
+    ) -> None:
+        """Initialization of MNE-LSL offline player.
+        Either a filename (PathLike) is provided,
+        or data and sampling frequency to initialize an example mock-up stream.
+
+
+        Parameters
+        ----------
+        stream_name : str, optional
+            LSL stream name, by default "example_stream"
+        f_name : str | None, optional
+            file name used for streaming, by default None
+        sfreq : int | float | None, optional
+            sampling rate, by default None
+        data : np.ndarray | None, optional
+            data used for streaming, by default None
+        ch_type: str | None, optional
+            channel type to select for streaming, by default "dbs"
+
+        Raises
+        ------
+        ValueError
+            _description_
+        """
+        self.sfreq = sfreq
+        self.stream_name = stream_name
+        got_raw = raw is not None
+        got_fname = f_name is not None
+        got_sfreq_data = sfreq is not None and data is not None
+
+        if not (got_fname or got_sfreq_data or got_raw):
+            error_msg = "Either f_name or raw or sfreq and data must be provided."
+            logger.critical(error_msg)
+            raise ValueError(error_msg)
+
+        if got_fname: 
+            (self._path_raw, data, sfreq, line_noise, coord_list, coord_names) = nm_IO.read_BIDS_data(f_name)
+
+        elif got_raw: 
+            self._path_raw = raw
+
+        elif got_sfreq_data:
+            info = mne.create_info(
+                ch_names=[f"ch{i}" for i in range(data.shape[0])],
+                ch_types=[ch_type for _ in range(data.shape[0])],
+                sfreq=sfreq,
+            )
+            raw = mne.io.RawArray(data, info)
+            self._path_raw = Path.cwd() / "temp_raw.fif"
+            raw.save(self._path_raw, overwrite=True)
+
+    def start_player(self, chunk_size: int = 10, n_repeat: int = 1):
+        """Start MNE-LSL Player
+
+        Parameters
+        ----------
+        chunk_size : int, optional
+            _description_, by default 1
+        n_repeat : int, optional
+            _description_, by default 1
+        """
+        from mne_lsl.player import PlayerLSL
+
+        self.player = PlayerLSL(
+            self._path_raw,
+            name=self.stream_name,
+            chunk_size=chunk_size,
+            n_repeat=n_repeat,
+        )
+        self.player = self.player.start()
+
+    def stop_player(self):
+        """Stop MNE-LSL Player"""
+        self.player.stop()

py_neuromodulation/nm_mnelsl_stream.py

@@ -0,0 +1,116 @@
+from collections.abc import Iterator
+import time
+from typing import TYPE_CHECKING
+import numpy as np
+from py_neuromodulation import logger
+from mne_lsl.lsl import resolve_streams
+import os
+
+if TYPE_CHECKING:
+    from py_neuromodulation import NMSettings
+    
+
+class LSLStream:
+    """
+    Class is used to create and connect to a LSL stream and pull data from it.
+    """
+
+    def __init__(self, settings: "NMSettings", stream_name: str | None = None) -> None:
+        """
+        Initialize the LSL stream.
+
+        Parameters:
+        -----------
+        settings : nm_settings.NMSettings object
+        stream_name : str, optional
+            Name of the stream to connect to. If not provided, the first available stream is used.
+
+        Raises:
+        -------
+        RuntimeError
+            If no stream is running under the provided name or if there are multiple streams running
+            under the same name.
+        """
+        from mne_lsl.stream import StreamLSL
+        self.stream: StreamLSL
+        
+        self.settings = settings
+        self._n_seconds_wait_before_disconnect = 3
+        try:
+            if stream_name is None:
+                stream_name = resolve_streams()[0].name
+                logger.info(
+                    f"Stream name not provided. Using first available stream: {stream_name}"
+                )
+            self.stream = StreamLSL(name=stream_name, bufsize=2).connect(timeout=2)
+        except Exception as e:
+            msg = f"Could not connect to stream: {e}. Either no stream is running under the name {stream_name} or there is several streams under this name."
+            logger.exception(msg)
+            raise RuntimeError(msg)
+
+        if self.stream.sinfo is None:
+            raise RuntimeError("Stream info is None. Check if the stream is running.")
+        
+        self.winsize = settings.segment_length_features_ms / self.stream.sinfo.sfreq
+        self.sampling_interval = 1 / self.settings.sampling_rate_features_hz
+
+        # If not running the generator when the escape key is pressed.
+        self.headless: bool = not os.environ.get("DISPLAY")
+        if not self.headless:
+            from pynput import keyboard
+
+            self.listener = keyboard.Listener(
+                on_press=lambda key: key != keyboard.Key.esc  # type: ignore
+            )
+            self.listener.start()
+
+    def get_next_batch(self) -> Iterator[tuple[np.ndarray, np.ndarray]]:
+        self.last_time = time.time()
+        check_data = None
+        data = None
+        stream_start_time = None
+
+        while self.stream.connected:
+            time_diff = time.time() - self.last_time  # in s
+            time.sleep(0.005)
+            if time_diff >= self.sampling_interval:
+                self.last_time = time.time()
+
+                logger.debug(f"Pull data - current time: {self.last_time}")
+                logger.debug(f"time since last data pull {time_diff} seconds")
+
+                if time_diff >= 2 * self.sampling_interval:
+                    logger.warning(
+                        "Feature computation time between two consecutive samples"
+                        "was twice the feature sampling interval"
+                    )
+                if data is not None:
+                    check_data = data
+
+                data, timestamp = self.stream.get_data(winsize=self.winsize)
+                if stream_start_time is None:
+                    stream_start_time = timestamp[0]
+
+                for i in range(self._n_seconds_wait_before_disconnect):
+                    if (
+                        data is not None
+                        and check_data is not None
+                        and np.allclose(data, check_data, atol=1e-7, rtol=1e-7)
+                    ):
+                        logger.warning(
+                            f"No new data incoming. Disconnecting stream in {3-i} seconds."
+                        )
+                        time.sleep(1)
+                        i += 1
+                        if i == self._n_seconds_wait_before_disconnect:
+                            self.stream.disconnect()
+                            logger.warning("Stream disconnected.")
+                            break
+
+                yield timestamp, data
+
+                logger.info(f"Stream time: {timestamp[-1] - stream_start_time}")
+
+                if not self.headless and not self.listener.running:
+                    logger.info("Keyboard interrupt")
+                    self.stream.disconnect()