py-neuromodulation 0.0.6__py3-none-any.whl → 0.1.0__py3-none-any.whl

py_neuromodulation/stream/mnelsl_player.py

@@ -1,13 +1,20 @@
 import numpy as np
 import mne
 from pathlib import Path
+import multiprocessing as mp
+import atexit
+import time
+import signal
 
 from py_neuromodulation.utils.types import _PathLike
-from py_neuromodulation.utils import io
+from py_neuromodulation.utils.io import read_BIDS_data
 from py_neuromodulation import logger
 
 
 class LSLOfflinePlayer:
+    _instances: set["LSLOfflinePlayer"] = set()  # Keep track of initialized players
+    _atexit_registered: bool = False  # Flag to register atexit
+
     def __init__(
         self,
         stream_name: str | None = "lsl_offline_player",
@@ -16,6 +23,8 @@ class LSLOfflinePlayer:
         sfreq: int | float | None = None,
         data: np.ndarray | None = None,
         ch_type: str | None = "dbs",
+        chunk_size: int = 10,
+        n_repeat: int = 1,
     ) -> None:
         """Initialization of MNE-LSL offline player.
         Either a filename (PathLike) is provided,
@@ -42,24 +51,16 @@ class LSLOfflinePlayer:
         """
         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
+        self.chunk_size = chunk_size
+        self.n_repeat = n_repeat
 
-        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:
+        if f_name:
             (self._path_raw, data, sfreq, line_noise, coord_list, coord_names) = (
-                io.read_BIDS_data(f_name)
+                read_BIDS_data(f_name)
             )
-
-        elif got_raw:
+        elif raw:
             self._path_raw = raw
-
-        elif got_sfreq_data:
+        elif sfreq and 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])],
@@ -68,27 +69,126 @@ class LSLOfflinePlayer:
             raw = mne.io.RawArray(data, info)
             self._path_raw = Path.cwd() / "temp_raw.fif"
             raw.save(self._path_raw, overwrite=True)
+        else:
+            error_msg = "Either f_name or raw or sfreq and data must be provided."
+            logger.critical(error_msg)
+            raise ValueError(error_msg)
 
-    def start_player(self, chunk_size: int = 10, n_repeat: int = 1):
+        # Flags to control the player subprocess
+        self._streaming_complete = mp.Event()
+        self._player_process = None
+        self._stop_flag = mp.Event()
+
+        LSLOfflinePlayer._instances.add(self)  # Register instancwe
+        if LSLOfflinePlayer._atexit_registered:
+            atexit.register(LSLOfflinePlayer._stop_all_players)
+            LSLOfflinePlayer._atexit_registered = True
+
+    def start_player(
+        self,
+        chunk_size: int | None = None,
+        n_repeat: int | None = None,
+        block: bool = False,
+    ):
         """Start MNE-LSL Player
 
         Parameters
         ----------
         chunk_size : int, optional
-            _description_, by default 1
+            Number of samples to stream at once, by default 10
         n_repeat : int, optional
-            _description_, by default 1
+            Number of times to repeat the stream, by default 1
+        block : bool, optional
+            If True, block until streaming is complete, by default False
         """
+
+        if chunk_size:
+            self.chunk_size = chunk_size
+        if n_repeat:
+            self.n_repeat = n_repeat
+
+        self._stop_flag.clear()
+        self._streaming_complete.clear()
+
+        self._player_process = mp.Process(
+            target=self._run_player,
+            args=(
+                self.chunk_size,
+                self.n_repeat,
+                self._stop_flag,
+                self._streaming_complete,
+            ),
+        )
+        self._player_process.start()
+
+        if block:
+            try:
+                self.wait_for_completion()
+            except KeyboardInterrupt:
+                logger.info("\nKeyboard interrupt received. Stopping the player...")
+                self.stop_player()
+
+    def _run_player(self, chunk_size, n_repeat, stop_flag, streaming_complete):
         from mne_lsl.player import PlayerLSL
 
-        self.player = PlayerLSL(
+        signal.signal(signal.SIGINT, lambda: stop_flag.set())
+
+        player = PlayerLSL(
             self._path_raw,
             name=self.stream_name,
             chunk_size=chunk_size,
             n_repeat=n_repeat,
         )
-        self.player = self.player.start()
+        player = player.start()
+
+        try:
+            while not stop_flag.is_set() and not player._end_streaming:
+                time.sleep(0.1)
+        finally:
+            try:
+                player.stop()
+            except RuntimeError:
+                # player already stopped
+                pass
+            streaming_complete.set()
+
+    def wait_for_completion(self):
+        """Block until streaming is complete"""
+        while self._player_process and self._player_process.is_alive():
+            try:
+                self._streaming_complete.wait(timeout=1.0)
+                if self._streaming_complete.is_set():
+                    break
+            except KeyboardInterrupt:
+                logger.info("\nKeyboard interrupt received. Stopping the player...")
+                self.stop_player()
+                break
 
     def stop_player(self):
         """Stop MNE-LSL Player"""
-        self.player.stop()
+        if self._player_process and self._player_process.is_alive():
+            self._stop_flag.set()
+            self._player_process.join(timeout=5)
+            if self._player_process.is_alive():
+                self._player_process.terminate()
+                self._player_process.join(timeout=1)
+            if self._player_process.is_alive():
+                self._player_process.kill()
+            self._player_process = None
+
+        logger.info(f"Player stopped: {self.stream_name}")
+        LSLOfflinePlayer._instances.discard(self)
+
+    @classmethod
+    def _stop_all_players(cls):
+        """Stop all player instances (used for atexit)"""
+        for player in cls._instances:
+            player.stop_player()
+
+    # Enable use as a context manager
+    def __enter__(self):
+        self.start_player()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop_player()

py_neuromodulation/stream/mnelsl_stream.py

@@ -34,7 +34,7 @@ class LSLStream:
         from mne_lsl.stream import StreamLSL
 
         self.stream: StreamLSL
-        self.keyboard_interrupt = False
+        # self.keyboard_interrupt = False
 
         self.settings = settings
         self._n_seconds_wait_before_disconnect = 3
@@ -58,11 +58,11 @@ class LSLStream:
 
         # If not running the generator when the escape key is pressed.
         self.headless: bool = not os.environ.get("DISPLAY")
-        if not self.headless:
-            from py_neuromodulation.utils.keyboard import KeyboardListener
+        # if not self.headless:
+        # from py_neuromodulation.utils.keyboard import KeyboardListener
 
-            self.listener = KeyboardListener(("esc", self.set_keyboard_interrupt))
-            self.listener.start()
+        # self.listener = KeyboardListener(("esc", self.set_keyboard_interrupt))
+        # self.listener.start()
 
     def get_next_batch(self) -> Iterator[tuple[np.ndarray, np.ndarray]]:
         self.last_time = time.time()
@@ -91,6 +91,8 @@ class LSLStream:
                 if stream_start_time is None:
                     stream_start_time = timestamp[0]
 
+                logger.info(f"Stream time: {timestamp[-1] - stream_start_time}")
+
                 for i in range(self._n_seconds_wait_before_disconnect):
                     if (
                         data is not None
@@ -98,7 +100,7 @@ class LSLStream:
                         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."
+                            f"No new data incoming. Disconnecting stream in {3 - i} seconds."
                         )
                         time.sleep(1)
                         i += 1
@@ -107,14 +109,4 @@ class LSLStream:
                             logger.warning("Stream disconnected.")
                             break
 
-                yield timestamp, data
-
-                logger.info(f"Stream time: {timestamp[-1] - stream_start_time}")
-
-                if not self.headless and self.keyboard_interrupt:
-                    logger.info("Keyboard interrupt")
-                    self.listener.stop()
-                    self.stream.disconnect()
-
-    def set_keyboard_interrupt(self):
-        self.keyboard_interrupt = True
+                yield timestamp - stream_start_time, data

py_neuromodulation/stream/settings.py

@@ -1,22 +1,27 @@
 """Module for handling settings."""
 
 from pathlib import Path
-from typing import ClassVar
-from pydantic import Field, model_validator
+from typing import Any, ClassVar, get_args
+from pydantic import model_validator, ValidationError
+from pydantic.functional_validators import ModelWrapValidatorHandler
 
-from py_neuromodulation import PYNM_DIR, logger, user_features
+from py_neuromodulation import logger, user_features, PYNM_DIR
 
 from py_neuromodulation.utils.types import (
     BoolSelector,
     FrequencyRange,
-    PreprocessorName,
     _PathLike,
     NMBaseModel,
-    NormMethod,
+    NORM_METHOD,
+    PREPROCESSOR_NAME,
 )
+from py_neuromodulation.utils.pydantic_extensions import NMErrorList, NMField
 
 from py_neuromodulation.processing.filter_preprocessing import FilterSettings
-from py_neuromodulation.processing.normalization import NormalizationSettings
+from py_neuromodulation.processing.normalization import (
+    FeatureNormalizationSettings,
+    NormalizationSettings,
+)
 from py_neuromodulation.processing.resample import ResamplerSettings
 from py_neuromodulation.processing.projection import ProjectionSettings
 
@@ -31,7 +36,9 @@ from py_neuromodulation.features import OscillatorySettings, BandPowerSettings
 from py_neuromodulation.features import BurstsSettings
 
 
-class FeatureSelection(BoolSelector):
+# TONI: this class has the proble that if a feature is absent,
+# it won't default to False but to whatever is defined here as default
+class FeatureSelector(BoolSelector):
     raw_hjorth: bool = True
     return_raw: bool = True
     bandpass_filter: bool = False
@@ -54,13 +61,24 @@ class PostprocessingSettings(BoolSelector):
     project_subcortex: bool = False
 
 
+DEFAULT_PREPROCESSORS: list[PREPROCESSOR_NAME] = [
+    "raw_resampling",
+    "notch_filter",
+    "re_referencing",
+]
+
+
 class NMSettings(NMBaseModel):
     # Class variable to store instances
     _instances: ClassVar[list["NMSettings"]] = []
 
     # General settings
-    sampling_rate_features_hz: float = Field(default=10, gt=0)
-    segment_length_features_ms: float = Field(default=1000, gt=0)
+    sampling_rate_features_hz: float = NMField(
+        default=10, gt=0, custom_metadata={"unit": "Hz"}
+    )
+    segment_length_features_ms: float = NMField(
+        default=1000, gt=0, custom_metadata={"unit": "ms"}
+    )
     frequency_ranges_hz: dict[str, FrequencyRange] = {
         "theta": FrequencyRange(4, 8),
         "alpha": FrequencyRange(8, 12),
@@ -72,35 +90,40 @@ class NMSettings(NMBaseModel):
     }
 
     # Preproceessing settings
-    preprocessing: list[PreprocessorName] = [
-        "raw_resampling",
-        "notch_filter",
-        "re_referencing",
-    ]
+    preprocessing: list[PREPROCESSOR_NAME] = NMField(
+        default=DEFAULT_PREPROCESSORS,
+        custom_metadata={
+            "field_type": "PreprocessorList",
+            "valid_values": list(get_args(PREPROCESSOR_NAME)),
+        },
+    )
+
     raw_resampling_settings: ResamplerSettings = ResamplerSettings()
     preprocessing_filter: FilterSettings = FilterSettings()
     raw_normalization_settings: NormalizationSettings = NormalizationSettings()
 
     # Postprocessing settings
     postprocessing: PostprocessingSettings = PostprocessingSettings()
-    feature_normalization_settings: NormalizationSettings = NormalizationSettings()
+    feature_normalization_settings: FeatureNormalizationSettings = (
+        FeatureNormalizationSettings()
+    )
     project_cortex_settings: ProjectionSettings = ProjectionSettings(max_dist_mm=20)
     project_subcortex_settings: ProjectionSettings = ProjectionSettings(max_dist_mm=5)
 
     # Feature settings
-    features: FeatureSelection = FeatureSelection()
+    features: FeatureSelector = FeatureSelector()
 
     fft_settings: OscillatorySettings = OscillatorySettings()
     welch_settings: OscillatorySettings = OscillatorySettings()
     stft_settings: OscillatorySettings = OscillatorySettings()
     bandpass_filter_settings: BandPowerSettings = BandPowerSettings()
     kalman_filter_settings: KalmanSettings = KalmanSettings()
-    burst_settings: BurstsSettings = BurstsSettings()
+    bursts_settings: BurstsSettings = BurstsSettings()
     sharpwave_analysis_settings: SharpwaveSettings = SharpwaveSettings()
     mne_connectivity_settings: MNEConnectivitySettings = MNEConnectivitySettings()
     coherence_settings: CoherenceSettings = CoherenceSettings()
     fooof_settings: FooofSettings = FooofSettings()
-    nolds_settings: NoldsSettings = NoldsSettings()
+    nolds_features: NoldsSettings = NoldsSettings()
     bispectrum_settings: BispectraSettings = BispectraSettings()
 
     def __init__(self, *args, **kwargs) -> None:
@@ -126,10 +149,38 @@ class NMSettings(NMBaseModel):
         for instance in cls._instances:
             delattr(instance.features, feature)
 
-    @model_validator(mode="after")
-    def validate_settings(self):
+    @model_validator(mode="wrap")  # type: ignore[reportIncompatibleMethodOverride]
+    def validate_settings(self, handler: ModelWrapValidatorHandler) -> Any:
+        # Perform all necessary custom validations in the settings class and also
+        # all validations in the feature classes that need additional information from
+        # the settings class
+        errors: NMErrorList = NMErrorList()
+
+        def remove_private_keys(data):
+            if isinstance(data, dict):
+                if "__value__" in data:
+                    return data["__value__"]
+                else:
+                    return {
+                        key: remove_private_keys(value)
+                        for key, value in data.items()
+                        if not key.startswith("__")
+                    }
+            elif isinstance(data, (list, tuple, set)):
+                return type(data)(remove_private_keys(item) for item in data)
+            else:
+                return data
+
+        self = remove_private_keys(self)
+
+        try:
+            self = handler(self)  # validate the model
+        except ValidationError as e:
+            self = NMSettings.unvalidated(**self)  # type: ignore
+            errors.extend(NMErrorList(e.errors()))
+
         if len(self.features.get_enabled()) == 0:
-            raise ValueError("At least one feature must be selected.")
+            errors.add_error("At least one feature must be selected.")
 
         # Replace spaces with underscores in frequency band names
         self.frequency_ranges_hz = {
@@ -138,32 +189,27 @@ class NMSettings(NMBaseModel):
 
         if self.features.bandpass_filter:
             # Check BandPass settings frequency bands
-            self.bandpass_filter_settings.validate_fbands(self)
+            errors.extend(self.bandpass_filter_settings.validate_fbands(self))
 
             # Check Kalman filter frequency bands
             if self.bandpass_filter_settings.kalman_filter:
-                self.kalman_filter_settings.validate_fbands(self)
+                errors.extend(self.kalman_filter_settings.validate_fbands(self))
 
-        for k, v in self.frequency_ranges_hz.items():
-            if not isinstance(v, FrequencyRange):
-                self.frequency_ranges_hz[k] = FrequencyRange.create_from(v)
+        if len(errors) > 0:
+            raise errors.create_error()
 
         return self
 
     def reset(self) -> "NMSettings":
         self.features.disable_all()
-        self.preprocessing = []
+        self.preprocessing = DEFAULT_PREPROCESSORS
         self.postprocessing.disable_all()
         return self
 
     def set_fast_compute(self) -> "NMSettings":
         self.reset()
         self.features.fft = True
-        self.preprocessing = [
-            "raw_resampling",
-            "notch_filter",
-            "re_referencing",
-        ]
+        self.preprocessing = DEFAULT_PREPROCESSORS
         self.postprocessing.feature_normalization = True
         self.postprocessing.project_cortex = False
         self.postprocessing.project_subcortex = False
@@ -253,7 +299,7 @@ class NMSettings(NMBaseModel):
         return NMSettings.from_file(PYNM_DIR / "default_settings.yaml")
 
     @staticmethod
-    def list_normalization_methods() -> list[NormMethod]:
+    def list_normalization_methods() -> list[NORM_METHOD]:
         return NormalizationSettings.list_normalization_methods()
 
     def save(
@@ -261,7 +307,7 @@ class NMSettings(NMBaseModel):
     ) -> None:
         filename = f"{prefix}_SETTINGS.{format}" if prefix else f"SETTINGS.{format}"
 
-        path_out = Path(out_dir) / filename
+        path_out = Path(out_dir) / prefix / filename
 
         with open(path_out, "w") as f:
             match format: