Merge pull request #46 from TGSAI/enh/optimize_zfp

Improved ZFP compression quality using `FIXED-ACCURACY` mode

Author: tasansal

src/mdio/commands/segy.py

@@ -123,12 +123,12 @@
     show_default=True,
 )
 @click.option(
-    "-ratio",
-    "--compression-ratio",
+    "-tolerance",
+    "--compression-tolerance",
     required=False,
-    default=4,
-    help="Lossy compression ratio.",
-    type=click.INT,
+    default=0.01,
+    help="Lossy compression tolerance in ZFP.",
+    type=click.FLOAT,
     show_default=True,
 )
 @click.option(
@@ -156,7 +156,7 @@ def segy_import(
     chunk_size,
     endian,
     lossless,
-    compression_ratio,
+    compression_tolerance,
     storage_options,
     overwrite,
 ):
@@ -213,13 +213,18 @@ def segy_import(
 
     By default, the data is ingested with LOSSLESS compression. This
     saves disk space in the range of 20% to 40%. MDIO also allows
-    data to be compressed using the ZFP compressor's fixed rate lossy
-    compression. If lossless parameter is set to False and MDIO was
-    installed using the lossy extra; then the data will be compressed
+    data to be compressed using the ZFP compressor's fixed accuracy
+    lossy compression. If lossless parameter is set to False and MDIO
+    was installed using the lossy extra; then the data will be compressed
     to approximately 30% of its original size and will be perceptually
-    lossless. The compression ratio can be adjusted using the option
-    compression_ratio (integer). Higher values will compress more, but
-    will introduce artifacts.
+    lossless. The compression amount can be adjusted using the option
+    compression_tolerance (float). Values less than 1 gives good results.
+    The higher the value, the more compression, but will introduce artifacts.
+    The default value is 0.01 tolerance, however we get good results
+    up to 0.5; where data is almost compressed to 10% of its original size.
+    NOTE: This assumes data has amplitudes normalized to have approximately
+    standard deviation of 1. If dataset has values smaller than this
+    tolerance, a lot of loss may occur.
 
     Usage:
 
@@ -265,7 +270,7 @@ def segy_import(
         chunksize=chunk_size,
         endian=endian,
         lossless=lossless,
-        compression_ratio=compression_ratio,
+        compression_tolerance=compression_tolerance,
         storage_options=storage_options,
         overwrite=overwrite,
     )

src/mdio/converters/segy.py

@@ -36,7 +36,7 @@ def segy_to_mdio(
     chunksize: Sequence[int] | None = None,
     endian: str = "big",
     lossless: bool = True,
-    compression_ratio: int | float = 4,
+    compression_tolerance: float = 0.01,
     storage_options: dict[str, Any] | None = None,
     overwrite: bool = False,
 ) -> None:
@@ -84,8 +84,10 @@ def segy_to_mdio(
         endian: Endianness of the input SEG-Y. Rev.2 allows little endian.
             Default is 'big'. Must be in `{"big", "little"}`
         lossless: Lossless Blosc with zstandard, or ZFP with fixed precision.
-        compression_ratio: Approximate compression ratio for ZFP compression.
-            Will be ignored if `lossless=True`
+        compression_tolerance: Tolerance ZFP compression, optional. The fixed
+            accuracy mode in ZFP guarantees there won't be any errors larger
+            than this value. The default is 0.01, which gives about 70%
+            reduction in size. Will be ignored if `lossless=True`.
         storage_options: Storage options for the cloud storage backend.
             Default is `None` (will assume anonymous)
         overwrite: Toggle for overwriting existing store
@@ -253,7 +255,7 @@ def segy_to_mdio(
         dtype="float32",
         chunks=chunksize,
         lossless=lossless,
-        compression_ratio=compression_ratio,
+        compression_tolerance=compression_tolerance,
     )
 
     for key, value in stats.items():

src/mdio/segy/blocked_io.py

@@ -19,9 +19,12 @@
 
 
 try:
-    from zarr import ZFPY
+    import zfpy  # Base library
+    from zarr import ZFPY  # Codec
+
 except ImportError:
     ZFPY = None
+    zfpy = None
 
 # Globals
 NUM_CORES = cpu_count(logical=False)
@@ -36,7 +39,7 @@ def to_zarr(
     name: str,
     chunks: tuple[int, ...],
     lossless: bool,
-    compression_ratio: int | float,
+    compression_tolerance: float = 0.01,
     **kwargs,
 ) -> dict:
     """Blocked I/O from SEG-Y to chunked `zarr.core.Array`.
@@ -50,8 +53,10 @@ def to_zarr(
         name: Name of the zarr.Array
         chunks: Chunk sizes for trace data
         lossless: Lossless Blosc with zstandard, or ZFP with fixed precision.
-        compression_ratio: Approximate compression ratio for ZFP compression.
-            Will be ignored for `lossless=True`
+        compression_tolerance: Tolerance ZFP compression, optional. The fixed
+            accuracy mode in ZFP guarantees there won't be any errors larger
+            than this value. The default is 0.01, which gives about 70%
+            reduction in size.
         **kwargs: Additional keyword arguments passed to zarr.core.Array  # noqa: RST210
 
     Returns:
@@ -63,13 +68,11 @@ def to_zarr(
     if lossless is True:
         trace_compressor = Blosc("zstd")
         header_compressor = trace_compressor
-    elif ZFPY is not None:
-        # Compression precision is 32 bit float divided by ratio.
-        # We round it to the nearest integer since ZFP expects an
-        # integer precision for `mode=3`. This will approximately
-        # give the user the ratio they asked for.
-        zfp_precision = np.round(32 / compression_ratio)
-        trace_compressor = ZFPY(mode=3, precision=zfp_precision)
+    elif ZFPY is not None or zfpy is not None:
+        trace_compressor = ZFPY(
+            mode=zfpy.mode_fixed_accuracy,
+            tolerance=compression_tolerance,
+        )
         header_compressor = Blosc("zstd")
     else:
         raise ImportError(