Cleans up ann2rr; adds rr2ann

Lucas-Mc · Lucas-Mc · commit f3c081ad63df · 2021-03-31T13:37:28.000-04:00
diff --git a/wfdb/__init__.py b/wfdb/__init__.py
@@ -2,7 +2,7 @@
                             wrsamp, dl_database, edf2mit, mit2edf, wav2mit, mit2wav,
                             wfdb2mat, csv2mit, sampfreq, signame)
 from wfdb.io.annotation import (Annotation, rdann, wrann, show_ann_labels,
-                                show_ann_classes, ann2rr)
+                                show_ann_classes, ann2rr, rr2ann)
 from wfdb.io.download import get_dbs, get_record_list, dl_files, set_db_index_url
 from wfdb.plot.plot import plot_items, plot_wfdb, plot_all_records
 
diff --git a/wfdb/io/__init__.py b/wfdb/io/__init__.py
@@ -3,6 +3,6 @@
                             csv2mit, sampfreq, signame, SIGNAL_CLASSES)
 from wfdb.io._signal import est_res, wr_dat_file
 from wfdb.io.annotation import (Annotation, rdann, wrann, show_ann_labels,
-                                show_ann_classes, ann2rr)
+                                show_ann_classes, ann2rr, rr2ann)
 from wfdb.io.download import get_dbs, get_record_list, dl_files, set_db_index_url
 from wfdb.io.tff import rdtff
diff --git a/wfdb/io/annotation.py b/wfdb/io/annotation.py
@@ -2212,7 +2212,7 @@ def rm_last(*args):
 
 
 def ann2rr(record_name, extension, pn_dir=None, start_time=None,
-           stop_time=None, format=None):
+           stop_time=None, format=None, as_array=True):
     from wfdb.processing import hr
     """
     Obtain RR interval series from ECG annotation files.
@@ -2225,26 +2225,28 @@ def ann2rr(record_name, extension, pn_dir=None, start_time=None,
     extension : str
         The annotatator extension of the annotation file. ie. for  file
         '100.atr', extension='atr'.
-    pn_dir : str
+    pn_dir : str, optional
         Option used to stream data from Physionet. The PhysioNet database
         directory from which to find the required annotation file. eg. For
         record '100' in 'http://physionet.org/content/mitdb': pn_dir='mitdb'.
-    start_time : float
+    start_time : float, optional
         The time to start the intervals in seconds.
-    stop_time : float
+    stop_time : float, optional
         The time to stop the intervals in seconds.
-    format : str
+    format : str, optional
         Print intervals in the specified format. By default, intervals are
         printed in units of sample intervals. Other formats include
         's' (seconds), 'm' (minutes), 'h' (hours). Set to 'None' for samples.
+    as_array : bool, optional
+        If True, return an an 'ndarray', else print the output.
 
     Returns
     -------
     N/A
 
     Examples
     --------
-    >>> wfdb.ann2rr('sample-data/100', 'atr')
+    >>> wfdb.ann2rr('sample-data/100', 'atr', as_array=False)
     >>> 18
     >>> 59
     >>> ...
@@ -2278,7 +2280,78 @@ def ann2rr(record_name, extension, pn_dir=None, start_time=None,
     else:
         out_interval = np.around(time_interval * ann.fs).astype(np.int)
 
-    print(*out_interval, sep='\n')
+    if as_array:
+        return out_interval
+    else:
+        print(*out_interval, sep='\n')
+
+
+def rr2ann(rr_array, record_name, extension, fs=250, as_time=False):
+    """
+    Creates an annotation file from the standard input, which should usually
+    be a Numpy array of intervals in the format produced by `ann2rr`. (For
+    exceptions, see the `as_time` parameter below.). An optional second column
+    may be provided which gives the respective annotation mnemonic.
+
+    Parameters
+    ----------
+    rr_array : ndarray
+        A Numpy array consisting of the input RR intervals. If `as_time` is
+        set to True, then the input should consist of times of occurences. If,
+        the shape of the input array is '(n_annot,2)', then treat the second
+        column as the annotation mnemonic ('N', 'V', etc.). If a second column
+        is not specified, then the default annotation will the '"' which
+        specifies a comment.
+    record_name : str
+        The record name of the WFDB annotation file. ie. for file '100.atr',
+        record_name='100'.
+    extension : str
+        The annotatator extension of the annotation file. ie. for  file
+        '100.atr', extension='atr'.
+    fs : float, int, optional
+        Assume the specified sampling frequency. This option has no effect
+        unless the `as_time` parameter is set to convert to samples; in this
+        case, a sampling frequency of 250 Hz is assumed if this option is
+        omitted.
+    as_time : bool
+        Interpret the input as times of occurrence (if True), rather than as
+        samples (if False). There is not currently a way to input RR intervals
+        in time format between beats. For example, 0.2 seconds between beats
+        1->2, 0.3 seconds between beats 2->3, etc.
+
+    Returns
+    -------
+    N/A
+
+    Examples
+    --------
+    Using time of occurence as input:
+    >>> import numpy as np
+    >>> rr_array = np.array([[0.2, 0.6, 1.3], ['V', 'N', 'V']]).T
+    >>> wfdb.rr2ann(rr_array, 'test_ann', 'atr', fs=100, as_time=True)
+
+    Using samples as input:
+    >>> import numpy as np
+    >>> rr_array = np.array([4, 17, 18, 16])
+    >>> wfdb.rr2ann(rr_array, 'test_ann', 'atr')
+
+    """
+    try:
+        ann_sample = rr_array[:,0]
+    except IndexError:
+        ann_sample = rr_array
+
+    if as_time:
+        ann_sample = (fs * ann_sample.astype(np.float64)).astype(np.int64)
+    else:
+        ann_sample = np.cumsum(ann_sample).astype(np.int64)
+
+    try:
+        ann_symbol = rr_array[:,1].tolist()
+    except IndexError:
+        ann_symbol = rr_array.shape[0] * ['"']
+
+    wrann(record_name, extension, ann_sample, symbol=ann_symbol)
 
 
 ## ------------- Annotation Field Specifications ------------- ##
