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

py_neuromodulation/features/mne_connectivity.py

@@ -1,18 +1,44 @@
 from collections.abc import Iterable
 import numpy as np
-from typing import TYPE_CHECKING
+
+from typing import TYPE_CHECKING, Annotated, Literal
+from pydantic import Field
 
 from py_neuromodulation.utils.types import NMFeature, NMBaseModel
+from py_neuromodulation.utils.pydantic_extensions import NMField
 
 if TYPE_CHECKING:
     from py_neuromodulation import NMSettings
-    from mne.io import RawArray
-    from mne import Epochs
+
+
+ListOfTwoStr = Annotated[list[str], Field(min_length=2, max_length=2)]
+
+
+MNE_CONNECTIVITY_METHOD = Literal[
+    "coh",
+    "cohy",
+    "imcoh",
+    "cacoh",
+    "mic",
+    "mim",
+    "plv",
+    "ciplv",
+    "ppc",
+    "pli",
+    "dpli",
+    "wpli",
+    "wpli2_debiased",
+    "gc",
+    "gc_tr",
+]
+
+MNE_CONNECTIVITY_MODE = Literal["multitaper", "fourier", "cwt_morlet"]
 
 
 class MNEConnectivitySettings(NMBaseModel):
-    method: str = "plv"
-    mode: str = "multitaper"
+    method: MNE_CONNECTIVITY_METHOD = NMField(default="plv")
+    mode: MNE_CONNECTIVITY_MODE = NMField(default="multitaper")
+    channels: list[ListOfTwoStr] = []
 
 
 class MNEConnectivity(NMFeature):
@@ -22,102 +48,42 @@ class MNEConnectivity(NMFeature):
         ch_names: Iterable[str],
         sfreq: float,
     ) -> None:
-        from mne import create_info
-
         self.settings = settings
 
         self.ch_names = ch_names
         self.sfreq = sfreq
 
+        self.channels = settings.mne_connectivity_settings.channels
+
         # Params used by spectral_connectivity_epochs
         self.mode = settings.mne_connectivity_settings.mode
         self.method = settings.mne_connectivity_settings.method
+        self.indices = ([], [])  # convert channel names to channel indices in data
+        for con_idx in range(len(self.channels)):
+            seed_name = self.channels[con_idx][0]
+            target_name = self.channels[con_idx][1]
+            seed_name_reref = [ch for ch in self.ch_names if ch.startswith(seed_name)][0]
+            target_name_reref = [ch for ch in self.ch_names if ch.startswith(target_name)][0]
+            self.indices[0].append(self.ch_names.index(seed_name_reref))
+            self.indices[1].append(self.ch_names.index(target_name_reref))
 
         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,
+            data=np.expand_dims(data, axis=0),  # add singleton epoch dimension
             sfreq=self.sfreq,
             method=self.method,
             mode=self.mode,
-            indices=(np.array([0, 0, 1, 1]), np.array([2, 3, 2, 3])),
+            indices=self.indices,
             verbose=False,
         )
         dat_conn: np.ndarray = spec_out.get_data()
@@ -127,20 +93,27 @@ class MNEConnectivity(NMFeature):
             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])
+                        (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]])
-                )
+        for con_idx in np.arange(dat_conn.shape[0]):
+            for fband_idx, fband_name in enumerate(self.fbands):
+                # TODO: Add support for max_fband and max_allfbands
+                feature_results[
+                    "_".join(
+                        [
+                            self.method,
+                            self.channels[con_idx][0],  # seed channel name
+                            "to",
+                            self.channels[con_idx][1],  # target channel name
+                            "mean_fband",
+                            fband_name,
+                        ]
+                    )
+                ] = np.mean(dat_conn[con_idx, self.fband_ranges[fband_idx]])
 
         # Store current experiment parameters to check if re-initialization is needed
         self.prev_batch_shape = data.shape

py_neuromodulation/features/oscillatory.py

@@ -3,6 +3,7 @@ import numpy as np
 from itertools import product
 
 from py_neuromodulation.utils.types import NMBaseModel, BoolSelector, NMFeature
+from py_neuromodulation.utils.pydantic_extensions import NMField
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -17,12 +18,12 @@ class OscillatoryFeatures(BoolSelector):
 
 
 class OscillatorySettings(NMBaseModel):
-    windowlength_ms: int = 1000
+    windowlength_ms: int = NMField(1000, gt=0, custom_metadata={"unit": "ms"})
     log_transform: bool = True
     features: OscillatoryFeatures = OscillatoryFeatures(
         mean=True, median=False, std=False, max=False
     )
-    return_spectrum: bool = False
+    return_spectrum: bool = True
 
 
 ESTIMATOR_DICT = {
@@ -176,7 +177,7 @@ class Welch(OscillatoryFeature):
         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]
+                feature_results[f"{ch_name}_welch_psd_{int(f)}"] = Z[ch_idx][idx]
 
         return feature_results
 
@@ -242,7 +243,7 @@ class STFT(OscillatoryFeature):
         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(
+                feature_results[f"{ch_name}_stft_psd_{int(f)}"] = Z[ch_idx].mean(
                     axis=1
                 )[idx]
 

py_neuromodulation/features/sharpwaves.py

@@ -267,6 +267,14 @@ class SharpwaveAnalyzer(NMFeature):
 
             # for each feature take the respective fun.
             for feature_name, estimator_name, estimator in estimator_combinations:
+                if feature_name == "num_peaks":
+                    key_name = f"{ch_name}_Sharpwave_{feature_name}_{filter_name}"
+                    if len(waveform_results[feature_name]) == 1:
+                        dict_ch_features[key_name][key_name_pt] = waveform_results[feature_name][0]
+                        continue
+                    else:
+                        raise ValueError("num_peaks should be a list with length 1")
+                    # there can be only one num_peak in each batch
                 feature_data = waveform_results[feature_name]
                 key_name = f"{ch_name}_Sharpwave_{estimator_name.title()}_{feature_name}_{filter_name}"
 
@@ -280,12 +288,25 @@ class SharpwaveAnalyzer(NMFeature):
 
             # the key_name stays, since the estimator function stays between peaks and troughs
             for key_name, estimator in self.estimator_key_map.items():
+                if len(dict_ch_features[key_name]) == 0:
+                    # might happen if num_peaks was written in estimator
+                    # e.g. estimator["mean"] = ["num_peaks"]
+                    # for conveniance this doesn't raise an exception
+                    continue
+
                 feature_results[key_name] = estimator(
                     [
                         list(dict_ch_features[key_name].values())[0],
                         list(dict_ch_features[key_name].values())[1],
                     ]
                 )
+            # add here also the num_peaks features
+            if self.sw_settings.sharpwave_features.num_peaks:
+                for ch_name in self.ch_names:
+                    for filter_name in self.filter_names:
+                        key_name = f"{ch_name}_Sharpwave_num_peaks_{filter_name}"
+                        feature_results[key_name] = np_mean([dict_ch_features[key_name]["Peak"],
+                                                            dict_ch_features[key_name]["Trough"]])
         else:
             # otherwise, save all write all "flattened" key value pairs in feature_results
             for key, subdict in dict_ch_features.items():

py_neuromodulation/filter/kalman_filter.py

@@ -1,7 +1,9 @@
 import numpy as np
 from typing import TYPE_CHECKING
 
+
 from py_neuromodulation.utils.types import NMBaseModel
+from py_neuromodulation.utils.pydantic_extensions import NMErrorList
 
 
 if TYPE_CHECKING:
@@ -22,13 +24,22 @@ class KalmanSettings(NMBaseModel):
         "HFA",
     ]
 
-    def validate_fbands(self, settings: "NMSettings") -> None:
-        assert all(
+    def validate_fbands(self, settings: "NMSettings") -> NMErrorList:
+        errors: NMErrorList = NMErrorList()
+
+        if not 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."
-        )
+        ):
+            errors.add_error(
+                "Frequency bands for Kalman filter must also be specified in "
+                "frequency_ranges_hz.",
+                location=[
+                    "kalman_filter_settings",
+                    "frequency_bands",
+                ],
+            )
+
+        return errors
 
 
 def define_KF(Tp, sigma_w, sigma_v):

py_neuromodulation/gui/__init__.py

@@ -0,0 +1,3 @@
+from .backend.app_manager import AppManager as App
+
+__all__ = ["App"]