diff --git a/neurokit2/ppg/ppg_intervalrelated.py b/neurokit2/ppg/ppg_intervalrelated.py index 1bb0221894..7f91483140 100644 --- a/neurokit2/ppg/ppg_intervalrelated.py +++ b/neurokit2/ppg/ppg_intervalrelated.py @@ -61,7 +61,7 @@ def ppg_intervalrelated(data, sampling_rate=1000): rate_cols = [col for col in data.columns if "PPG_Rate" in col] if len(rate_cols) == 1: intervals.update(_ppg_intervalrelated_formatinput(data)) - intervals.update(_ppg_intervalrelated_hrv(data, sampling_rate)) + intervals.update(*_ppg_intervalrelated_hrv(data, sampling_rate)) else: raise ValueError( "NeuroKit error: ppg_intervalrelated(): Wrong input," @@ -71,20 +71,18 @@ def ppg_intervalrelated(data, sampling_rate=1000): ppg_intervals = pd.DataFrame.from_dict(intervals, orient="index").T elif isinstance(data, dict): - for index in data: - intervals[index] = {} # Initialize empty container + for epoch in data: + intervals[epoch] = {} # Initialize empty container # Add label info - intervals[index]["Label"] = data[index]["Label"].iloc[0] + intervals[epoch]["Label"] = data[epoch]["Label"].iloc[0] # Rate - intervals[index] = _ppg_intervalrelated_formatinput( - data[index], intervals[index] - ) + intervals[epoch].update(_ppg_intervalrelated_formatinput(data[epoch])) # HRV - intervals[index] = _ppg_intervalrelated_hrv( - data[index], sampling_rate, intervals[index] + intervals[epoch].update( + *_ppg_intervalrelated_hrv(data[epoch], sampling_rate) ) ppg_intervals = pd.DataFrame.from_dict(intervals, orient="index") @@ -97,27 +95,28 @@ def ppg_intervalrelated(data, sampling_rate=1000): # ============================================================================= -def _ppg_intervalrelated_formatinput(data, output={}): - +def _ppg_intervalrelated_formatinput(data): # Sanitize input colnames = data.columns.values - if len([i for i in colnames if "PPG_Rate" in i]) == 0: + + if "PPG_Rate" not in colnames: raise ValueError( "NeuroKit error: ppg_intervalrelated(): Wrong input," "we couldn't extract heart rate. Please make sure" "your DataFrame contains a `PPG_Rate` column." ) signal = data["PPG_Rate"].values - output["PPG_Rate_Mean"] = np.mean(signal) - return output + PPG_Rate_Mean = np.mean(signal) + return {"PPG_Rate_Mean": PPG_Rate_Mean} -def _ppg_intervalrelated_hrv(data, sampling_rate, output={}): +def _ppg_intervalrelated_hrv(data, sampling_rate): # Sanitize input colnames = data.columns.values - if len([i for i in colnames if "PPG_Peaks" in i]) == 0: + + if "PPG_Peaks" not in colnames: raise ValueError( "NeuroKit error: ppg_intervalrelated(): Wrong input," "we couldn't extract peaks. Please make sure" @@ -129,7 +128,5 @@ def _ppg_intervalrelated_hrv(data, sampling_rate, output={}): peaks = {"PPG_Peaks": peaks} results = hrv(peaks, sampling_rate=sampling_rate) - for column in results.columns: - output[column] = results[column].values.astype("float") - return output + return results.astype("float").to_dict("records") diff --git a/tests/tests_ppg.py b/tests/tests_ppg.py index 01c1c6f974..7f43e32015 100644 --- a/tests/tests_ppg.py +++ b/tests/tests_ppg.py @@ -18,9 +18,10 @@ params_combis = list(itertools.product(*params)) -@pytest.mark.parametrize("duration, sampling_rate, heart_rate, freq_modulation", params_combis) +@pytest.mark.parametrize( + "duration, sampling_rate, heart_rate, freq_modulation", params_combis +) def test_ppg_simulate(duration, sampling_rate, heart_rate, freq_modulation): - ppg = nk.ppg_simulate( duration=duration, sampling_rate=sampling_rate, @@ -42,7 +43,6 @@ def test_ppg_simulate(duration, sampling_rate, heart_rate, freq_modulation): signals, _ = nk.ppg_process(ppg, sampling_rate=sampling_rate) if sampling_rate > 25: assert np.allclose(signals["PPG_Rate"].mean(), heart_rate, atol=1) - # Ensure that the heart rate fluctuates in the requested range. groundtruth_range = freq_modulation * heart_rate observed_range = np.percentile(signals["PPG_Rate"], 90) - np.percentile(signals["PPG_Rate"], 10) @@ -56,7 +56,6 @@ def test_ppg_simulate(duration, sampling_rate, heart_rate, freq_modulation): [(0.1, 3), (0.2, 5), (0.3, 8), (0.4, 11), (0.5, 14), (0.6, 19)], ) def test_ppg_simulate_ibi(ibi_randomness, std_heart_rate): - ppg = nk.ppg_simulate( duration=20, sampling_rate=50, @@ -109,7 +108,6 @@ def test_ppg_simulate_legacy_rng(): def test_ppg_clean(): - sampling_rate = 500 ppg = nk.ppg_simulate( @@ -126,7 +124,9 @@ def test_ppg_clean(): random_state=42, show=False, ) - ppg_cleaned_elgendi = nk.ppg_clean(ppg, sampling_rate=sampling_rate, method="elgendi") + ppg_cleaned_elgendi = nk.ppg_clean( + ppg, sampling_rate=sampling_rate, method="elgendi" + ) assert ppg.size == ppg_cleaned_elgendi.size @@ -141,7 +141,6 @@ def test_ppg_clean(): def test_ppg_findpeaks(): - sampling_rate = 500 # Test Elgendi method @@ -159,9 +158,13 @@ def test_ppg_findpeaks(): random_state=42, show=True, ) - ppg_cleaned_elgendi = nk.ppg_clean(ppg, sampling_rate=sampling_rate, method="elgendi") + ppg_cleaned_elgendi = nk.ppg_clean( + ppg, sampling_rate=sampling_rate, method="elgendi" + ) - info_elgendi = nk.ppg_findpeaks(ppg_cleaned_elgendi, sampling_rate=sampling_rate, show=True) + info_elgendi = nk.ppg_findpeaks( + ppg_cleaned_elgendi, sampling_rate=sampling_rate, show=True + ) peaks = info_elgendi["PPG_Peaks"] @@ -169,7 +172,9 @@ def test_ppg_findpeaks(): assert np.abs(peaks.sum() - 219764) < 5 # off by no more than 5 samples in total # Test MSPTD method - info_msptd = nk.ppg_findpeaks(ppg, sampling_rate=sampling_rate, method="bishop", show=True) + info_msptd = nk.ppg_findpeaks( + ppg, sampling_rate=sampling_rate, method="bishop", show=True + ) peaks = info_msptd["PPG_Peaks"] @@ -182,7 +187,6 @@ def test_ppg_findpeaks(): [("elgendi", "elgendi"), ("nabian2018", "elgendi"), ("elgendi", "bishop")], ) def test_ppg_report(tmp_path, method_cleaning, method_peaks): - sampling_rate = 100 ppg = nk.ppg_simulate( @@ -212,3 +216,32 @@ def test_ppg_report(tmp_path, method_cleaning, method_peaks): method_peaks=method_peaks, ) assert p.is_file() + + +def test_ppg_intervalrelated(): + sampling_rate = 100 + + ppg = nk.ppg_simulate( + duration=500, + sampling_rate=sampling_rate, + heart_rate=70, + frequency_modulation=0.025, + ibi_randomness=0.15, + drift=0.5, + motion_amplitude=0.25, + powerline_amplitude=0.25, + burst_amplitude=0.5, + burst_number=3, + random_state=0, + show=True, + ) + # Process the data + df, info = nk.ppg_process(ppg, sampling_rate=sampling_rate) + epochs = nk.epochs_create( + df, events=[0, 15000], sampling_rate=sampling_rate, epochs_end=150 + ) + epochs_ppg_intervals = nk.ppg_intervalrelated(epochs) + assert "PPG_Rate_Mean" in epochs_ppg_intervals.columns + + ppg_intervals = nk.ppg_intervalrelated(df) + assert "PPG_Rate_Mean" in ppg_intervals.columns