Add documentation and heart rate computing function.

Dubrzr · Dubrzr · commit 84f35795714c · 2017-06-26T15:57:24.000Z
diff --git a/README.rst b/README.rst
@@ -579,7 +579,7 @@ Example Usage:
     tf = 20000
     sig, fields = wfdb.srdsamp('sampledata/100', sampfrom=t0, sampto=tf, channels=[0])
     record = wfdb.rdsamp("sampledata/100", sampfrom=t0, sampto=tf, channels=[0], physical=False)
-    peak_indexes = wfdb.processing.gqrs_detect(x=sig[:,0], freq=fields['fs'], gain=record.adcgain[0], adczero=record.adczero[0], threshold=1.0)
+    peaks_indexes = wfdb.processing.gqrs_detect(x=sig[:,0], frequency=fields['fs'], adcgain=record.adcgain[0], adczero=record.adczero[0], threshold=1.0)
 
 Input arguments:
 
@@ -599,6 +599,10 @@ Input arguments:
 - ``QRSa`` (default=750): Typical QRS peak-to-peak amplitude, in microvolts.
 - ``QRSamin`` (default=130): Minimum QRS peak-to-peak amplitude, in microvolts.
 
+Output Arguments:
+
+- ``peaks_indexes``: A python list containing the peaks indexes.
+
 
 **correct_peaks** - A post-processing algorithm to correct peaks position.
 
@@ -624,7 +628,7 @@ Example Usage:
     max_bpm = 350
     min_gap = fs*60/min_bpm
     max_gap = fs*60/max_bpm
-    y_idxs = wfdb.processing.correct_peaks(x=sig[:,0], peak_indexes=peak_indexes, min_gap=min_gap, max_gap=max_gap, smooth_window=150)
+    new_indexes = wfdb.processing.correct_peaks(x=sig[:,0], peak_indexes=peak_indexes, min_gap=min_gap, max_gap=max_gap, smooth_window=150)
 
 Input arguments:
 peaks_indexes, min_gap, max_gap, smooth_window
@@ -634,6 +638,45 @@ peaks_indexes, min_gap, max_gap, smooth_window
 - ``max_gap`` (required): The maximum gap in samples between two peaks.
 - ``smooth_window`` (required): The size of the smoothing window.
 
+Output Arguments:
+
+- ``new_indexes``: A python list containing the new peaks indexes.
+
+
+Heart rate
+~~~~~~~~~~~~~~
+
+**compute_hr** - Compute heart rate from peak indexes and signal frequency.
+
+::
+
+  compute_hr(length, peaks_indexes, fs)
+
+Example Usage:
+
+::
+
+    import wfdb
+    t0 = 10000
+    tf = 20000
+    sig, fields = wfdb.srdsamp('sampledata/100', sampfrom=t0, sampto=tf, channels=[0])
+    record = wfdb.rdsamp("sampledata/100", sampfrom=t0, sampto=tf, channels=[0], physical=False)
+    peaks_indexes = wfdb.processing.gqrs_detect(x=sig[:,0], frequency=fields['fs'], adcgain=record.adcgain[0], adczero=record.adczero[0], threshold=1.0)
+    hr = compute_hr(length=tf-t0, peaks_indexes=peaks_indexes, fs=fields['fs'])
+
+Input arguments:
+
+- ``length`` (required): The length of the corresponding signal.
+- ``peaks_indexes`` (required): The peak indexes.
+- ``fs`` (required): The signal frequency.
+
+
+Output Arguments:
+
+- ``hr``: A numpy.array containing heart rate for each sample. Contains numpy.nan where heart rate could not be computed.
+
+
+
 Based on the original WFDB software package specifications
 ----------------------------------------------------------
 
diff --git a/tests/test_processing.py b/tests/test_processing.py
@@ -61,7 +61,7 @@ def test_6(self):
         peaks = wfdb.processing.gqrs_detect(x, frequency, adcgain, adczero, threshold=1.0)
         print(peaks)
         print(expecting)
-        assert peaks == expecting
+        assert numpy.array_equal(peaks, expecting)
 
     def test_7(self):
         sig, fields = wfdb.srdsamp('sampledata/100', channels = [0, 1])
diff --git a/wfdb/processing/__init__.py b/wfdb/processing/__init__.py
@@ -1,3 +1,4 @@
 from .basic import resample_ann, resample_sig, resample_singlechan, resample_multichan, normalize
 from .gqrs import gqrs_detect
+from .hr import compute_hr
 from .peaks import find_peaks, correct_peaks
diff --git a/wfdb/processing/hr.py b/wfdb/processing/hr.py
@@ -0,0 +1,21 @@
+import numpy
+
+
+def compute_hr(length, peaks_indexes, fs):
+    result = numpy.full(length, numpy.nan, dtype='float32')
+
+    if len(peaks_indexes) < 2:
+        return result
+
+    current_hr = numpy.nan
+
+    for i in range(0, len(peaks_indexes)-2):
+        a = peaks_indexes[i]
+        b = peaks_indexes[i+1]
+        c = peaks_indexes[i+2]
+        RR = (b-a) * (1.0 / fs) * 1000
+        hr = 60000.0 / RR
+        result[b+1:c+1] = hr
+    result[peaks_indexes[-1]:] = result[peaks_indexes[-1]]
+
+    return result
diff --git a/wfdb/processing/peaks.py b/wfdb/processing/peaks.py
@@ -82,11 +82,9 @@ def correct_peaks(x, peaks_indexes, min_gap, max_gap, smooth_window):
     rpeaks = numpy.zeros(N)
     rpeaks[peaks_indexes] = 1.0
 
-    # Multiple ones '1' side by side:
-    # in order to prevent that, the following code computes the best peak
     rpeaks = rpeaks.astype('int32')
 
-    # 1- Extract ranges where we have one or many ones side by side (rpeaks locations predicted by NN)
+    # 1- Extract ranges where we have one or many ones side by side
     rpeaks_ranges = []
     tmp_idx = 0
     for i in range(1, len(rpeaks)):
@@ -144,4 +142,4 @@ def correct_peaks(x, peaks_indexes, min_gap, max_gap, smooth_window):
     for v, _ in to_remove.items():
         rpeaks_indexes.remove(v)
 
-    return rpeaks_indexes
+    return sorted(rpeaks_indexes)
