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

py_neuromodulation/nm_fooof.py

@@ -1,159 +1,139 @@
-import logging
-from typing import Iterable
-
-import numpy as np
-from fooof import FOOOF
-from scipy import fft
-
-from py_neuromodulation import nm_features_abc
-
-
-class FooofAnalyzer(nm_features_abc.Feature):
-    def __init__(
-        self, settings: dict, ch_names: Iterable[str], sfreq: float
-    ) -> None:
-        self.settings_fooof = settings["fooof"]
-        self.sfreq = sfreq
-        self.ch_names = ch_names
-
-        self.freq_range = self.settings_fooof["freq_range_hz"]
-        self.ap_mode = "knee" if self.settings_fooof["knee"] else "fixed"
-        self.max_n_peaks = self.settings_fooof["max_n_peaks"]
-
-        self.num_samples = int(
-            self.settings_fooof["windowlength_ms"] * sfreq / 1000
-        )
-
-        self.f_vec = np.arange(0, int(self.num_samples / 2) + 1, 1)
-
-    def test_settings(
-        s: dict,
-        ch_names: Iterable[str],
-        sfreq: int | float,
-    ):
-        assert isinstance(s["fooof"]["aperiodic"]["exponent"], bool)
-        assert isinstance(s["fooof"]["aperiodic"]["offset"], bool)
-        assert isinstance(s["fooof"]["aperiodic"]["knee"], bool)
-        assert isinstance(s["fooof"]["periodic"]["center_frequency"], bool)
-        assert isinstance(s["fooof"]["periodic"]["band_width"], bool)
-        assert isinstance(s["fooof"]["periodic"]["height_over_ap"], bool)
-        assert isinstance(s["fooof"]["knee"], bool)
-        assert isinstance(s["fooof"]["windowlength_ms"], (int, float))
-        assert (
-            s["fooof"]["windowlength_ms"] <= s["segment_length_features_ms"]
-        ), (
-            "fooof windowlength_ms needs to be smaller equal than segment_length_features_ms "
-            f"got windowlength_ms: {s['fooof']['windowlength_ms']} and {s['segment_length_features_ms']}"
-        )
-
-        assert (
-            s["fooof"]["freq_range_hz"][0] < sfreq
-            and s["fooof"]["freq_range_hz"][1] < sfreq
-        ), f"fooof frequency range needs to be below sfreq, got {s['fooof']['freq_range_hz']}"
-
-    def _get_spectrum(self, data: np.array):
-        """return absolute value fft spectrum"""
-
-        data = data[-self.num_samples :]
-        Z = np.abs(fft.rfft(data))
-
-        return Z
-
-    def calc_feature(
-        self,
-        data: np.array,
-        features_compute: dict,
-    ) -> dict:
-        for ch_idx, ch_name in enumerate(self.ch_names):
-            spectrum = self._get_spectrum(data[ch_idx, :])
-
-            try:
-                fm = FOOOF(
-                    aperiodic_mode=self.ap_mode,
-                    peak_width_limits=self.settings_fooof["peak_width_limits"],
-                    max_n_peaks=self.settings_fooof["max_n_peaks"],
-                    min_peak_height=self.settings_fooof["min_peak_height"],
-                    peak_threshold=self.settings_fooof["peak_threshold"],
-                    verbose=False,
-                )
-                fm.fit(self.f_vec, spectrum, self.freq_range)
-            except Exception as e:
-                logging.critical(e, exc_info=True)
-
-            if fm.fooofed_spectrum_ is None:
-                FIT_PASSED = False
-            else:
-                FIT_PASSED = True
-
-            if self.settings_fooof["aperiodic"]["exponent"]:
-                features_compute[f"{ch_name}_fooof_a_exp"] = (
-                    np.nan_to_num(fm.get_params("aperiodic_params", "exponent"))
-                    if FIT_PASSED is True
-                    else None
-                )
-
-            if self.settings_fooof["aperiodic"]["offset"]:
-                features_compute[f"{ch_name}_fooof_a_offset"] = (
-                    np.nan_to_num(fm.get_params("aperiodic_params", "offset"))
-                    if FIT_PASSED is True
-                    else None
-                )
-
-            if self.settings_fooof["aperiodic"]["knee"]:
-                if FIT_PASSED is False:
-                    knee_freq = None
-                else:
-                    if fm.get_params("aperiodic_params", "exponent") != 0:
-                        knee_fooof = fm.get_params("aperiodic_params", "knee")
-                        knee_freq = np.nan_to_num(
-                            knee_fooof
-                            ** (
-                                1
-                                / fm.get_params("aperiodic_params", "exponent")
-                            )
-                        )
-                    else:
-                        knee_freq = None
-
-                features_compute[f"{ch_name}_fooof_a_knee_frequency"] = (
-                    knee_freq
-                )
-
-            peaks_bw = (
-                fm.get_params("peak_params", "BW")
-                if FIT_PASSED is True
-                else None
-            )
-            peaks_cf = (
-                fm.get_params("peak_params", "CF")
-                if FIT_PASSED is True
-                else None
-            )
-            peaks_pw = (
-                fm.get_params("peak_params", "PW")
-                if FIT_PASSED is True
-                else None
-            )
-
-            if type(peaks_bw) is np.float64 or peaks_bw is None:
-                peaks_bw = [peaks_bw]
-                peaks_cf = [peaks_cf]
-                peaks_pw = [peaks_pw]
-
-            for peak_idx in range(self.max_n_peaks):
-                if self.settings_fooof["periodic"]["band_width"]:
-                    features_compute[f"{ch_name}_fooof_p_{peak_idx}_bw"] = (
-                        peaks_bw[peak_idx] if peak_idx < len(peaks_bw) else None
-                    )
-
-                if self.settings_fooof["periodic"]["center_frequency"]:
-                    features_compute[f"{ch_name}_fooof_p_{peak_idx}_cf"] = (
-                        peaks_cf[peak_idx] if peak_idx < len(peaks_bw) else None
-                    )
-
-                if self.settings_fooof["periodic"]["height_over_ap"]:
-                    features_compute[f"{ch_name}_fooof_p_{peak_idx}_pw"] = (
-                        peaks_pw[peak_idx] if peak_idx < len(peaks_bw) else None
-                    )
-
-        return features_compute
+from collections.abc import Iterable
+import numpy as np
+
+from typing import TYPE_CHECKING
+from py_neuromodulation.nm_types import NMBaseModel
+
+from py_neuromodulation.nm_features import NMFeature
+from py_neuromodulation.nm_types import BoolSelector, FrequencyRange
+
+if TYPE_CHECKING:
+    from py_neuromodulation.nm_settings import NMSettings
+
+
+class FooofAperiodicSettings(BoolSelector):
+    exponent: bool = True
+    offset: bool = True
+    knee: bool = True
+
+
+class FooofPeriodicSettings(BoolSelector):
+    center_frequency: bool = False
+    band_width: bool = False
+    height_over_ap: bool = False
+
+
+class FooofSettings(NMBaseModel):
+    aperiodic: FooofAperiodicSettings = FooofAperiodicSettings()
+    periodic: FooofPeriodicSettings = FooofPeriodicSettings()
+    windowlength_ms: float = 800
+    peak_width_limits: FrequencyRange = FrequencyRange(0.5, 12)
+    max_n_peaks: int = 3
+    min_peak_height: float = 0
+    peak_threshold: float = 2
+    freq_range_hz: FrequencyRange = FrequencyRange(2, 40)
+    knee: bool = True
+
+
+class FooofAnalyzer(NMFeature):
+    feat_name_map = {
+        "exponent": "exp",
+        "offset": "offset",
+        "knee": "knee_frequency",
+        "center_frequency": "cf",
+        "band_width": "bw",
+        "height_over_ap": "pw",
+    }
+
+    def __init__(
+        self, settings: "NMSettings", ch_names: Iterable[str], sfreq: float
+    ) -> None:
+        self.settings = settings.fooof
+        self.sfreq = sfreq
+        self.ch_names = ch_names
+
+        self.ap_mode = "knee" if self.settings.knee else "fixed"
+
+        self.num_samples = int(self.settings.windowlength_ms * sfreq / 1000)
+
+        self.f_vec = np.arange(0, int(self.num_samples / 2) + 1, 1)
+
+        assert (
+            settings.fooof.windowlength_ms <= settings.segment_length_features_ms
+        ), f"fooof windowlength_ms ({settings.fooof.windowlength_ms}) needs to be smaller equal than segment_length_features_ms ({settings.segment_length_features_ms})."
+
+        assert (
+            settings.fooof.freq_range_hz[0] < sfreq
+            and settings.fooof.freq_range_hz[1] < sfreq
+        ), f"fooof frequency range needs to be below sfreq, got {settings.fooof.freq_range_hz}"
+
+        from fooof import FOOOFGroup
+
+        self.fm = FOOOFGroup(
+            aperiodic_mode=self.ap_mode,
+            peak_width_limits=tuple(self.settings.peak_width_limits),
+            max_n_peaks=self.settings.max_n_peaks,
+            min_peak_height=self.settings.min_peak_height,
+            peak_threshold=self.settings.peak_threshold,
+            verbose=False,
+        )
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        from scipy.fft import rfft
+
+        spectra = np.abs(rfft(data[:, -self.num_samples :]))  # type: ignore
+
+        self.fm.fit(self.f_vec, spectra, self.settings.freq_range_hz)
+
+        if not self.fm.has_model or self.fm.null_inds_ is None:
+            raise RuntimeError("FOOOF failed to fit model to data.")
+
+        failed_fits: list[int] = self.fm.null_inds_
+
+        feature_results = {}
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            FIT_PASSED = ch_idx not in failed_fits
+            exp = self.fm.get_params("aperiodic_params", "exponent")[ch_idx]
+
+            for feat in self.settings.aperiodic.get_enabled():
+                f_name = f"{ch_name}_fooof_a_{self.feat_name_map[feat]}"
+
+                if not FIT_PASSED:
+                    feature_results[f_name] = None
+
+                elif feat == "knee" and exp == 0:
+                    feature_results[f_name] = None
+
+                else:
+                    params = self.fm.get_params("aperiodic_params", feat)[ch_idx]
+                    if feat == "knee":
+                        # If knee parameter is negative, set knee frequency to 0
+                        if params < 0:
+                            params = 0
+                        else:
+                            params = params ** (1 / exp)
+
+                    feature_results[f_name] = np.nan_to_num(params)
+
+            peaks_dict: dict[str, np.ndarray | None] = {
+                "bw": self.fm.get_params("peak_params", "BW") if FIT_PASSED else None,
+                "cf": self.fm.get_params("peak_params", "CF") if FIT_PASSED else None,
+                "pw": self.fm.get_params("peak_params", "PW") if FIT_PASSED else None,
+            }
+
+            if type(peaks_dict["bw"]) is np.float64 or peaks_dict["bw"] is None:
+                peaks_dict["bw"] = [peaks_dict["bw"]]
+                peaks_dict["cf"] = [peaks_dict["cf"]]
+                peaks_dict["pw"] = [peaks_dict["pw"]]
+
+            for peak_idx in range(self.settings.max_n_peaks):
+                for feat in self.settings.periodic.get_enabled():
+                    f_name = f"{ch_name}_fooof_p_{peak_idx}_{self.feat_name_map[feat]}"
+
+                    feature_results[f_name] = (
+                        peaks_dict[self.feat_name_map[feat]][peak_idx]
+                        if peak_idx < len(peaks_dict[self.feat_name_map[feat]])
+                        else None
+                    )
+
+        return feature_results

py_neuromodulation/nm_generator.py

@@ -1,37 +1,45 @@
-from typing import Iterator
-
-import numpy as np
-
-
-def raw_data_generator(
-    data: np.ndarray,
-    settings: dict,
-    sfreq: int,
-) -> Iterator[np.ndarray]:
-    """
-    This generator function mimics online data acquisition.
-    The data are iteratively sampled with sfreq_new.
-    Arguments
-    ---------
-        ieeg_raw (np array): shape (channels, time)
-        sfreq: int
-        sfreq_new: int
-        offset_time: int | float
-    Returns
-    -------
-        np.array: new batch for run function of full segment length shape
-    """
-    sfreq_new = settings["sampling_rate_features_hz"]
-    offset_time = settings["segment_length_features_ms"]
-    offset_start = offset_time / 1000 * sfreq
-
-    ratio_samples_features = sfreq / sfreq_new
-
-    ratio_counter = 0
-    for cnt in range(data.shape[1]+1):  # shape + 1 guarantees that the last sample is also included
-
-        if (cnt - offset_start) >= ratio_samples_features * ratio_counter:
-
-            ratio_counter += 1
-
-            yield data[:, np.floor(cnt-offset_start).astype(int) : cnt]
+from collections.abc import Iterator
+from typing import TYPE_CHECKING
+import numpy as np
+
+if TYPE_CHECKING:
+    from py_neuromodulation.nm_settings import NMSettings
+
+
+def raw_data_generator(
+    data: np.ndarray,
+    settings: "NMSettings",
+    sfreq: float,
+) -> Iterator[tuple[np.ndarray, np.ndarray]]:
+    """
+    This generator function mimics online data acquisition.
+    The data are iteratively sampled with settings.sampling_rate_features_hz
+
+    Arguments
+    ---------
+        data (np array): shape (channels, time)
+        settings (nm_settings.NMSettings): settings object
+        sfreq (float): sampling frequency of the data
+
+    Returns
+    -------
+        np.array: 1D array of time stamps
+        np.array: new batch for run function of full segment length shape
+    """
+    sfreq_new = settings.sampling_rate_features_hz
+    offset_time = settings.segment_length_features_ms
+    offset_start = offset_time / 1000 * sfreq
+
+    ratio_samples_features = sfreq / sfreq_new
+
+    ratio_counter = 0
+    for cnt in range(
+        data.shape[1] + 1
+    ):  # shape + 1 guarantees that the last sample is also included
+        if (cnt - offset_start) >= ratio_samples_features * ratio_counter:
+            ratio_counter += 1
+
+            yield (
+                np.arange(cnt - offset_start, cnt) / sfreq,
+                data[:, np.floor(cnt - offset_start).astype(int) : cnt],
+            )

py_neuromodulation/nm_hjorth_raw.py

@@ -1,73 +1,52 @@
-import enum
-import numpy as np
-from typing import Iterable
-
-from py_neuromodulation import nm_features_abc
-
-
-class Hjorth(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 test_settings(
-        settings: dict,
-        ch_names: Iterable[str],
-        sfreq: int | float,
-    ):
-        # no settings to test
-        pass
-
-    def calc_feature(self, data: np.array, features_compute: dict) -> dict:
-        for ch_idx, ch_name in enumerate(self.ch_names):
-            features_compute[
-                "_".join([ch_name, "RawHjorth_Activity"])
-            ] = np.nan_to_num(np.var(data[ch_idx, :]))
-            deriv_variance = np.nan_to_num(np.var(np.diff(data[ch_idx, :])))
-            mobility = np.nan_to_num(
-                np.sqrt(deriv_variance / np.var(data[ch_idx, :]))
-            )
-            features_compute[
-                "_".join([ch_name, "RawHjorth_Mobility"])
-            ] = mobility
-
-            dat_deriv_2_var = np.nan_to_num(
-                np.var(np.diff(np.diff(data[ch_idx, :])))
-            )
-            deriv_mobility = np.nan_to_num(
-                np.sqrt(dat_deriv_2_var / deriv_variance)
-            )
-            features_compute[
-                "_".join([ch_name, "RawHjorth_Complexity"])
-            ] = np.nan_to_num(deriv_mobility / mobility)
-
-        return features_compute
-
-
-class Raw(nm_features_abc.Feature):
-    def __init__(
-        self, settings: dict, ch_names: Iterable[str], sfreq: float
-    ) -> None:
-        self.ch_names = ch_names
-
-    def calc_feature(
-        self,
-        data: np.array,
-        features_compute: dict,
-    ) -> dict:
-        for ch_idx, ch_name in enumerate(self.ch_names):
-            features_compute["_".join([ch_name, "raw"])] = data[ch_idx, -1]
-
-        return features_compute
-
-    @staticmethod
-    def test_settings(
-        settings: dict,
-        ch_names: Iterable[str],
-        sfreq: int | float,
-    ):
-        # no settings to test
-        pass
+"""
+Reference:  B Hjorth
+            EEG analysis based on time domain properties
+            Electroencephalogr Clin Neurophysiol. 1970 Sep;29(3):306-10.
+            DOI: 10.1016/0013-4694(70)90143-4
+"""
+
+import numpy as np
+from collections.abc import Iterable
+
+from py_neuromodulation.nm_features import NMFeature
+from py_neuromodulation.nm_settings import NMSettings
+
+
+class Hjorth(NMFeature):
+    def __init__(
+        self, settings: NMSettings, ch_names: Iterable[str], sfreq: float
+    ) -> None:
+        self.ch_names = ch_names
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        var = np.var(data, axis=-1)
+        deriv1 = np.diff(data, axis=-1)
+        deriv2 = np.diff(deriv1, axis=-1)
+        deriv1_var = np.var(deriv1, axis=-1)
+        deriv2_var = np.var(deriv2, axis=-1)
+        deriv1_mobility = np.sqrt(deriv2_var / deriv1_var)
+
+        activity = np.nan_to_num(var)
+        mobility = np.nan_to_num(np.sqrt(deriv1_var / var))
+        complexity = np.nan_to_num(deriv1_mobility / mobility)
+
+        feature_results = {}
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            feature_results[f"{ch_name}_RawHjorth_Activity"] = activity[ch_idx]
+            feature_results[f"{ch_name}_RawHjorth_Mobility"] = mobility[ch_idx]
+            feature_results[f"{ch_name}_RawHjorth_Complexity"] = complexity[ch_idx]
+
+        return feature_results
+
+
+class Raw(NMFeature):
+    def __init__(self, settings: dict, ch_names: Iterable[str], sfreq: float) -> None:
+        self.ch_names = ch_names
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        feature_results = {}
+
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            feature_results["_".join([ch_name, "raw"])] = data[ch_idx, -1]
+
+        return feature_results

py_neuromodulation/nm_kalmanfilter.py

@@ -1,58 +1,71 @@
-from numpy import array, cov
-from typing import Iterable
-
-from filterpy.kalman import KalmanFilter
-
-
-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
-    """
-
-    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
-
-def test_kf_settings(
-        s: dict,
-        ch_names: Iterable[str],
-        sfreq: int | float,
-    ):
-        assert isinstance(s["kalman_filter_settings"]["Tp"], (float, int))
-        assert isinstance(s["kalman_filter_settings"]["sigma_w"], (float, int))
-        assert isinstance(s["kalman_filter_settings"]["sigma_v"], (float, int))
-        assert s["kalman_filter_settings"][
-            "frequency_bands"
-        ], "No frequency bands specified for Kalman filter."
-        assert isinstance(
-            s["kalman_filter_settings"]["frequency_bands"], list
-        ), "Frequency bands for Kalman filter must be specified as a list."
-        assert (
-            item
-            in s["frequency_ranges_hz"].values()
-            for item in s["kalman_filter_settings"]["frequency_bands"]
-        ), (
-            "Frequency bands for Kalman filter must also be specified in "
-            "bandpass_filter_settings."
-        )
+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