add bad flac tests

cx1111 · cx1111 · commit 18a496cabbff · 2018-05-23T17:14:24.000-04:00
diff --git a/wfdb/io/__init__.py b/wfdb/io/__init__.py
@@ -4,3 +4,4 @@
 from .annotation import (Annotation, rdann, wrann, show_ann_labels,
                          show_ann_classes)
 from .download import get_dbs, get_record_list, dl_files
+from .compress import compress_file, test_compression, compare_compressions
diff --git a/wfdb/io/compress.py b/wfdb/io/compress.py
@@ -1,24 +1,44 @@
 import bz2
+from datetime import timedelta
 import gzip
 from multiprocessing import Pool
 import os
 from shutil import copyfileobj
+import subprocess
 import time
 
 import cxutils as cx
 import lz4.frame
 import numpy as np
 import pandas as pd
+import wfdb
+from wfdb.io._signal import wfdbfmtres
 import zstd
 
 
 def compress_file(file, fmt, level):
     """
-    Compress and decompress a file with a particular fmt. Return
-    compressed size and time for compression/decompression.
+    Compress and decompress a single file.
+
+    Parameters
+    ----------
+    file : str
+        Full file path
+    fmt : str
+        The compression format
+    level : int
+        The compression level
+
+    Returns
+    -------
+    compressed_size : int
+       Compressed file size in bytes
+    compression_time : float
+        Time taken to compress, in seconds.
+    decompression_time : float
+        Time taken to decompress, in seconds.
 
     """
-
     with open(file, 'rb') as f_in:
         u_data = f_in.read()
 
@@ -42,34 +62,39 @@ def compress_file(file, fmt, level):
             c_data = lz4.frame.compress(u_data, compression_level=level)
             t1 = time.time()
             u_data = lz4.frame.decompress(c_data)
-
-        compressed_size = len(c_data)
+        elif fmt == 'flac':
+            # command line processing
+            record = wfdb.rdheader(file[:-4])
+            out_file = os.path.join('/home/cx1111/Downloads/writedir/', os.path.basename(file).strip('.dat') + '.flac')
+            # Write the file since we need to decompress it
+            compress_command = "flac %s --endian=little --channels=%d --sample-rate=%d --bps=%d --sign=signed -%d -o %s" % (
+                file, record.n_sig, record.fs, wfdbfmtres(record.fmt[0]), level, out_file)
+            subprocess.run(compress_command, shell=True)
+            t1 = time.time()
+            decompress_command = "flac -d %s -c" % out_file
+            subprocess.run(decompress_command, shell=True)
 
         t2 = time.time()
-        t_compress = t1 - t0
-        t_decompress = t2 - t1
 
-    return compressed_size, t_compress, t_decompress
+        if fmt == 'flac':
+            compressed_size = os.path.getsize(out_file)
+            os.remove(out_file)
+        else:
+            compressed_size = len(c_data)
 
+        compression_time = t1 - t0
+        decompression_time = t2 - t1
 
-# can change header
-# flac, wabpack
-def test_compression(fmt, compress_level):
-    """
-    Test compression on target dat files.
-
-    From mitdb and first 50 patient records of mimic3wdb/matched/
+    return compressed_size, compression_time, decompression_time
 
-    Total size is about 10 Gb.
 
+def test_compression(fmt, compress_level, test_dat_files):
     """
-    data_dirs = (['/home/cx1111/Downloads/data/mitdb']
-                 + cx.list_dirs('/home/cx1111/Downloads/data/mimic3wdb/matched'))
-
-    test_dat_files = cx.list_files(data_dirs)
+    Test a type of compression of a specified level, on all target dat
+    files.
 
+    """
     n_files = len(test_dat_files)
-
     uncompressed_sizes = [os.path.getsize(file) for file in test_dat_files]
 
     with Pool(os.cpu_count() - 1) as pool:
@@ -79,92 +104,65 @@ def test_compression(fmt, compress_level):
                                             n_files * [compress_level]))
     compressed_sizes, compression_times, decompression_times = zip(*output)
 
-    uncompressed_sizes = np.array(uncompressed_sizes)
-    compressed_sizes = np.array(compressed_sizes)
-    decompression_times = np.array(decompression_times)
-    compression_ratios = uncompressed_sizes / compressed_sizes
+    # Calculate performance summary
+    compression_ratio = np.sum(uncompressed_sizes) / np.sum(compressed_sizes)
+    compression_time = np.sum(compression_times)
+    decompression_time = np.sum(decompression_times)
 
-    # Return the compression ratios and time taken
-    return (uncompressed_sizes, compressed_sizes, compression_ratios,
-            compression_times, decompression_times)
+    return compression_ratio, compression_time, decompression_time
 
 
-def summarize_compression(uncompressed_sizes, compressed_sizes,
-                          compression_ratios, compression_times,
-                          decompression_times, mode='print'):
+def compare_compressions(fmts, compress_levels):
     """
-    Print or return a summary of the compression
+    For each compression format/level pair, run the full compression
+    test. Return the aggregate results. Rounds to nearest second.
 
-    Input parameters are outputs of `test_compression`.
+    The data is the waveforms of the first 100 patients
+    mimic3wdb/matched/ Total size is about 22 Gb.
+
+    Returns
+    -------
+    compression_results : pandas dataframe
+      Dataframe of results for each compression format/level combination.
+      Results include compression ratio, compression time, and decompression
+      time.
+    dataset_info : dict
+      Dictionary of
 
     """
-    n_files = len(uncompressed_sizes)
-    uncompressed_total = np.sum(uncompressed_sizes)
-    compressed_total = np.sum(compressed_sizes)
-
-    overall_compression_ratio = uncompressed_total / compressed_total
-
-    # Sum of min(compressed, uncompressed) for all files
-    smallest_total = np.sum([min(uncompressed_sizes[i], compressed_sizes[i])
-                             for i in range(n_files)])
-    smallest_overall_compression_ratio = uncompressed_total / smallest_total
-
-    # Total times
-    t_compress = np.sum(compression_times)
-    t_decompress = np.sum(decompression_times)
-
-    if mode == 'print':
-        print('Number of files compressed: %d' % n_files)
-        print('Total size of uncompressed files: %s'
-              % cx.readable_size(uncompressed_total, 'string'))
-        print('Total size of compressed files: %s'
-              % cx.readable_size(compressed_total, 'string'))
-        print('Overall compression ratio: %.2f'
-              % overall_compression_ratio)
-        print('Overall compression ratio without compressing inflated files: %.2f'
-              % smallest_overall_compression_ratio)
-        print('Total compression time: %.2f' % t_compress)
-        print('Total compression time: %.2f' % t_decompress)
-    else:
-        return (n_files, uncompressed_total, compressed_total,
-                overall_compression_ratio, t_compress, t_decompress)
+    # Files to be compressed
+    data_dirs = cx.list_dirs('/home/cx1111/Downloads/data/mimic3wdb/matched')
+    test_dat_files = cx.list_files(data_dirs, extensions=['dat'])
 
+    # kloogy inaccurate fix for flac files
+    if 'flac' in fmts:
+        test_dat_files = [file for file in test_dat_files if not file.endswith('n.dat')]
 
-def compare_compressions(fmts, compress_levels):
-    """
-    Run the compression tests and summarize the results of multiple
-    formats/compress pairs.
+    n_files = len(test_dat_files)
+    uncompressed_sizes = [os.path.getsize(file) for file in test_dat_files]
+    uncompressed_total = cx.readable_size(np.sum(uncompressed_sizes))
+    dataset_info = {'n_files':n_files, 'uncompressed_total':uncompressed_total}
 
-    """
-    df = pd.DataFrame(columns=['fmt', 'compress_level', 'n_files',
-                               'uncompressed_total', 'compressed_total',
-                               'compression_ratio', 'time_compress',
-                               'time_decompress'])
+    # Compression results
+    compression_results = pd.DataFrame(columns=['fmt', 'compress_level',
+                                       'compression_ratio', 'time_compress',
+                                       'time_decompress'])
 
     # Iterate through formats and compress levels
     for i in range(len(fmts)):
         fmt = fmts[i]
         compress_level = compress_levels[i]
-        print('Testing fmt: %s, compress level: %d' % (fmt, compress_level))
-
-        (uncompressed_sizes, compressed_sizes, compression_ratios,
-         compression_times, decompression_times) = test_compression(fmt=fmt,
-                                          compress_level=compress_level)
-
-        (n_files, uncompressed_total,
-         compressed_total,
-         overall_compression_ratio,
-         t_compress, t_decompress) = summarize_compression(uncompressed_sizes,
-                                                           compressed_sizes,
-                                                           compression_ratios,
-                                                           compression_times,
-                                                           decompression_times,
-                                                           mode='return')
-
-        df.loc[i] = [fmt, compress_level, n_files,
-                     cx.readable_size(uncompressed_total, 'string'),
-                     cx.readable_size(compressed_total, 'string'),
-                     '%.2f' % overall_compression_ratio, '%.2f' % t_compress,
-                     '%.2f' % t_decompress]
-
-    return df
+        print('Testing %s with compress level=%d ...' % (fmt, compress_level))
+
+        (compression_ratio, compression_time,
+         decompression_time) = test_compression(fmt=fmt,
+                                                compress_level=compress_level,
+                                                test_dat_files=test_dat_files)
+
+        compression_results.loc[i] = [fmt, compress_level,
+                                      '%.2f' % compression_ratio,
+                                      str(timedelta(seconds=int(compression_time))),
+                                      str(timedelta(seconds=int(decompression_time)))]
+
+    print('Full benchmark complete')
+    return compression_results, dataset_info
