diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 169bc5441e..e53622bc24 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -12,7 +12,7 @@ jobs: - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v4 with: - python-version: "3.10" + python-version: "3.11" - name: Install dependencies run: | diff --git a/.github/workflows/releasePR.yml b/.github/workflows/releasePR.yml index 0c284d8416..40b0fc5b47 100644 --- a/.github/workflows/releasePR.yml +++ b/.github/workflows/releasePR.yml @@ -15,7 +15,7 @@ jobs: - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v4 with: - python-version: "3.7" + python-version: "3.11" - name: Install dependencies run: | @@ -50,7 +50,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v4 with: - python-version: "3.x" + python-version: "3.11" - name: Install pep517 run: >- diff --git a/README.rst b/README.rst index 5b9881b6b6..1fbea5c8f3 100644 --- a/README.rst +++ b/README.rst @@ -248,7 +248,7 @@ Electrodermal Activity (EDA/GSR) signals, info = nk.eda_process(eda, sampling_rate=250) # Visualise the processing - nk.eda_plot(signals, sampling_rate=250) + nk.eda_plot(signals, info) .. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_eda.png :target: https://neuropsychology.github.io/NeuroKit/examples/eda_peaks/eda_peaks.html @@ -266,7 +266,7 @@ Cardiac activity (ECG) signals, info = nk.ecg_process(ecg, sampling_rate=250) # Visualise the processing - nk.ecg_plot(signals, sampling_rate=250) + nk.ecg_plot(signals, info) .. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_ecg.png @@ -285,47 +285,48 @@ Respiration (RSP) signals, info = nk.rsp_process(rsp, sampling_rate=250) # Visualise the processing - nk.rsp_plot(signals, sampling_rate=250) + nk.rsp_plot(signals, info) .. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_rsp.png :target: https://neuropsychology.github.io/NeuroKit/examples/rsp_rrv/rsp_rrv.html -Electromyography (EMG) -^^^^^^^^^^^^^^^^^^^^^^^ +Photoplethysmography (PPG/BVP) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. code-block:: python - # Generate 10 seconds of EMG signal (recorded at 250 samples/second) - emg = nk.emg_simulate(duration=10, sampling_rate=250, burst_number=3) + # Generate 15 seconds of PPG signal (recorded at 250 samples/second) + ppg = nk.ppg_simulate(duration=15, sampling_rate=250, heart_rate=70) # Process it - signals, info = nk.emg_process(emg, sampling_rate=250) + signals, info = nk.ppg_process(ppg, sampling_rate=250) - # Visualise the processing - nk.emg_plot(signals, sampling_rate=250) + # Visualize the processing + nk.ppg_plot(signals, info) -.. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_emg.png +.. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_ppg.png -Photoplethysmography (PPG/BVP) -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Electromyography (EMG) +^^^^^^^^^^^^^^^^^^^^^^^ .. code-block:: python - # Generate 15 seconds of PPG signal (recorded at 250 samples/second) - ppg = nk.ppg_simulate(duration=15, sampling_rate=250, heart_rate=70) + # Generate 10 seconds of EMG signal (recorded at 250 samples/second) + emg = nk.emg_simulate(duration=10, sampling_rate=250, burst_number=3) # Process it - signals, info = nk.ppg_process(ppg, sampling_rate=250) + signals, info = nk.emg_process(emg, sampling_rate=250) - # Visualize the processing - nk.ppg_plot(signals, sampling_rate=250) + # Visualise the processing + nk.emg_plot(signals, info) -.. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_ppg.png +.. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_emg.png + Electrooculography (EOG) @@ -340,7 +341,7 @@ Electrooculography (EOG) signals, info = nk.eog_process(eog_signal, sampling_rate=100) # Plot - nk.eog_plot(signals, info, sampling_rate=100) + nk.eog_plot(signals, info) .. image:: https://raw.github.com/neuropsychology/NeuroKit/master/docs/readme/README_eog.png diff --git a/docs/readme/README_ecg.png b/docs/readme/README_ecg.png index 31f6436c0d..695b5685e4 100644 Binary files a/docs/readme/README_ecg.png and b/docs/readme/README_ecg.png differ diff --git a/docs/readme/README_eda.png b/docs/readme/README_eda.png index 2a7c101d03..db2e3daf4b 100644 Binary files a/docs/readme/README_eda.png and b/docs/readme/README_eda.png differ diff --git a/docs/readme/README_emg.png b/docs/readme/README_emg.png index 1bd480cdc4..391647277e 100644 Binary files a/docs/readme/README_emg.png and b/docs/readme/README_emg.png differ diff --git a/docs/readme/README_eog.png b/docs/readme/README_eog.png index bc39d38120..348c4296e6 100644 Binary files a/docs/readme/README_eog.png and b/docs/readme/README_eog.png differ diff --git a/docs/readme/README_examples.py b/docs/readme/README_examples.py index ce1295d4c3..a3a6770e49 100644 --- a/docs/readme/README_examples.py +++ b/docs/readme/README_examples.py @@ -9,7 +9,7 @@ # Setup matplotlib with Agg to run on server matplotlib.use("Agg") -plt.rcParams["figure.figsize"] = (10, 6.5) +plt.rcParams["figure.figsize"] = (2**0.5 * 10, 10) plt.rcParams["savefig.facecolor"] = "white" # ============================================================================= @@ -55,13 +55,15 @@ } ) plot = data.plot( - subplots=True, layout=(5, 1), color=["#f44336", "#E91E63", "#2196F3", "#9C27B0", "#FF9800"] + subplots=True, + layout=(5, 1), + color=["#f44336", "#E91E63", "#2196F3", "#9C27B0", "#FF9800"], ) fig = plt.gcf() fig.set_size_inches(10, 6, forward=True) [ax.legend(loc=1) for ax in plt.gcf().axes] plt.tight_layout() -fig.savefig("README_simulation.png", dpi=300) +fig.savefig("README_simulation.png", dpi=150) # ============================================================================= # Electrodermal Activity (EDA) processing @@ -73,13 +75,10 @@ # Process it signals, info = nk.eda_process(eda, sampling_rate=250) -# Visualise the processing -nk.eda_plot(signals, sampling_rate=None) - # Save it -nk.eda_plot(signals, sampling_rate=None) +nk.eda_plot(signals, info) plt.tight_layout() -plt.savefig("README_eda.png", dpi=300) +plt.savefig("README_eda.png", dpi=150) # ============================================================================= # Cardiac activity (ECG) processing @@ -91,13 +90,10 @@ # Process it signals, info = nk.ecg_process(ecg, sampling_rate=250) -# Visualise the processing -nk.ecg_plot(signals, sampling_rate=250) - # Save it -nk.ecg_plot(signals, sampling_rate=250) +nk.ecg_plot(signals, info) plt.tight_layout() -plt.savefig("README_ecg.png", dpi=300) +plt.savefig("README_ecg.png", dpi=150) # ============================================================================= # Respiration (RSP) processing @@ -109,51 +105,43 @@ # Process it signals, info = nk.rsp_process(rsp, sampling_rate=250) -# Visualise the processing -nk.rsp_plot(signals, sampling_rate=250) - # Save it -nk.rsp_plot(signals, sampling_rate=250) +nk.rsp_plot(signals, info) fig = plt.gcf() fig.set_size_inches(10, 12, forward=True) plt.tight_layout() -plt.savefig("README_rsp.png", dpi=300) +plt.savefig("README_rsp.png", dpi=150) # ============================================================================= -# Electromyography (EMG) processing +# Photoplethysmography (PPG/BVP) # ============================================================================= -# Generate 10 seconds of EMG signal (recorded at 250 samples / second) -emg = nk.emg_simulate(duration=10, sampling_rate=250, burst_number=3) +# Generate 15 seconds of PPG signal (recorded at 250 samples / second) +ppg = nk.ppg_simulate(duration=15, sampling_rate=250, heart_rate=70, random_state=333) # Process it -signals, info = nk.emg_process(emg, sampling_rate=250) - -# Visualise the processing -nk.emg_plot(signals, sampling_rate=250) +signals, info = nk.ppg_process(ppg, sampling_rate=250) # Save it -nk.emg_plot(signals, sampling_rate=250) +nk.ppg_plot(signals, info) plt.tight_layout() -plt.savefig("README_emg.png", dpi=300) +plt.savefig("README_ppg.png", dpi=150) # ============================================================================= -# Photoplethysmography (PPG/BVP) +# Electromyography (EMG) processing # ============================================================================= -# Generate 15 seconds of PPG signal (recorded at 250 samples / second) -ppg = nk.ppg_simulate(duration=15, sampling_rate=250, heart_rate=70, random_state=333) +# Generate 10 seconds of EMG signal (recorded at 250 samples / second) +emg = nk.emg_simulate(duration=10, sampling_rate=250, burst_number=3) # Process it -signals, info = nk.ppg_process(ppg, sampling_rate=250) - -# Visualize the processing -nk.ppg_plot(signals, sampling_rate=250) +signals, info = nk.emg_process(emg, sampling_rate=250) # Save it -nk.ppg_plot(signals, sampling_rate=250) +nk.emg_plot(signals, info) plt.tight_layout() -plt.savefig("README_ppg.png", dpi=300) +plt.savefig("README_emg.png", dpi=150) + # ============================================================================= # Electrooculography (EOG) @@ -166,9 +154,9 @@ signals, info = nk.eog_process(eog_signal, sampling_rate=100) # Plot -nk.eog_plot(signals, peaks=info, sampling_rate=100) +nk.eog_plot(signals, info) plt.tight_layout() -plt.savefig("README_eog.png", dpi=300) +plt.savefig("README_eog.png", dpi=150) # ============================================================================= # Signal Processing @@ -280,7 +268,9 @@ signal = nk.signal_simulate(duration=10, frequency=1) # High freq signal += 3 * nk.signal_simulate(duration=10, frequency=3) # Higher freq signal += 3 * np.linspace(0, 2, len(signal)) # Add baseline and linear trend -signal += 2 * nk.signal_simulate(duration=10, frequency=0.1, noise=0) # Non-linear trend +signal += 2 * nk.signal_simulate( + duration=10, frequency=0.1, noise=0 +) # Non-linear trend signal += np.random.normal(0, 0.02, len(signal)) # Add noise # Decompose signal using Empirical Mode Decomposition (EMD) @@ -323,10 +313,16 @@ ) # Get the PSD using different methods -welch = nk.signal_psd(signal, method="welch", min_frequency=1, max_frequency=20, show=True) +welch = nk.signal_psd( + signal, method="welch", min_frequency=1, max_frequency=20, show=True +) multitaper = nk.signal_psd(signal, method="multitapers", max_frequency=20, show=True) -lomb = nk.signal_psd(signal, method="lomb", min_frequency=1, max_frequency=20, show=True) -burg = nk.signal_psd(signal, method="burg", min_frequency=1, max_frequency=20, order=10, show=True) +lomb = nk.signal_psd( + signal, method="lomb", min_frequency=1, max_frequency=20, show=True +) +burg = nk.signal_psd( + signal, method="burg", min_frequency=1, max_frequency=20, order=10, show=True +) # Visualize the different methods together @@ -337,7 +333,12 @@ axes[0].set_xlabel("Time (s)") axes[1].plot( - welch["Frequency"], welch["Power"], label="Welch", color="#E91E63", linewidth=2, zorder=1 + welch["Frequency"], + welch["Power"], + label="Welch", + color="#E91E63", + linewidth=2, + zorder=1, ) axes[1].plot( multitaper["Frequency"], @@ -347,9 +348,21 @@ linewidth=2, zorder=2, ) -axes[1].plot(burg["Frequency"], burg["Power"], label="Burg", color="#4CAF50", linewidth=2, zorder=3) axes[1].plot( - lomb["Frequency"], lomb["Power"], label="Lomb", color="#FFC107", linewidth=0.5, zorder=4 + burg["Frequency"], + burg["Power"], + label="Burg", + color="#4CAF50", + linewidth=2, + zorder=3, +) +axes[1].plot( + lomb["Frequency"], + lomb["Power"], + label="Lomb", + color="#FFC107", + linewidth=0.5, + zorder=4, ) axes[1].set_title("Power Spectrum Density (PSD)") diff --git a/docs/readme/README_ppg.png b/docs/readme/README_ppg.png index 99b081461f..f4efe726bc 100644 Binary files a/docs/readme/README_ppg.png and b/docs/readme/README_ppg.png differ diff --git a/docs/readme/README_rsp.png b/docs/readme/README_rsp.png index 56660de397..77431b49e1 100644 Binary files a/docs/readme/README_rsp.png and b/docs/readme/README_rsp.png differ diff --git a/docs/readme/README_simulation.png b/docs/readme/README_simulation.png index 3c31f54975..79c3c0e855 100644 Binary files a/docs/readme/README_simulation.png and b/docs/readme/README_simulation.png differ diff --git a/neurokit2/ecg/ecg_clean.py b/neurokit2/ecg/ecg_clean.py index ebf11fead8..c1d2830cf6 100644 --- a/neurokit2/ecg/ecg_clean.py +++ b/neurokit2/ecg/ecg_clean.py @@ -6,7 +6,7 @@ import scipy.signal from ..misc import NeuroKitWarning, as_vector -from ..signal import signal_filter, signal_resample +from ..signal import signal_filter def ecg_clean(ecg_signal, sampling_rate=1000, method="neurokit", **kwargs): @@ -335,7 +335,9 @@ def _ecg_clean_templateconvolution(ecg_signal, sampling_rate=1000): """Filter and Convolve ECG signal with QRS complex template. Totally exploratory method by Dominique Makowski, use at your own risks. - The idea is to use a QRS template to convolve the signal with, in order to magnify the QRS features. However, it doens't work well and creates a lot of artifacts. If you have ideas for improvement please let me know! + The idea is to use a QRS template to convolve the signal with, in order to magnify the QRS + features. However, it doens't work well and creates a lot of artifacts. If you have ideas for + improvement please let me know! """ window_size = int(np.round(sampling_rate / 4)) diff --git a/neurokit2/ecg/ecg_peaks.py b/neurokit2/ecg/ecg_peaks.py index 27125b0f85..325fd435f3 100644 --- a/neurokit2/ecg/ecg_peaks.py +++ b/neurokit2/ecg/ecg_peaks.py @@ -1,4 +1,8 @@ +import matplotlib.pyplot as plt +import numpy as np + from ..signal import signal_fixpeaks, signal_formatpeaks +from ..stats import rescale from .ecg_findpeaks import ecg_findpeaks @@ -7,6 +11,7 @@ def ecg_peaks( sampling_rate=1000, method="neurokit", correct_artifacts=False, + show=False, **kwargs ): """**Find R-peaks in an ECG signal** @@ -59,6 +64,8 @@ def ecg_peaks( correct_artifacts : bool Whether or not to first identify and fix artifacts, using the method by Lipponen & Tarvainen (2019). + show : bool + If ``True``, will show a plot of the signal with peaks. Defaults to ``False``. **kwargs Additional keyword arguments, usually specific for each method. @@ -85,11 +92,10 @@ def ecg_peaks( import neurokit2 as nk - ecg = nk.ecg_simulate(duration=10, sampling_rate=1000) - signals, info = nk.ecg_peaks(ecg, correct_artifacts=True) + ecg = nk.ecg_simulate(duration=10, sampling_rate=250) @savefig p_ecg_peaks1.png scale=100% - nk.events_plot(info["ECG_R_Peaks"], ecg) + signals, info = nk.ecg_peaks(ecg, sampling_rate=250, correct_artifacts=True, show=True) @suppress plt.close() @@ -98,50 +104,50 @@ def ecg_peaks( .. ipython:: python # neurokit (default) - cleaned = nk.ecg_clean(ecg, method="neurokit") - _, neurokit = nk.ecg_peaks(cleaned, method="neurokit") + cleaned = nk.ecg_clean(ecg, sampling_rate=250, method="neurokit") + _, neurokit = nk.ecg_peaks(cleaned, sampling_rate=250, method="neurokit") # pantompkins1985 - cleaned = nk.ecg_clean(ecg, method="pantompkins1985") - _, pantompkins1985 = nk.ecg_peaks(cleaned, method="pantompkins1985") + cleaned = nk.ecg_clean(ecg, sampling_rate=250, method="pantompkins1985") + _, pantompkins1985 = nk.ecg_peaks(cleaned, sampling_rate=250, method="pantompkins1985") # nabian2018 - _, nabian2018 = nk.ecg_peaks(ecg, method="nabian2018") + _, nabian2018 = nk.ecg_peaks(ecg, sampling_rate=250, method="nabian2018") # hamilton2002 - cleaned = nk.ecg_clean(ecg, method="hamilton2002") - _, hamilton2002 = nk.ecg_peaks(cleaned, method="hamilton2002") + cleaned = nk.ecg_clean(ecg, sampling_rate=250, method="hamilton2002") + _, hamilton2002 = nk.ecg_peaks(cleaned, sampling_rate=250, method="hamilton2002") # martinez2004 - _, martinez2004 = nk.ecg_peaks(ecg, method="martinez2004") + _, martinez2004 = nk.ecg_peaks(ecg, sampling_rate=250, method="martinez2004") # zong2003 - _, zong2003 = nk.ecg_peaks(ecg, method="zong2003") + _, zong2003 = nk.ecg_peaks(ecg, sampling_rate=250, method="zong2003") # christov2004 - _, christov2004 = nk.ecg_peaks(cleaned, method="christov2004") + _, christov2004 = nk.ecg_peaks(cleaned, sampling_rate=250, method="christov2004") # gamboa2008 - cleaned = nk.ecg_clean(ecg, method="gamboa2008") - _, gamboa2008 = nk.ecg_peaks(cleaned, method="gamboa2008") + cleaned = nk.ecg_clean(ecg, sampling_rate=250, method="gamboa2008") + _, gamboa2008 = nk.ecg_peaks(cleaned, sampling_rate=250, method="gamboa2008") # elgendi2010 - cleaned = nk.ecg_clean(ecg, method="elgendi2010") - _, elgendi2010 = nk.ecg_peaks(cleaned, method="elgendi2010") + cleaned = nk.ecg_clean(ecg, sampling_rate=250, method="elgendi2010") + _, elgendi2010 = nk.ecg_peaks(cleaned, sampling_rate=250, method="elgendi2010") # engzeemod2012 - cleaned = nk.ecg_clean(ecg, method="engzeemod2012") - _, engzeemod2012 = nk.ecg_peaks(cleaned, method="engzeemod2012") + cleaned = nk.ecg_clean(ecg, sampling_rate=250, method="engzeemod2012") + _, engzeemod2012 = nk.ecg_peaks(cleaned, sampling_rate=250, method="engzeemod2012") # kalidas2017 - cleaned = nk.ecg_clean(ecg, method="kalidas2017") - _, kalidas2017 = nk.ecg_peaks(cleaned, method="kalidas2017") + cleaned = nk.ecg_clean(ecg, sampling_rate=250, method="kalidas2017") + _, kalidas2017 = nk.ecg_peaks(cleaned, sampling_rate=250, method="kalidas2017") # rodrigues2021 - _, rodrigues2021 = nk.ecg_peaks(ecg, method="rodrigues2021") + _, rodrigues2021 = nk.ecg_peaks(ecg, sampling_rate=250, method="rodrigues2021") # koka2022 - _, koka2022 = nk.ecg_peaks(ecg, method="koka2022") + _, koka2022 = nk.ecg_peaks(ecg, sampling_rate=250, method="koka2022") # Collect all R-peak lists by iterating through the result dicts rpeaks = [ @@ -177,7 +183,7 @@ def ecg_peaks( noise_amplitude=0.05, noise_frequency=[25, 50], artifacts_amplitude=0.05, artifacts_frequency=50) @savefig p_ecg_peaks3.png scale=100% - info = nk.ecg_findpeaks(ecg, sampling_rate=500, method="promac", show=True) + info = nk.ecg_findpeaks(ecg, sampling_rate=250, method="promac", show=True) @suppress plt.close() @@ -224,27 +230,6 @@ def ecg_peaks( * T. Koka and M. Muma, "Fast and Sample Accurate R-Peak Detection for Noisy ECG Using Visibility Graphs," 2022 44th Annual International Conference of the IEEE Engineering in Medicine & Biology Society (EMBC), 2022, pp. 121-126. - - * ``nabian2018`` - - * ``gamboa2008`` - - - * ``hamilton2002`` - - * ``christov2004`` - - * ``engzeemod2012`` - - * ``elgendi2010`` - - * ``kalidas2017`` - - - * ``rodrigues2021`` - - * ``koka2022`` - * ``promac`` * **Unpublished.** It runs different methods and derives a probability index using convolution. See this discussion for more information on the method: @@ -254,22 +239,168 @@ def ecg_peaks( engineering & technology, 43(3), 173-181. """ - rpeaks = ecg_findpeaks( - ecg_cleaned, sampling_rate=sampling_rate, method=method, **kwargs + # Store info + info = {"method_peaks": method.lower(), "method_fixpeaks": "None"} + + # First peak detection + info.update( + ecg_findpeaks( + ecg_cleaned, + sampling_rate=sampling_rate, + method=info["method_peaks"], + **kwargs + ) ) + # Peak correction if correct_artifacts: - _, rpeaks = signal_fixpeaks( - rpeaks, sampling_rate=sampling_rate, iterative=True, method="Kubios" + info["ECG_R_Peaks_Uncorrected"] = info["ECG_R_Peaks"].copy() + + fixpeaks, info["ECG_R_Peaks"] = signal_fixpeaks( + info["ECG_R_Peaks"], sampling_rate=sampling_rate, method="Kubios" ) - rpeaks = {"ECG_R_Peaks": rpeaks} + # Add prefix and merge + fixpeaks = {"ECG_fixpeaks_" + str(key): val for key, val in fixpeaks.items()} + info.update(fixpeaks) - instant_peaks = signal_formatpeaks( - rpeaks, desired_length=len(ecg_cleaned), peak_indices=rpeaks + # Format output + signals = signal_formatpeaks( + dict(ECG_R_Peaks=info["ECG_R_Peaks"]), # Takes a dict as input + desired_length=len(ecg_cleaned), + peak_indices=info["ECG_R_Peaks"], ) - signals = instant_peaks - info = rpeaks + info["sampling_rate"] = sampling_rate # Add sampling rate in dict info + if show is True: + _ecg_peaks_plot(ecg_cleaned, info, sampling_rate) + return signals, info + + +# ============================================================================= +# Internals +# ============================================================================= +def _ecg_peaks_plot( + ecg_cleaned, + info=None, + sampling_rate=1000, + raw=None, + quality=None, + phase=None, + ax=None, +): + x_axis = np.linspace(0, len(ecg_cleaned) / sampling_rate, len(ecg_cleaned)) + + # Prepare plot + if ax is None: + _, ax = plt.subplots() + + ax.set_xlabel("Time (seconds)") + ax.set_title("ECG signal and peaks") + + # Quality Area ------------------------------------------------------------- + if quality is not None: + quality = rescale( + quality, + to=[ + np.min([np.min(raw), np.min(ecg_cleaned)]), + np.max([np.max(raw), np.max(ecg_cleaned)]), + ], + ) + minimum_line = np.full(len(x_axis), quality.min()) + + # Plot quality area first + ax.fill_between( + x_axis, + minimum_line, + quality, + alpha=0.12, + zorder=0, + interpolate=True, + facecolor="#4CAF50", + label="Signal quality", + ) + + # Raw Signal --------------------------------------------------------------- + if raw is not None: + ax.plot(x_axis, raw, color="#B0BEC5", label="Raw signal", zorder=1) + label_clean = "Cleaned signal" + else: + label_clean = "Signal" + + # Peaks ------------------------------------------------------------------- + ax.scatter( + x_axis[info["ECG_R_Peaks"]], + ecg_cleaned[info["ECG_R_Peaks"]], + color="#FFC107", + label="R-peaks", + zorder=2, + ) + + # TODO + # # Artifacts --------------------------------------------------------------- + # def _plot_artifact(artifact, color, label, ax): + # if artifact in info.keys() and len(info[artifact]) > 0: + # ax.scatter( + # x_axis[info[artifact]], + # ecg_cleaned[info[artifact]], + # color=color, + # label=label, + # marker="x", + # zorder=2, + # ) + + # _plot_artifact("ECG_fixpeaks_missed", "#1E88E5", "Missed Peaks", ax) + # _plot_artifact("ECG_fixpeaks_longshort", "#1E88E5", "Long/Short", ax) + + # Clean Signal ------------------------------------------------------------ + if phase is not None: + mask = (phase == 0) | (np.isnan(phase)) + diastole = ecg_cleaned.copy() + diastole[~mask] = np.nan + + # Create overlap to avoid interuptions in signal + mask[np.where(np.diff(mask))[0] + 1] = True + systole = ecg_cleaned.copy() + systole[mask] = np.nan + + ax.plot( + x_axis, + diastole, + color="#B71C1C", + label=label_clean, + zorder=3, + linewidth=1, + ) + ax.plot( + x_axis, + systole, + color="#F44336", + zorder=3, + linewidth=1, + ) + else: + ax.plot( + x_axis, + ecg_cleaned, + color="#F44336", + label=label_clean, + zorder=3, + linewidth=1, + ) + + # Optimize legend + if raw is not None: + handles, labels = ax.get_legend_handles_labels() + order = [2, 0, 1, 3] + ax.legend( + [handles[idx] for idx in order], + [labels[idx] for idx in order], + loc="upper right", + ) + else: + ax.legend(loc="upper right") + + return ax diff --git a/neurokit2/ecg/ecg_phase.py b/neurokit2/ecg/ecg_phase.py index c30312ab8b..4848802956 100644 --- a/neurokit2/ecg/ecg_phase.py +++ b/neurokit2/ecg/ecg_phase.py @@ -55,7 +55,18 @@ def ecg_phase(ecg_cleaned, rpeaks=None, delineate_info=None, sampling_rate=None) cardiac_phase = nk.ecg_phase(ecg_cleaned=ecg, rpeaks=rpeaks, delineate_info=waves, sampling_rate=1000) @savefig p_ecg_phase.png scale=100% - nk.signal_plot([ecg, cardiac_phase], standardize=True) + _, ax = plt.subplots(nrows=2) + ax[0].plot(nk.rescale(ecg), label="ECG", color="red", alpha=0.3) + ax[0].plot(cardiac_phase["ECG_Phase_Atrial"], label="Atrial Phase", color="orange") + ax[0].plot(cardiac_phase["ECG_Phase_Completion_Atrial"], + label="Atrial Phase Completion", linestyle="dotted") + ax[0].legend(loc="upper right") + + ax[1].plot(nk.rescale(ecg), label="ECG", color="red", alpha=0.3) + ax[1].plot(cardiac_phase["ECG_Phase_Ventricular"], label="Ventricular Phase", color="green") + ax[1].plot(cardiac_phase["ECG_Phase_Completion_Ventricular"], + label="Ventricular Phase Completion", linestyle="dotted") + ax[1].legend(loc="upper right") @suppress plt.close() @@ -79,9 +90,10 @@ def ecg_phase(ecg_cleaned, rpeaks=None, delineate_info=None, sampling_rate=None) if isinstance( delineate_info, dict ): # FIXME: if this evaluates to False, toffsets and ppeaks are not instantiated - toffsets = np.full(len(ecg_cleaned), False, dtype=bool) - toffsets_idcs = [int(x) for x in delineate_info["ECG_T_Offsets"] if ~np.isnan(x)] + toffsets_idcs = [ + int(x) for x in delineate_info["ECG_T_Offsets"] if ~np.isnan(x) + ] toffsets[toffsets_idcs] = True ppeaks = np.full(len(ecg_cleaned), False, dtype=bool) @@ -93,8 +105,12 @@ def ecg_phase(ecg_cleaned, rpeaks=None, delineate_info=None, sampling_rate=None) atrial[rpeaks] = 0.0 atrial[ppeaks] = 1.0 - last_element = np.where(~np.isnan(atrial))[0][-1] # Avoid filling beyond the last peak/trough - atrial[0:last_element] = pd.Series(atrial).fillna(method="ffill").values[0:last_element] + last_element = np.where(~np.isnan(atrial))[0][ + -1 + ] # Avoid filling beyond the last peak/trough + atrial[0:last_element] = ( + pd.Series(atrial).fillna(method="ffill").values[0:last_element] + ) # Atrial Phase Completion atrial_completion = signal_phase(atrial, method="percent") diff --git a/neurokit2/ecg/ecg_plot.py b/neurokit2/ecg/ecg_plot.py index 9bc589266c..754900af7f 100644 --- a/neurokit2/ecg/ecg_plot.py +++ b/neurokit2/ecg/ecg_plot.py @@ -1,17 +1,18 @@ # -*- coding: utf-8 -*- +from warnings import warn + import matplotlib.gridspec import matplotlib.pyplot as plt import numpy as np import pandas as pd -from ..ecg import ecg_peaks -from ..signal import signal_fixpeaks +from ..misc import NeuroKitWarning from ..signal.signal_rate import _signal_rate_plot -from ..stats import rescale +from .ecg_peaks import _ecg_peaks_plot from .ecg_segment import ecg_segment -def ecg_plot(ecg_signals, rpeaks=None, sampling_rate=1000, show_type="default"): +def ecg_plot(ecg_signals, info=None): """**Visualize ECG data** Plot ECG signals and R-peaks. @@ -20,14 +21,8 @@ def ecg_plot(ecg_signals, rpeaks=None, sampling_rate=1000, show_type="default"): ---------- ecg_signals : DataFrame DataFrame obtained from ``ecg_process()``. - rpeaks : dict - The samples at which the R-peak occur. Dict returned by - ``ecg_process()``. Defaults to ``None``. - sampling_rate : int - The sampling frequency of ``ecg_cleaned`` (in Hz, i.e., samples/second). Defaults to 1000. - show_type : str - Visualize the ECG data with ``"default"`` or visualize artifacts thresholds with - ``"artifacts"`` produced by ``ecg_fixpeaks()``, or ``"full"`` to visualize both. + info : dict + The information Dict returned by ``ecg_process()``. Defaults to ``None``. See Also -------- @@ -37,10 +32,11 @@ def ecg_plot(ecg_signals, rpeaks=None, sampling_rate=1000, show_type="default"): ------- Though the function returns nothing, the figure can be retrieved and saved as follows: - .. code-block:: console + .. code-block:: python # To be run after ecg_plot() fig = plt.gcf() + fig.set_size_inches(10, 12, forward=True) fig.savefig("myfig.png") Examples @@ -57,7 +53,7 @@ def ecg_plot(ecg_signals, rpeaks=None, sampling_rate=1000, show_type="default"): # Plot @savefig p_ecg_plot.png scale=100% - nk.ecg_plot(signals, sampling_rate=1000, show_type='default') + nk.ecg_plot(signals, info) @suppress plt.close() @@ -70,102 +66,57 @@ def ecg_plot(ecg_signals, rpeaks=None, sampling_rate=1000, show_type="default"): ) # Extract R-peaks. - peaks = np.where(ecg_signals["ECG_R_Peaks"] == 1)[0] - - # Prepare figure and set axes. - if show_type in ["default", "full"]: - x_axis = np.linspace(0, len(ecg_signals) / sampling_rate, len(ecg_signals)) - gs = matplotlib.gridspec.GridSpec(2, 2, width_ratios=[2 / 3, 1 / 3]) - fig = plt.figure(constrained_layout=False) - ax0 = fig.add_subplot(gs[0, :-1]) - ax0.set_xlabel("Time (seconds)") - - ax1 = fig.add_subplot(gs[1, :-1], sharex=ax0) - ax2 = fig.add_subplot(gs[:, -1]) - - fig.suptitle("Electrocardiogram (ECG)", fontweight="bold") - - # Plot cleaned, raw ECG, R-peaks and signal quality. - ax0.set_title("Raw and Cleaned Signal") - - quality = rescale( - ecg_signals["ECG_Quality"], - to=[np.min(ecg_signals["ECG_Clean"]), np.max(ecg_signals["ECG_Clean"])], - ) - minimum_line = np.full(len(x_axis), quality.min()) - - # Plot quality area first - ax0.fill_between( - x_axis, - minimum_line, - quality, - alpha=0.12, - zorder=0, - interpolate=True, - facecolor="#4CAF50", - label="Quality", - ) - - # Plot signals - ax0.plot(x_axis, ecg_signals["ECG_Raw"], color="#B0BEC5", label="Raw", zorder=1) - ax0.plot( - x_axis, - ecg_signals["ECG_Clean"], - color="#F44336", - label="Cleaned", - zorder=1, - linewidth=1.5, - ) - ax0.scatter( - x_axis[peaks], - ecg_signals["ECG_Clean"][peaks], - color="#FFC107", - label="R-peaks", - zorder=2, + if info is None: + warn( + "'info' dict not provided. Some information might be missing." + + " Sampling rate will be set to 1000 Hz.", + category=NeuroKitWarning, ) - # Optimize legend - handles, labels = ax0.get_legend_handles_labels() - order = [2, 0, 1, 3] - ax0.legend( - [handles[idx] for idx in order], - [labels[idx] for idx in order], - loc="upper right", - ) - - # Plot Heart Rate - ax1 = _signal_rate_plot( - ecg_signals["ECG_Rate"].values, - peaks, - sampling_rate=sampling_rate, - title="Heart Rate", - ytitle="Beats per minute (bpm)", - color="#FF5722", - color_mean="#FF9800", - color_points="red", - ax=ax1, - ) + info = { + "ECG_R_Peaks": np.where(ecg_signals["ECG_R_Peaks"] == 1)[0], + "sampling_rate": 1000, + } - # Plot individual heart beats - ax2 = ecg_segment( - ecg_signals["ECG_Clean"], peaks, sampling_rate, show="return", ax=ax2 - ) - - # Plot artifacts - if show_type in ["artifacts", "full"]: - if sampling_rate is None: - raise ValueError( - "NeuroKit error: ecg_plot(): Sampling rate must be specified for artifacts" - " to be plotted." - ) - - if rpeaks is None: - _, rpeaks = ecg_peaks(ecg_signals["ECG_Clean"], sampling_rate=sampling_rate) - - fig = signal_fixpeaks( - rpeaks, - sampling_rate=sampling_rate, - iterative=True, - show=True, - method="Kubios", - ) + # Prepare figure and set axes. + gs = matplotlib.gridspec.GridSpec(2, 2, width_ratios=[2 / 3, 1 / 3]) + + fig = plt.figure(constrained_layout=False) + fig.suptitle("Electrocardiogram (ECG)", fontweight="bold") + + ax0 = fig.add_subplot(gs[0, :-1]) + ax1 = fig.add_subplot(gs[1, :-1], sharex=ax0) + ax2 = fig.add_subplot(gs[:, -1]) + + # Plot signals + ax0 = _ecg_peaks_plot( + ecg_signals["ECG_Clean"].values, + info=info, + sampling_rate=info["sampling_rate"], + raw=ecg_signals["ECG_Raw"].values, + quality=ecg_signals["ECG_Quality"].values, + phase=ecg_signals["ECG_Phase_Ventricular"].values, + ax=ax0, + ) + + # Plot Heart Rate + ax1 = _signal_rate_plot( + ecg_signals["ECG_Rate"].values, + info["ECG_R_Peaks"], + sampling_rate=info["sampling_rate"], + title="Heart Rate", + ytitle="Beats per minute (bpm)", + color="#FF5722", + color_mean="#FF9800", + color_points="#FFC107", + ax=ax1, + ) + + # Plot individual heart beats + ax2 = ecg_segment( + ecg_signals, + info["ECG_R_Peaks"], + info["sampling_rate"], + show="return", + ax=ax2, + ) diff --git a/neurokit2/ecg/ecg_process.py b/neurokit2/ecg/ecg_process.py index aacf017c44..39a6274a0a 100644 --- a/neurokit2/ecg/ecg_process.py +++ b/neurokit2/ecg/ecg_process.py @@ -84,7 +84,7 @@ def ecg_process(ecg_signal, sampling_rate=1000, method="neurokit"): # Visualize @savefig p_ecg_process.png scale=100% - nk.ecg_plot(signals) + nk.ecg_plot(signals, info) @suppress plt.close() @@ -95,29 +95,50 @@ def ecg_process(ecg_signal, sampling_rate=1000, method="neurokit"): ecg_cleaned = ecg_clean(ecg_signal, sampling_rate=sampling_rate, method=method) # Detect R-peaks - instant_peaks, rpeaks = ecg_peaks( - ecg_cleaned=ecg_cleaned, sampling_rate=sampling_rate, method=method, correct_artifacts=True + instant_peaks, info = ecg_peaks( + ecg_cleaned=ecg_cleaned, + sampling_rate=sampling_rate, + method=method, + correct_artifacts=True, ) # Calculate heart rate - rate = signal_rate(rpeaks, sampling_rate=sampling_rate, desired_length=len(ecg_cleaned)) + rate = signal_rate( + info, sampling_rate=sampling_rate, desired_length=len(ecg_cleaned) + ) # Assess signal quality - quality = ecg_quality(ecg_cleaned, rpeaks=rpeaks["ECG_R_Peaks"], sampling_rate=sampling_rate) + quality = ecg_quality( + ecg_cleaned, rpeaks=info["ECG_R_Peaks"], sampling_rate=sampling_rate + ) # Merge signals in a DataFrame - signals = pd.DataFrame({"ECG_Raw": ecg_signal, "ECG_Clean": ecg_cleaned, "ECG_Rate": rate, "ECG_Quality": quality}) + signals = pd.DataFrame( + { + "ECG_Raw": ecg_signal, + "ECG_Clean": ecg_cleaned, + "ECG_Rate": rate, + "ECG_Quality": quality, + } + ) # Delineate QRS complex delineate_signal, delineate_info = ecg_delineate( - ecg_cleaned=ecg_cleaned, rpeaks=rpeaks, sampling_rate=sampling_rate + ecg_cleaned=ecg_cleaned, rpeaks=info["ECG_R_Peaks"], sampling_rate=sampling_rate ) + info.update(delineate_info) # Merge waves indices dict with info dict # Determine cardiac phases - cardiac_phase = ecg_phase(ecg_cleaned=ecg_cleaned, rpeaks=rpeaks, delineate_info=delineate_info) + cardiac_phase = ecg_phase( + ecg_cleaned=ecg_cleaned, + rpeaks=info["ECG_R_Peaks"], + delineate_info=delineate_info, + ) # Add additional information to signals DataFrame - signals = pd.concat([signals, instant_peaks, delineate_signal, cardiac_phase], axis=1) + signals = pd.concat( + [signals, instant_peaks, delineate_signal, cardiac_phase], axis=1 + ) # return signals DataFrame and R-peak locations - return signals, rpeaks + return signals, info diff --git a/neurokit2/ecg/ecg_segment.py b/neurokit2/ecg/ecg_segment.py index 1cb3e5aec7..0c177753f3 100644 --- a/neurokit2/ecg/ecg_segment.py +++ b/neurokit2/ecg/ecg_segment.py @@ -66,11 +66,14 @@ def ecg_segment(ecg_cleaned, rpeaks=None, sampling_rate=1000, show=False, **kwar epochs_end=epochs_end, ) - # pad last heartbeat with nan so that segments are equal length + # Pad last heartbeats with nan so that segments are equal length last_heartbeat_key = str(np.max(np.array(list(heartbeats.keys()), dtype=int))) after_last_index = heartbeats[last_heartbeat_key]["Index"] < len(ecg_cleaned) - heartbeats[last_heartbeat_key].loc[after_last_index, "Signal"] = np.nan + for col in ["Signal", "ECG_Raw", "ECG_Clean"]: + if col in heartbeats[last_heartbeat_key].columns: + heartbeats[last_heartbeat_key].loc[after_last_index, col] = np.nan + # Plot or return plot axis (feature meant to be used internally in ecg_plot) if show is not False: ax = _ecg_segment_plot(heartbeats, heartrate=average_hr, ytitle="ECG", **kwargs) if show == "return": @@ -84,13 +87,17 @@ def ecg_segment(ecg_cleaned, rpeaks=None, sampling_rate=1000, show=False, **kwar # ============================================================================= def _ecg_segment_plot(heartbeats, heartrate=0, ytitle="ECG", color="#F44336", ax=None): df = epochs_to_df(heartbeats) + + # Get main signal column name + col = [c for c in ["Signal", "ECG_Raw", "ECG_Clean"] if c in df.columns][-1] + # Average heartbeat - mean_heartbeat = df.drop(["Index", "Label"], axis=1).groupby("Time").mean() - df_pivoted = df.pivot(index="Time", columns="Label", values="Signal") + mean_heartbeat = df.groupby("Time")[[col]].mean() + df_pivoted = df.pivot(index="Time", columns="Label", values=col) # Prepare plot if ax is None: - fig, ax = plt.subplots() + _, ax = plt.subplots() ax.set_title(f"Individual Heart Beats (average heart rate: {heartrate:0.1f} bpm)") ax.set_xlabel("Time (seconds)") @@ -101,19 +108,49 @@ def _ecg_segment_plot(heartbeats, heartrate=0, ytitle="ECG", color="#F44336", ax # Plot average heartbeat ax.plot( - mean_heartbeat.index, mean_heartbeat, color=color, linewidth=10, label="Average" + mean_heartbeat.index, + mean_heartbeat, + color=color, + linewidth=7, + label="Average beat shape", + zorder=1, ) + # Alpha of individual beats decreases with more heartbeats + alpha = 1 / np.log1p(np.log2(1 + df_pivoted.shape[1])) + # Plot all heartbeats - alpha = 1 / np.log2(1 + df_pivoted.shape[1]) # alpha decreases with more heartbeats - ax.plot(df_pivoted, color="grey", linewidth=alpha) - ax.legend(loc="upper right") + ax.plot(df_pivoted, color="grey", linewidth=alpha, zorder=2) + + # Plot individual waves + for wave in [ + ("P", "#3949AB"), + ("Q", "#1E88E5"), + ("S", "#039BE5"), + ("T", "#00ACC1"), + ]: + wave_col = f"ECG_{wave[0]}_Peaks" + if wave_col in df.columns: + ax.scatter( + df["Time"][df[wave_col] == 1], + df[col][df[wave_col] == 1], + color=wave[1], + marker="+", + label=f"{wave[0]}-waves", + zorder=3, + ) + # Legend + ax.legend(loc="upper right") return ax def _ecg_segment_window( - heart_rate=None, rpeaks=None, sampling_rate=1000, desired_length=None + heart_rate=None, + rpeaks=None, + sampling_rate=1000, + desired_length=None, + ratio_pre=0.35, ): # Extract heart rate if heart_rate is not None: @@ -130,7 +167,7 @@ def _ecg_segment_window( window_size = 60 / heart_rate # Beats per second # Window - epochs_start = 1 / 3 * window_size - epochs_end = 2 / 3 * window_size + epochs_start = ratio_pre * window_size + epochs_end = (1 - ratio_pre) * window_size return -epochs_start, epochs_end, heart_rate diff --git a/neurokit2/eda/eda_plot.py b/neurokit2/eda/eda_plot.py index e0a5fb4536..d61274e435 100644 --- a/neurokit2/eda/eda_plot.py +++ b/neurokit2/eda/eda_plot.py @@ -1,19 +1,23 @@ # -*- coding: utf-8 -*- +from warnings import warn + import matplotlib.collections import matplotlib.pyplot as plt import numpy as np import pandas as pd +from ..misc import NeuroKitWarning + -def eda_plot(eda_signals, sampling_rate=None, static=True): +def eda_plot(eda_signals, info=None, static=True): """**Visualize electrodermal activity (EDA) data** Parameters ---------- eda_signals : DataFrame DataFrame obtained from :func:`eda_process()`. - sampling_rate : int - The desired sampling rate (in Hz, i.e., samples/second). Defaults to None. + info : dict + The information Dict returned by ``eda_process()``. Defaults to ``None``. static : bool If True, a static plot will be generated with matplotlib. If False, an interactive plot will be generated with plotly. @@ -21,8 +25,7 @@ def eda_plot(eda_signals, sampling_rate=None, static=True): Returns ------- - fig - Figure representing a plot of the processed EDA signals. + See :func:`.ecg_plot` for details on how to access the figure, modify the size and save it. Examples -------- @@ -32,8 +35,9 @@ def eda_plot(eda_signals, sampling_rate=None, static=True): eda_signal = nk.eda_simulate(duration=30, scr_number=5, drift=0.1, noise=0, sampling_rate=250) eda_signals, info = nk.eda_process(eda_signal, sampling_rate=250) + @savefig p_eda_plot1.png scale=100% - nk.eda_plot(eda_signals) + nk.eda_plot(eda_signals, info) @suppress plt.close() @@ -42,25 +46,31 @@ def eda_plot(eda_signals, sampling_rate=None, static=True): eda_process """ + if info is None: + warn( + "'info' dict not provided. Some information might be missing." + + " Sampling rate will be set to 1000 Hz.", + category=NeuroKitWarning, + ) + + info = { + "sampling_rate": 1000, + } + # Determine peaks, onsets, and half recovery. peaks = np.where(eda_signals["SCR_Peaks"] == 1)[0] onsets = np.where(eda_signals["SCR_Onsets"] == 1)[0] half_recovery = np.where(eda_signals["SCR_Recovery"] == 1)[0] # Determine unit of x-axis. - if sampling_rate is not None: - x_label = "Seconds" - x_axis = np.linspace(0, len(eda_signals) / sampling_rate, len(eda_signals)) - else: - x_label = "Samples" - x_axis = np.arange(0, len(eda_signals)) + x_label = "Time (seconds)" + x_axis = np.linspace(0, len(eda_signals) / info["sampling_rate"], len(eda_signals)) if static: fig, (ax0, ax1, ax2) = plt.subplots(nrows=3, ncols=1, sharex=True) last_ax = fig.get_axes()[-1] last_ax.set_xlabel(x_label) - plt.tight_layout(h_pad=0.2) # Plot cleaned and raw electrodermal activity. ax0.set_title("Raw and Cleaned Signal") @@ -121,7 +131,7 @@ def eda_plot(eda_signals, sampling_rate=None, static=True): linewidth=1.5, ) ax2.legend(loc="upper right") - return fig + else: # Create interactive plot with plotly. try: diff --git a/neurokit2/eda/eda_process.py b/neurokit2/eda/eda_process.py index abff6d7e28..862dd5cd2d 100644 --- a/neurokit2/eda/eda_process.py +++ b/neurokit2/eda/eda_process.py @@ -4,13 +4,15 @@ from ..misc.report import create_report from ..signal import signal_sanitize from .eda_clean import eda_clean +from .eda_methods import eda_methods from .eda_peaks import eda_peaks from .eda_phasic import eda_phasic -from .eda_methods import eda_methods from .eda_plot import eda_plot -def eda_process(eda_signal, sampling_rate=1000, method="neurokit", report=None, **kwargs): +def eda_process( + eda_signal, sampling_rate=1000, method="neurokit", report=None, **kwargs +): """**Process Electrodermal Activity (EDA)** Convenience function that automatically processes electrodermal activity (EDA) signal. @@ -79,8 +81,9 @@ def eda_process(eda_signal, sampling_rate=1000, method="neurokit", report=None, eda_signal = nk.eda_simulate(duration=30, scr_number=5, drift=0.1, noise=0) signals, info = nk.eda_process(eda_signal, sampling_rate=1000) + @savefig p_eda_process.png scale=100% - nk.eda_plot(signals) + nk.eda_plot(signals, info) @suppress plt.close() @@ -91,17 +94,21 @@ def eda_process(eda_signal, sampling_rate=1000, method="neurokit", report=None, # Preprocess # Clean signal - eda_cleaned = eda_clean(eda_signal, - sampling_rate=sampling_rate, - method=methods["method_cleaning"], - **methods["kwargs_cleaning"]) + eda_cleaned = eda_clean( + eda_signal, + sampling_rate=sampling_rate, + method=methods["method_cleaning"], + **methods["kwargs_cleaning"], + ) if methods["method_phasic"] is None or methods["method_phasic"].lower() == "none": eda_decomposed = pd.DataFrame({"EDA_Phasic": eda_cleaned}) else: - eda_decomposed = eda_phasic(eda_cleaned, - sampling_rate=sampling_rate, - method=methods["method_phasic"], - **methods["kwargs_phasic"]) + eda_decomposed = eda_phasic( + eda_cleaned, + sampling_rate=sampling_rate, + method=methods["method_phasic"], + **methods["kwargs_phasic"], + ) # Find peaks peak_signal, info = eda_peaks( @@ -121,7 +128,7 @@ def eda_process(eda_signal, sampling_rate=1000, method="neurokit", report=None, if report is not None: # Generate report containing description and figures of processing if ".html" in str(report): - fig = eda_plot(signals, sampling_rate=sampling_rate, static=False) + fig = eda_plot(signals, info, static=False) else: fig = None create_report(file=report, signals=signals, info=methods, fig=fig) diff --git a/neurokit2/emg/emg_plot.py b/neurokit2/emg/emg_plot.py index 9b493202d0..f5c7cb489b 100644 --- a/neurokit2/emg/emg_plot.py +++ b/neurokit2/emg/emg_plot.py @@ -1,10 +1,14 @@ # -*- coding: utf-8 -*- +from warnings import warn + import matplotlib.pyplot as plt import numpy as np import pandas as pd +from ..misc import NeuroKitWarning + -def emg_plot(emg_signals, sampling_rate=None, static=True): +def emg_plot(emg_signals, info=None, static=True): """**EMG Graph** Visualize electromyography (EMG) data. @@ -13,10 +17,8 @@ def emg_plot(emg_signals, sampling_rate=None, static=True): ---------- emg_signals : DataFrame DataFrame obtained from ``emg_process()``. - sampling_rate : int - The sampling frequency of the EMG (in Hz, i.e., samples/second). Needs to be supplied if the - data should be plotted over time in seconds. Otherwise the data is plotted over samples. - Defaults to ``None``. + info : dict + The information Dict returned by ``emg_process()``. Defaults to ``None``. static : bool If True, a static plot will be generated with matplotlib. If False, an interactive plot will be generated with plotly. @@ -28,13 +30,7 @@ def emg_plot(emg_signals, sampling_rate=None, static=True): Returns ------- - Though the function returns nothing, the figure can be retrieved and saved as follows: - - .. code-block:: console - - # To be run after emg_plot() - fig = plt.gcf() - fig.savefig("myfig.png") + See :func:`.ecg_plot` for details on how to access the figure, modify the size and save it. Examples -------- @@ -46,17 +42,26 @@ def emg_plot(emg_signals, sampling_rate=None, static=True): emg = nk.emg_simulate(duration=10, sampling_rate=1000, burst_number=3) # Process signal - emg_signals, _ = nk.emg_process(emg, sampling_rate=1000) + emg_signals, info = nk.emg_process(emg, sampling_rate=1000) # Plot @savefig p_emg_plot.png scale=100% - nk.emg_plot(emg_signals) + nk.emg_plot(emg_signals, info) @suppress plt.close() + """ + if info is None: + warn( + "'info' dict not provided. Some information might be missing." + + " Sampling rate will be set to 1000 Hz.", + category=NeuroKitWarning, + ) + info = { + "sampling_rate": 1000, + } - """ # Mark onsets, offsets, activity onsets = np.where(emg_signals["EMG_Onsets"] == 1)[0] offsets = np.where(emg_signals["EMG_Offsets"] == 1)[0] @@ -69,22 +74,22 @@ def emg_plot(emg_signals, sampling_rate=None, static=True): ) # Determine what to display on the x-axis, mark activity. - if sampling_rate is not None: - x_axis = np.linspace(0, emg_signals.shape[0] / sampling_rate, emg_signals.shape[0]) - else: - x_axis = np.arange(0, emg_signals.shape[0]) + x_axis = np.linspace( + 0, emg_signals.shape[0] / info["sampling_rate"], emg_signals.shape[0] + ) if static is True: - return _emg_plot_static(emg_signals, x_axis, onsets, offsets, sampling_rate) + _emg_plot_static(emg_signals, x_axis, onsets, offsets, info["sampling_rate"]) else: - return _emg_plot_interactive(emg_signals, x_axis, onsets, offsets, sampling_rate) + return _emg_plot_interactive( + emg_signals, x_axis, onsets, offsets, info["sampling_rate"] + ) # ============================================================================= # Internals # ============================================================================= def _emg_plot_activity(emg_signals, onsets, offsets): - activity_signal = pd.Series(np.full(len(emg_signals), np.nan)) activity_signal[onsets] = emg_signals["EMG_Amplitude"][onsets].values activity_signal[offsets] = emg_signals["EMG_Amplitude"][offsets].values @@ -113,14 +118,23 @@ def _emg_plot_static(emg_signals, x_axis, onsets, offsets, sampling_rate): ax0.set_title("Raw and Cleaned Signal") ax0.plot(x_axis, emg_signals["EMG_Raw"], color="#B0BEC5", label="Raw", zorder=1) ax0.plot( - x_axis, emg_signals["EMG_Clean"], color="#FFC107", label="Cleaned", zorder=1, linewidth=1.5 + x_axis, + emg_signals["EMG_Clean"], + color="#FFC107", + label="Cleaned", + zorder=1, + linewidth=1.5, ) ax0.legend(loc="upper right") # Plot Amplitude. ax1.set_title("Muscle Activation") ax1.plot( - x_axis, emg_signals["EMG_Amplitude"], color="#FF9800", label="Amplitude", linewidth=1.5 + x_axis, + emg_signals["EMG_Amplitude"], + color="#FF9800", + label="Amplitude", + linewidth=1.5, ) # Shade activity regions. @@ -137,7 +151,11 @@ def _emg_plot_static(emg_signals, x_axis, onsets, offsets, sampling_rate): # Mark onsets and offsets. ax1.scatter( - x_axis[onsets], emg_signals["EMG_Amplitude"][onsets], color="#f03e65", label=None, zorder=3 + x_axis[onsets], + emg_signals["EMG_Amplitude"][onsets], + color="#f03e65", + label=None, + zorder=3, ) ax1.scatter( x_axis[offsets], @@ -155,8 +173,6 @@ def _emg_plot_static(emg_signals, x_axis, onsets, offsets, sampling_rate): ax1.axvline(i, color="#4a4a4a", linestyle="--", label=None, zorder=2) ax1.axvline(j, color="#4a4a4a", linestyle="--", label=None, zorder=2) ax1.legend(loc="upper right") - plt.close() - return fig def _emg_plot_interactive(emg_signals, x_axis, onsets, offsets, sampling_rate): diff --git a/neurokit2/emg/emg_process.py b/neurokit2/emg/emg_process.py index 4742770216..91229ae890 100644 --- a/neurokit2/emg/emg_process.py +++ b/neurokit2/emg/emg_process.py @@ -61,7 +61,7 @@ def emg_process(emg_signal, sampling_rate=1000, report=None, **kwargs): signals, info = nk.emg_process(emg, sampling_rate=1000) @savefig p_emg_process1.png scale=100% - nk.emg_plot(signals) + nk.emg_plot(signals, info) @suppress plt.close() @@ -98,7 +98,7 @@ def emg_process(emg_signal, sampling_rate=1000, report=None, **kwargs): if report is not None: # Generate report containing description and figures of processing if ".html" in str(report): - fig = emg_plot(signals, sampling_rate=sampling_rate, static=False) + fig = emg_plot(signals, info, static=False) else: fig = None create_report(file=report, signals=signals, info=methods, fig=fig) diff --git a/neurokit2/eog/eog_findpeaks.py b/neurokit2/eog/eog_findpeaks.py index 739ce6a067..2dfa5aeba8 100644 --- a/neurokit2/eog/eog_findpeaks.py +++ b/neurokit2/eog/eog_findpeaks.py @@ -113,7 +113,9 @@ def eog_findpeaks(veog_cleaned, sampling_rate=None, method="mne", **kwargs): elif method in ["blinker"]: peaks = _eog_findpeaks_blinker(eog_cleaned, sampling_rate=sampling_rate) elif method in ["neurokit", "nk"]: - peaks = _eog_findpeaks_neurokit(eog_cleaned, sampling_rate=sampling_rate, **kwargs) + peaks = _eog_findpeaks_neurokit( + eog_cleaned, sampling_rate=sampling_rate, **kwargs + ) # elif method in ["jammes2008", "jammes"]: # peaks = _eog_findpeaks_jammes2008(eog_cleaned, sampling_rate=sampling_rate) else: @@ -131,7 +133,7 @@ def eog_findpeaks(veog_cleaned, sampling_rate=None, method="mne", **kwargs): def _eog_findpeaks_neurokit(eog_cleaned, sampling_rate=1000, threshold=0.33, show=True): """In-house EOG blink detection.""" peaks = signal_findpeaks(eog_cleaned, relative_height_min=1.25)["Peaks"] - peaks = signal_fixpeaks( + _, peaks = signal_fixpeaks( peaks=peaks, sampling_rate=sampling_rate, interval_min=0.2, method="neurokit" ) peaks = _eog_findpeaks_neurokit_filterblinks( @@ -146,7 +148,11 @@ def _eog_findpeaks_neurokit_filterblinks( """Compare each detected event to blink template and reject it if too different.""" # Get epoch around each blink events = epochs_create( - eog_cleaned, peaks, sampling_rate=sampling_rate, epochs_start=-0.4, epochs_end=0.6 + eog_cleaned, + peaks, + sampling_rate=sampling_rate, + epochs_start=-0.4, + epochs_end=0.6, ) events = epochs_to_array(events) # Convert to 2D array diff --git a/neurokit2/eog/eog_plot.py b/neurokit2/eog/eog_plot.py index c4a78578a3..d754c07607 100644 --- a/neurokit2/eog/eog_plot.py +++ b/neurokit2/eog/eog_plot.py @@ -1,27 +1,25 @@ # -*- coding: utf-8 -*- +from warnings import warn + import matplotlib.gridspec import matplotlib.pyplot as plt import numpy as np import pandas as pd from ..epochs import epochs_create, epochs_to_array, epochs_to_df +from ..misc import NeuroKitWarning from ..stats import standardize -def eog_plot(eog_signals, peaks=None, sampling_rate=None): +def eog_plot(eog_signals, info=None): """**Visualize EOG data** Parameters ---------- eog_signals : DataFrame DataFrame obtained from :func:`.eog_process`. - peaks : dict - The samples at which the blink peaks occur. Dict returned by - :func:`.eog_process`. Defaults to ``None``. Must be specified to plot individual blinks. - sampling_rate : int - The sampling frequency of the EOG (in Hz, i.e., samples/second). Needs to be supplied if - the data should be plotted over time in seconds. Otherwise the data is plotted over - samples. Defaults to ``None``. Must be specified to plot individual blinks. + info : dict + The information Dict returned by ``eog_process()``. Defaults to ``None``. See Also -------- @@ -47,15 +45,25 @@ def eog_plot(eog_signals, peaks=None, sampling_rate=None): eog_signal = nk.data('eog_100hz') # Process signal - eog_signals, peaks = nk.eog_process(eog_signal, sampling_rate=100) + eog_signals, info = nk.eog_process(eog_signal, sampling_rate=100) # Plot - @savefig p.eog_plot.png scale=100% - nk.eog_plot(eog_signals, peaks, sampling_rate=100) + @savefig p_eog_plot.png scale=100% + nk.eog_plot(eog_signals, info) @suppress plt.close() """ + if info is None: + warn( + "'info' dict not provided. Some information might be missing." + + " Sampling rate will be set to 1000 Hz.", + category=NeuroKitWarning, + ) + + info = { + "sampling_rate": 1000, + } # Sanity-check input. if not isinstance(eog_signals, pd.DataFrame): @@ -65,21 +73,17 @@ def eog_plot(eog_signals, peaks=None, sampling_rate=None): ) # Prepare figure - if sampling_rate is not None: - x_axis = np.linspace(0, eog_signals.shape[0] / sampling_rate, eog_signals.shape[0]) - gs = matplotlib.gridspec.GridSpec(2, 2, width_ratios=[1 - 1 / np.pi, 1 / np.pi]) - fig = plt.figure(constrained_layout=False) - ax0 = fig.add_subplot(gs[0, :-1]) - ax1 = fig.add_subplot(gs[1, :-1]) - ax2 = fig.add_subplot(gs[:, -1]) - ax0.set_xlabel("Time (seconds)") - ax1.set_xlabel("Time (seconds)") - ax2.set_xlabel("Time (seconds)") - else: - x_axis = np.arange(0, eog_signals.shape[0]) - fig, (ax0, ax1) = plt.subplots(nrows=2, ncols=1, sharex=True) - ax0.set_xlabel("Samples") - ax1.set_xlabel("Samples") + x_axis = np.linspace( + 0, eog_signals.shape[0] / info["sampling_rate"], eog_signals.shape[0] + ) + gs = matplotlib.gridspec.GridSpec(2, 2, width_ratios=[1 - 1 / np.pi, 1 / np.pi]) + fig = plt.figure(constrained_layout=False) + ax0 = fig.add_subplot(gs[0, :-1]) + ax1 = fig.add_subplot(gs[1, :-1]) + ax2 = fig.add_subplot(gs[:, -1]) + ax0.set_xlabel("Time (seconds)") + ax1.set_xlabel("Time (seconds)") + ax2.set_xlabel("Time (seconds)") fig.suptitle("Electrooculography (EOG)", fontweight="bold") plt.tight_layout(h_pad=0.3, w_pad=0.2) @@ -88,14 +92,23 @@ def eog_plot(eog_signals, peaks=None, sampling_rate=None): ax0.set_title("Raw and Cleaned Signal") ax0.plot(x_axis, eog_signals["EOG_Raw"], color="#B0BEC5", label="Raw", zorder=1) ax0.plot( - x_axis, eog_signals["EOG_Clean"], color="#49A4FD", label="Cleaned", zorder=1, linewidth=1.5 + x_axis, + eog_signals["EOG_Clean"], + color="#49A4FD", + label="Cleaned", + zorder=1, + linewidth=1.5, ) ax0.set_ylabel("Amplitude (mV)") # Plot blinks blinks = np.where(eog_signals["EOG_Blinks"] == 1)[0] ax0.scatter( - x_axis[blinks], eog_signals["EOG_Clean"][blinks], color="#0146D7", label="Blinks", zorder=2 + x_axis[blinks], + eog_signals["EOG_Clean"][blinks], + color="#0146D7", + label="Blinks", + zorder=2, ) ax0.legend(loc="upper right") @@ -103,42 +116,43 @@ def eog_plot(eog_signals, peaks=None, sampling_rate=None): ax1.set_title("Blink Rate") ax1.set_ylabel("Blinks per minute") blink_rate_mean = eog_signals["EOG_Rate"].mean() - ax1.plot(x_axis, eog_signals["EOG_Rate"], color="#9C5AFF", label="Rate", linewidth=1.5) + ax1.plot( + x_axis, eog_signals["EOG_Rate"], color="#9C5AFF", label="Rate", linewidth=1.5 + ) ax1.axhline(y=blink_rate_mean, label="Mean", linestyle="--", color="#CEAFFF") ax1.legend(loc="upper right") # Plot individual blinks - if sampling_rate is not None: - ax2.set_title("Individual Blinks") - - # Create epochs - events = epochs_create( - eog_signals["EOG_Clean"], - peaks["EOG_Blinks"], - sampling_rate=sampling_rate, - epochs_start=-0.3, - epochs_end=0.7, - ) - events_array = epochs_to_array(events) # Convert to 2D array - events_array = standardize( - events_array - ) # Rescale so that all the blinks are on the same scale - - blinks_df = epochs_to_df(events) - blinks_wide = blinks_df.pivot(index="Time", columns="Label", values="Signal") - blinks_wide = standardize(blinks_wide) - - cmap = iter(plt.cm.RdBu(np.linspace(0, 1, num=len(events)))) - for x, color in zip(blinks_wide, cmap): - ax2.plot(blinks_wide[x], color=color, linewidth=0.4, zorder=1) - - # Plot with their median (used here as a robust average) - ax2.plot( - np.array(blinks_wide.index), - np.median(events_array, axis=1), - linewidth=2, - linestyle="--", - color="black", - label="Median", - ) - ax2.legend(loc="upper right") + ax2.set_title("Individual Blinks") + + # Create epochs + events = epochs_create( + eog_signals["EOG_Clean"], + info["EOG_Blinks"], + sampling_rate=info["sampling_rate"], + epochs_start=-0.3, + epochs_end=0.7, + ) + events_array = epochs_to_array(events) # Convert to 2D array + events_array = standardize( + events_array + ) # Rescale so that all the blinks are on the same scale + + blinks_df = epochs_to_df(events) + blinks_wide = blinks_df.pivot(index="Time", columns="Label", values="Signal") + blinks_wide = standardize(blinks_wide) + + cmap = iter(plt.cm.RdBu(np.linspace(0, 1, num=len(events)))) + for x, color in zip(blinks_wide, cmap): + ax2.plot(blinks_wide[x], color=color, linewidth=0.4, zorder=1) + + # Plot with their median (used here as a robust average) + ax2.plot( + np.array(blinks_wide.index), + np.median(events_array, axis=1), + linewidth=2, + linestyle="--", + color="black", + label="Median", + ) + ax2.legend(loc="upper right") diff --git a/neurokit2/eog/eog_process.py b/neurokit2/eog/eog_process.py index 9f9b5555ec..13ec77a858 100644 --- a/neurokit2/eog/eog_process.py +++ b/neurokit2/eog/eog_process.py @@ -50,9 +50,15 @@ def eog_process(veog_signal, sampling_rate=1000, **kwargs): import neurokit2 as nk # Get data - eog_signal = nk.data('eog_100hz') + eog = nk.data('eog_100hz') - signals, info = nk.eog_process(eog_signal, sampling_rate=100) + eog_signals, info = nk.eog_process(eog, sampling_rate=100) + + # Plot + @savefig p_eog_process.png scale=100% + nk.eog_plot(eog_signals, info) + @suppress + plt.close() References ---------- @@ -71,17 +77,24 @@ def eog_process(veog_signal, sampling_rate=1000, **kwargs): peaks = eog_findpeaks(eog_cleaned, sampling_rate=sampling_rate, **kwargs) info = {"EOG_Blinks": peaks} - info['sampling_rate'] = sampling_rate # Add sampling rate in dict info + info["sampling_rate"] = sampling_rate # Add sampling rate in dict info # Mark (potential) blink events signal_blinks = _signal_from_indices(peaks, desired_length=len(eog_cleaned)) # Rate computation - rate = signal_rate(peaks, sampling_rate=sampling_rate, desired_length=len(eog_cleaned)) + rate = signal_rate( + peaks, sampling_rate=sampling_rate, desired_length=len(eog_cleaned) + ) # Prepare output signals = pd.DataFrame( - {"EOG_Raw": eog_signal, "EOG_Clean": eog_cleaned, "EOG_Blinks": signal_blinks, "EOG_Rate": rate} + { + "EOG_Raw": eog_signal, + "EOG_Clean": eog_cleaned, + "EOG_Blinks": signal_blinks, + "EOG_Rate": rate, + } ) return signals, info diff --git a/neurokit2/ppg/ppg_peaks.py b/neurokit2/ppg/ppg_peaks.py index c09251ffac..ae0cccf727 100644 --- a/neurokit2/ppg/ppg_peaks.py +++ b/neurokit2/ppg/ppg_peaks.py @@ -1,8 +1,11 @@ +import matplotlib.pyplot as plt +import numpy as np + from ..signal import signal_formatpeaks from .ppg_findpeaks import ppg_findpeaks -def ppg_peaks(ppg_cleaned, sampling_rate=1000, method="elgendi", **kwargs): +def ppg_peaks(ppg_cleaned, sampling_rate=1000, method="elgendi", show=False, **kwargs): """**Find systolic peaks in a photoplethysmogram (PPG) signal** Find the peaks in an PPG signal using the specified method. You can pass an unfiltered PPG @@ -23,6 +26,8 @@ def ppg_peaks(ppg_cleaned, sampling_rate=1000, method="elgendi", **kwargs): method : str The processing pipeline to apply. Can be one of ``"elgendi"``, ``"bishop"``. The default is ``"elgendi"``. + show : bool + If ``True``, will show a plot of the signal with peaks. Defaults to ``False``. **kwargs Additional keyword arguments, usually specific for each method. @@ -75,7 +80,7 @@ def ppg_peaks(ppg_cleaned, sampling_rate=1000, method="elgendi", **kwargs): """ peaks = ppg_findpeaks( - ppg_cleaned, sampling_rate=sampling_rate, method=method, **kwargs + ppg_cleaned, sampling_rate=sampling_rate, method=method, show=False, **kwargs ) instant_peaks = signal_formatpeaks( @@ -85,4 +90,58 @@ def ppg_peaks(ppg_cleaned, sampling_rate=1000, method="elgendi", **kwargs): info = peaks info["sampling_rate"] = sampling_rate # Add sampling rate in dict info + if show is True: + _ppg_peaks_plot(ppg_cleaned, info, sampling_rate) + return signals, info + + +# ============================================================================= +# Internals +# ============================================================================= +def _ppg_peaks_plot( + ppg_cleaned, + info=None, + sampling_rate=1000, + raw=None, + ax=None, +): + x_axis = np.linspace(0, len(ppg_cleaned) / sampling_rate, len(ppg_cleaned)) + + # Prepare plot + if ax is None: + _, ax = plt.subplots() + + ax.set_xlabel("Time (seconds)") + ax.set_title("PPG signal and peaks") + + # Raw Signal --------------------------------------------------------------- + if raw is not None: + ax.plot(x_axis, raw, color="#B0BEC5", label="Raw signal", zorder=1) + label_clean = "Cleaned signal" + else: + label_clean = "Signal" + + # Peaks ------------------------------------------------------------------- + ax.scatter( + x_axis[info["PPG_Peaks"]], + ppg_cleaned[info["PPG_Peaks"]], + color="#FFC107", + label="Systolic peaks", + zorder=2, + ) + + # Clean Signal ------------------------------------------------------------ + ax.plot( + x_axis, + ppg_cleaned, + color="#E91E63", + label=label_clean, + zorder=3, + linewidth=1, + ) + + # Optimize legend + ax.legend(loc="upper right") + + return ax diff --git a/neurokit2/ppg/ppg_plot.py b/neurokit2/ppg/ppg_plot.py index 694d00890c..e9efb4c9e4 100644 --- a/neurokit2/ppg/ppg_plot.py +++ b/neurokit2/ppg/ppg_plot.py @@ -1,14 +1,18 @@ # -*- coding: utf-8 -*- +from warnings import warn + import matplotlib import matplotlib.pyplot as plt import numpy as np import pandas as pd +from ..misc import NeuroKitWarning from ..signal.signal_rate import _signal_rate_plot +from .ppg_peaks import _ppg_peaks_plot from .ppg_segment import ppg_segment -def ppg_plot(ppg_signals, sampling_rate=1000, static=True): +def ppg_plot(ppg_signals, info=None, static=True): """**Visualize photoplethysmogram (PPG) data** Visualize the PPG signal processing. @@ -17,8 +21,8 @@ def ppg_plot(ppg_signals, sampling_rate=1000, static=True): ---------- ppg_signals : DataFrame DataFrame obtained from :func:`.ppg_process`. - sampling_rate : int - The sampling frequency of the PPG (in Hz, i.e., samples/second). Defaults to 1000. + info : dict + The information Dict returned by ``ppg_process()``. Defaults to ``None``. static : bool If True, a static plot will be generated with matplotlib. If False, an interactive plot will be generated with plotly. @@ -26,8 +30,7 @@ def ppg_plot(ppg_signals, sampling_rate=1000, static=True): Returns ------- - fig - Figure representing a plot of the processed PPG signals. + See :func:`.ecg_plot` for details on how to access the figure, modify the size and save it. See Also -------- @@ -46,7 +49,7 @@ def ppg_plot(ppg_signals, sampling_rate=1000, static=True): # Plot @savefig p_ppg_plot1.png scale=100% - nk.ppg_plot(signals, sampling_rate=100) + nk.ppg_plot(signals, info) @suppress plt.close() @@ -60,63 +63,58 @@ def ppg_plot(ppg_signals, sampling_rate=1000, static=True): ) # Extract Peaks. - peaks = None - if "PPG_Peaks" in ppg_signals.columns: - peaks = np.where(ppg_signals["PPG_Peaks"] == 1)[0] + if info is None: + warn( + "'info' dict not provided. Some information might be missing." + + " Sampling rate will be set to 1000 Hz.", + category=NeuroKitWarning, + ) - # X-axis - x_axis = np.linspace(0, len(ppg_signals) / sampling_rate, len(ppg_signals)) + info = { + "PPG_Peaks": np.where(ppg_signals["PPG_Peaks"] == 1)[0], + "sampling_rate": 1000, + } if static: # Prepare figure gs = matplotlib.gridspec.GridSpec(2, 2, width_ratios=[2 / 3, 1 / 3]) fig = plt.figure(constrained_layout=False) - ax0 = fig.add_subplot(gs[0, :-1]) - ax0.set_xlabel("Time (seconds)") + ax0 = fig.add_subplot(gs[0, :-1]) ax1 = fig.add_subplot(gs[1, :-1], sharex=ax0) ax2 = fig.add_subplot(gs[:, -1]) fig.suptitle("Photoplethysmogram (PPG)", fontweight="bold") # Plot cleaned and raw PPG - ax0.set_title("Raw and Cleaned Signal") - ax0.plot(x_axis, ppg_signals["PPG_Raw"], color="#B0BEC5", label="Raw", zorder=1) - ax0.plot( - x_axis, - ppg_signals["PPG_Clean"], - color="#9C27B0", - label="Cleaned", - zorder=1, - linewidth=1.5, + ax0 = _ppg_peaks_plot( + ppg_signals["PPG_Clean"].values, + info=info, + sampling_rate=info["sampling_rate"], + raw=ppg_signals["PPG_Raw"].values, + ax=ax0, ) - # Plot peaks - ax0.scatter( - x_axis[peaks], - ppg_signals["PPG_Clean"][peaks], - color="#FFC107", - label="Peaks", - zorder=2, - ) - ax0.legend(loc="upper right") - # Plot Heart Rate ax1 = _signal_rate_plot( ppg_signals["PPG_Rate"].values, - peaks, - sampling_rate=sampling_rate, + info["PPG_Peaks"], + sampling_rate=info["sampling_rate"], title="Heart Rate", ytitle="Beats per minute (bpm)", color="#FB661C", color_mean="#FBB41C", - color_points="red", + color_points="#FF9800", ax=ax1, ) # Plot individual heart beats ax2 = ppg_segment( - ppg_signals["PPG_Clean"].values, peaks, sampling_rate, show="return", ax=ax2 + ppg_signals["PPG_Clean"].values, + info["PPG_Peaks"], + info["sampling_rate"], + show="return", + ax=ax2, ) else: @@ -131,6 +129,11 @@ def ppg_plot(ppg_signals, sampling_rate=1000, static=True): " Please install it first (`pip install plotly`).", ) from e + # X-axis + x_axis = np.linspace( + 0, len(ppg_signals) / info["sampling_rate"], len(ppg_signals) + ) + fig = make_subplots( rows=2, cols=1, @@ -156,8 +159,8 @@ def ppg_plot(ppg_signals, sampling_rate=1000, static=True): # Plot peaks fig.add_trace( go.Scatter( - x=x_axis[peaks], - y=ppg_signals["PPG_Clean"][peaks], + x=x_axis[info["PPG_Peaks"]], + y=ppg_signals["PPG_Clean"][info["PPG_Peaks"]], name="Peaks", mode="markers", marker_color="#D60574", @@ -188,10 +191,10 @@ def ppg_plot(ppg_signals, sampling_rate=1000, static=True): col=1, ) fig.update_layout(title_text="Photoplethysmogram (PPG)", height=500, width=750) - if sampling_rate is not None: + if info["sampling_rate"] is not None: fig.update_xaxes(title_text="Time (seconds)", row=1, col=1) fig.update_xaxes(title_text="Time (seconds)", row=2, col=1) - elif sampling_rate is None: + elif info["sampling_rate"] is None: fig.update_xaxes(title_text="Samples", row=1, col=1) fig.update_xaxes(title_text="Samples", row=2, col=1) return fig diff --git a/neurokit2/ppg/ppg_process.py b/neurokit2/ppg/ppg_process.py index f87fac0538..e90a5d5880 100644 --- a/neurokit2/ppg/ppg_process.py +++ b/neurokit2/ppg/ppg_process.py @@ -60,8 +60,9 @@ def ppg_process( ppg = nk.ppg_simulate(duration=10, sampling_rate=1000, heart_rate=70) signals, info = nk.ppg_process(ppg, sampling_rate=1000) + @savefig p_ppg_process1.png scale=100% - nk.ppg_plot(signals) + nk.ppg_plot(signals, info) @suppress plt.close() @@ -106,7 +107,7 @@ def ppg_process( if report is not None: # Generate report containing description and figures of processing if ".html" in str(report): - fig = ppg_plot(signals, sampling_rate=sampling_rate) + fig = ppg_plot(signals, info) else: fig = None create_report(file=report, signals=signals, info=methods, fig=fig) diff --git a/neurokit2/ppg/ppg_segment.py b/neurokit2/ppg/ppg_segment.py index 15572a4bc1..5d7c6435f8 100644 --- a/neurokit2/ppg/ppg_segment.py +++ b/neurokit2/ppg/ppg_segment.py @@ -52,7 +52,10 @@ def ppg_segment(ppg_cleaned, peaks=None, sampling_rate=1000, show=False, **kwarg peaks = peaks["PPG_Peaks"] epochs_start, epochs_end, average_hr = _ecg_segment_window( - rpeaks=peaks, sampling_rate=sampling_rate, desired_length=len(ppg_cleaned) + rpeaks=peaks, + sampling_rate=sampling_rate, + desired_length=len(ppg_cleaned), + ratio_pre=0.3, ) heartbeats = epochs_create( ppg_cleaned, @@ -69,7 +72,7 @@ def ppg_segment(ppg_cleaned, peaks=None, sampling_rate=1000, show=False, **kwarg if show is not False: ax = _ecg_segment_plot( - heartbeats, heartrate=average_hr, ytitle="PPG", color="#9C27B0", **kwargs + heartbeats, heartrate=average_hr, ytitle="PPG", color="#E91E63", **kwargs ) if show == "return": return ax diff --git a/neurokit2/rsp/rsp_plot.py b/neurokit2/rsp/rsp_plot.py index 7c7378b151..8505c53641 100644 --- a/neurokit2/rsp/rsp_plot.py +++ b/neurokit2/rsp/rsp_plot.py @@ -1,20 +1,22 @@ # -*- coding: utf-8 -*- +from warnings import warn + import matplotlib.pyplot as plt import numpy as np import pandas as pd +from ..misc import NeuroKitWarning + -def rsp_plot(rsp_signals, sampling_rate=None, figsize=(10, 10), static=True): +def rsp_plot(rsp_signals, info=None, static=True): """**Visualize respiration (RSP) data** Parameters ---------- rsp_signals : DataFrame DataFrame obtained from :func:`.rsp_process`. - sampling_rate : int - The desired sampling rate (in Hz, i.e., samples/second). - figsize : tuple - The size of the figure (width, height) in inches. + info : dict + The information Dict returned by ``rsp_process()``. Defaults to ``None``. static : bool If True, a static plot will be generated with matplotlib. If False, an interactive plot will be generated with plotly. @@ -26,8 +28,7 @@ def rsp_plot(rsp_signals, sampling_rate=None, figsize=(10, 10), static=True): Returns ------- - fig - Figure representing a plot of the processed RSP signals. + See :func:`.ecg_plot` for details on how to access the figure, modify the size and save it. Examples -------- @@ -36,18 +37,29 @@ def rsp_plot(rsp_signals, sampling_rate=None, figsize=(10, 10), static=True): import neurokit2 as nk # Simulate data - rsp = nk.rsp_simulate(duration=90, respiratory_rate=15) + rsp = nk.rsp_simulate(duration=90, respiratory_rate=15, sampling_rate=100) # Process signal - rsp_signals, info = nk.rsp_process(rsp, sampling_rate=1000) + rsp_signals, info = nk.rsp_process(rsp, sampling_rate=100) # Plot @savefig p_rsp_plot1.png scale=100% - nk.rsp_plot(rsp_signals, sampling_rate=1000) + nk.rsp_plot(rsp_signals, info) @suppress plt.close() """ + if info is None: + warn( + "'info' dict not provided. Some information might be missing." + + " Sampling rate will be set to 1000 Hz.", + category=NeuroKitWarning, + ) + + info = { + "sampling_rate": 1000, + } + # Mark peaks, troughs and phases. peaks = np.where(rsp_signals["RSP_Peaks"] == 1)[0] troughs = np.where(rsp_signals["RSP_Troughs"] == 1)[0] @@ -74,15 +86,11 @@ def rsp_plot(rsp_signals, sampling_rate=None, figsize=(10, 10), static=True): exhale_signal, inhale_signal = _rsp_plot_phase(rsp_signals, troughs, peaks) # Determine unit of x-axis. - if sampling_rate is not None: - x_label = "Time (seconds)" - x_axis = np.linspace(0, len(rsp_signals) / sampling_rate, len(rsp_signals)) - else: - x_label = "Samples" - x_axis = np.arange(0, len(rsp_signals)) + x_label = "Time (seconds)" + x_axis = np.linspace(0, len(rsp_signals) / info["sampling_rate"], len(rsp_signals)) if static: - fig, ax = plt.subplots(nrows=nrow, ncols=1, sharex=True, figsize=figsize) + fig, ax = plt.subplots(nrows=nrow, ncols=1, sharex=True) last_ax = fig.get_axes()[-1] last_ax.set_xlabel(x_label) @@ -198,7 +206,6 @@ def rsp_plot(rsp_signals, sampling_rate=None, figsize=(10, 10), static=True): linewidth=1.5, ) ax[4].legend(loc="upper right") - return fig else: # Generate interactive plot with plotly. try: diff --git a/neurokit2/rsp/rsp_process.py b/neurokit2/rsp/rsp_process.py index 2751218a14..61acd3ac58 100644 --- a/neurokit2/rsp/rsp_process.py +++ b/neurokit2/rsp/rsp_process.py @@ -88,7 +88,7 @@ def rsp_process( signals, info = nk.rsp_process(rsp, sampling_rate=1000, report="text") @savefig p_rsp_process_1.png scale=100% - fig = nk.rsp_plot(signals, sampling_rate=1000) + fig = nk.rsp_plot(signals, info) @suppress plt.close() @@ -146,7 +146,7 @@ def rsp_process( if report is not None: # Generate report containing description and figures of processing if ".html" in str(report): - fig = rsp_plot(signals, sampling_rate=sampling_rate) + fig = rsp_plot(signals, info) else: fig = None create_report(file=report, signals=signals, info=methods, fig=fig) diff --git a/neurokit2/signal/signal_fixpeaks.py b/neurokit2/signal/signal_fixpeaks.py index 664a62f4ad..cd0ae5241d 100644 --- a/neurokit2/signal/signal_fixpeaks.py +++ b/neurokit2/signal/signal_fixpeaks.py @@ -45,9 +45,9 @@ def signal_fixpeaks( show : bool Whether or not to visualize artifacts and artifact thresholds. interval_min : float - Only when ``method = "neurokit"``. The minimum interval between the peaks. + Only when ``method = "neurokit"``. The minimum interval between the peaks (in seconds). interval_max : float - Only when ``method = "neurokit"``. The maximum interval between the peaks. + Only when ``method = "neurokit"``. The maximum interval between the peaks (in seconds). relative_interval_min : float Only when ``method = "neurokit"``. The minimum interval between the peaks as relative to the sample (expressed in standard deviation from the mean). @@ -83,15 +83,17 @@ def signal_fixpeaks( import neurokit2 as nk - # Simulate ECG data - ecg = nk.ecg_simulate(duration=240, noise=0.25, heart_rate=70, random_state=42) + # Simulate ECG data and add noisy period + ecg = nk.ecg_simulate(duration=240, sampling_rate=250, noise=2, random_state=42) + ecg[20000:30000] += np.random.uniform(size=10000) + ecg[40000:43000] = 0 # Identify and Correct Peaks using "Kubios" Method - rpeaks_uncorrected = nk.ecg_findpeaks(ecg) + rpeaks_uncorrected = nk.ecg_findpeaks(ecg, method="pantompkins", sampling_rate=250) @savefig p_signal_fixpeaks1.png scale=100% - artifacts, rpeaks_corrected = nk.signal_fixpeaks( - rpeaks_uncorrected, iterative=True, method="Kubios", show=True + info, rpeaks_corrected = nk.signal_fixpeaks( + rpeaks_uncorrected, sampling_rate=250, iterative=True, method="Kubios", show=True ) @suppress plt.close() @@ -125,7 +127,7 @@ def signal_fixpeaks( peaks = np.sort(np.append(peaks, [1350, 11350, 18350])) # add artifacts # Identify and Correct Peaks using 'NeuroKit' Method - peaks_corrected = nk.signal_fixpeaks( + info, peaks_corrected = nk.signal_fixpeaks( peaks=peaks, interval_min=0.5, interval_max=1.5, method="neurokit" ) @@ -148,10 +150,12 @@ def signal_fixpeaks( # If method Kubios if method.lower() == "kubios": - return _signal_fixpeaks_kubios(peaks, sampling_rate=sampling_rate, iterative=iterative, show=show, **kwargs) + info, peaks_clean = _signal_fixpeaks_kubios( + peaks, sampling_rate=sampling_rate, iterative=iterative, show=show, **kwargs + ) else: # Else method is NeuroKit - return _signal_fixpeaks_neurokit( + info, peaks_clean = _signal_fixpeaks_neurokit( peaks, sampling_rate=sampling_rate, interval_min=interval_min, @@ -160,6 +164,7 @@ def signal_fixpeaks( relative_interval_max=relative_interval_max, robust=robust, ) + return info, peaks_clean # ============================================================================= @@ -176,21 +181,38 @@ def _signal_fixpeaks_neurokit( ): """NeuroKit method.""" - peaks_clean = _remove_small(peaks, sampling_rate, interval_min, relative_interval_min, robust) - peaks_clean = _interpolate_big(peaks_clean, sampling_rate, interval_max, relative_interval_max, robust,) + peaks_clean = _remove_small( + peaks, sampling_rate, interval_min, relative_interval_min, robust + ) + peaks_clean = _interpolate_big( + peaks_clean, + sampling_rate, + interval_max, + relative_interval_max, + robust, + ) valid_peaks = peaks_clean[peaks_clean >= 0] n_invalid_idcs = len(peaks_clean) - len(valid_peaks) if n_invalid_idcs > 0: warn( - f" Negative peak indices detected in output. " f" Removing {n_invalid_idcs} invalid peaks. ", + f" Negative peak indices detected in output. " + f" Removing {n_invalid_idcs} invalid peaks. ", category=NeuroKitWarning, ) peaks_clean = valid_peaks - return peaks_clean + + info = { + "method": "neurokit", + "extra": [i for i in peaks if i not in peaks_clean], + "missed": [i for i in peaks_clean if i not in peaks], + } + return info, peaks_clean -def _signal_fixpeaks_kubios(peaks, sampling_rate=1000, iterative=True, show=False, **kwargs): +def _signal_fixpeaks_kubios( + peaks, sampling_rate=1000, iterative=True, show=False, **kwargs +): """kubios method.""" # Get corrected peaks and normal-to-normal intervals. @@ -198,14 +220,14 @@ def _signal_fixpeaks_kubios(peaks, sampling_rate=1000, iterative=True, show=Fals peaks_clean = _correct_artifacts(artifacts, peaks) if iterative: - # Iteratively apply the artifact correction until the number # of artifacts stops decreasing. n_artifacts_current = sum([len(i) for i in artifacts.values()]) while True: - - new_artifacts, new_subspaces = _find_artifacts(peaks_clean, sampling_rate=sampling_rate, **kwargs) + new_artifacts, new_subspaces = _find_artifacts( + peaks_clean, sampling_rate=sampling_rate, **kwargs + ) n_artifacts_previous = n_artifacts_current n_artifacts_current = sum([len(i) for i in new_artifacts.values()]) @@ -214,8 +236,13 @@ def _signal_fixpeaks_kubios(peaks, sampling_rate=1000, iterative=True, show=Fals artifacts = new_artifacts subspaces = new_subspaces peaks_clean = _correct_artifacts(artifacts, peaks_clean) + + artifacts["method"] = "kubios" + artifacts.update(subspaces) + if show: - _plot_artifacts_lipponen2019(artifacts, subspaces) + _plot_artifacts_lipponen2019(artifacts) + return artifacts, peaks_clean @@ -223,9 +250,14 @@ def _signal_fixpeaks_kubios(peaks, sampling_rate=1000, iterative=True, show=Fals # Kubios: Lipponen & Tarvainen (2019). # ============================================================================= def _find_artifacts( - peaks, c1=0.13, c2=0.17, alpha=5.2, window_width=91, medfilt_order=11, sampling_rate=1000, + peaks, + c1=0.13, + c2=0.17, + alpha=5.2, + window_width=91, + medfilt_order=11, + sampling_rate=1000, ): - # Compute period series (make sure it has same numer of elements as peaks); # peaks are in samples, convert to seconds. rr = np.ediff1d(peaks, to_begin=0) / sampling_rate @@ -254,7 +286,6 @@ def _find_artifacts( s12 = np.zeros(drrs.size) for d in np.arange(padding, padding + drrs.size): - if drrs_pad[d] > 0: s12[d - padding] = np.max([drrs_pad[d - 1], drrs_pad[d + 1]]) elif drrs_pad[d] < 0: @@ -262,7 +293,6 @@ def _find_artifacts( # Cast dRRs to subspace s22. s22 = np.zeros(drrs.size) for d in np.arange(padding, padding + drrs.size): - if drrs_pad[d] >= 0: s22[d - padding] = np.min([drrs_pad[d + 1], drrs_pad[d + 2]]) elif drrs_pad[d] < 0: @@ -287,12 +317,15 @@ def _find_artifacts( i = 0 while i < rr.size - 2: # The flow control is implemented based on Figure 1 - if np.abs(drrs[i]) <= 1: # Figure 1 i += 1 continue - eq1 = np.logical_and(drrs[i] > 1, s12[i] < (-c1 * drrs[i] - c2)) # pylint: disable=E1111 - eq2 = np.logical_and(drrs[i] < -1, s12[i] > (-c1 * drrs[i] + c2)) # pylint: disable=E1111 + eq1 = np.logical_and( + drrs[i] > 1, s12[i] < (-c1 * drrs[i] - c2) + ) # pylint: disable=E1111 + eq2 = np.logical_and( + drrs[i] < -1, s12[i] > (-c1 * drrs[i] + c2) + ) # pylint: disable=E1111 if np.any([eq1, eq2]): # If any of the two equations is true. @@ -363,17 +396,23 @@ def _find_artifacts( def _compute_threshold(signal, alpha, window_width): - df = pd.DataFrame({"signal": np.abs(signal)}) - q1 = df.rolling(window_width, center=True, min_periods=1).quantile(0.25).signal.values - q3 = df.rolling(window_width, center=True, min_periods=1).quantile(0.75).signal.values + q1 = ( + df.rolling(window_width, center=True, min_periods=1) + .quantile(0.25) + .signal.values + ) + q3 = ( + df.rolling(window_width, center=True, min_periods=1) + .quantile(0.75) + .signal.values + ) th = alpha * ((q3 - q1) / 2) return th def _correct_artifacts(artifacts, peaks): - # Artifact correction ##################### # The integrity of indices must be maintained if peaks are inserted or @@ -406,7 +445,6 @@ def _correct_artifacts(artifacts, peaks): def _correct_extra(extra_idcs, peaks): - corrected_peaks = peaks.copy() corrected_peaks = np.delete(corrected_peaks, extra_idcs) @@ -414,13 +452,14 @@ def _correct_extra(extra_idcs, peaks): def _correct_missed(missed_idcs, peaks): - corrected_peaks = peaks.copy() missed_idcs = np.array(missed_idcs) # Calculate the position(s) of new beat(s). Make sure to not generate # negative indices. prev_peaks and next_peaks must have the same # number of elements. - valid_idcs = np.logical_and(missed_idcs > 1, missed_idcs < len(corrected_peaks)) # pylint: disable=E1111 + valid_idcs = np.logical_and( + missed_idcs > 1, missed_idcs < len(corrected_peaks) + ) # pylint: disable=E1111 missed_idcs = missed_idcs[valid_idcs] prev_peaks = corrected_peaks[[i - 1 for i in missed_idcs]] next_peaks = corrected_peaks[missed_idcs] @@ -432,14 +471,14 @@ def _correct_missed(missed_idcs, peaks): def _correct_misaligned(misaligned_idcs, peaks): - corrected_peaks = peaks.copy() misaligned_idcs = np.array(misaligned_idcs) # Make sure to not generate negative indices, or indices that exceed # the total number of peaks. prev_peaks and next_peaks must have the # same number of elements. valid_idcs = np.logical_and( - misaligned_idcs > 1, misaligned_idcs < len(corrected_peaks) - 1, # pylint: disable=E1111 + misaligned_idcs > 1, + misaligned_idcs < len(corrected_peaks) - 1, # pylint: disable=E1111 ) misaligned_idcs = misaligned_idcs[valid_idcs] prev_peaks = corrected_peaks[[i - 1 for i in misaligned_idcs]] @@ -465,65 +504,94 @@ def _update_indices(source_idcs, update_idcs, update): return list(np.unique(update_idcs)) -def _plot_artifacts_lipponen2019(artifacts, info): +def _plot_artifacts_lipponen2019(info): + # Covnenience function to extract relevant stuff. + def _get_which_endswith(info, string): + return [s for key, s in info.items() if key.endswith(string)][0] # Extract parameters - longshort_idcs = artifacts["longshort"] - ectopic_idcs = artifacts["ectopic"] - extra_idcs = artifacts["extra"] - missed_idcs = artifacts["missed"] - - rr = info["rr"] - drrs = info["drrs"] - mrrs = info["mrrs"] - s12 = info["s12"] - s22 = info["s22"] - c1 = info["c1"] - c2 = info["c2"] + longshort_idcs = _get_which_endswith(info, "longshort") + ectopic_idcs = _get_which_endswith(info, "ectopic") + extra_idcs = _get_which_endswith(info, "extra") + missed_idcs = _get_which_endswith(info, "missed") + + # Extract subspace info + rr = _get_which_endswith(info, "rr") + drrs = _get_which_endswith(info, "drrs") + mrrs = _get_which_endswith(info, "mrrs") + s12 = _get_which_endswith(info, "s12") + s22 = _get_which_endswith(info, "s22") + c1 = _get_which_endswith(info, "c1") + c2 = _get_which_endswith(info, "c2") # Visualize artifact type indices. # Set grids - gs = matplotlib.gridspec.GridSpec(ncols=4, nrows=3, width_ratios=[1, 2, 2, 2]) - fig = plt.figure(constrained_layout=False, figsize=(15, 10)) - ax0 = fig.add_subplot(gs[0, :-2]) - ax1 = fig.add_subplot(gs[1, :-2]) - ax2 = fig.add_subplot(gs[2, :-2]) - ax3 = fig.add_subplot(gs[:, -1]) - ax4 = fig.add_subplot(gs[:, -2]) - - ax0.set_title("Artifact types", fontweight="bold") + gs = matplotlib.gridspec.GridSpec(ncols=2, nrows=6) + fig = plt.figure(constrained_layout=False, figsize=(17, 12)) + fig.suptitle("Peak Correction", fontweight="bold") + ax0 = fig.add_subplot(gs[0:2, 0]) + ax1 = fig.add_subplot(gs[2:4, 0]) + ax2 = fig.add_subplot(gs[4:6, 0]) + ax3 = fig.add_subplot(gs[0:3:, 1]) + ax4 = fig.add_subplot(gs[3:6, 1]) + + ax0.set_title("Artifact types") ax0.plot(rr, label="heart period") ax0.scatter( - longshort_idcs, rr[longshort_idcs], marker="x", c="m", s=100, zorder=3, label="long/short", + longshort_idcs, + rr[longshort_idcs], + marker="x", + c="m", + s=100, + zorder=3, + label="long/short", ) ax0.scatter( - ectopic_idcs, rr[ectopic_idcs], marker="x", c="g", s=100, zorder=3, label="ectopic", + ectopic_idcs, + rr[ectopic_idcs], + marker="x", + c="g", + s=100, + zorder=3, + label="ectopic", ) ax0.scatter( - extra_idcs, rr[extra_idcs], marker="x", c="y", s=100, zorder=3, label="false positive", + extra_idcs, + rr[extra_idcs], + marker="x", + c="y", + s=100, + zorder=3, + label="false positive", ) ax0.scatter( - missed_idcs, rr[missed_idcs], marker="x", c="r", s=100, zorder=3, label="false negative", + missed_idcs, + rr[missed_idcs], + marker="x", + c="r", + s=100, + zorder=3, + label="false negative", ) ax0.legend(loc="upper right") # Visualize first threshold. - ax1.set_title("Consecutive-difference criterion", fontweight="bold") + ax1.set_title("Consecutive-difference criterion") ax1.plot(np.abs(drrs), label="normalized difference consecutive heart periods") ax1.axhline(1, c="r", label="artifact threshold") ax1.legend(loc="upper right") ax1.set_ylim(0, 5) # Visualize second threshold. - ax2.set_title("Difference-from-median criterion", fontweight="bold") + ax2.set_title("Difference-from-median criterion") ax2.plot(np.abs(mrrs), label="difference from median over 11 periods") ax2.axhline(3, c="r", label="artifact threshold") ax2.legend(loc="upper right") ax2.set_ylim(0, 5) # Visualize subspaces. - ax4.set_title("Subspace 1", fontweight="bold") + ax4.set_title("Subspace 1") ax4.set_xlabel("S11") ax4.set_ylabel("S12") ax4.scatter(drrs, s12, marker="x", label="heart periods") @@ -531,52 +599,73 @@ def _plot_artifacts_lipponen2019(artifacts, info): ax4.set_xlim(-10, 10) verts0 = [(-10, 5), (-10, -c1 * -10 + c2), (-1, -c1 * -1 + c2), (-1, 5)] - poly0 = matplotlib.patches.Polygon(verts0, alpha=0.3, facecolor="r", edgecolor=None, label="ectopic periods") + poly0 = matplotlib.patches.Polygon( + verts0, alpha=0.3, facecolor="r", edgecolor=None, label="ectopic periods" + ) ax4.add_patch(poly0) verts1 = [(1, -c1 * 1 - c2), (1, -5), (10, -5), (10, -c1 * 10 - c2)] poly1 = matplotlib.patches.Polygon(verts1, alpha=0.3, facecolor="r", edgecolor=None) ax4.add_patch(poly1) ax4.legend(loc="upper right") - ax3.set_title("Subspace 2", fontweight="bold") + ax3.set_title("Subspace 2") ax3.set_xlabel("S21") ax3.set_ylabel("S22") ax3.scatter(drrs, s22, marker="x", label="heart periods") ax3.set_xlim(-10, 10) ax3.set_ylim(-10, 10) verts2 = [(-10, 10), (-10, 1), (-1, 1), (-1, 10)] - poly2 = matplotlib.patches.Polygon(verts2, alpha=0.3, facecolor="r", edgecolor=None, label="short periods") + poly2 = matplotlib.patches.Polygon( + verts2, alpha=0.3, facecolor="r", edgecolor=None, label="short periods" + ) ax3.add_patch(poly2) verts3 = [(1, -1), (1, -10), (10, -10), (10, -1)] - poly3 = matplotlib.patches.Polygon(verts3, alpha=0.3, facecolor="y", edgecolor=None, label="long periods") + poly3 = matplotlib.patches.Polygon( + verts3, alpha=0.3, facecolor="y", edgecolor=None, label="long periods" + ) ax3.add_patch(poly3) ax3.legend(loc="upper right") + plt.tight_layout() # ============================================================================= # NeuroKit # ============================================================================= def _remove_small( - peaks, sampling_rate=1000, interval_min=None, relative_interval_min=None, robust=False, + peaks, + sampling_rate=1000, + interval_min=None, + relative_interval_min=None, + robust=False, ): if interval_min is None and relative_interval_min is None: return peaks if interval_min is not None: - interval = signal_period(peaks, sampling_rate=sampling_rate, desired_length=None) + interval = signal_period( + peaks, sampling_rate=sampling_rate, desired_length=None + ) peaks = peaks[interval > interval_min] if relative_interval_min is not None: - interval = signal_period(peaks, sampling_rate=sampling_rate, desired_length=None) + interval = signal_period( + peaks, sampling_rate=sampling_rate, desired_length=None + ) peaks = peaks[standardize(interval, robust=robust) > relative_interval_min] return peaks def _interpolate_big( - peaks, sampling_rate=1000, interval_max=None, relative_interval_max=None, robust=False, + peaks, + sampling_rate=1000, + interval_max=None, + relative_interval_max=None, + robust=False, ): if interval_max is None and relative_interval_max is None: return peaks else: - interval = signal_period(peaks, sampling_rate=sampling_rate, desired_length=None) + interval = signal_period( + peaks, sampling_rate=sampling_rate, desired_length=None + ) if relative_interval_max is not None: outliers = standardize(interval, robust=robust) > relative_interval_max else: @@ -605,7 +694,11 @@ def _interpolate_big( peaks_to_correct[loc] = np.nan peaks_to_correct = np.insert(peaks_to_correct, loc, [np.nan] * (n_nan - 1)) # Interpolate values - interpolated_peaks = pd.Series(peaks_to_correct).interpolate(limit_area="inside").values + interpolated_peaks = ( + pd.Series(peaks_to_correct).interpolate(limit_area="inside").values + ) # If there are missing values remaining, remove - peaks = interpolated_peaks[np.invert(np.isnan(interpolated_peaks))].astype(peaks.dtype) + peaks = interpolated_peaks[np.invert(np.isnan(interpolated_peaks))].astype( + peaks.dtype + ) return peaks diff --git a/tests/tests_ecg.py b/tests/tests_ecg.py index 43968687fd..32ae86e6a1 100644 --- a/tests/tests_ecg.py +++ b/tests/tests_ecg.py @@ -50,27 +50,33 @@ def test_ecg_clean(): def test_ecg_peaks(): - sampling_rate = 1000 - noise = 0.15 + sampling_rate = 200 + noise = 1 ecg = nk.ecg_simulate( duration=120, sampling_rate=sampling_rate, noise=noise, random_state=42 ) - ecg_cleaned_nk = nk.ecg_clean(ecg, sampling_rate=sampling_rate, method="neurokit") + ecg[3000:3600] = 0 # Test without request to correct artifacts. signals, _ = nk.ecg_peaks( - ecg_cleaned_nk, correct_artifacts=False, method="neurokit" + ecg, sampling_rate=sampling_rate, method="neurokit", correct_artifacts=False ) - assert signals.shape == (120000, 1) - assert np.allclose(signals["ECG_R_Peaks"].values.sum(dtype=np.int64), 139, atol=1) + assert signals.shape == (24000, 1) + assert np.allclose(signals["ECG_R_Peaks"].values.sum(dtype=np.int64), 137, atol=1) # Test with request to correct artifacts. - signals, _ = nk.ecg_peaks(ecg_cleaned_nk, correct_artifacts=True, method="neurokit") + signals, info = nk.ecg_peaks( + ecg, + sampling_rate=sampling_rate, + correct_artifacts=True, + method="neurokit", + ) - assert signals.shape == (120000, 1) - assert np.allclose(signals["ECG_R_Peaks"].values.sum(dtype=np.int64), 139, atol=1) + assert signals.shape == (24000, 1) + assert np.allclose(signals["ECG_R_Peaks"].values.sum(dtype=np.int64), 136, atol=1) + assert 17 in info["ECG_fixpeaks_longshort"] def test_ecg_process(): @@ -84,10 +90,10 @@ def test_ecg_process(): def test_ecg_plot(): ecg = nk.ecg_simulate(duration=60, heart_rate=70, noise=0.05, random_state=5) - ecg_summary, _ = nk.ecg_process(ecg, sampling_rate=1000, method="neurokit") + ecg_summary, info = nk.ecg_process(ecg, sampling_rate=1000, method="neurokit") # Plot data over seconds. - nk.ecg_plot(ecg_summary, sampling_rate=1000) + nk.ecg_plot(ecg_summary, info) fig = plt.gcf() # Extract the latest figure. assert len(fig.axes) == 3 assert fig.get_axes()[1].get_xlabel() == "Time (seconds)" diff --git a/tests/tests_eda.py b/tests/tests_eda.py index 490e271773..028dc43cf9 100644 --- a/tests/tests_eda.py +++ b/tests/tests_eda.py @@ -159,10 +159,10 @@ def test_eda_plot(): drift=0.01, random_state=42, ) - eda_summary, _ = nk.eda_process(eda, sampling_rate=sampling_rate) + eda_summary, info = nk.eda_process(eda, sampling_rate=sampling_rate) # Plot data over samples. - nk.eda_plot(eda_summary) + nk.eda_plot(eda_summary, info) # This will identify the latest figure. fig = plt.gcf() assert len(fig.axes) == 3 @@ -173,18 +173,12 @@ def test_eda_plot(): ] for ax, title in zip(fig.get_axes(), titles): assert ax.get_title() == title - assert fig.get_axes()[2].get_xlabel() == "Samples" + assert fig.get_axes()[2].get_xlabel() == "Time (seconds)" np.testing.assert_array_equal( fig.axes[0].get_xticks(), fig.axes[1].get_xticks(), fig.axes[2].get_xticks() ) plt.close(fig) - # Plot data over seconds. - nk.eda_plot(eda_summary, sampling_rate=sampling_rate) - # This will identify the latest figure. - fig = plt.gcf() - assert fig.get_axes()[2].get_xlabel() == "Seconds" - def test_eda_eventrelated(): eda = nk.eda_simulate(duration=15, scr_number=3) diff --git a/tests/tests_emg.py b/tests/tests_emg.py index b757c6ded6..8fd49ee26b 100644 --- a/tests/tests_emg.py +++ b/tests/tests_emg.py @@ -13,7 +13,6 @@ def test_emg_simulate(): - emg1 = nk.emg_simulate(duration=20, length=5000, burst_number=1) assert len(emg1) == 5000 @@ -30,7 +29,6 @@ def test_emg_simulate(): def test_emg_activation(): - emg = nk.emg_simulate(duration=10, burst_number=3) cleaned = nk.emg_clean(emg) emg_amplitude = nk.emg_amplitude(cleaned) @@ -44,7 +42,6 @@ def test_emg_activation(): def test_emg_clean(): - sampling_rate = 1000 emg = nk.emg_simulate(duration=20, sampling_rate=sampling_rate) @@ -66,29 +63,24 @@ def test_emg_clean(): def test_emg_plot(): - sampling_rate = 1000 emg = nk.emg_simulate(duration=10, sampling_rate=1000, burst_number=3) - emg_summary, _ = nk.emg_process(emg, sampling_rate=sampling_rate) + emg_summary, info = nk.emg_process(emg, sampling_rate=sampling_rate) # Plot data over samples. - fig = nk.emg_plot(emg_summary) + nk.emg_plot(emg_summary, info) + fig = plt.gcf() assert len(fig.axes) == 2 titles = ["Raw and Cleaned Signal", "Muscle Activation"] - for (ax, title) in zip(fig.get_axes(), titles): + for ax, title in zip(fig.get_axes(), titles): assert ax.get_title() == title - assert fig.get_axes()[1].get_xlabel() == "Samples" + assert fig.get_axes()[1].get_xlabel() == "Time (seconds)" np.testing.assert_array_equal(fig.axes[0].get_xticks(), fig.axes[1].get_xticks()) plt.close(fig) - # Plot data over time. - fig = nk.emg_plot(emg_summary, sampling_rate=sampling_rate) - assert fig.get_axes()[1].get_xlabel() == "Time (seconds)" - def test_emg_eventrelated(): - emg = nk.emg_simulate(duration=20, sampling_rate=1000, burst_number=3) emg_signals, info = nk.emg_process(emg, sampling_rate=1000) epochs = nk.epochs_create( @@ -138,7 +130,6 @@ def test_emg_eventrelated(): def test_emg_intervalrelated(): - emg = nk.emg_simulate(duration=40, sampling_rate=1000, burst_number=3) emg_signals, info = nk.emg_process(emg, sampling_rate=1000) columns = ["EMG_Activation_N", "EMG_Amplitude_Mean"] @@ -163,16 +154,18 @@ def test_emg_intervalrelated(): ) assert features_dict.shape[0] == 2 # Number of rows + @pytest.mark.parametrize( "method_cleaning, method_activation, threshold", - [("none", "threshold", "default"), - ("biosppy", "pelt", 0.5), - ("biosppy", "mixture", 0.05), - ("biosppy", "biosppy", "default"), - ("biosppy", "silva", "default")], + [ + ("none", "threshold", "default"), + ("biosppy", "pelt", 0.5), + ("biosppy", "mixture", 0.05), + ("biosppy", "biosppy", "default"), + ("biosppy", "silva", "default"), + ], ) def test_emg_report(tmp_path, method_cleaning, method_activation, threshold): - sampling_rate = 250 emg = nk.emg_simulate( @@ -191,7 +184,7 @@ def test_emg_report(tmp_path, method_cleaning, method_activation, threshold): report=str(p), method_cleaning=method_cleaning, method_activation=method_activation, - threshold=threshold + threshold=threshold, ) assert p.is_file() assert "EMG_Activity" in signals.columns diff --git a/tests/tests_eog.py b/tests/tests_eog.py index d052d182eb..a4c7e7bc49 100644 --- a/tests/tests_eog.py +++ b/tests/tests_eog.py @@ -8,7 +8,6 @@ def test_eog_clean(): - # test with exported csv eog_signal = nk.data("eog_200hz")["vEOG"] eog_cleaned = nk.eog_clean(eog_signal, sampling_rate=200) @@ -16,7 +15,8 @@ def test_eog_clean(): # test with mne.io.Raw raw = mne.io.read_raw_fif( - str(mne.datasets.sample.data_path()) + "/MEG/sample/sample_audvis_raw.fif", preload=True + str(mne.datasets.sample.data_path()) + "/MEG/sample/sample_audvis_raw.fif", + preload=True, ) sampling_rate = raw.info["sfreq"] @@ -43,7 +43,6 @@ def test_eog_clean(): def test_eog_findpeaks(): - eog_signal = nk.data("eog_100hz") eog_cleaned = nk.eog_clean(eog_signal, sampling_rate=100) @@ -66,7 +65,6 @@ def test_eog_findpeaks(): def test_eog_process(): - eog_signal = nk.data("eog_200hz")["vEOG"] signals, info = nk.eog_process(eog_signal, sampling_rate=200) @@ -76,12 +74,11 @@ def test_eog_process(): def test_eog_plot(): - eog_signal = nk.data("eog_100hz") signals, info = nk.eog_process(eog_signal, sampling_rate=100) # Plot - nk.eog_plot(signals, peaks=info, sampling_rate=100) + nk.eog_plot(signals, info) fig = plt.gcf() assert len(fig.axes) == 3 @@ -89,7 +86,7 @@ def test_eog_plot(): legends = [["Raw", "Cleaned", "Blinks"], ["Rate", "Mean"], ["Median"]] ylabels = ["Amplitude (mV)", "Blinks per minute"] - for (ax, title, legend, ylabel) in zip(fig.get_axes(), titles, legends, ylabels): + for ax, title, legend, ylabel in zip(fig.get_axes(), titles, legends, ylabels): assert ax.get_title() == title subplot = ax.get_legend_handles_labels() assert subplot[1] == legend @@ -104,7 +101,6 @@ def test_eog_plot(): def test_eog_eventrelated(): - eog = nk.data("eog_200hz")["vEOG"] eog_signals, info = nk.eog_process(eog, sampling_rate=200) epochs = nk.epochs_create( @@ -114,24 +110,32 @@ def test_eog_eventrelated(): # Test rate features assert np.alltrue( - np.array(eog_eventrelated["EOG_Rate_Min"]) < np.array(eog_eventrelated["EOG_Rate_Mean"]) + np.array(eog_eventrelated["EOG_Rate_Min"]) + < np.array(eog_eventrelated["EOG_Rate_Mean"]) ) assert np.alltrue( - np.array(eog_eventrelated["EOG_Rate_Mean"]) < np.array(eog_eventrelated["EOG_Rate_Max"]) + np.array(eog_eventrelated["EOG_Rate_Mean"]) + < np.array(eog_eventrelated["EOG_Rate_Max"]) ) # Test blink presence - assert np.alltrue(np.array(eog_eventrelated["EOG_Blinks_Presence"]) == np.array([1, 0, 0])) + assert np.alltrue( + np.array(eog_eventrelated["EOG_Blinks_Presence"]) == np.array([1, 0, 0]) + ) # Test warning on missing columns - with pytest.warns(nk.misc.NeuroKitWarning, match=r".*does not have an `EOG_Blinks`.*"): + with pytest.warns( + nk.misc.NeuroKitWarning, match=r".*does not have an `EOG_Blinks`.*" + ): first_epoch_key = list(epochs.keys())[0] first_epoch_copy = epochs[first_epoch_key].copy() del first_epoch_copy["EOG_Blinks"] nk.eog_eventrelated({**epochs, first_epoch_key: first_epoch_copy}) - with pytest.warns(nk.misc.NeuroKitWarning, match=r".*does not have an `EOG_Rate`.*"): + with pytest.warns( + nk.misc.NeuroKitWarning, match=r".*does not have an `EOG_Rate`.*" + ): first_epoch_key = list(epochs.keys())[0] first_epoch_copy = epochs[first_epoch_key].copy() del first_epoch_copy["EOG_Rate"] @@ -139,7 +143,6 @@ def test_eog_eventrelated(): def test_eog_intervalrelated(): - eog = nk.data("eog_200hz")["vEOG"] eog_signals, info = nk.eog_process(eog, sampling_rate=200) @@ -158,5 +161,7 @@ def test_eog_intervalrelated(): ) epochs_dict = nk.eog_intervalrelated(epochs) - assert all(elem in columns for elem in np.array(epochs_dict.columns.values, dtype=str)) + assert all( + elem in columns for elem in np.array(epochs_dict.columns.values, dtype=str) + ) assert epochs_dict.shape[0] == len(epochs) # Number of rows diff --git a/tests/tests_rsp.py b/tests/tests_rsp.py index 1175c70390..3a7f13378d 100644 --- a/tests/tests_rsp.py +++ b/tests/tests_rsp.py @@ -243,8 +243,8 @@ def test_rsp_plot(): rsp = nk.rsp_simulate( duration=120, sampling_rate=1000, respiratory_rate=15, random_state=3 ) - rsp_summary, _ = nk.rsp_process(rsp, sampling_rate=1000) - nk.rsp_plot(rsp_summary) + rsp_summary, info = nk.rsp_process(rsp, sampling_rate=1000) + nk.rsp_plot(rsp_summary, info) # This will identify the latest figure. fig = plt.gcf() assert len(fig.axes) == 5 diff --git a/tests/tests_signal_fixpeaks.py b/tests/tests_signal_fixpeaks.py index 791136d70f..3e8c1658ab 100644 --- a/tests/tests_signal_fixpeaks.py +++ b/tests/tests_signal_fixpeaks.py @@ -3,15 +3,19 @@ import numpy as np import numpy.random import pytest + import neurokit2 as nk -from neurokit2.signal.signal_fixpeaks import _correct_artifacts, _find_artifacts, signal_fixpeaks +from neurokit2.signal.signal_fixpeaks import ( + _correct_artifacts, + _find_artifacts, + signal_fixpeaks, +) def compute_rmssd(peaks): - rr = np.ediff1d(peaks, to_begin=0) rr[0] = np.mean(rr[1:]) - rmssd = np.sqrt(np.mean(rr ** 2)) + rmssd = np.sqrt(np.mean(rr**2)) return rmssd @@ -50,7 +54,6 @@ def peaks_correct(n_peaks): @pytest.fixture def peaks_misaligned(request, peaks_correct, artifact_idcs): - rmssd = compute_rmssd(peaks_correct) displacement = request.param * rmssd @@ -62,7 +65,6 @@ def peaks_misaligned(request, peaks_correct, artifact_idcs): @pytest.fixture def peaks_missed(peaks_correct, artifact_idcs): - peaks_missed = peaks_correct.copy() peaks_missed = np.delete(peaks_missed, artifact_idcs) @@ -71,8 +73,9 @@ def peaks_missed(peaks_correct, artifact_idcs): @pytest.fixture def peaks_extra(peaks_correct, artifact_idcs): - - extra_peaks = (peaks_correct[artifact_idcs + 1] - peaks_correct[artifact_idcs]) / 15 + peaks_correct[artifact_idcs] + extra_peaks = ( + peaks_correct[artifact_idcs + 1] - peaks_correct[artifact_idcs] + ) / 15 + peaks_correct[artifact_idcs] peaks_extra = peaks_correct.copy() peaks_extra = np.insert(peaks_extra, artifact_idcs, extra_peaks) @@ -82,7 +85,12 @@ def peaks_extra(peaks_correct, artifact_idcs): @pytest.fixture def artifacts_misaligned(artifact_idcs): - artifacts = {"ectopic": list(artifact_idcs + 1), "missed": [], "extra": [], "longshort": list(artifact_idcs)} + artifacts = { + "ectopic": list(artifact_idcs + 1), + "missed": [], + "extra": [], + "longshort": list(artifact_idcs), + } return artifacts @@ -106,26 +114,22 @@ def artifacts_extra(artifact_idcs): @pytest.mark.parametrize("peaks_misaligned", [2, 4, 8], indirect=["peaks_misaligned"]) def test_misaligned_detection(peaks_misaligned, artifacts_misaligned): - artifacts, _ = _find_artifacts(peaks_misaligned, sampling_rate=1) assert artifacts == artifacts_misaligned # check for identical key-value pairs def test_missed_detection(peaks_missed, artifacts_missed): - artifacts, _ = _find_artifacts(peaks_missed, sampling_rate=1) assert artifacts == artifacts_missed def test_extra_detection(peaks_extra, artifacts_extra): - artifacts, _ = _find_artifacts(peaks_extra, sampling_rate=1) assert artifacts == artifacts_extra @pytest.mark.parametrize("peaks_misaligned", [2, 4, 8], indirect=["peaks_misaligned"]) def test_misaligned_correction(peaks_misaligned, artifacts_misaligned): - peaks_corrected = _correct_artifacts(artifacts_misaligned, peaks_misaligned) assert ( @@ -134,17 +138,19 @@ def test_misaligned_correction(peaks_misaligned, artifacts_misaligned): def test_missed_correction(peaks_missed, artifacts_missed): - peaks_corrected = _correct_artifacts(artifacts_missed, peaks_missed) - assert np.unique(peaks_corrected).size == (peaks_missed.size + len(artifacts_missed["missed"])) + assert np.unique(peaks_corrected).size == ( + peaks_missed.size + len(artifacts_missed["missed"]) + ) def test_extra_correction(peaks_extra, artifacts_extra): - peaks_corrected = _correct_artifacts(artifacts_extra, peaks_extra) - assert np.unique(peaks_corrected).size == (peaks_extra.size - len(artifacts_extra["extra"])) + assert np.unique(peaks_corrected).size == ( + peaks_extra.size - len(artifacts_extra["extra"]) + ) def idfn(val): @@ -154,13 +160,23 @@ def idfn(val): @pytest.mark.parametrize( "peaks_misaligned, iterative, rmssd_diff", - [(2, True, 27), (2, False, 27), (4, True, 113), (4, False, 113), (8, True, 444), (8, False, 444)], + [ + (2, True, 27), + (2, False, 27), + (4, True, 113), + (4, False, 113), + (8, True, 444), + (8, False, 444), + ], indirect=["peaks_misaligned"], ids=idfn, ) -def test_misaligned_correction_wrapper(peaks_correct, peaks_misaligned, iterative, rmssd_diff): - - _, peaks_corrected = signal_fixpeaks(peaks_misaligned, sampling_rate=1, iterative=iterative) +def test_misaligned_correction_wrapper( + peaks_correct, peaks_misaligned, iterative, rmssd_diff +): + _, peaks_corrected = signal_fixpeaks( + peaks_misaligned, sampling_rate=1, iterative=iterative + ) rmssd_correct = compute_rmssd(peaks_correct) rmssd_corrected = compute_rmssd(peaks_corrected) @@ -181,8 +197,9 @@ def test_misaligned_correction_wrapper(peaks_correct, peaks_misaligned, iterativ @pytest.mark.parametrize("iterative, rmssd_diff", [(True, 3), (False, 3)], ids=idfn) def test_extra_correction_wrapper(peaks_correct, peaks_extra, iterative, rmssd_diff): - - _, peaks_corrected = signal_fixpeaks(peaks_extra, sampling_rate=1, iterative=iterative) + _, peaks_corrected = signal_fixpeaks( + peaks_extra, sampling_rate=1, iterative=iterative + ) rmssd_correct = compute_rmssd(peaks_correct) rmssd_corrected = compute_rmssd(peaks_corrected) @@ -204,8 +221,9 @@ def test_extra_correction_wrapper(peaks_correct, peaks_extra, iterative, rmssd_d @pytest.mark.parametrize("iterative, rmssd_diff", [(True, 13), (False, 13)], ids=idfn) def test_missed_correction_wrapper(peaks_correct, peaks_missed, iterative, rmssd_diff): - - _, peaks_corrected = signal_fixpeaks(peaks_missed, sampling_rate=1, iterative=iterative) + _, peaks_corrected = signal_fixpeaks( + peaks_missed, sampling_rate=1, iterative=iterative + ) rmssd_correct = compute_rmssd(peaks_correct) rmssd_corrected = compute_rmssd(peaks_corrected) @@ -229,31 +247,46 @@ def test_missed_correction_wrapper(peaks_correct, peaks_missed, iterative, rmssd def testpeaks_for_neurokit_method(): signal = nk.signal_simulate(duration=20, sampling_rate=1000, frequency=1) peaks_true = nk.signal_findpeaks(signal)["Peaks"] - peaks = np.delete(peaks_true, [5,6,7,8,9,10,15,16,17,19]) # create gaps + peaks = np.delete(peaks_true, [5, 6, 7, 8, 9, 10, 15, 16, 17, 19]) # create gaps # (I added more than in the example in the function docstring) peaks = np.sort(np.append(peaks, [1350, 11350, 18350])) # add artifacts return peaks @pytest.mark.parametrize("interval_max", [None, 1.5, 2.0]) -def test_neurokit_method_returns_only_positive_indices(testpeaks_for_neurokit_method, interval_max): - peaks_corrected = nk.signal_fixpeaks(peaks=testpeaks_for_neurokit_method, interval_min=0.5, - interval_max=interval_max, - method="neurokit") +def test_neurokit_method_returns_only_positive_indices( + testpeaks_for_neurokit_method, interval_max +): + _, peaks_corrected = nk.signal_fixpeaks( + peaks=testpeaks_for_neurokit_method, + interval_min=0.5, + interval_max=interval_max, + method="neurokit", + ) assert np.all(peaks_corrected >= 0) @pytest.mark.parametrize("interval_max", [None, 1.5, 2.0]) -def test_neurokit_method_returns_no_duplicates(testpeaks_for_neurokit_method, interval_max): - peaks_corrected = nk.signal_fixpeaks(peaks=testpeaks_for_neurokit_method, interval_min=0.5, - interval_max=interval_max, - method="neurokit") +def test_neurokit_method_returns_no_duplicates( + testpeaks_for_neurokit_method, interval_max +): + _, peaks_corrected = nk.signal_fixpeaks( + peaks=testpeaks_for_neurokit_method, + interval_min=0.5, + interval_max=interval_max, + method="neurokit", + ) assert np.unique(peaks_corrected).size == peaks_corrected.size @pytest.mark.parametrize("interval_max", [None, 1.5, 2.0]) -def test_neurokit_method_returns_strictly_increasing_indices(testpeaks_for_neurokit_method, interval_max): - peaks_corrected = nk.signal_fixpeaks(peaks=testpeaks_for_neurokit_method, interval_min=0.5, - interval_max=interval_max, - method="neurokit") - assert np.all(np.diff(peaks_corrected) > 0) \ No newline at end of file +def test_neurokit_method_returns_strictly_increasing_indices( + testpeaks_for_neurokit_method, interval_max +): + _, peaks_corrected = nk.signal_fixpeaks( + peaks=testpeaks_for_neurokit_method, + interval_min=0.5, + interval_max=interval_max, + method="neurokit", + ) + assert np.all(np.diff(peaks_corrected) > 0)