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

py_neuromodulation/{nm_oscillatory.py → features/oscillatory.py}

@@ -1,17 +1,12 @@
-from collections.abc import Iterable
+from collections.abc import Sequence
 import numpy as np
 from itertools import product
 
-from py_neuromodulation.nm_types import NMBaseModel
-from pydantic import field_validator
+from py_neuromodulation.utils.types import NMBaseModel, BoolSelector, NMFeature
 from typing import TYPE_CHECKING
 
-from py_neuromodulation.nm_features import NMFeature
-from py_neuromodulation.nm_types import BoolSelector
-
 if TYPE_CHECKING:
-    from py_neuromodulation.nm_settings import NMSettings
-    from py_neuromodulation.nm_kalmanfilter import KalmanSettings
+    from py_neuromodulation.stream.settings import NMSettings
 
 
 class OscillatoryFeatures(BoolSelector):
@@ -40,7 +35,7 @@ ESTIMATOR_DICT = {
 
 class OscillatoryFeature(NMFeature):
     def __init__(
-        self, settings: "NMSettings", ch_names: Iterable[str], sfreq: int
+        self, settings: "NMSettings", ch_names: Sequence[str], sfreq: int
     ) -> None:
         settings.validate()
         self.settings: OscillatorySettings  # Assignment in subclass __init__
@@ -63,7 +58,7 @@ class FFT(OscillatoryFeature):
     def __init__(
         self,
         settings: "NMSettings",
-        ch_names: Iterable[str],
+        ch_names: Sequence[str],
         sfreq: int,
     ) -> None:
         from scipy.fft import rfftfreq
@@ -127,7 +122,7 @@ class Welch(OscillatoryFeature):
     def __init__(
         self,
         settings: "NMSettings",
-        ch_names: Iterable[str],
+        ch_names: Sequence[str],
         sfreq: int,
     ) -> None:
         from scipy.fft import rfftfreq
@@ -190,7 +185,7 @@ class STFT(OscillatoryFeature):
     def __init__(
         self,
         settings: "NMSettings",
-        ch_names: Iterable[str],
+        ch_names: Sequence[str],
         sfreq: int,
     ) -> None:
         from scipy.fft import rfftfreq
@@ -252,172 +247,3 @@ class STFT(OscillatoryFeature):
                 )[idx]
 
         return feature_results
-
-
-class BandpowerFeatures(BoolSelector):
-    activity: bool = True
-    mobility: bool = False
-    complexity: bool = False
-
-
-###################################
-######## BANDPOWER FEATURE ########
-###################################
-
-
-class BandpassSettings(NMBaseModel):
-    segment_lengths_ms: dict[str, int] = {
-        "theta": 1000,
-        "alpha": 500,
-        "low_beta": 333,
-        "high_beta": 333,
-        "low_gamma": 100,
-        "high_gamma": 100,
-        "HFA": 100,
-    }
-    bandpower_features: BandpowerFeatures = BandpowerFeatures()
-    log_transform: bool = True
-    kalman_filter: bool = False
-
-    @field_validator("segment_lengths_ms")
-    @classmethod
-    # Replace spaces with underscores in frequency band names
-    def fbands_spaces_to_underscores(cls, segment_lengths_ms: dict[str, int]):
-        return {k.replace(" ", "_"): v for k, v in segment_lengths_ms.items()}
-
-    @field_validator("bandpower_features")
-    @classmethod
-    def bandpower_features_validator(cls, bandpower_features: BandpowerFeatures):
-        assert (
-            len(bandpower_features.get_enabled()) > 0
-        ), "Set at least one bandpower_feature to True."
-
-        return bandpower_features
-
-    def validate_fbands(self, settings: "NMSettings") -> None:
-        # Ensure that each freq-band is defined in the global settings
-        for fband_name in settings.frequency_ranges_hz.keys():
-            assert fband_name in self.segment_lengths_ms, (
-                f"frequency range {fband_name} "
-                "needs to be defined in settings.bandpass_filter_settings.segment_lengths_ms]"
-            )
-
-        # Ensure that segment length for each freq-band is smaller than the feature segment length setting
-        for fband_name, seg_length_fband in self.segment_lengths_ms.items():
-            assert seg_length_fband <= settings.segment_length_features_ms, (
-                f"segment length {seg_length_fband} needs to be smaller than "
-                f" settings['segment_length_features_ms'] = {settings.segment_length_features_ms}"
-            )
-
-
-class BandPower(NMFeature):
-    def __init__(
-        self,
-        settings: "NMSettings",
-        ch_names: Iterable[str],
-        sfreq: float,
-        use_kf: bool | None = None,
-    ) -> None:
-        settings.validate()
-
-        self.bp_settings: BandpassSettings = settings.bandpass_filter_settings
-        self.kalman_filter_settings: KalmanSettings = settings.kalman_filter_settings
-        self.sfreq = sfreq
-        self.ch_names = ch_names
-        self.KF_dict: dict = {}
-
-        from py_neuromodulation.nm_filter import MNEFilter
-
-        self.bandpass_filter = MNEFilter(
-            f_ranges=[
-                tuple(frange) for frange in settings.frequency_ranges_hz.values()
-            ],
-            sfreq=self.sfreq,
-            filter_length=self.sfreq - 1,
-            verbose=False,
-        )
-
-        if use_kf or (use_kf is None and self.bp_settings.kalman_filter):
-            self.init_KF("bandpass_activity")
-
-        seglengths = self.bp_settings.segment_lengths_ms
-
-        self.feature_params = []
-        for ch_idx, ch_name in enumerate(self.ch_names):
-            for f_band_idx, f_band in enumerate(settings.frequency_ranges_hz.keys()):
-                seglength_ms = seglengths[f_band]
-                seglen = int(np.floor(self.sfreq / 1000 * seglength_ms))
-                for bp_feature in self.bp_settings.bandpower_features.get_enabled():
-                    feature_name = "_".join([ch_name, "bandpass", bp_feature, f_band])
-                    self.feature_params.append(
-                        (
-                            ch_idx,
-                            f_band_idx,
-                            seglen,
-                            bp_feature,
-                            feature_name,
-                        )
-                    )
-
-    def init_KF(self, feature: str) -> None:
-        from py_neuromodulation.nm_kalmanfilter import define_KF
-
-        for f_band in self.kalman_filter_settings.frequency_bands:
-            for channel in self.ch_names:
-                self.KF_dict["_".join([channel, feature, f_band])] = define_KF(
-                    self.kalman_filter_settings.Tp,
-                    self.kalman_filter_settings.sigma_w,
-                    self.kalman_filter_settings.sigma_v,
-                )
-
-    def update_KF(self, feature_calc: np.floating, KF_name: str) -> np.floating:
-        if KF_name in self.KF_dict:
-            self.KF_dict[KF_name].predict()
-            self.KF_dict[KF_name].update(feature_calc)
-            feature_calc = self.KF_dict[KF_name].x[0]
-        return feature_calc
-
-    def calc_feature(self, data: np.ndarray) -> dict:
-        data = self.bandpass_filter.filter_data(data)
-
-        feature_results = {}
-
-        for (
-            ch_idx,
-            f_band_idx,
-            seglen,
-            bp_feature,
-            feature_name,
-        ) in self.feature_params:
-            feature_results[feature_name] = self.calc_bp_feature(
-                bp_feature, feature_name, data[ch_idx, f_band_idx, -seglen:]
-            )
-
-        return feature_results
-
-    def calc_bp_feature(self, bp_feature, feature_name, data):
-        match bp_feature:
-            case "activity":
-                feature_calc = np.var(data)
-                if self.bp_settings.log_transform:
-                    feature_calc = np.log10(feature_calc)
-                if self.KF_dict:
-                    feature_calc = self.update_KF(feature_calc, feature_name)
-            case "mobility":
-                feature_calc = np.sqrt(np.var(np.diff(data)) / np.var(data))
-            case "complexity":
-                feature_calc = self.calc_complexity(data)
-            case _:
-                raise ValueError(f"Unknown bandpower feature: {bp_feature}")
-
-        return np.nan_to_num(feature_calc)
-
-    @staticmethod
-    def calc_complexity(data: np.ndarray) -> float:
-        dat_deriv = np.diff(data)
-        deriv_variance = np.var(dat_deriv)
-        mobility = np.sqrt(deriv_variance / np.var(data))
-        dat_deriv_2_var = np.var(np.diff(dat_deriv))
-        deriv_mobility = np.sqrt(dat_deriv_2_var / deriv_variance)
-
-        return deriv_mobility / mobility

py_neuromodulation/{nm_sharpwaves.py → features/sharpwaves.py}

@@ -2,7 +2,6 @@ from collections.abc import Sequence
 from collections import defaultdict
 from itertools import product
 
-from py_neuromodulation.nm_types import NMBaseModel
 from pydantic import model_validator
 from typing import TYPE_CHECKING, Any, Callable
 
@@ -13,11 +12,15 @@ if np.__version__ >= "2.0.0":
 else:
     from numpy.core._methods import _mean as np_mean
 
-from py_neuromodulation.nm_features import NMFeature
-from py_neuromodulation.nm_types import BoolSelector, FrequencyRange
+from py_neuromodulation.utils.types import (
+    NMFeature,
+    NMBaseModel,
+    BoolSelector,
+    FrequencyRange,
+)
 
 if TYPE_CHECKING:
-    from py_neuromodulation.nm_settings import NMSettings
+    from py_neuromodulation import NMSettings
 
 # Using low-level numpy mean function for performance, could do the same for the other estimators
 ESTIMATOR_DICT = {
@@ -38,6 +41,7 @@ class PeakDetectionSettings(NMBaseModel):
 class SharpwaveFeatures(BoolSelector):
     peak_left: bool = False
     peak_right: bool = False
+    num_peaks: bool = False
     trough: bool = False
     width: bool = False
     prominence: bool = True
@@ -368,6 +372,11 @@ class SharpwaveAnalyzer(NMFeature):
             # results["sharpness"] = ((trough_height - left_height) + (trough_height - right_height)) / 2
             results["sharpness"] = trough_height - 0.5 * (left_height + right_height)
 
+        if self.sw_settings.sharpwave_features.num_peaks:
+            results["num_peaks"] = [
+                trough_idx.shape[0]
+            ]  # keep list to the estimator can be applied
+
         if self.need_steepness:
             # steepness is calculated as the first derivative
             steepness: np.ndarray = np.concatenate((np.zeros(1), np.diff(data)))

py_neuromodulation/filter/__init__.py

@@ -0,0 +1,3 @@
+from .kalman_filter import define_KF, KalmanSettings
+from .notch_filter import NotchFilter
+from .mne_filter import MNEFilter

py_neuromodulation/{nm_kalmanfilter.py → filter/kalman_filter.py}

@@ -1,71 +1,67 @@
-from numpy import array, cov
-from py_neuromodulation.nm_types import NMBaseModel
-from typing import TYPE_CHECKING
-
-from pydantic import field_validator
-
-if TYPE_CHECKING:
-    from py_neuromodulation.nm_settings import NMSettings
-
-
-class KalmanSettings(NMBaseModel):
-    Tp: float = 0.1
-    sigma_w: float = 0.7
-    sigma_v: float = 1.0
-    frequency_bands: list[str] = [
-        "theta",
-        "alpha",
-        "low_beta",
-        "high_beta",
-        "low_gamma",
-        "high_gamma",
-        "HFA",
-    ]
-
-    @field_validator("frequency_bands")
-    def fbands_spaces_to_underscores(cls, frequency_bands):
-        return [f.replace(" ", "_") for f in frequency_bands]
-
-    def validate_fbands(self, settings: "NMSettings") -> None:
-        assert all(
-            (item in settings.frequency_ranges_hz for item in self.frequency_bands)
-        ), (
-            "Frequency bands for Kalman filter must also be specified in "
-            "bandpass_filter_settings."
-        )
-
-
-def define_KF(Tp, sigma_w, sigma_v):
-    """Define Kalman filter according to white noise acceleration model.
-    See DOI: 10.1109/TBME.2009.2038990  for explanation
-    See https://filterpy.readthedocs.io/en/latest/kalman/KalmanFilter.html#r64ca38088676-2 for implementation details
-
-    Parameters
-    ----------
-    Tp : float
-        prediction interval
-    sigma_w : float
-        process noise
-    sigma_v : float
-        measurement noise
-
-    Returns
-    -------
-    filterpy.KalmanFilter
-        initialized KalmanFilter object
-    """
-    from filterpy.kalman import KalmanFilter
-
-    f = KalmanFilter(dim_x=2, dim_z=1)
-    f.x = array([0, 1])  # x here sensor signal and it's first derivative
-    f.F = array([[1, Tp], [0, 1]])
-    f.H = array([[1, 0]])
-    f.R = sigma_v
-    f.Q = array(
-        [
-            [(sigma_w**2) * (Tp**3) / 3, (sigma_w**2) * (Tp**2) / 2],
-            [(sigma_w**2) * (Tp**2) / 2, (sigma_w**2) * Tp],
-        ]
-    )
-    f.P = cov([[1, 0], [0, 1]])
-    return f
+import numpy as np
+from typing import TYPE_CHECKING
+
+from py_neuromodulation.utils.types import NMBaseModel
+
+
+if TYPE_CHECKING:
+    from py_neuromodulation.stream.settings import NMSettings
+
+
+class KalmanSettings(NMBaseModel):
+    Tp: float = 0.1
+    sigma_w: float = 0.7
+    sigma_v: float = 1.0
+    frequency_bands: list[str] = [
+        "theta",
+        "alpha",
+        "low_beta",
+        "high_beta",
+        "low_gamma",
+        "high_gamma",
+        "HFA",
+    ]
+
+    def validate_fbands(self, settings: "NMSettings") -> None:
+        assert all(
+            (item in settings.frequency_ranges_hz for item in self.frequency_bands)
+        ), (
+            "Frequency bands for Kalman filter must also be specified in "
+            "bandpass_filter_settings."
+        )
+
+
+def define_KF(Tp, sigma_w, sigma_v):
+    """Define Kalman filter according to white noise acceleration model.
+    See DOI: 10.1109/TBME.2009.2038990  for explanation
+    See https://filterpy.readthedocs.io/en/latest/kalman/KalmanFilter.html#r64ca38088676-2 for implementation details
+
+    Parameters
+    ----------
+    Tp : float
+        prediction interval
+    sigma_w : float
+        process noise
+    sigma_v : float
+        measurement noise
+
+    Returns
+    -------
+    filterpy.KalmanFilter
+        initialized KalmanFilter object
+    """
+    from .kalman_filter_external import KalmanFilter
+
+    f = KalmanFilter(dim_x=2, dim_z=1)
+    f.x = np.array([0, 1])  # x here sensor signal and it's first derivative
+    f.F = np.array([[1, Tp], [0, 1]])
+    f.H = np.array([[1, 0]])
+    f.R = sigma_v
+    f.Q = np.array(
+        [
+            [(sigma_w**2) * (Tp**3) / 3, (sigma_w**2) * (Tp**2) / 2],
+            [(sigma_w**2) * (Tp**2) / 2, (sigma_w**2) * Tp],
+        ]
+    )
+    f.P = np.cov([[1, 0], [0, 1]])
+    return f