diff --git a/src/climatebenchpress/compressor/compressors/utils.py b/src/climatebenchpress/compressor/compressors/utils.py index bee70dc..efc149b 100644 --- a/src/climatebenchpress/compressor/compressors/utils.py +++ b/src/climatebenchpress/compressor/compressors/utils.py @@ -12,6 +12,12 @@ np.dtype("float16"): 10, } +NONMANTISSA_BITS = { + np.dtype("float32"): 32 - MANTISSA_BITS[np.dtype("float32")], + np.dtype("float64"): 64 - MANTISSA_BITS[np.dtype("float64")], + np.dtype("float16"): 16 - MANTISSA_BITS[np.dtype("float16")], +} + def compute_keepbits(dtype: np.dtype, rel_error: float) -> int: """ diff --git a/src/climatebenchpress/compressor/compressors/zfp.py b/src/climatebenchpress/compressor/compressors/zfp.py index 9e28cb8..7ef5ef9 100644 --- a/src/climatebenchpress/compressor/compressors/zfp.py +++ b/src/climatebenchpress/compressor/compressors/zfp.py @@ -3,22 +3,37 @@ import numcodecs_wasm_zfp_classic from .abc import Compressor +from .utils import NONMANTISSA_BITS, compute_keepbits + +# From: https://github.com/LLNL/zfp/blob/4baa4c7eeae8e0b6a7ace4dde242ac165bcd59d9/include/zfp.h#L18 +ZFP_MIN_BITS = 1 +ZFP_MAX_BITS = 16658 class Zfp(Compressor): name = "zfp" description = "ZFP" + @staticmethod + def abs_bound_codec(error_bound, **kwargs): + return numcodecs_wasm_zfp_classic.ZfpClassic( + mode="fixed-accuracy", tolerance=error_bound + ) + # NOTE: # ZFP mechanism for strictly supporting relative error bounds is to - # truncate the floating point bit representation and then use ZFP's lossless - # mode for compression. This is essentially equivalent to the BitRound - # compressors we are already implementing (with a difference what the lossless - # compression algorithm is). + # apply bitshaving and then use ZFP's lossless mode for compression. # See https://zfp.readthedocs.io/en/release1.0.1/faq.html#q-relerr for more details. - @staticmethod - def abs_bound_codec(error_bound, **kwargs): + def rel_bound_codec(error_bound, *, dtype=None, **kwargs): + assert dtype is not None, "dtype must be provided" + + mantissa_keepbits = compute_keepbits(dtype, error_bound) + total_keepbits = mantissa_keepbits + NONMANTISSA_BITS[dtype] return numcodecs_wasm_zfp_classic.ZfpClassic( - mode="fixed-accuracy", tolerance=error_bound + mode="expert", + min_bits=ZFP_MIN_BITS, + max_bits=ZFP_MAX_BITS, + max_prec=total_keepbits, + min_exp=-1075, ) diff --git a/src/climatebenchpress/compressor/compressors/zfp_round.py b/src/climatebenchpress/compressor/compressors/zfp_round.py index bc55e81..df0f877 100644 --- a/src/climatebenchpress/compressor/compressors/zfp_round.py +++ b/src/climatebenchpress/compressor/compressors/zfp_round.py @@ -3,20 +3,32 @@ import numcodecs_wasm_zfp from .abc import Compressor +from .utils import NONMANTISSA_BITS, compute_keepbits +from .zfp import ZFP_MAX_BITS, ZFP_MIN_BITS class ZfpRound(Compressor): name = "zfp-round" description = "ZFP-ROUND" + @staticmethod + def abs_bound_codec(error_bound, **kwargs): + return numcodecs_wasm_zfp.Zfp(mode="fixed-accuracy", tolerance=error_bound) + # NOTE: # ZFP mechanism for strictly supporting relative error bounds is to - # truncate the floating point bit representation and then use ZFP's lossless - # mode for compression. This is essentially equivalent to the BitRound - # compressors we are already implementing (with a difference what the lossless - # compression algorithm is). + # apply bitshaving and then use ZFP's lossless mode for compression. # See https://zfp.readthedocs.io/en/release1.0.1/faq.html#q-relerr for more details. - @staticmethod - def abs_bound_codec(error_bound, **kwargs): - return numcodecs_wasm_zfp.Zfp(mode="fixed-accuracy", tolerance=error_bound) + def rel_bound_codec(error_bound, *, dtype=None, **kwargs): + assert dtype is not None, "dtype must be provided" + + mantissa_keepbits = compute_keepbits(dtype, error_bound) + total_keepbits = mantissa_keepbits + NONMANTISSA_BITS[dtype] + return numcodecs_wasm_zfp.Zfp( + mode="expert", + min_bits=ZFP_MIN_BITS, + max_bits=ZFP_MAX_BITS, + max_prec=total_keepbits, + min_exp=-1075, + )