cbrnr / sleepecg

Sleep stage detection using ECG
BSD 3-Clause "New" or "Revised" License
96 stars 25 forks source link

Add more tests for detect_heartbeats #88

Closed raphaelvallat closed 2 years ago

raphaelvallat commented 2 years ago

Hi @cbrnr and @hofaflo,

Really awesome work on this package! I have been using the detect_heartbeats function as my default QRS detector in many projects, and I love it. It is super fast and robust.

This issue is related to https://github.com/cbrnr/sleepecg/issues/87. I think that the unit tests in test_heartbeats.py could be more comprehensive. Specifically, the following scenarios should be tested:

I have included below some quick-and-dirty code to check these three scenarios, using SciPy's ECG dataset. Let me know if you think this would be a valuable addition.

import mne
import sleepecg
import numpy as np
from scipy.misc import electrocardiogram

ecg = electrocardiogram()
sf = 360

# The ground-truth
pks_360hz = sleepecg.detect_heartbeats(ecg, sf)

def f1_score(y_pred, y_true):
    tp, fp, fn = sleepecg.compare_heartbeats(y_pred, y_true, max_distance=5)
    tp, fp, fn = len(tp), len(fp), len(fn)
    f1 = tp / (tp + 0.5 * (fp + fn))
    return f1

Resampling

for new_sf in [72, 90, 100, 120, 180, 256, 360, 720, 1080, 1800, 3600, 7200]:
    # Resample
    res_factor = new_sf / sf
    ecg_res = mne.filter.resample(ecg, up=res_factor)
    pks = sleepecg.detect_heartbeats(ecg_res, new_sf)
    # Compare to the original SF
    pks = np.round(pks / res_factor).astype(int)
    f1 = f1_score(pks, pks_360hz)
    print(f"{new_sf} Hz (x{res_factor:.2f}) » F1 = {f1:.3f}")

The F1-score is 0.96 at 72 Hz, 0.97 at 120 Hz and 0.99 at 256 Hz and higher frequencies.

Rescaling

for scale in [1e-5, 1e-4, 1e-3, 1e-2, 0.1, 1, 10, 1e2, 1e3, 1e4, 1e5]:
    pks = sleepecg.detect_heartbeats(ecg * scale, sf)
    # Compare to the original SF
    f1 = f1_score(pks, pks_360hz)
    print(f"{scale} » F1 = {f1:.3f}")

The F1-score is always 1, so the algorithm is indeed unit-independent.

Padding with flat data

ecg_pad_after = np.pad(ecg, pad_width=(0, ecg.size), mode="edge")
pks = sleepecg.detect_heartbeats(ecg_pad_after, sf)
f1_score(pks, pks_360hz)  # F1 = 1

ecg_pad_before = np.pad(ecg, pad_width=(ecg.size, 0), mode="edge")
pks = sleepecg.detect_heartbeats(ecg_pad_before, sf)
f1_score(pks, pks_360hz + ecg.size)  # F1 = 0.36

Cheers, Raphael

cbrnr commented 2 years ago

Hi @raphaelvallat! I think these are very interesting additional properties of the detector, and we should definitely include this somewhere. I'm just not sure if a unit test is the right place. It seems like if these are properties of the algorithm, there's no point in testing that in our test suite, because we will not be able to change that anyway. For example, it seems intuitive that the lower the sampling frequency, the lower the detector performance will get. Is this important information? Yes, definitely? Should we include tests for that? I'm not sure. My feeling is that this would very nicely fit into our docs. WDYT?

raphaelvallat commented 2 years ago

Hi @cbrnr,

Absolutely, I think it would be great to have this explicit in the documentation, as I'm sure other people will be looking for these properties. As for unit testing, I guess that it only makes sense if at some point in the future you anticipate that there will be changes in the algorithm. Then, unit tests will ensure that the updates are not breaking these fundamental properties of the algorithm. To be conservative, I would probably add these tests both in the documentation and in the unit tests I think.

Thanks! Raphael

cbrnr commented 2 years ago

OK then, let's add it to the docs and tests. However, I wouldn't add the padding issue to the tests. We know that the algorithm cannot deal with flat segments, and this will hopefully be addressed in the future. So it is a known issue, and there's no use in testing for a known limitation that will get resolved.

cbrnr commented 2 years ago

Fixed by #89.