work on evaluator

cx1111 · cx1111 · commit b7f69bcb493a · 2018-01-19T22:05:17.000-05:00
diff --git a/wfdb/processing/__init__.py b/wfdb/processing/__init__.py
@@ -3,7 +3,9 @@
 """
 from .basic import (resample_ann, resample_sig, resample_singlechan,
     resample_multichan, normalize)
+from .evaluate import 
 from .gqrs import gqrs_detect
 from .hr import compute_hr
 from .peaks import find_peaks, correct_peaks
 from .qrs import Conf, XQRS, xqrs_detect
+
diff --git a/wfdb/processing/evaluate.py b/wfdb/processing/evaluate.py
@@ -1,8 +1,8 @@
 import numpy as np
-
-class Comparitor(object):
+import matplotlib.pyplot as plt
 
 
+class Comparitor(object):
 
     def __init__(self, ref_sample, test_sample, window_width):
         """
@@ -19,65 +19,103 @@ def __init__(self, ref_sample, test_sample, window_width):
         
         self.ref_sample = ref_sample
         self.test_sample = test_sample
-
-        self.fp = 0
-        self.tp = 0
-
-        self.n_missed = 0
-        self.n_detected = 0
-
-        # How many there are
         self.n_ref = len(ref_sample)
         self.n_comp = len(test_sample)
 
-
-        # # Just derive these 4 at the end?
-        # # Index info about the reference samples
-        # self.detected_inds = []
-        # self.missed_inds = []
-        # # About the testing samples
-        # self.correct_test_inds = []
-        # self.wrong_test_inds = []
-
-
         # The matching test sample numbers. -1 for indices with no match
         self.matching_sample_nums = -1 * np.ones(n_ref)
 
         # TODO: rdann return annotations.where
 
-    def compare(self):
+    def calc_stats(self):
+        """
+        Calculate performance statistics after the two sets of annotations
+        are compared.
+
+        Example:
+        -------------------
+         ref=500  test=480
+        {  30 { 470 } 10  }
+        -------------------
+        
+        tp = 470
+        fp = 10
+        fn = 30
+
+        specificity = 470 / 500
+        positive_predictivity = 470 / 480
+        false_positive_rate = 10 / 480
+
+        """
+        self.detected_ref_inds = np.where(self.matching_sample_nums != -1)
+        self.missed_ref_inds = np.where(self.matching_sample_nums == -1)
+        self.matched_test_inds = self.matching_sample_nums(
+            self.matching_sample_nums != -1)
+        self.unmached_test_inds = np.setdiff1d(np.array(range(self.n_test)),
+            self.matched_test_inds, assume_unique=True)
+
+        # True positives = matched reference samples
+        self.tp = len(detected_ref_inds)
+        # False positives = extra test samples not matched
+        self.fp = self.n_test - self.tp
+        # False negatives = undetected reference samples
+        self.fn = self.n_ref - self.tp
+        # No tn attribute
+
+        self.specificity = self.tp / self.n_ref
+        self.positive_predictivity = self.tp / self.n_test
+        self.false_positive_rate = self.fp / self.n_test
 
 
+    def compare(self):
 
         test_samp_num = 0
         ref_samp_num = 0
         
-        while ref_samp_num < n_ref:
+        # Why can't this just be a for loop of ref_samp_num?
+        while ref_samp_num < n_ref and test_samp_num < n_test:
 
             closest_samp_num, smallest_samp_diff = (
-                self.get_closest_samp_num(ref_samp_num, test_samp_num))
+                self.get_closest_samp_num(ref_samp_num, test_samp_num,
+                                          self.n_test))
+            # This needs to work for last index
             closest_samp_num_next, smallest_samp_diff_next = (
-                self.get_closest_samp_num(ref_samp_num + 1, test_samp_num))
+                self.get_closest_samp_num(ref_samp_num + 1, test_samp_num,
+                                          self.n_test))
 
             # Found a contested test sample number. Decide which reference
             # sample it belongs to.
             if closest_samp_num == closest_samp_num_next:
-                pass
-            # No clash. Assign the reference-test pair
-            else:
+                # If the sample is closer to the next reference sample, get
+                # the next closest sample for this reference sample.
+                if smallest_samp_diff_next < smallest_samp_diff:
+                    # Get the next closest sample.
+                    # Can this be empty? Need to catch case where nothing left?
+                    closest_samp_num, smallest_samp_diff = (
+                        self.get_closest_samp_num(ref_samp_num, test_samp_num,
+                                                  closest_samp_num))
+
+
                 self.matching_sample_nums[ref_samp_num] = closest_samp_num
 
-                ref_samp_num += 1
-                test_samp_num = closest_samp_num + 1
+            # If no clash, it is straightforward.
+            
+            # Assign the reference-test pair if close enough
+            if smallest_sample_diff < self.window_width:
+                self.matching_sample_nums[ref_samp_num] = closest_samp_num
 
+            ref_samp_num += 1
+            test_samp_num = closest_samp_num + 1
 
         self.calc_stats()
 
             
-        def get_closest_samp_num(self, ref_samp_num, start_test_samp_num):
+        def get_closest_samp_num(self, ref_samp_num, start_test_samp_num,
+                                 stop_test_samp_num):
             """
             Return the closest testing sample number for the given reference
-            sample number. Begin the search from start_test_samp_num.
+            sample number. Limit the search between start_test_samp_num and
+            stop_test_samp_num.
             """
 
             if start_test_samp_num >= self.n_test:
@@ -92,7 +130,7 @@ def get_closest_samp_num(self, ref_samp_num, start_test_samp_num):
             smallest_samp_diff = abs(samp_diff)
 
             # Iterate through the testing samples
-            for test_samp_num in range(start_test_samp_num, self.n_test):
+            for test_samp_num in range(start_test_samp_num, stop_test_samp_num):
                 test_samp = self.test_sample[test_samp_num]
                 samp_diff = ref_samp - test_samp
                 abs_samp_diff = abs(samp_diff)
@@ -109,32 +147,22 @@ def get_closest_samp_num(self, ref_samp_num, start_test_samp_num):
             return closest_samp_num, smallest_samp_diff
 
 
-
-
-
-
-
-def compare_annotations(ind_ref, ind_comp):
+def compare_annotations(ref_sample, test_sample, window_width):
     """
+    
+    Parameters
+    ----------
 
-    """
-
-
-    detected_inds
-    missed_inds
-
-
-    tp
-    tn
-    fp
-    fn
+    Returns
+    -------
+    comparitor : Comparitor object
+        Object containing parameters about the two sets of annotations
 
-    tpr
-    tnr
-    fpr
-    fnr
+    """
+    comparitor = Comparitor(ref_sample, test_sample, window_width)
+    comparitor.compare()
 
-    return evaluation
+    return comparitor
 
 
 def plot_comparitor(comparitor, sig=None):
