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

py_neuromodulation/features/mne_connectivity.py

@@ -0,0 +1,148 @@
+from collections.abc import Iterable
+import numpy as np
+from typing import TYPE_CHECKING
+
+from py_neuromodulation.utils.types import NMFeature, NMBaseModel
+
+if TYPE_CHECKING:
+    from py_neuromodulation import NMSettings
+    from mne.io import RawArray
+    from mne import Epochs
+
+
+class MNEConnectivitySettings(NMBaseModel):
+    method: str = "plv"
+    mode: str = "multitaper"
+
+
+class MNEConnectivity(NMFeature):
+    def __init__(
+        self,
+        settings: "NMSettings",
+        ch_names: Iterable[str],
+        sfreq: float,
+    ) -> None:
+        from mne import create_info
+
+        self.settings = settings
+
+        self.ch_names = ch_names
+        self.sfreq = sfreq
+
+        # Params used by spectral_connectivity_epochs
+        self.mode = settings.mne_connectivity_settings.mode
+        self.method = settings.mne_connectivity_settings.method
+
+        self.fbands = settings.frequency_ranges_hz
+        self.fband_ranges: list = []
+        self.result_keys = []
+
+        self.raw_info = create_info(ch_names=self.ch_names, sfreq=self.sfreq)
+        self.raw_array: "RawArray"
+        self.epochs: "Epochs"
+        self.prev_batch_shape: tuple = (-1, -1)  # sentinel value
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        from mne.io import RawArray
+        from mne import Epochs
+        from mne_connectivity import spectral_connectivity_epochs
+        import pandas as pd
+
+        time_samples_s = data.shape[1] / self.sfreq
+        epoch_length: float = 1  # TODO: Make this a parameter?
+
+        if epoch_length > time_samples_s:
+            raise ValueError(
+                f"the intended epoch length for mne connectivity: {epoch_length}s"
+                f" are longer than the passed data array {np.round(time_samples_s, 2)}s"
+            )
+
+        # Only reinitialize the raw_array and epochs object if the data shape has changed
+        # That could mean that the channels have been re-selected, or we're in the last batch
+        # TODO: If sfreq or channels change, do we re-initialize the whole Stream object?
+        if data.shape != self.prev_batch_shape:
+            self.raw_array = RawArray(
+                data=data,
+                info=self.raw_info,
+                copy=None,  # type: ignore
+                verbose=False,
+            )
+
+            # self.events = make_fixed_length_events(self.raw_array, duration=epoch_length)
+            # Equivalent code for those parameters:
+            event_times = np.arange(
+                0, data.shape[-1], self.sfreq * epoch_length, dtype=int
+            )
+            events = np.column_stack(
+                (
+                    event_times,
+                    np.zeros_like(event_times, dtype=int),
+                    np.ones_like(event_times, dtype=int),
+                )
+            )
+
+            # there need to be minimum 2 of two epochs, otherwise mne_connectivity
+            # is not correctly initialized
+            if events.shape[0] < 2:
+                raise RuntimeError(
+                    f"A minimum of 2 epochs is required for mne_connectivity,"
+                    f" got only {events.shape[0]}. Increase settings['segment_length_features_ms']"
+                )
+
+            self.epochs = Epochs(
+                self.raw_array,
+                events=events,
+                event_id={"rest": 1},
+                tmin=0,
+                tmax=epoch_length,
+                baseline=None,
+                reject_by_annotation=True,
+                verbose=False,
+            )
+
+            # Trick the function "spectral_connectivity_epochs" into not calling "add_annotations_to_metadata"
+            # TODO: This is a hack, and maybe needs a fix in the mne_connectivity library
+            self.epochs._metadata = pd.DataFrame(index=np.arange(events.shape[0]))
+
+        else:
+            # As long as the initialization parameters, channels, sfreq and batch size are the same
+            # We can re-use the existing epochs object by updating the raw data
+            self.raw_array._data = data
+            self.epochs._raw = self.raw_array
+
+        # n_jobs is here kept to 1, since setup of the multiprocessing Pool
+        # takes longer than most batch computing sizes
+        spec_out = spectral_connectivity_epochs(
+            data=self.epochs,
+            sfreq=self.sfreq,
+            method=self.method,
+            mode=self.mode,
+            indices=(np.array([0, 0, 1, 1]), np.array([2, 3, 2, 3])),
+            verbose=False,
+        )
+        dat_conn: np.ndarray = spec_out.get_data()
+
+        # Get frequency band ranges only for the first batch, it's already the same
+        if len(self.fband_ranges) == 0:
+            for fband_range in self.fbands.values():
+                self.fband_ranges.append(
+                    np.where(
+                        (np.array(spec_out.freqs) > fband_range[0])
+                        & (np.array(spec_out.freqs) < fband_range[1])
+                    )[0]
+                )
+
+        # TODO: If I compute the mean for the entire fband, results are almost the same before
+        # normalization (0.9999999... vs 1.0), but some change wildly after normalization (-3 vs 0)
+        # Investigate why, is this a bug in normalization?
+        feature_results = {}
+        for conn in np.arange(dat_conn.shape[0]):
+            for fband_idx, fband in enumerate(self.fbands):
+                feature_results["_".join(["ch1", self.method, str(conn), fband])] = (
+                    np.mean(dat_conn[conn, self.fband_ranges[fband_idx]])
+                )
+
+        # Store current experiment parameters to check if re-initialization is needed
+        self.prev_batch_shape = data.shape
+
+        return feature_results

py_neuromodulation/features/nolds.py

@@ -0,0 +1,94 @@
+import numpy as np
+from collections.abc import Iterable
+
+from typing import TYPE_CHECKING
+
+from py_neuromodulation.utils.types import NMFeature, BoolSelector, NMBaseModel
+
+from pydantic import field_validator
+
+if TYPE_CHECKING:
+    from py_neuromodulation import NMSettings
+
+
+class NoldsFeatures(BoolSelector):
+    sample_entropy: bool = False
+    correlation_dimension: bool = False
+    lyapunov_exponent: bool = True
+    hurst_exponent: bool = False
+    detrended_fluctuation_analysis: bool = False
+
+
+class NoldsSettings(NMBaseModel):
+    raw: bool = True
+    frequency_bands: list[str] = ["low_beta"]
+    features: NoldsFeatures = NoldsFeatures()
+
+    @field_validator("frequency_bands")
+    def fbands_spaces_to_underscores(cls, frequency_bands):
+        return [f.replace(" ", "_") for f in frequency_bands]
+
+
+class Nolds(NMFeature):
+    def __init__(
+        self, settings: "NMSettings", ch_names: Iterable[str], sfreq: float
+    ) -> None:
+        self.settings = settings.nolds_settings
+        self.ch_names = ch_names
+
+        if len(self.settings.frequency_bands) > 0:
+            from py_neuromodulation.features.bandpower import BandPower
+
+            self.bp_filter = BandPower(settings, ch_names, sfreq, use_kf=False)
+
+        # Check if the selected frequency bands are defined in the global settings
+        for fb in settings.nolds_settings.frequency_bands:
+            assert (
+                fb in settings.frequency_ranges_hz
+            ), f"{fb} selected in nolds_features, but not defined in s['frequency_ranges_hz']"
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        feature_results = {}
+        data = np.nan_to_num(data)
+        if self.settings.raw:
+            feature_results = self.calc_nolds(data, feature_results)
+        if len(self.settings.frequency_bands) > 0:
+            data_filt = self.bp_filter.bandpass_filter.filter_data(data)
+
+            for f_band_idx, f_band in enumerate(self.settings.frequency_bands):
+                # filter data now for a specific fband and pass to calc_nolds
+                feature_results = self.calc_nolds(
+                    data_filt[:, f_band_idx, :], feature_results, f_band
+                )  # ch, bands, samples
+        return feature_results
+
+    def calc_nolds(
+        self, data: np.ndarray, feature_results: dict, data_str: str = "raw"
+    ) -> dict:
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            for f_name in self.settings.features.get_enabled():
+                feature_results[f"{ch_name}_nolds_{f_name}_{data_str}"] = (
+                    self.calc_nolds_feature(f_name, data[ch_idx, :])
+                    if data[ch_idx, :].sum()
+                    else 0
+                )
+
+        return feature_results
+
+    @staticmethod
+    def calc_nolds_feature(f_name: str, dat: np.ndarray):
+        import nolds
+
+        match f_name:
+            case "sample_entropy":
+                return nolds.sampen(dat)
+            case "correlation_dimension":
+                return nolds.corr_dim(dat, emb_dim=2)
+            case "lyapunov_exponent":
+                return nolds.lyap_r(dat)
+            case "hurst_exponent":
+                return nolds.hurst_rs(dat)
+            case "detrended_fluctuation_analysis":
+                return nolds.dfa(dat)
+            case _:
+                raise ValueError(f"Invalid nolds feature name: {f_name}")

py_neuromodulation/features/oscillatory.py

@@ -0,0 +1,249 @@
+from collections.abc import Sequence
+import numpy as np
+from itertools import product
+
+from py_neuromodulation.utils.types import NMBaseModel, BoolSelector, NMFeature
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from py_neuromodulation.stream.settings import NMSettings
+
+
+class OscillatoryFeatures(BoolSelector):
+    mean: bool = True
+    median: bool = False
+    std: bool = False
+    max: bool = False
+
+
+class OscillatorySettings(NMBaseModel):
+    windowlength_ms: int = 1000
+    log_transform: bool = True
+    features: OscillatoryFeatures = OscillatoryFeatures(
+        mean=True, median=False, std=False, max=False
+    )
+    return_spectrum: bool = False
+
+
+ESTIMATOR_DICT = {
+    "mean": np.nanmean,
+    "median": np.nanmedian,
+    "std": np.nanstd,
+    "max": np.nanmax,
+}
+
+
+class OscillatoryFeature(NMFeature):
+    def __init__(
+        self, settings: "NMSettings", ch_names: Sequence[str], sfreq: int
+    ) -> None:
+        settings.validate()
+        self.settings: OscillatorySettings  # Assignment in subclass __init__
+        self.osc_feature_name: str  # Required for output
+
+        self.sfreq = int(sfreq)
+        self.ch_names = ch_names
+
+        self.frequency_ranges = settings.frequency_ranges_hz
+
+        # Test settings
+        assert self.settings.windowlength_ms <= settings.segment_length_features_ms, (
+            f"oscillatory feature windowlength_ms = ({self.settings.windowlength_ms})"
+            f"needs to be smaller than"
+            f"settings['segment_length_features_ms'] = {settings.segment_length_features_ms}",
+        )
+
+
+class FFT(OscillatoryFeature):
+    def __init__(
+        self,
+        settings: "NMSettings",
+        ch_names: Sequence[str],
+        sfreq: int,
+    ) -> None:
+        from scipy.fft import rfftfreq
+
+        self.osc_feature_name = "fft"
+        self.settings = settings.fft_settings
+        # super.__init__ needs osc_feature_name and settings
+        super().__init__(settings, ch_names, sfreq)
+
+        window_ms = self.settings.windowlength_ms
+
+        self.window_samples = int(-np.floor(window_ms / 1000 * sfreq))
+        self.freqs = rfftfreq(-self.window_samples, 1 / np.floor(self.sfreq))
+
+        # Pre-calculate frequency ranges
+        self.idx_range = [
+            (
+                f_band,
+                np.where((self.freqs >= f_range[0]) & (self.freqs < f_range[1]))[0],
+            )
+            for f_band, f_range in self.frequency_ranges.items()
+        ]
+
+        self.estimators = [
+            (est, ESTIMATOR_DICT[est]) for est in self.settings.features.get_enabled()
+        ]
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        data = data[:, self.window_samples :]
+
+        from scipy.fft import rfft
+
+        Z = np.abs(rfft(data))  # type: ignore
+
+        if self.settings.log_transform:
+            Z = np.log10(Z)
+
+        feature_results = {}
+
+        for f_band_name, idx_range in self.idx_range:
+            # TODO Can we get rid of this for-loop? Hard to vectorize windows of different lengths...
+            Z_band = Z[:, idx_range]  # Data for all channels
+
+            for est_name, est_fun in self.estimators:
+                result = est_fun(Z_band, axis=1)
+
+                for ch_idx, ch_name in enumerate(self.ch_names):
+                    feature_results[
+                        f"{ch_name}_{self.osc_feature_name}_{f_band_name}_{est_name}"
+                    ] = result[ch_idx]
+
+        if self.settings.return_spectrum:
+            combinations = product(enumerate(self.ch_names), enumerate(self.freqs))
+            for (ch_idx, ch_name), (idx, f) in combinations:
+                feature_results[f"{ch_name}_fft_psd_{int(f)}"] = Z[ch_idx][idx]
+
+        return feature_results
+
+
+class Welch(OscillatoryFeature):
+    def __init__(
+        self,
+        settings: "NMSettings",
+        ch_names: Sequence[str],
+        sfreq: int,
+    ) -> None:
+        from scipy.fft import rfftfreq
+
+        self.osc_feature_name = "welch"
+        self.settings = settings.welch_settings
+        # super.__init__ needs osc_feature_name and settings
+        super().__init__(settings, ch_names, sfreq)
+
+        self.freqs = rfftfreq(self.sfreq, 1 / self.sfreq)
+
+        self.idx_range = [
+            (
+                f_band,
+                np.where((self.freqs >= f_range[0]) & (self.freqs < f_range[1]))[0],
+            )
+            for f_band, f_range in self.frequency_ranges.items()
+        ]
+
+        self.estimators = [
+            (est, ESTIMATOR_DICT[est]) for est in self.settings.features.get_enabled()
+        ]
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        from scipy.signal import welch
+
+        _, Z = welch(
+            data,
+            fs=self.sfreq,
+            window="hann",
+            nperseg=self.sfreq,
+            noverlap=None,
+        )
+
+        if self.settings.log_transform:
+            Z = np.log10(Z)
+
+        feature_results = {}
+
+        for f_band_name, idx_range in self.idx_range:
+            Z_band = Z[:, idx_range]
+
+            for est_name, est_fun in self.estimators:
+                result = est_fun(Z_band, axis=1)
+
+                for ch_idx, ch_name in enumerate(self.ch_names):
+                    feature_results[
+                        f"{ch_name}_{self.osc_feature_name}_{f_band_name}_{est_name}"
+                    ] = result[ch_idx]
+
+        if self.settings.return_spectrum:
+            combinations = product(enumerate(self.ch_names), enumerate(self.freqs))
+            for (ch_idx, ch_name), (idx, f) in combinations:
+                feature_results[f"{ch_name}_welch_psd_{str(f)}"] = Z[ch_idx][idx]
+
+        return feature_results
+
+
+class STFT(OscillatoryFeature):
+    def __init__(
+        self,
+        settings: "NMSettings",
+        ch_names: Sequence[str],
+        sfreq: int,
+    ) -> None:
+        from scipy.fft import rfftfreq
+
+        self.osc_feature_name = "stft"
+        self.settings = settings.stft_settings
+        # super.__init__ needs osc_feature_name and settings
+        super().__init__(settings, ch_names, sfreq)
+
+        self.nperseg = self.settings.windowlength_ms
+
+        self.freqs = rfftfreq(self.nperseg, 1 / self.sfreq)
+
+        self.idx_range = [
+            (
+                f_band,
+                np.where((self.freqs >= f_range[0]) & (self.freqs <= f_range[1]))[0],
+            )
+            for f_band, f_range in self.frequency_ranges.items()
+        ]
+
+        self.estimators = [
+            (est, ESTIMATOR_DICT[est]) for est in self.settings.features.get_enabled()
+        ]
+
+    def calc_feature(self, data: np.ndarray) -> dict:
+        from scipy.signal import stft
+
+        _, _, Zxx = stft(
+            data,
+            fs=self.sfreq,
+            window="hamming",
+            nperseg=self.nperseg,
+            boundary="even",
+        )
+
+        Z = np.abs(Zxx)
+        if self.settings.log_transform:
+            Z = np.log10(Z)
+
+        feature_results = {}
+
+        for f_band_name, idx_range in self.idx_range:
+            Z_band = Z[:, idx_range, :]
+
+            for est_name, est_fun in self.estimators:
+                result = est_fun(Z_band, axis=(1, 2))
+
+                for ch_idx, ch_name in enumerate(self.ch_names):
+                    feature_results[
+                        f"{ch_name}_{self.osc_feature_name}_{f_band_name}_{est_name}"
+                    ] = result[ch_idx]
+
+        if self.settings.return_spectrum:
+            combinations = product(enumerate(self.ch_names), enumerate(self.freqs))
+            for (ch_idx, ch_name), (idx, f) in combinations:
+                feature_results[f"{ch_name}_stft_psd_{str(f)}"] = Z[ch_idx].mean(
+                    axis=1
+                )[idx]
+
+        return feature_results