py-neuromodulation 0.0.7__py3-none-any.whl → 0.1.1__py3-none-any.whl

py_neuromodulation/stream/stream.py

@@ -1,16 +1,19 @@
 """Module for generic and offline data streams."""
 
-from typing import TYPE_CHECKING
+import time
+from typing import TYPE_CHECKING, Any
 from collections.abc import Iterator
 import numpy as np
 from pathlib import Path
 
 import py_neuromodulation as nm
-from contextlib import suppress
 
 from py_neuromodulation.stream.data_processor import DataProcessor
-from py_neuromodulation.utils.types import _PathLike, FeatureName
+from py_neuromodulation.utils.types import _PathLike, FEATURE_NAME
+from py_neuromodulation.utils.file_writer import MsgPackFileWriter
 from py_neuromodulation.stream.settings import NMSettings
+from py_neuromodulation.analysis.decode import RealTimeDecoder
+from py_neuromodulation.stream.backend_interface import StreamBackendInterface
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -35,9 +38,6 @@ class Stream:
         sampling_rate_features_hz: float | None = None,
         path_grids: _PathLike | None = None,
         coord_names: list | None = None,
-        stream_name: str
-        | None = "example_stream",  # Timon: do we need those in the nmstream_abc?
-        is_stream_lsl: bool = False,
         coord_list: list | None = None,
         verbose: bool = True,
     ) -> None:
@@ -67,6 +67,7 @@ class Stream:
         verbose : bool, optional
             print out stream computation time information, by default True
         """
+        # This is calling NMSettings.validate() which is making a copy
         self.settings: NMSettings = NMSettings.load(settings)
 
         if channels is None and data is not None:
@@ -84,7 +85,7 @@ class Stream:
             raise ValueError("Either `channels` or `data` must be passed to `Stream`.")
 
         # If features that use frequency ranges are on, test them against nyquist frequency
-        use_freq_ranges: list[FeatureName] = [
+        use_freq_ranges: list[FEATURE_NAME] = [
             "bandpass_filter",
             "stft",
             "fft",
@@ -124,8 +125,8 @@ class Stream:
         self.sess_right = None
         self.projection = None
         self.model = None
+        self.is_running = False
 
-        # TODO(toni): is it necessary to initialize the DataProcessor on stream init?
         self.data_processor = DataProcessor(
             sfreq=self.sfreq,
             settings=self.settings,
@@ -201,13 +202,20 @@ class Stream:
         experiment_name: str = "sub",
         is_stream_lsl: bool = False,
         stream_lsl_name: str | None = None,
-        plot_lsl: bool = False,
-        save_csv: bool = False,
+        save_csv: bool = True,
         save_interval: int = 10,
         return_df: bool = True,
+        simulate_real_time: bool = False,
+        decoder: RealTimeDecoder | None = None,
+        backend_interface: StreamBackendInterface | None = None,
     ) -> "pd.DataFrame":
         self.is_stream_lsl = is_stream_lsl
         self.stream_lsl_name = stream_lsl_name
+        self.save_csv = save_csv
+        self.save_interval = save_interval
+        self.return_df = return_df
+        self.out_dir = Path.cwd() if not out_dir else Path(out_dir)
+        self.experiment_name = experiment_name
 
         # Validate input data
         if data is not None:
@@ -217,24 +225,10 @@ class Stream:
         elif self.data is None and data is None and self.is_stream_lsl is False:
             raise ValueError("No data passed to run function.")
 
-        # Generate output dirs
-        self.out_dir_root = Path.cwd() if not out_dir else Path(out_dir)
-        self.out_dir = self.out_dir_root / experiment_name
-        # TONI: Need better default experiment name
-        self.experiment_name = experiment_name if experiment_name else "sub"
-
-        self.out_dir.mkdir(parents=True, exist_ok=True)
-
-        # Open database connection
-        # TONI: we should give the user control over the save format
-        from py_neuromodulation.utils.database import NMDatabase
-
-        db = NMDatabase(experiment_name, out_dir)  # Create output database
+        file_writer = MsgPackFileWriter(name=experiment_name, out_dir=out_dir)
 
         self.batch_count: int = 0  # Keep track of the number of batches processed
 
-        # Reinitialize the data processor in case the nm_channels or nm_settings changed between runs of the same Stream
-        # TONI: then I think we can just not initialize the data processor in the init function
         self.data_processor = DataProcessor(
             sfreq=self.sfreq,
             settings=self.settings,
@@ -248,9 +242,8 @@ class Stream:
 
         nm.logger.log_to_file(out_dir)
 
-        # Initialize generator
         self.generator: Iterator
-        if not is_stream_lsl:
+        if not is_stream_lsl and data is not None:
             from py_neuromodulation.stream.generator import RawDataGenerator
 
             self.generator = RawDataGenerator(
@@ -259,6 +252,7 @@ class Stream:
                 self.settings.sampling_rate_features_hz,
                 self.settings.segment_length_features_ms,
             )
+            nm.logger.info("Initializing RawDataGenerator")
         else:
             from py_neuromodulation.stream.mnelsl_stream import LSLStream
 
@@ -266,13 +260,10 @@ class Stream:
                 settings=self.settings, stream_name=stream_lsl_name
             )
 
-            if plot_lsl:
-                from mne_lsl.stream_viewer import StreamViewer
-
-                viewer = StreamViewer(stream_name=stream_lsl_name)
-                viewer.start()
-
-            if self.sfreq != self.lsl_stream.stream.sinfo.sfreq:
+            if (
+                self.lsl_stream.stream.sinfo is not None
+                and self.sfreq != self.lsl_stream.stream.sinfo.sfreq
+            ):
                 error_msg = (
                     f"Sampling frequency of the lsl-stream ({self.lsl_stream.stream.sinfo.sfreq}) "
                     f"does not match the settings ({self.sfreq})."
@@ -285,9 +276,21 @@ class Stream:
 
         prev_batch_end = 0
         for timestamps, data_batch in self.generator:
+            self.is_running = True
+            if backend_interface:
+                # Only simulate real-time if connected to GUI
+                if simulate_real_time:
+                    time.sleep(1 / self.settings.sampling_rate_features_hz)
+
+                signal = backend_interface.check_control_signals()
+                if signal == "stop":
+                    break
+
             if data_batch is None:
+                nm.logger.info("Data batch is None, stopping run function")
                 break
 
+            nm.logger.info("Processing new data batch")
             feature_dict = self.data_processor.process(data_batch)
 
             this_batch_end = timestamps[-1]
@@ -296,10 +299,13 @@ class Stream:
                 f"{batch_length:.3f} seconds of new data processed",
             )
 
-            feature_dict["time"] = (
-                batch_length if is_stream_lsl else np.ceil(this_batch_end * 1000 + 1)
-            )
+            if decoder is not None:
+                ch_to_decode = self.channels.query("used == 1").iloc[0]["name"]
+                feature_dict = decoder.predict(
+                    feature_dict, ch_to_decode, fft_bands_only=True
+                )
 
+            feature_dict["time"] = np.ceil(this_batch_end * 1000 + 1)
             prev_batch_end = this_batch_end
 
             if self.verbose:
@@ -307,29 +313,42 @@ class Stream:
 
             self._add_target(feature_dict, data_batch)
 
-            # We should ensure that feature output is always either float64 or None and remove this
-            with suppress(TypeError):  # Need this because some features output None
-                for key, value in feature_dict.items():
-                    feature_dict[key] = np.float64(value)
+            # Push data to file writer
+            file_writer.insert_data(feature_dict)
 
-            db.insert_data(feature_dict)
+            # Send data to frontend
+            if backend_interface:
+                nm.logger.debug("stream.run: Sending features to frontend")
+                backend_interface.send_features(feature_dict)
+                backend_interface.send_raw_data(self._prepare_raw_data_dict(data_batch))
 
+            # Save features to file in intervals
             self.batch_count += 1
-            if self.batch_count % save_interval == 0:
-                db.commit()
+            if self.batch_count % self.save_interval == 0:
+                file_writer.save()
 
-        db.commit()  # Save last batches
+        file_writer.save()
 
-        # If save_csv is False, still save the first row to get the column names
-        feature_df: "pd.DataFrame" = (
-            db.fetch_all() if (save_csv or return_df) else db.head()
-        )
+        if self.save_csv:
+            file_writer.save_as_csv(save_all_combined=True)
+
+        feature_df = file_writer.load_all() if self.return_df else {}
 
-        db.close()  # Close the database connection
+        self._save_after_stream()
+        self.is_running = False
 
-        self._save_after_stream(feature_arr=feature_df)
+        return feature_df  # Timon: We could think of returnader instead
 
-        return feature_df  # TONI: Not sure if this makes sense anymore
+    def _prepare_raw_data_dict(self, data_batch: np.ndarray) -> dict[str, Any]:
+        """Prepare raw data dictionary for sending through queue"""
+        new_time_ms = 1000 / self.settings.sampling_rate_features_hz
+        new_samples = int(new_time_ms * self.sfreq / 1000)
+        return {
+            "raw_data": {
+                ch: list(data_batch[i, -new_samples:])
+                for i, ch in enumerate(self.channels["name"])
+            }
+        }
 
     def plot_raw_signal(
         self,
@@ -365,11 +384,15 @@ class Stream:
         ValueError
             raise Exception when no data is passed
         """
-        if self.data is None and data is None:
-            raise ValueError("No data passed to plot_raw_signal function.")
-
-        if data is None and self.data is not None:
-            data = self.data
+        if data is None:
+            if self.data is None:
+                raise ValueError("No data passed to plot_raw_signal function.")
+            else:
+                data = (
+                    self.data.to_numpy()
+                    if isinstance(self.data, pd.DataFrame)
+                    else self.data
+                )
 
         if sfreq is None:
             sfreq = self.sfreq
@@ -384,7 +407,7 @@ class Stream:
         from mne import create_info
         from mne.io import RawArray
 
-        info = create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
+        info = create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)  # type: ignore
         raw = RawArray(data, info)
 
         if picks is not None:
@@ -397,12 +420,9 @@ class Stream:
 
     def _save_after_stream(
         self,
-        feature_arr: "pd.DataFrame | None" = None,
     ) -> None:
-        """Save features, settings, nm_channels and sidecar after run"""
+        """Save settings, nm_channels and sidecar after run"""
         self._save_sidecar()
-        if feature_arr is not None:
-            self._save_features(feature_arr)
         self._save_settings()
         self._save_channels()
 
@@ -422,6 +442,7 @@ class Stream:
         """Save sidecar incduing fs, coords, sess_right to
         out_path_root and subfolder 'folder_name'"""
         additional_args = {"sess_right": self.sess_right}
+
         self.data_processor.save_sidecar(
             self.out_dir, self.experiment_name, additional_args
         )

py_neuromodulation/utils/channels.py

@@ -251,7 +251,7 @@ def _get_default_references(
 
 
 def get_default_channels_from_data(
-    data: np.ndarray,
+    data: "np.ndarray | pd.DataFrame",
     car_rereferencing: bool = True,
 ):
     """Return default channels dataframe with

py_neuromodulation/utils/file_writer.py

@@ -0,0 +1,110 @@
+import msgpack
+from abc import ABC, abstractmethod
+from pathlib import Path
+import pandas as pd
+import numpy as np
+from py_neuromodulation.utils.types import _PathLike
+from py_neuromodulation import logger
+
+class AbstractFileWriter(ABC):
+
+    @abstractmethod
+    def insert_data(self, feature_dict: dict):
+        pass
+
+    @abstractmethod
+    def save(self):
+        pass
+
+    @abstractmethod
+    def load_all(self):
+        pass
+
+    @abstractmethod
+    def save_as_csv(self, save_all_combined: bool = False):
+        pass
+
+class MsgPackFileWriter(AbstractFileWriter):
+    """
+    Class to store data in a serialized MessagePack file and load it back efficiently.
+    Parameters
+    ----------
+    out_dir : _PathLike
+        The directory to save the MessagePack database.
+    """
+
+    def __init__(
+        self,
+        name: str = "sub",
+        out_dir: _PathLike = "",
+    ):
+        # Make sure out_dir exists
+
+        self.out_dir = Path.cwd() if not out_dir else Path(out_dir)
+        self.out_dir = self.out_dir / name
+
+        Path(self.out_dir).mkdir(parents=True, exist_ok=True)
+
+        self.idx = 0
+        self.name = name
+        self.csv_path = Path(self.out_dir, f"{name}_FEATURES.csv")
+        self.data_l = []
+
+    def insert_data(self, feature_dict: dict):
+        """
+        Insert data into the MessagePack database.
+        Parameters
+        ----------
+        feature_dict : dict
+            The dictionary with the feature names and values.
+        """
+        # transform every value to float s.t. msgpack can serialize the data
+        for key, value in feature_dict.items():
+            feature_dict[key] = float(value) if value is not None else 0
+
+        self.data_l.append(feature_dict)
+
+    def save(self):
+        """
+        Save the current data to the MessagePack file.
+        """
+        if len(self.data_l) == 0:
+            return
+        with open(self.out_dir / f"{self.name}-{self.idx}.msgpack", "wb") as f:
+            msgpack.pack(self.data_l, f)
+        self.idx += 1
+        self.data_l = []
+
+    def load_all(self):
+        """
+        Load data from the MessagePack file into memory.
+        """
+        data_l = []
+        for i in range(self.idx):
+            with open(self.out_dir / f"{self.name}-{i}.msgpack", "rb") as f:
+                data_l.append(msgpack.unpack(f))
+        if len(data_l) == 0:
+            raise ValueError("No data to load")
+        data = pd.DataFrame(list(np.concatenate(data_l)))
+        return data
+
+    def save_as_csv(self, save_all_combined: bool = False):
+        """
+        Save the data as a CSV file.
+        """
+
+        if save_all_combined:
+            try:
+                data = self.load_all()
+            except ValueError as e:
+                logger.error(e)
+                return
+            data.to_csv(self.csv_path, index=False)
+        else:
+            if len(self.data_l) > 0:
+                self.data_l[-1].to_csv(self.csv_path, index=False)
+            else:
+                outpath =self.out_dir / f"{self.name}-0.msgpack"
+                with open(outpath, "rb") as f:
+                    data = msgpack.unpack(f)
+                data.to_csv(self.csv_path, index=False)

py_neuromodulation/utils/io.py

@@ -1,6 +1,6 @@
 import json
 from pathlib import PurePath, Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 
@@ -79,7 +79,7 @@ def read_BIDS_data(
 def read_mne_data(
     PATH_RUN: "_PathLike | BIDSPath",
     line_noise: int = 50,
-):
+) -> tuple[np.ndarray, float, list[str], list[str], list[str]]:
     """Read data in the mne.io.read_raw supported format.
 
     Parameters
@@ -117,7 +117,8 @@ def read_mne_data(
             f"Line noise is not available in the data, using value of {line_noise} Hz."
         )
 
-    return raw_arr.get_data(), sfreq, ch_names, ch_types, bads
+    data = cast(np.ndarray, raw_arr.get_data())
+    return data, sfreq, ch_names, ch_types, bads
 
 
 def get_coord_list(
@@ -190,6 +191,46 @@ def get_annotations(PATH_ANNOTATIONS: str, PATH_RUN: str, raw_arr: "mne_io.RawAr
     return annot, annot_data, raw_arr
 
 
+def read_plot_modules(
+    PATH_PLOT: _PathLike = PYNM_DIR / "plots",
+):
+    """Read required .mat files for plotting
+
+    Parameters
+    ----------
+    PATH_PLOT : regexp, optional
+        path to plotting files, by default
+    """
+
+    faces = loadmat(PurePath(PATH_PLOT, "faces.mat"))
+    vertices = loadmat(PurePath(PATH_PLOT, "Vertices.mat"))
+    grid = loadmat(PurePath(PATH_PLOT, "grid.mat"))["grid"]
+    stn_surf = loadmat(PurePath(PATH_PLOT, "STN_surf.mat"))
+    x_ver = stn_surf["vertices"][::2, 0]
+    y_ver = stn_surf["vertices"][::2, 1]
+    x_ecog = vertices["Vertices"][::1, 0]
+    y_ecog = vertices["Vertices"][::1, 1]
+    z_ecog = vertices["Vertices"][::1, 2]
+    x_stn = stn_surf["vertices"][::1, 0]
+    y_stn = stn_surf["vertices"][::1, 1]
+    z_stn = stn_surf["vertices"][::1, 2]
+
+    return (
+        faces,
+        vertices,
+        grid,
+        stn_surf,
+        x_ver,
+        y_ver,
+        x_ecog,
+        y_ecog,
+        z_ecog,
+        x_stn,
+        y_stn,
+        z_stn,
+    )
+
+
 def write_csv(df, path_out):
     """
     Function to save Pandas dataframes to disk as CSV using
@@ -209,7 +250,7 @@ def save_channels(
 ) -> None:
     out_dir = Path.cwd() if not out_dir else Path(out_dir)
     filename = "channels.csv" if not prefix else prefix + "_channels.csv"
-    write_csv(nmchannels, out_dir / filename)
+    write_csv(nmchannels, out_dir / prefix / filename)
     logger.info(f"{filename} saved to {out_dir}")
 
 
@@ -241,7 +282,7 @@ def save_general_dict(
     out_dir = Path.cwd() if not out_dir else Path(out_dir)
     filename = f"{prefix}{str_add}"
 
-    with open(out_dir / filename, "w") as f:
+    with open(out_dir / prefix / filename, "w") as f:
         json.dump(
             dict_,
             f,

py_neuromodulation/utils/perf.py

@@ -0,0 +1,156 @@
+from typing import Any
+import time
+import threading
+import logging
+from dataclasses import dataclass, field
+from collections import deque
+import statistics
+
+
+@dataclass
+class MetricPoint:
+    timestamp: float
+    value: float
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class MetricBuffer:
+    """Maintains a rolling buffer of metric values with timestamps."""
+
+    def __init__(self, max_size: int = 1000):
+        self.buffer: deque[MetricPoint] = deque(maxlen=max_size)
+        self.lock = threading.Lock()
+
+    def add(self, value: float, metadata: dict[str, Any] | None = None):
+        with self.lock:
+            self.buffer.append(
+                MetricPoint(timestamp=time.time(), value=value, metadata=metadata or {})
+            )
+
+    def get_stats(self, window_seconds: float | None = None) -> dict[str, float]:
+        with self.lock:
+            if not self.buffer:
+                return {"count": 0, "mean": 0.0, "min": 0.0, "max": 0.0, "std_dev": 0.0}
+
+            current_time = time.time()
+            values = [
+                point.value
+                for point in self.buffer
+                if window_seconds is None
+                or (current_time - point.timestamp) <= window_seconds
+            ]
+
+            if not values:
+                return {"count": 0, "mean": 0.0, "min": 0.0, "max": 0.0, "std_dev": 0.0}
+
+            return {
+                "count": len(values),
+                "mean": statistics.mean(values),
+                "min": min(values),
+                "max": max(values),
+                "std_dev": statistics.stdev(values) if len(values) > 1 else 0.0,
+            }
+
+
+class PerformanceMonitor:
+    """Centralized system for tracking performance metrics across the application."""
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super().__new__(cls)
+            return cls._instance
+
+    def __init__(self):
+        if not hasattr(self, "initialized"):
+            self.logger = logging.getLogger("PyNM.Performance")
+            self.metrics: dict[str, MetricBuffer] = {}
+            self.timers: dict[str, float] = {}
+            self.counters: dict[str, int] = {}
+            self.metrics_lock = threading.Lock()
+            self.initialized = True
+
+    def record_metric(
+        self, name: str, value: float, metadata: dict[str, Any] | None = None
+    ):
+        """Record a metric value with optional metadata."""
+        with self.metrics_lock:
+            if name not in self.metrics:
+                self.metrics[name] = MetricBuffer()
+            self.metrics[name].add(value, metadata)
+
+    def start_timer(self, name: str):
+        """Start a timer for measuring operation duration."""
+        self.timers[name] = time.time()
+
+    def stop_timer(self, name: str, record: bool = True) -> float:
+        """Stop a timer and optionally record its duration as a metric."""
+        if name not in self.timers:
+            raise KeyError(f"Timer '{name}' was never started")
+
+        duration = time.time() - self.timers[name]
+        if record:
+            self.record_metric(f"{name}_duration", duration)
+
+        del self.timers[name]
+        return duration
+
+    def increment_counter(self, name: str, amount: int = 1):
+        """Increment a counter by the specified amount."""
+        with self.metrics_lock:
+            self.counters[name] = self.counters.get(name, 0) + amount
+
+    def get_counter(self, name: str) -> int:
+        """Get the current value of a counter."""
+        return self.counters.get(name, 0)
+
+    def get_metric_stats(
+        self, name: str, window_seconds: float | None = None
+    ) -> dict[str, float]:
+        """Get statistics for a metric over the specified time window."""
+        if name not in self.metrics:
+            return {"count": 0, "mean": 0.0, "min": 0.0, "max": 0.0, "std_dev": 0.0}
+        return self.metrics[name].get_stats(window_seconds)
+
+    def get_all_metrics(
+        self, window_seconds: float | None = None
+    ) -> dict[str, dict[str, float]]:
+        """Get statistics for all metrics."""
+        return {
+            name: self.get_metric_stats(name, window_seconds) for name in self.metrics
+        }
+
+    def log_summary(self, window_seconds: float | None = None):
+        """Log a summary of all metrics and counters."""
+        stats = self.get_all_metrics(window_seconds)
+        self.logger.info("Performance Summary:")
+
+        for name, metric_stats in stats.items():
+            self.logger.info(f"{name}:")
+            for stat_name, value in metric_stats.items():
+                self.logger.info(f"  {stat_name}: {value:.3f}")
+
+        self.logger.info("Counters:")
+        for name, value in self.counters.items():
+            self.logger.info(f"  {name}: {value}")
+
+
+# Example usage:
+# monitor = PerformanceMonitor()
+#
+# # Record individual metrics
+# monitor.record_metric("queue_size", queue.qsize())
+#
+# # Time operations
+# monitor.start_timer("websocket_send")
+# await websocket.send_bytes(data)
+# duration = monitor.stop_timer("websocket_send")
+#
+# # Track message counts
+# monitor.increment_counter("messages_received")
+#
+# # Get stats for the last 5 minutes
+# stats = monitor.get_metric_stats("websocket_send", window_seconds=300)