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

py_neuromodulation/features/coherence.py

@@ -0,0 +1,255 @@
+import numpy as np
+from collections.abc import Iterable
+
+
+from typing import TYPE_CHECKING, Annotated
+from pydantic import Field, field_validator
+
+from py_neuromodulation.utils.types import (
+    NMFeature,
+    BoolSelector,
+    FrequencyRange,
+    NMBaseModel,
+)
+from py_neuromodulation import logger
+
+if TYPE_CHECKING:
+    from py_neuromodulation import NMSettings
+
+
+class CoherenceMethods(BoolSelector):
+    coh: bool = True
+    icoh: bool = True
+
+
+class CoherenceFeatures(BoolSelector):
+    mean_fband: bool = True
+    max_fband: bool = True
+    max_allfbands: bool = True
+
+
+ListOfTwoStr = Annotated[list[str], Field(min_length=2, max_length=2)]
+
+
+class CoherenceSettings(NMBaseModel):
+    features: CoherenceFeatures = CoherenceFeatures()
+    method: CoherenceMethods = CoherenceMethods()
+    channels: list[ListOfTwoStr] = []
+    frequency_bands: list[str] = Field(default=["high_beta"], min_length=1)
+
+    @field_validator("frequency_bands")
+    def fbands_spaces_to_underscores(cls, frequency_bands):
+        return [f.replace(" ", "_") for f in frequency_bands]
+
+
+class CoherenceObject:
+    def __init__(
+        self,
+        sfreq: float,
+        window: str,
+        fbands: list[FrequencyRange],
+        fband_names: list[str],
+        ch_1_name: str,
+        ch_2_name: str,
+        ch_1_idx: int,
+        ch_2_idx: int,
+        coh: bool,
+        icoh: bool,
+        features_coh: CoherenceFeatures,
+    ) -> None:
+        self.sfreq = sfreq
+        self.window = window
+        self.fbands = fbands
+        self.fband_names = fband_names
+        self.ch_1 = ch_1_name
+        self.ch_2 = ch_2_name
+        self.ch_1_idx = ch_1_idx
+        self.ch_2_idx = ch_2_idx
+        self.coh = coh
+        self.icoh = icoh
+        self.features_coh = features_coh
+
+        self.Pxx = None
+        self.Pyy = None
+        self.Pxy = None
+        self.f = None
+        self.coh_val = None
+        self.icoh_val = None
+
+    def get_coh(self, feature_results, x, y):
+        from scipy.signal import welch, csd
+
+        self.f, self.Pxx = welch(x, self.sfreq, self.window, nperseg=128)
+        self.Pyy = welch(y, self.sfreq, self.window, nperseg=128)[1]
+        self.Pxy = csd(x, y, self.sfreq, self.window, nperseg=128)[1]
+
+        if self.coh:
+            self.coh_val = np.abs(self.Pxy**2) / (self.Pxx * self.Pyy)
+        if self.icoh:
+            self.icoh_val = np.array(self.Pxy / (self.Pxx * self.Pyy)).imag
+
+        for coh_idx, coh_type in enumerate([self.coh, self.icoh]):
+            if coh_type:
+                if coh_idx == 0:
+                    coh_val = self.coh_val
+                    coh_name = "coh"
+                else:
+                    coh_val = self.icoh_val
+                    coh_name = "icoh"
+
+            for idx, fband in enumerate(self.fbands):
+                if self.features_coh.mean_fband:
+                    feature_calc = np.mean(
+                        coh_val[np.bitwise_and(self.f > fband[0], self.f < fband[1])]
+                    )
+                    feature_name = "_".join(
+                        [
+                            coh_name,
+                            self.ch_1,
+                            "to",
+                            self.ch_2,
+                            "mean_fband",
+                            self.fband_names[idx],
+                        ]
+                    )
+                    feature_results[feature_name] = feature_calc
+                if self.features_coh.max_fband:
+                    feature_calc = np.max(
+                        coh_val[np.bitwise_and(self.f > fband[0], self.f < fband[1])]
+                    )
+                    feature_name = "_".join(
+                        [
+                            coh_name,
+                            self.ch_1,
+                            "to",
+                            self.ch_2,
+                            "max_fband",
+                            self.fband_names[idx],
+                        ]
+                    )
+                    feature_results[feature_name] = feature_calc
+            if self.features_coh.max_allfbands:
+                feature_calc = self.f[np.argmax(coh_val)]
+                feature_name = "_".join(
+                    [
+                        coh_name,
+                        self.ch_1,
+                        "to",
+                        self.ch_2,
+                        "max_allfbands",
+                        self.fband_names[idx],
+                    ]
+                )
+                feature_results[feature_name] = feature_calc
+        return feature_results
+
+
+class Coherence(NMFeature):
+    def __init__(
+        self, settings: "NMSettings", ch_names: list[str], sfreq: float
+    ) -> None:
+        self.settings = settings.coherence_settings
+        self.frequency_ranges_hz = settings.frequency_ranges_hz
+        self.sfreq = sfreq
+        self.ch_names = ch_names
+        self.coherence_objects: Iterable[CoherenceObject] = []
+
+        self.test_settings(settings, ch_names, sfreq)
+
+        for idx_coh in range(len(self.settings.channels)):
+            fband_names = self.settings.frequency_bands
+            fband_specs = []
+            for band_name in fband_names:
+                fband_specs.append(self.frequency_ranges_hz[band_name])
+
+            ch_1_name = self.settings.channels[idx_coh][0]
+            ch_1_name_reref = [ch for ch in self.ch_names if ch.startswith(ch_1_name)][
+                0
+            ]
+            ch_1_idx = self.ch_names.index(ch_1_name_reref)
+
+            ch_2_name = self.settings.channels[idx_coh][1]
+            ch_2_name_reref = [ch for ch in self.ch_names if ch.startswith(ch_2_name)][
+                0
+            ]
+            ch_2_idx = self.ch_names.index(ch_2_name_reref)
+
+            self.coherence_objects.append(
+                CoherenceObject(
+                    sfreq,
+                    "hann",
+                    fband_specs,
+                    fband_names,
+                    ch_1_name,
+                    ch_2_name,
+                    ch_1_idx,
+                    ch_2_idx,
+                    self.settings.method.coh,
+                    self.settings.method.icoh,
+                    self.settings.features,
+                )
+            )
+
+    @staticmethod
+    def test_settings(
+        settings: "NMSettings",
+        ch_names: Iterable[str],
+        sfreq: float,
+    ):
+        flat_channels = [
+            ch for ch_pair in settings.coherence_settings.channels for ch in ch_pair
+        ]
+
+        valid_coh_channel = [
+            sum(ch.startswith(ch_coh) for ch in ch_names) for ch_coh in flat_channels
+        ]
+        for ch_idx, ch_coh in enumerate(flat_channels):
+            if valid_coh_channel[ch_idx] == 0:
+                raise RuntimeError(
+                    f"Coherence selected channel {ch_coh} does not match any channel name: \n"
+                    f"  - settings.coherence_settings.channels: {settings.coherence_settings.channels}\n"
+                    f"  - ch_names: {ch_names} \n"
+                )
+
+            if valid_coh_channel[ch_idx] > 1:
+                raise RuntimeError(
+                    f"Coherence selected channel {ch_coh} is ambigous and matches more than one channel name: \n"
+                    f"  - settings.coherence_settings.channels: {settings.coherence_settings.channels}\n"
+                    f"  - ch_names: {ch_names} \n"
+                )
+
+        assert all(
+            f_band_coh in settings.frequency_ranges_hz
+            for f_band_coh in settings.coherence_settings.frequency_bands
+        ), (
+            "coherence selected frequency bands don't match the ones"
+            "specified in s['frequency_ranges_hz']"
+            f"coherence frequency bands: {settings.coherence_settings.frequency_bands}"
+            f"specified frequency_ranges_hz: {settings.frequency_ranges_hz}"
+        )
+
+        assert all(
+            settings.frequency_ranges_hz[fb][0] < sfreq / 2
+            and settings.frequency_ranges_hz[fb][1] < sfreq / 2
+            for fb in settings.coherence_settings.frequency_bands
+        ), (
+            "the coherence frequency band ranges need to be smaller than the Nyquist frequency"
+            f"got sfreq = {sfreq} and fband ranges {settings.coherence_settings.frequency_bands}"
+        )
+
+        if not settings.coherence_settings.method.get_enabled():
+            logger.warn(
+                "feature coherence enabled, but no coherence['method'] selected"
+            )
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        feature_results = {}
+
+        for coh_obj in self.coherence_objects:
+            feature_results = coh_obj.get_coh(
+                feature_results,
+                data[coh_obj.ch_1_idx, :],
+                data[coh_obj.ch_2_idx, :],
+            )
+
+        return feature_results

py_neuromodulation/features/feature_processor.py

@@ -0,0 +1,121 @@
+from typing import Type, TYPE_CHECKING
+
+from py_neuromodulation.utils.types import NMFeature, FeatureName
+
+if TYPE_CHECKING:
+    import numpy as np
+    from py_neuromodulation import NMSettings
+
+
+FEATURE_DICT: dict[FeatureName | str, str] = {
+    "raw_hjorth": "Hjorth",
+    "return_raw": "Raw",
+    "bandpass_filter": "BandPower",
+    "stft": "STFT",
+    "fft": "FFT",
+    "welch": "Welch",
+    "sharpwave_analysis": "SharpwaveAnalyzer",
+    "fooof": "FooofAnalyzer",
+    "nolds": "Nolds",
+    "coherence": "Coherence",
+    "bursts": "Bursts",
+    "linelength": "LineLength",
+    "mne_connectivity": "MNEConnectivity",
+    "bispectrum": "Bispectra",
+}
+
+
+class FeatureProcessors:
+    """Class for storing NMFeature objects and calculating features during processing"""
+
+    def __init__(
+        self, settings: "NMSettings", ch_names: list[str], sfreq: float
+    ) -> None:
+        """Initialize FeatureProcessors object with settings, channel names and sampling frequency.
+
+        Args:
+            settings (NMSettings): PyNM settings object
+            ch_names (list[str]): list of channel names
+            sfreq (float): sampling frequency in Hz
+        """
+        from py_neuromodulation import user_features
+        from importlib import import_module
+
+        # Accept 'str' for custom features
+        self.features: dict[FeatureName | str, NMFeature] = {
+            feature_name: getattr(
+                import_module("py_neuromodulation.features"), FEATURE_DICT[feature_name]
+            )(settings, ch_names, sfreq)
+            for feature_name in settings.features.get_enabled()
+        }
+
+        for feature_name, feature in user_features.items():
+            self.features[feature_name] = feature(settings, ch_names, sfreq)
+
+    def register_new_feature(self, feature_name: str, feature: NMFeature) -> None:
+        """Register new feature.
+
+        Parameters
+        ----------
+        feature : features_abc.Feature
+            New feature to add to feature list
+        """
+        self.features[feature_name] = feature  # type: ignore
+
+    def estimate_features(self, data: "np.ndarray") -> dict:
+        """Calculate features, as defined in settings.json
+        Features are based on bandpower, raw Hjorth parameters and sharp wave
+        characteristics.
+
+        Parameters
+        ----------
+        data (np array) : (channels, time)
+
+        Returns
+        -------
+        dat (dict): naming convention : channel_method_feature_(f_band)
+        """
+
+        feature_results: dict = {}
+
+        for feature in self.features.values():
+            feature_results.update(feature.calc_feature(data))
+
+        return feature_results
+
+    def get_feature(self, fname: FeatureName) -> NMFeature:
+        return self.features[fname]
+
+
+def add_custom_feature(feature_name: str, new_feature: Type[NMFeature]):
+    """Add a custom feature to the dictionary of user-defined features.
+        The feature will be automatically enabled in the settings,
+        and computed when the Stream.run() method is called.
+
+    Args:
+        feature_name (str): A name for the feature that will be used to
+        enable/disable the feature in settings and to store the feature
+        class instance in the DataProcessor
+
+        new_feature (NMFeature): Class that implements the user-defined
+        feature. It should implement the NMSettings protocol (defined
+        in this file).
+    """
+    from py_neuromodulation import user_features
+    from py_neuromodulation import NMSettings
+
+    user_features[feature_name] = new_feature
+    NMSettings._add_feature(feature_name)
+
+
+def remove_custom_feature(feature_name: str):
+    """Remove a custom feature from the dictionary of user-defined features.
+
+    Args:
+        feature_name (str): Name of the feature to remove
+    """
+    from py_neuromodulation import user_features
+    from py_neuromodulation import NMSettings
+
+    user_features.pop(feature_name)
+    NMSettings._remove_feature(feature_name)

py_neuromodulation/features/fooof.py

@@ -0,0 +1,142 @@
+from collections.abc import Iterable
+import numpy as np
+
+from typing import TYPE_CHECKING
+
+from py_neuromodulation.utils.types import (
+    NMBaseModel,
+    NMFeature,
+    BoolSelector,
+    FrequencyRange,
+)
+
+if TYPE_CHECKING:
+    from py_neuromodulation 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_settings
+        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 (
+            self.settings.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 (
+            self.settings.freq_range_hz[0] < sfreq
+            and self.settings.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/features/hjorth_raw.py

@@ -0,0 +1,57 @@
+"""
+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
+"""
+
+from typing import TYPE_CHECKING
+import numpy as np
+from collections.abc import Sequence
+
+from py_neuromodulation.utils.types import NMFeature
+
+if TYPE_CHECKING:
+    from py_neuromodulation import NMSettings
+
+
+class Hjorth(NMFeature):
+    def __init__(
+        self, settings: "NMSettings", ch_names: Sequence[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: "NMSettings", ch_names: Sequence[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/features/linelength.py

@@ -0,0 +1,21 @@
+import numpy as np
+from collections.abc import Sequence
+
+from py_neuromodulation.features.feature_processor 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