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

py_neuromodulation/nm_fooof.py

@@ -1,7 +1,9 @@
+import logging
+from typing import Iterable
+
 import numpy as np
 from fooof import FOOOF
 from scipy import fft
-from typing import Iterable
 
 from py_neuromodulation import nm_features_abc
 
@@ -18,15 +20,6 @@ class FooofAnalyzer(nm_features_abc.Feature):
         self.ap_mode = "knee" if self.settings_fooof["knee"] else "fixed"
         self.max_n_peaks = self.settings_fooof["max_n_peaks"]
 
-        self.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,
-        )
-
         self.num_samples = int(
             self.settings_fooof["windowlength_ms"] * sfreq / 1000
         )
@@ -75,44 +68,70 @@ class FooofAnalyzer(nm_features_abc.Feature):
             spectrum = self._get_spectrum(data[ch_idx, :])
 
             try:
-                self.fm.fit(self.f_vec, spectrum, self.freq_range)
+                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:
-                print(e)
-                print(f"failing spectrum: {spectrum}")
+                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(self.fm.get_params("aperiodic_params", "exponent"))
-                    if self.fm.fooofed_spectrum_ is not None
+                    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(self.fm.get_params("aperiodic_params", "offset"))
-                    if self.fm.fooofed_spectrum_ is not None
+                    np.nan_to_num(fm.get_params("aperiodic_params", "offset"))
+                    if FIT_PASSED is True
                     else None
                 )
-            
+
             if self.settings_fooof["aperiodic"]["knee"]:
-                features_compute[f"{ch_name}_fooof_a_knee"] = (
-                    np.nan_to_num(self.fm.get_params("aperiodic_params", "knee"))
-                    if self.fm.fooofed_spectrum_ is not None
-                    else None
+                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 = (
-                self.fm.get_params("peak_params", "BW")
-                if self.fm.fooofed_spectrum_ is not None
+                fm.get_params("peak_params", "BW")
+                if FIT_PASSED is True
                 else None
             )
             peaks_cf = (
-                self.fm.get_params("peak_params", "CF")
-                if self.fm.fooofed_spectrum_ is not None
+                fm.get_params("peak_params", "CF")
+                if FIT_PASSED is True
                 else None
             )
             peaks_pw = (
-                self.fm.get_params("peak_params", "PW")
-                if self.fm.fooofed_spectrum_ is not None
+                fm.get_params("peak_params", "PW")
+                if FIT_PASSED is True
                 else None
             )
 
@@ -122,7 +141,6 @@ class FooofAnalyzer(nm_features_abc.Feature):
                 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

py_neuromodulation/nm_mne_connectivity.py

@@ -59,7 +59,7 @@ class MNEConnectivity(nm_features_abc.Feature):
         if epochs.events.shape[0] < 2:
             raise Exception(
                 f"A minimum of 2 epochs is required for mne_connectivity,"
-                f" got only {epochs.events.shape[0]}. Increase settings['segment_length']"
+                f" got only {epochs.events.shape[0]}. Increase settings['segment_length_features_ms']"
             )
         return epochs
 

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()