from multiprocessing import cpu_count, Pool
import numpy as np
from wfdb.io.annotation import rdann
from wfdb.io.download import get_record_list
from wfdb.io.record import rdsamp
class Comparitor(object):
"""
The class to implement and hold comparisons between two sets of
annotations. See methods `compare`, `print_summary` and `plot`.
Attributes
----------
ref_sample : ndarray
An array of the reference sample locations.
test_sample : ndarray
An array of the comparison sample locations.
window_width : int
The width of the window.
signal : 1d numpy array, optional
The signal array the annotation samples are labelling. Only used
for plotting.
Examples
--------
>>> import wfdb
>>> from wfdb import processing
>>> sig, fields = wfdb.rdsamp('sample-data/100', channels=[0])
>>> ann_ref = wfdb.rdann('sample-data/100','atr')
>>> xqrs = processing.XQRS(sig=sig[:,0], fs=fields['fs'])
>>> xqrs.detect()
>>> comparitor = processing.Comparitor(ann_ref.sample[1:],
xqrs.qrs_inds,
int(0.1 * fields['fs']),
sig[:,0])
>>> comparitor.compare()
>>> comparitor.print_summary()
>>> comparitor.plot()
"""
def __init__(self, ref_sample, test_sample, window_width, signal=None):
if len(ref_sample) > 1 and len(test_sample) > 1:
if min(np.diff(ref_sample)) < 0 or min(np.diff(test_sample)) < 0:
raise ValueError(
(
"The sample locations must be monotonically"
+ " increasing"
)
)
self.ref_sample = ref_sample
self.test_sample = test_sample
self.n_ref = len(ref_sample)
self.n_test = len(test_sample)
self.window_width = window_width
# The matching test sample number for each reference annotation.
# -1 for indices with no match
self.matching_sample_nums = np.full(self.n_ref, -1, dtype="int")
self.signal = signal
# TODO: rdann return annotations.where
def _calc_stats(self):
"""
Calculate performance statistics after the two sets of annotations
are compared.
Parameters
----------
N/A
Returns
-------
N/A
Example:
-------------------
ref=500 test=480
{ 30 { 470 } 10 }
-------------------
tp = 470
fp = 10
fn = 30
sensitivity = 470 / 500
positive_predictivity = 470 / 480
"""
# Reference annotation indices that were detected
self.matched_ref_inds = np.where(self.matching_sample_nums != -1)[0]
# Reference annotation indices that were missed
self.unmatched_ref_inds = np.where(self.matching_sample_nums == -1)[0]
# Test annotation indices that were matched to a reference annotation
self.matched_test_inds = self.matching_sample_nums[
self.matching_sample_nums != -1
]
# Test annotation indices that were unmatched to a reference annotation
self.unmatched_test_inds = np.setdiff1d(
np.array(range(self.n_test)),
self.matched_test_inds,
assume_unique=True,
)
# Sample numbers that were matched and unmatched
self.matched_ref_sample = self.ref_sample[self.matched_ref_inds]
self.unmatched_ref_sample = self.ref_sample[self.unmatched_ref_inds]
self.matched_test_sample = self.test_sample[self.matched_test_inds]
self.unmatched_test_sample = self.test_sample[self.unmatched_test_inds]
# True positives = matched reference samples
self.tp = len(self.matched_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.sensitivity = float(self.tp) / float(self.tp + self.fn)
self.positive_predictivity = float(self.tp) / self.n_test
def compare(self):
"""
Main comparison function. Note: Make sure to be able to handle
these ref/test scenarios:
Parameters
-------
N/A
Returns
-------
N/A
Example
-------
A:
o----o---o---o
x-------x----x
B:
o----o-----o---o
x--------x--x--x
C:
o------o-----o---o
x-x--------x--x--x
D:
o------o-----o---o
x-x--------x-----x
"""
test_samp_num = 0
ref_samp_num = 0
# Iterate through the reference sample numbers
while ref_samp_num < self.n_ref and test_samp_num < self.n_test:
# Get the closest testing sample number for this reference sample
closest_samp_num, smallest_samp_diff = self._get_closest_samp_num(
ref_samp_num, test_samp_num
)
# Get the closest testing sample number for the next reference
# sample. This doesn't need to be called for the last index.
if ref_samp_num < self.n_ref - 1:
(
closest_samp_num_next,
smallest_samp_diff_next,
) = self._get_closest_samp_num(ref_samp_num + 1, test_samp_num)
else:
# Set non-matching value if there is no next reference sample
# to compete for the test sample
closest_samp_num_next = -1
# Found a contested test sample number. Decide which
# reference sample it belongs to. If the sample is closer to
# the next reference sample, leave it to the next reference
# sample and label this reference sample as unmatched.
if (
closest_samp_num == closest_samp_num_next
and smallest_samp_diff_next < smallest_samp_diff
):
# Get the next closest sample for this reference sample,
# if not already assigned to a previous sample.
# It will be the previous testing sample number in any
# possible case (scenario D below), or nothing.
if closest_samp_num and (
not ref_samp_num
or closest_samp_num - 1
!= self.matching_sample_nums[ref_samp_num - 1]
):
# The previous test annotation is inspected
closest_samp_num = closest_samp_num - 1
smallest_samp_diff = abs(
self.ref_sample[ref_samp_num]
- self.test_sample[closest_samp_num]
)
# Assign the reference-test pair if close enough
if smallest_samp_diff < self.window_width:
self.matching_sample_nums[
ref_samp_num
] = closest_samp_num
# Set the starting test sample number to inspect
# for the next reference sample.
test_samp_num = closest_samp_num + 1
# Otherwise there is no matching test annotation
# If there is no clash, or the contested test sample is
# closer to the current reference, keep the test sample
# for this reference sample.
else:
# Assign the reference-test pair if close enough
if smallest_samp_diff < self.window_width:
self.matching_sample_nums[ref_samp_num] = closest_samp_num
# Increment the starting test sample number to inspect
# for the next reference sample.
test_samp_num = closest_samp_num + 1
ref_samp_num += 1
self._calc_stats()
def _get_closest_samp_num(self, ref_samp_num, start_test_samp_num):
"""
Return the closest testing sample number for the given reference
sample number. Limit the search from start_test_samp_num.
Parameters
----------
ref_samp_num : int
The desired reference sample number to get the closest result.
start_test_samp_num
The desired testing reference sample number to get the
closest result.
Returns
-------
closest_samp_num : int
The closest sample number to the reference sample number.
smallest_samp_diff : int
The smallest difference between the reference sample and
the testing sample.
"""
if start_test_samp_num >= self.n_test:
raise ValueError("Invalid starting test sample number.")
ref_samp = self.ref_sample[ref_samp_num]
test_samp = self.test_sample[start_test_samp_num]
samp_diff = ref_samp - test_samp
# Initialize running parameters
closest_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):
test_samp = self.test_sample[test_samp_num]
samp_diff = ref_samp - test_samp
abs_samp_diff = abs(samp_diff)
# Found a better match
if abs_samp_diff < smallest_samp_diff:
closest_samp_num = test_samp_num
smallest_samp_diff = abs_samp_diff
# Stop iterating when the ref sample is first passed or reached
if samp_diff <= 0:
break
return closest_samp_num, smallest_samp_diff
def print_summary(self):
"""
Print summary metrics of the annotation comparisons.
Parameters
----------
N/A
Returns
-------
N/A
"""
if not hasattr(self, "sensitivity"):
self._calc_stats()
print(
"%d reference annotations, %d test annotations\n"
% (self.n_ref, self.n_test)
)
print("True Positives (matched samples): %d" % self.tp)
print("False Positives (unmatched test samples): %d" % self.fp)
print("False Negatives (unmatched reference samples): %d\n" % self.fn)
print(
"Sensitivity: %.4f (%d/%d)"
% (self.sensitivity, self.tp, self.n_ref)
)
print(
"Positive Predictivity: %.4f (%d/%d)"
% (self.positive_predictivity, self.tp, self.n_test)
)
def plot(self, sig_style="", title=None, figsize=None, return_fig=False):
"""
Plot the comparison of two sets of annotations, possibly
overlaid on their original signal.
Parameters
----------
sig_style : str, optional
The matplotlib style of the signal
title : str, optional
The title of the plot
figsize: tuple, optional
Tuple pair specifying the width, and height of the figure.
It is the'figsize' argument passed into matplotlib.pyplot's
`figure` function.
return_fig : bool, optional
Whether the figure is to be returned as an output argument.
Returns
-------
fig : matplotlib figure object
The figure information for the plot.
ax : matplotlib axes object
The axes information for the plot.
"""
import matplotlib.pyplot as plt
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(1, 1, 1)
legend = [
"Signal",
"Matched Reference Annotations (%d/%d)" % (self.tp, self.n_ref),
"Unmatched Reference Annotations (%d/%d)" % (self.fn, self.n_ref),
"Matched Test Annotations (%d/%d)" % (self.tp, self.n_test),
"Unmatched Test Annotations (%d/%d)" % (self.fp, self.n_test),
]
# Plot the signal if any
if self.signal is not None:
ax.plot(self.signal, sig_style)
# Plot reference annotations
ax.plot(
self.matched_ref_sample,
self.signal[self.matched_ref_sample],
"ko",
)
ax.plot(
self.unmatched_ref_sample,
self.signal[self.unmatched_ref_sample],
"ko",
fillstyle="none",
)
# Plot test annotations
ax.plot(
self.matched_test_sample,
self.signal[self.matched_test_sample],
"g+",
)
ax.plot(
self.unmatched_test_sample,
self.signal[self.unmatched_test_sample],
"rx",
)
ax.legend(legend)
# Just plot annotations
else:
# Plot reference annotations
ax.plot(self.matched_ref_sample, np.ones(self.tp), "ko")
ax.plot(
self.unmatched_ref_sample,
np.ones(self.fn),
"ko",
fillstyle="none",
)
# Plot test annotations
ax.plot(self.matched_test_sample, 0.5 * np.ones(self.tp), "g+")
ax.plot(self.unmatched_test_sample, 0.5 * np.ones(self.fp), "rx")
ax.legend(legend[1:])
if title:
ax.set_title(title)
ax.set_xlabel("time/sample")
fig.show()
if return_fig:
return fig, ax
def compare_annotations(ref_sample, test_sample, window_width, signal=None):
"""
Compare a set of reference annotation locations against a set of
test annotation locations. See the Comparitor class docstring
for more information.
Parameters
----------
ref_sample : 1d numpy array
Array of reference sample locations.
test_sample : 1d numpy array
Array of test sample locations to compare.
window_width : int
The maximum absolute difference in sample numbers that is
permitted for matching annotations.
signal : 1d numpy array, optional
The original signal of the two annotations. Only used for
plotting.
Returns
-------
comparitor : Comparitor object
Object containing parameters about the two sets of annotations.
Examples
--------
>>> import wfdb
>>> from wfdb import processing
>>> sig, fields = wfdb.rdsamp('sample-data/100', channels=[0])
>>> ann_ref = wfdb.rdann('sample-data/100','atr')
>>> xqrs = processing.XQRS(sig=sig[:,0], fs=fields['fs'])
>>> xqrs.detect()
>>> comparitor = processing.compare_annotations(ann_ref.sample[1:],
xqrs.qrs_inds,
int(0.1 * fields['fs']),
sig[:,0])
>>> comparitor.print_summary()
>>> comparitor.plot()
"""
comparitor = Comparitor(
ref_sample=ref_sample,
test_sample=test_sample,
window_width=window_width,
signal=signal,
)
comparitor.compare()
return comparitor
def benchmark_mitdb(detector, verbose=False, print_results=False):
"""
Benchmark a QRS detector against mitdb's records.
Parameters
----------
detector : function
The detector function.
verbose : bool, optional
The verbose option of the detector function.
print_results : bool, optional
Whether to print the overall performance, and the results for
each record.
Returns
-------
comparitors : dictionary
Dictionary of Comparitor objects run on the records, keyed on
the record names.
sensitivity : float
Aggregate sensitivity.
positive_predictivity : float
Aggregate positive_predictivity.
Notes
-----
TODO:
- remove non-qrs detections from reference annotations
- allow kwargs
Examples
--------
>>> import wfdb
>> from wfdb.processing import benchmark_mitdb, xqrs_detect
>>> comparitors, spec, pp = benchmark_mitdb(xqrs_detect)
"""
record_list = get_record_list("mitdb")
n_records = len(record_list)
# Function arguments for starmap
args = zip(record_list, n_records * [detector], n_records * [verbose])
# Run detector and compare against reference annotations for all
# records
with Pool(cpu_count() - 1) as p:
comparitors = p.starmap(benchmark_mitdb_record, args)
# Calculate aggregate stats
sensitivity = np.mean([c.sensitivity for c in comparitors])
positive_predictivity = np.mean(
[c.positive_predictivity for c in comparitors]
)
comparitors = dict(zip(record_list, comparitors))
print("Benchmark complete")
if print_results:
print(
"\nOverall MITDB Performance - Sensitivity: %.4f, Positive Predictivity: %.4f\n"
% (sensitivity, positive_predictivity)
)
for record_name in record_list:
print("Record %s:" % record_name)
comparitors[record_name].print_summary()
print("\n\n")
return comparitors, sensitivity, positive_predictivity
def benchmark_mitdb_record(rec, detector, verbose):
"""
Benchmark a single mitdb record.
Parameters
----------
rec : str
The mitdb record to be benchmarked.
detector : function
The detector function.
verbose : bool
Whether to print the record names (True) or not (False).
Returns
-------
comparitor : Comparitor object
Object containing parameters about the two sets of annotations.
"""
sig, fields = rdsamp(rec, pn_dir="mitdb", channels=[0])
ann_ref = rdann(rec, pn_dir="mitdb", extension="atr")
qrs_inds = detector(sig=sig[:, 0], fs=fields["fs"], verbose=verbose)
comparitor = compare_annotations(
ref_sample=ann_ref.sample[1:],
test_sample=qrs_inds,
window_width=int(0.1 * fields["fs"]),
)
if verbose:
print("Finished record %s" % rec)
return comparitor