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

py_neuromodulation/nm_normalization.py

@@ -3,8 +3,6 @@ from enum import Enum
 
 from sklearn import preprocessing
 import numpy as np
-
-
 class NORM_METHODS(Enum):
     MEAN = "mean"
     MEDIAN = "median"
@@ -138,6 +136,17 @@ class FeatureNormalizer:
 
         return data
 
+"""
+Functions to check for NaN's before deciding which Numpy function to call
+"""
+def nan_mean(data, axis):
+    return np.nanmean(data, axis=axis) if np.any(np.isnan(sum(data))) else np.mean(data, axis=axis)
+
+def nan_std(data, axis):
+    return np.nanstd(data, axis=axis) if np.any(np.isnan(sum(data))) else np.std(data, axis=axis)
+
+def nan_median(data, axis):
+    return np.nanmedian(data, axis=axis) if np.any(np.isnan(sum(data))) else np.median(data, axis=axis)
 
 def _normalize_and_clip(
     current: np.ndarray,
@@ -147,82 +156,49 @@ def _normalize_and_clip(
     description: str,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Normalize data."""
-    if method == NORM_METHODS.MEAN.value:
-        mean = np.nanmean(previous, axis=0)
-        current = (current - mean) / mean
-    elif method == NORM_METHODS.MEDIAN.value:
-        median = np.nanmedian(previous, axis=0)
-        current = (current - median) / median
-    elif method == NORM_METHODS.ZSCORE.value:
-        mean = np.nanmean(previous, axis=0)
-        current = (current - mean) / np.nanstd(previous, axis=0)
-    elif method == NORM_METHODS.ZSCORE_MEDIAN.value:
-        current = (current - np.nanmedian(previous, axis=0)) / np.nanstd(
-            previous, axis=0
-        )
-    # For the following methods we check for the shape of current
-    # when current is a 1D array, then it is the post-processing normalization,
-    # and we need to expand, and take the [0, :] component
-    # When current is a 2D array, then it is pre-processing normalization, and
-    # there's no need for expanding.
-    elif method == NORM_METHODS.QUANTILE.value:
-        if len(current.shape) == 1:
-            current = (
-                preprocessing.QuantileTransformer(n_quantiles=300)
-                .fit(np.nan_to_num(previous))
-                .transform(np.expand_dims(current, axis=0))[0, :]
-            )
-        else:
-            current = (
-                preprocessing.QuantileTransformer(n_quantiles=300)
-                .fit(np.nan_to_num(previous))
-                .transform(current)
-            )
-    elif method == NORM_METHODS.ROBUST.value:
-        if len(current.shape) == 1:
-            current = (
-                preprocessing.RobustScaler()
-                .fit(np.nan_to_num(previous))
-                .transform(np.expand_dims(current, axis=0))[0, :]
-            )
-        else:
+    match method:
+        case NORM_METHODS.MEAN.value:
+            mean = nan_mean(previous, axis=0)
+            current = (current - mean) / mean
+        case NORM_METHODS.MEDIAN.value:
+            median = nan_median(previous, axis=0)
+            current = (current - median) / median
+        case NORM_METHODS.ZSCORE.value:
+            current = (current - nan_mean(previous, axis=0)) / nan_std(previous, axis=0)
+        case NORM_METHODS.ZSCORE_MEDIAN.value:
+            current = (current - nan_median(previous, axis=0)) / nan_std(previous, axis=0)
+        # For the following methods we check for the shape of current
+        # when current is a 1D array, then it is the post-processing normalization,
+        # and we need to expand, and remove the extra dimension afterwards
+        # When current is a 2D array, then it is pre-processing normalization, and
+        # there's no need for expanding.
+        case (NORM_METHODS.QUANTILE.value | 
+              NORM_METHODS.ROBUST.value | 
+              NORM_METHODS.MINMAX.value | 
+              NORM_METHODS.POWER.value):
+            
+            norm_methods = {
+                NORM_METHODS.QUANTILE.value : lambda: preprocessing.QuantileTransformer(n_quantiles=300),
+                NORM_METHODS.ROBUST.value : preprocessing.RobustScaler,
+                NORM_METHODS.MINMAX.value : preprocessing.MinMaxScaler,
+                NORM_METHODS.POWER.value : preprocessing.PowerTransformer
+            }
+                
             current = (
-                preprocessing.RobustScaler()
+                norm_methods[method]()
                 .fit(np.nan_to_num(previous))
-                .transform(current)
+                .transform(
+                    # if post-processing: pad dimensions to 2
+                    np.reshape(current, (2-len(current.shape))*(1,) + current.shape)
+                    )
+                .squeeze() # if post-processing: remove extra dimension
             )
-
-    elif method == NORM_METHODS.MINMAX.value:
-        if len(current.shape) == 1:
-            current = (
-                preprocessing.MinMaxScaler()
-                .fit(np.nan_to_num(previous))
-                .transform(np.expand_dims(current, axis=0))[0, :]
-            )
-        else:
-            current = (
-                preprocessing.MinMaxScaler()
-                .fit(np.nan_to_num(previous))
-                .transform(current)
-            )
-    elif method == NORM_METHODS.POWER.value:
-        if len(current.shape) == 1:
-            current = (
-                preprocessing.PowerTransformer()
-                .fit(np.nan_to_num(previous))
-                .transform(np.expand_dims(current, axis=0))[0, :]
+            
+        case _:
+            raise ValueError(
+                f"Only {[e.value for e in NORM_METHODS]} are supported as "
+                f"{description} normalization methods. Got {method}."
             )
-        else:
-            current = (
-                preprocessing.PowerTransformer()
-                .fit(np.nan_to_num(previous))
-                .transform(current)
-            )
-    else:
-        raise ValueError(
-            f"Only {[e.value for e in NORM_METHODS]} are supported as "
-            f"{description} normalization methods. Got {method}."
-        )
 
     if clip:
         current = _clip(data=current, clip=clip)

py_neuromodulation/nm_oscillatory.py

@@ -38,6 +38,15 @@ class OscillatoryFeature(nm_features_abc.Feature):
             assert isinstance(
                 s[osc_feature_name]["windowlength_ms"], int
             ), f"windowlength_ms needs to be type int, got {s[osc_feature_name]['windowlength_ms']}"
+
+            assert (
+                s[osc_feature_name]["windowlength_ms"]
+                <= s["segment_length_features_ms"]
+            ), (
+                f"oscillatory feature windowlength_ms = ({s[osc_feature_name]['windowlength_ms']})"
+                f"needs to be smaller than"
+                f"s['segment_length_features_ms'] = {s['segment_length_features_ms']}",
+            )
         else:
             for seg_length in s[osc_feature_name][
                 "segment_lengths_ms"
@@ -48,12 +57,6 @@ class OscillatoryFeature(nm_features_abc.Feature):
         assert isinstance(
             s[osc_feature_name]["log_transform"], bool
         ), f"log_transform needs to be type bool, got {s[osc_feature_name]['log_transform']}"
-        assert isinstance(
-            s[osc_feature_name]["kalman_filter"], bool
-        ), f"kalman_filter needs to be type bool, got {s[osc_feature_name]['kalman_filter']}"
-
-        if s[osc_feature_name]["kalman_filter"] is True:
-            nm_kalmanfilter.test_kf_settings(s, ch_names, sfreq)
 
         assert isinstance(s["frequency_ranges_hz"], dict)
 
@@ -95,6 +98,36 @@ class OscillatoryFeature(nm_features_abc.Feature):
             feature_calc = self.KF_dict[KF_name].x[0]
         return feature_calc
 
+    def estimate_osc_features(
+        self,
+        features_compute: dict,
+        data: np.ndarray,
+        feature_name: np.ndarray,
+        est_name: str,
+    ):
+        for feature_est_name in list(self.s[est_name]["features"].keys()):
+            if self.s[est_name]["features"][feature_est_name] is True:
+                # switch case for feature_est_name
+                match feature_est_name:
+                    case "mean":
+                        features_compute[
+                            f"{feature_name}_{feature_est_name}"
+                        ] = np.nanmean(data)
+                    case "median":
+                        features_compute[
+                            f"{feature_name}_{feature_est_name}"
+                        ] = np.nanmedian(data)
+                    case "std":
+                        features_compute[
+                            f"{feature_name}_{feature_est_name}"
+                        ] = np.nanstd(data)
+                    case "max":
+                        features_compute[
+                            f"{feature_name}_{feature_est_name}"
+                        ] = np.nanmax(data)
+
+        return features_compute
+
 
 class FFT(OscillatoryFeature):
     def __init__(
@@ -104,8 +137,6 @@ class FFT(OscillatoryFeature):
         sfreq: float,
     ) -> None:
         super().__init__(settings, ch_names, sfreq)
-        if self.s["fft_settings"]["kalman_filter"]:
-            self.init_KF("fft")
 
         if self.s["fft_settings"]["log_transform"]:
             self.log_transform = True
@@ -114,13 +145,15 @@ class FFT(OscillatoryFeature):
 
         window_ms = self.s["fft_settings"]["windowlength_ms"]
         self.window_samples = int(-np.floor(window_ms / 1000 * sfreq))
-        freqs = fft.rfftfreq(-self.window_samples, 1 / np.floor(self.sfreq))
+        self.freqs = fft.rfftfreq(
+            -self.window_samples, 1 / np.floor(self.sfreq)
+        )
 
         self.feature_params = []
         for ch_idx, ch_name in enumerate(self.ch_names):
             for fband, f_range in self.f_ranges_dict.items():
                 idx_range = np.where(
-                    (freqs >= f_range[0]) & (freqs < f_range[1])
+                    (self.freqs >= f_range[0]) & (self.freqs < f_range[1])
                 )[0]
                 feature_name = "_".join([ch_name, "fft", fband])
                 self.feature_params.append((ch_idx, feature_name, idx_range))
@@ -132,17 +165,87 @@ class FFT(OscillatoryFeature):
     def calc_feature(self, data: np.ndarray, features_compute: dict) -> dict:
         data = data[:, self.window_samples :]
         Z = np.abs(fft.rfft(data))
+
+        if self.log_transform:
+            Z = np.log10(Z)
+
         for ch_idx, feature_name, idx_range in self.feature_params:
             Z_ch = Z[ch_idx, idx_range]
-            feature_calc = np.mean(Z_ch)
 
-            if self.log_transform:
-                feature_calc = np.log(feature_calc)
+            features_compute = self.estimate_osc_features(
+                features_compute, Z_ch, feature_name, "fft_settings"
+            )
 
-            if self.KF_dict:
-                feature_calc = self.update_KF(feature_calc, feature_name)
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            if self.s["fft_settings"]["return_spectrum"]:
+                features_compute.update(
+                    {
+                        f"{ch_name}_fft_psd_{str(f)}": Z[ch_idx][idx]
+                        for idx, f in enumerate(self.freqs.astype(int))
+                    }
+                )
+
+        return features_compute
+
+
+class Welch(OscillatoryFeature):
+    def __init__(
+        self,
+        settings: dict,
+        ch_names: Iterable[str],
+        sfreq: float,
+    ) -> None:
+        super().__init__(settings, ch_names, sfreq)
+
+        self.log_transform = self.s["welch_settings"]["log_transform"]
+
+        self.feature_params = []
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            for fband, f_range in self.f_ranges_dict.items():
+                feature_name = "_".join([ch_name, "welch", fband])
+                self.feature_params.append((ch_idx, feature_name, f_range))
+
+    @staticmethod
+    def test_settings(s: dict, ch_names: Iterable[str], sfreq: int | float):
+        OscillatoryFeature.test_settings_osc(
+            s, ch_names, sfreq, "welch_settings"
+        )
+
+    def calc_feature(self, data: np.ndarray, features_compute: dict) -> dict:
+        freqs, Z = signal.welch(
+            data,
+            fs=self.sfreq,
+            window="hann",
+            nperseg=self.sfreq,
+            noverlap=None,
+        )
+
+        if self.log_transform:
+            Z = np.log10(Z)
+
+        for ch_idx, feature_name, f_range in self.feature_params:
+            Z_ch = Z[ch_idx]
+
+            idx_range = np.where((freqs >= f_range[0]) & (freqs <= f_range[1]))[
+                0
+            ]
+
+            features_compute = self.estimate_osc_features(
+                features_compute,
+                Z_ch[idx_range],
+                feature_name,
+                "welch_settings",
+            )
+
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            if self.s["welch_settings"]["return_spectrum"]:
+                features_compute.update(
+                    {
+                        f"{ch_name}_welch_psd_{str(f)}": Z[ch_idx][idx]
+                        for idx, f in enumerate(freqs.astype(int))
+                    }
+                )
 
-            features_compute[feature_name] = feature_calc
         return features_compute
 
 
@@ -154,10 +257,9 @@ class STFT(OscillatoryFeature):
         sfreq: float,
     ) -> None:
         super().__init__(settings, ch_names, sfreq)
-        if self.s["stft_settings"]["kalman_filter"]:
-            self.init_KF("stft")
 
         self.nperseg = int(self.s["stft_settings"]["windowlength_ms"])
+        self.log_transform = self.s["stft_settings"]["log_transform"]
 
         self.feature_params = []
         for ch_idx, ch_name in enumerate(self.ch_names):
@@ -172,7 +274,7 @@ class STFT(OscillatoryFeature):
         )
 
     def calc_feature(self, data: np.ndarray, features_compute: dict) -> dict:
-        f, _, Zxx = signal.stft(
+        freqs, _, Zxx = signal.stft(
             data,
             fs=self.sfreq,
             window="hamming",
@@ -180,15 +282,30 @@ class STFT(OscillatoryFeature):
             boundary="even",
         )
         Z = np.abs(Zxx)
+        if self.log_transform:
+            Z = np.log10(Z)
         for ch_idx, feature_name, f_range in self.feature_params:
             Z_ch = Z[ch_idx]
-            idx_range = np.where((f >= f_range[0]) & (f <= f_range[1]))[0]
-            feature_calc = np.mean(Z_ch[idx_range, :])  # 1. dim: f, 2. dim: t
-
-            if self.KF_dict:
-                feature_calc = self.update_KF(feature_calc, feature_name)
+            idx_range = np.where((freqs >= f_range[0]) & (freqs <= f_range[1]))[
+                0
+            ]
+
+            features_compute = self.estimate_osc_features(
+                features_compute,
+                Z_ch[idx_range, :],
+                feature_name,
+                "stft_settings",
+            )
 
-            features_compute[feature_name] = feature_calc
+        for ch_idx, ch_name in enumerate(self.ch_names):
+            if self.s["stft_settings"]["return_spectrum"]:
+                Z_ch_mean = Z[ch_idx].mean(axis=1)
+                features_compute.update(
+                    {
+                        f"{ch_name}_stft_psd_{str(f)}": Z_ch_mean[idx]
+                        for idx, f in enumerate(freqs.astype(int))
+                    }
+                )
 
         return features_compute
 
@@ -204,7 +321,7 @@ class BandPower(OscillatoryFeature):
         super().__init__(settings, ch_names, sfreq)
         bp_settings = self.s["bandpass_filter_settings"]
 
-        self.bandpass_filter = nm_filter.BandPassFilter(
+        self.bandpass_filter = nm_filter.MNEFilter(
             f_ranges=list(self.f_ranges_dict.values()),
             sfreq=self.sfreq,
             filter_length=self.sfreq - 1,
@@ -265,7 +382,9 @@ class BandPower(OscillatoryFeature):
             ].values()
         ), "Set at least one bandpower_feature to True."
 
-        for fband_name, seg_length_fband in s["bandpass_filter_settings"]["segment_lengths_ms"].items():
+        for fband_name, seg_length_fband in s["bandpass_filter_settings"][
+            "segment_lengths_ms"
+        ].items():
             assert isinstance(seg_length_fband, int), (
                 f"bandpass segment_lengths_ms for {fband_name} "
                 f"needs to be of type int, got {seg_length_fband}"
@@ -275,15 +394,16 @@ class BandPower(OscillatoryFeature):
                 f"segment length {seg_length_fband} needs to be smaller than "
                 f" s['segment_length_features_ms'] = {s['segment_length_features_ms']}"
             )
- 
+
         for fband_name in list(s["frequency_ranges_hz"].keys()):
-            assert fband_name in list(s["bandpass_filter_settings"]["segment_lengths_ms"].keys()), (
+            assert fband_name in list(
+                s["bandpass_filter_settings"]["segment_lengths_ms"].keys()
+            ), (
                 f"frequency range {fband_name} "
                 "needs to be defined in s['bandpass_filter_settings']['segment_lengths_ms']"
             )
 
     def calc_feature(self, data: np.ndarray, features_compute: dict) -> dict:
-
         data = self.bandpass_filter.filter_data(data)
 
         for (
@@ -297,7 +417,7 @@ class BandPower(OscillatoryFeature):
         ) in self.feature_params:
             if bp_feature == "activity":
                 if self.log_transform:
-                    feature_calc = np.log(
+                    feature_calc = np.log10(
                         np.var(data[ch_idx, f_band_idx, -seglen:])
                     )
                 else:

py_neuromodulation/nm_plots.py

@@ -6,6 +6,9 @@ from matplotlib import gridspec
 from typing import Optional
 import seaborn as sb
 import pandas as pd
+import logging
+
+logger = logging.getLogger("PynmLogger")
 
 from py_neuromodulation import nm_IO, nm_stats
 
@@ -87,7 +90,9 @@ def plot_epoch(
 ):
     if z_score is None:
         X_epoch = stats.zscore(
-            np.nan_to_num(np.nanmean(np.squeeze(X_epoch), axis=0)), axis=0
+            np.nan_to_num(np.nanmean(np.squeeze(X_epoch), axis=0)),
+            axis=0,
+            nan_policy="omit",
         ).T
     y_epoch = np.stack(np.array(y_epoch))
     plt.figure(figsize=(6, 6))
@@ -237,7 +242,7 @@ def plot_corr_matrix(
 
     if save_plot:
         plt.savefig(plt_path, bbox_inches="tight")
-        print("Correlation matrix figure saved to " + str(plt_path))
+        logger.info(f"Correlation matrix figure saved to {plt_path}")
 
     if show_plot is False:
         plt.close()
@@ -329,7 +334,7 @@ def plot_epochs_avg(
 
     if normalize_data:
         X_epoch_mean = stats.zscore(
-            np.nanmean(np.squeeze(X_epoch), axis=0), axis=0
+            np.nanmean(np.squeeze(X_epoch), axis=0), axis=0, nan_policy="omit"
         ).T
     else:
         X_epoch_mean = np.nanmean(np.squeeze(X_epoch), axis=0).T
@@ -385,7 +390,7 @@ def plot_epochs_avg(
             feature_name=feature_str_add,
         )
         plt.savefig(plt_path, bbox_inches="tight")
-        print("Feature epoch average figure saved to: " + str(plt_path))
+        logger.info(f"Feature epoch average figure saved to: {str(plt_path)}")
     if show_plot is False:
         plt.close()
 
@@ -441,7 +446,6 @@ def plot_all_features(
     OUT_PATH: str = None,
     feature_file: str = None,
 ):
-
     if time_limit_high_s is not None:
         df = df[df["time"] < time_limit_high_s * 1000]
     if time_limit_low_s is not None:
@@ -449,7 +453,7 @@ def plot_all_features(
 
     cols_plt = [c for c in df.columns if c != "time"]
     if normalize is True:
-        data_plt = stats.zscore(df[cols_plt])
+        data_plt = stats.zscore(df[cols_plt], nan_policy="omit")
     else:
         data_plt = df[cols_plt]
 
@@ -487,7 +491,6 @@ class NM_Plot:
         sess_right: Optional[bool] = False,
         proj_matrix_cortex: np.ndarray | None = None,
     ) -> None:
-
         self.grid_cortex = grid_cortex
         self.grid_subcortex = grid_subcortex
         self.ecog_strip = ecog_strip
@@ -510,7 +513,6 @@ class NM_Plot:
         ) = nm_IO.read_plot_modules()
 
     def plot_grid_elec_3d(self) -> None:
-
         plot_grid_elec_3d(np.array(self.grid_cortex), np.array(self.ecog_strip))
 
     def plot_cortex(
@@ -552,7 +554,6 @@ class NM_Plot:
         axes.axes.set_aspect("equal", anchor="C")
 
         if grid_cortex is not None:
-
             grid_color = (
                 np.ones(grid_cortex.shape[0])
                 if grid_color is None
@@ -604,6 +605,8 @@ class NM_Plot:
                 feature_name=feature_str_add,
             )
             plt.savefig(plt_path, bbox_inches="tight")
-            print("Feature epoch average figure saved to: " + str(plt_path))
+            logger.info(
+                f"Feature epoch average figure saved to: {str(plt_path)}"
+            )
         if show_plot is False:
             plt.close()

py_neuromodulation/nm_rereference.py

@@ -4,10 +4,8 @@ import pandas as pd
 
 
 class ReReferencer:
-
     ref_matrix: np.ndarray
 
-
     def __init__(
         self,
         sfreq: int | float,
@@ -28,11 +26,15 @@ class ReReferencer:
             ValueError: rereferencing using undefined channel
             ValueError: rereferencing to same channel
         """
-        (channels_used,) = np.where((nm_channels.used == 1))
+        nm_channels = nm_channels[nm_channels["used"] == 1].reset_index(
+            drop=True
+        )
+        # (channels_used,) = np.where((nm_channels.used == 1))
 
         ch_names = nm_channels["name"].tolist()
 
-        if len(ch_names) == 1:
+        # no re-referencing is being performed when there is a single channel present only
+        if nm_channels.shape[0] in (0, 1):
             self.ref_matrix = None
             return
 
@@ -48,8 +50,8 @@ class ReReferencer:
         ref_matrix = np.zeros((len(nm_channels), len(nm_channels)))
         for ind in range(len(nm_channels)):
             ref_matrix[ind, ind] = 1
-            if ind not in channels_used:
-                continue
+            # if ind not in channels_used:
+            #    continue
             ref = refs[ind]
             if ref.lower() == "none" or pd.isnull(ref):
                 ref_idx = None
@@ -84,10 +86,10 @@ class ReReferencer:
                 shape(n_channels, n_samples) - data to be rereferenced.
 
         Returns:
-            reref_data (numpy ndarray): 
+            reref_data (numpy ndarray):
             shape(n_channels, n_samples) - rereferenced data
         """
         if self.ref_matrix is not None:
             return self.ref_matrix @ data
         else:
-            return data
+            return data