diff --git a/pyproject.toml b/pyproject.toml index 42dcfe8..b1c69b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dependencies = [ "numcodecs>=0.13.0,<0.17", "numcodecs-combinators[xarray]~=0.2.10", "numcodecs-observers~=0.1.2", + "numcodecs-safeguards==0.1.0a1", "numcodecs-wasm~=0.2.1", "numcodecs-wasm-bit-round~=0.4.0", "numcodecs-wasm-fixed-offset-scale~=0.4.0", @@ -28,6 +29,7 @@ dependencies = [ "numcodecs-wasm-zfp~=0.6.0", "numcodecs-wasm-zfp-classic~=0.4.0", "numcodecs-wasm-zstd~=0.4.0", + "numcodecs-zero~=0.1.0", "pandas~=2.2", "scipy~=1.14", "seaborn~=0.13.2", diff --git a/src/climatebenchpress/compressor/compressors/__init__.py b/src/climatebenchpress/compressor/compressors/__init__.py index b523c75..b86813c 100644 --- a/src/climatebenchpress/compressor/compressors/__init__.py +++ b/src/climatebenchpress/compressor/compressors/__init__.py @@ -2,6 +2,11 @@ "BitRound", "BitRoundPco", "Jpeg2000", + "SafeguardsSperr", + "SafeguardsSz3", + "SafeguardsZero", + "SafeguardsZeroDssim", + "SafeguardsZfpRound", "Sperr", "StochRound", "StochRoundPco", @@ -15,6 +20,13 @@ from .bitround import BitRound from .bitround_pco import BitRoundPco from .jpeg2000 import Jpeg2000 +from .safeguards import ( + SafeguardsSperr, + SafeguardsSz3, + SafeguardsZero, + SafeguardsZeroDssim, + SafeguardsZfpRound, +) from .sperr import Sperr from .stochround import StochRound from .stochround_pco import StochRoundPco diff --git a/src/climatebenchpress/compressor/compressors/safeguards/__init__.py b/src/climatebenchpress/compressor/compressors/safeguards/__init__.py new file mode 100644 index 0000000..e75d448 --- /dev/null +++ b/src/climatebenchpress/compressor/compressors/safeguards/__init__.py @@ -0,0 +1,13 @@ +__all__ = [ + "SafeguardsSperr", + "SafeguardsSz3", + "SafeguardsZero", + "SafeguardsZeroDssim", + "SafeguardsZfpRound", +] + +from .sperr import SafeguardsSperr +from .sz3 import SafeguardsSz3 +from .zero import SafeguardsZero +from .zero_dssim import SafeguardsZeroDssim +from .zfp_round import SafeguardsZfpRound diff --git a/src/climatebenchpress/compressor/compressors/safeguards/sperr.py b/src/climatebenchpress/compressor/compressors/safeguards/sperr.py new file mode 100644 index 0000000..c49183c --- /dev/null +++ b/src/climatebenchpress/compressor/compressors/safeguards/sperr.py @@ -0,0 +1,22 @@ +__all__ = ["SafeguardsSperr"] + +import numcodecs_safeguards +import numcodecs_wasm_sperr + +from ..abc import Compressor + + +class SafeguardsSperr(Compressor): + """Safeguarded SPERR compressor.""" + + name = "safeguards-sperr" + description = "Safeguards(SPERR)" + + @staticmethod + def abs_bound_codec(error_bound, **kwargs): + return numcodecs_safeguards.SafeguardsCodec( + codec=numcodecs_wasm_sperr.Sperr(mode="pwe", pwe=error_bound), + safeguards=[ + dict(kind="eb", type="abs", eb=error_bound, equal_nan=True), + ], + ) diff --git a/src/climatebenchpress/compressor/compressors/safeguards/sz3.py b/src/climatebenchpress/compressor/compressors/safeguards/sz3.py new file mode 100644 index 0000000..480992c --- /dev/null +++ b/src/climatebenchpress/compressor/compressors/safeguards/sz3.py @@ -0,0 +1,31 @@ +__all__ = ["SafeguardsSz3"] + +import numcodecs_safeguards +import numcodecs_wasm_sz3 + +from ..abc import Compressor + + +class SafeguardsSz3(Compressor): + """Safeguarded SZ3 compressor.""" + + name = "safeguards-sz3" + description = "Safeguards(SZ3)" + + @staticmethod + def abs_bound_codec(error_bound, **kwargs): + return numcodecs_safeguards.SafeguardsCodec( + codec=numcodecs_wasm_sz3.Sz3(eb_mode="abs", eb_abs=error_bound), + safeguards=[ + dict(kind="eb", type="abs", eb=error_bound, equal_nan=True), + ], + ) + + @staticmethod + def rel_bound_codec(error_bound, **kwargs): + return numcodecs_safeguards.SafeguardsCodec( + codec=numcodecs_wasm_sz3.Sz3(eb_mode="rel", eb_rel=error_bound), + safeguards=[ + dict(kind="eb", type="rel", eb=error_bound, equal_nan=True), + ], + ) diff --git a/src/climatebenchpress/compressor/compressors/safeguards/zero.py b/src/climatebenchpress/compressor/compressors/safeguards/zero.py new file mode 100644 index 0000000..d156c8a --- /dev/null +++ b/src/climatebenchpress/compressor/compressors/safeguards/zero.py @@ -0,0 +1,31 @@ +__all__ = ["SafeguardsZero"] + +import numcodecs_safeguards +import numcodecs_zero + +from ..abc import Compressor + + +class SafeguardsZero(Compressor): + """Safeguarded all-zero compressor.""" + + name = "safeguards-zero" + description = "Safeguards(0)" + + @staticmethod + def abs_bound_codec(error_bound, **kwargs): + return numcodecs_safeguards.SafeguardsCodec( + codec=numcodecs_zero.ZeroCodec(), + safeguards=[ + dict(kind="eb", type="abs", eb=error_bound, equal_nan=True), + ], + ) + + @staticmethod + def rel_bound_codec(error_bound, **kwargs): + return numcodecs_safeguards.SafeguardsCodec( + codec=numcodecs_zero.ZeroCodec(), + safeguards=[ + dict(kind="eb", type="rel", eb=error_bound, equal_nan=True), + ], + ) diff --git a/src/climatebenchpress/compressor/compressors/safeguards/zero_dssim.py b/src/climatebenchpress/compressor/compressors/safeguards/zero_dssim.py new file mode 100644 index 0000000..d0d3c71 --- /dev/null +++ b/src/climatebenchpress/compressor/compressors/safeguards/zero_dssim.py @@ -0,0 +1,79 @@ +__all__ = ["SafeguardsZeroDssim"] + +import numcodecs_safeguards +import numcodecs_zero + +from ..abc import Compressor + + +class SafeguardsZeroDssim(Compressor): + """Safeguarded all-zero compressor that also safeguards the dSSIM score.""" + + name = "safeguards-zero-dssim" + description = "Safeguards(0, dSSIM)" + + @staticmethod + def abs_bound_codec(error_bound, **kwargs): + return numcodecs_safeguards.SafeguardsCodec( + codec=numcodecs_zero.ZeroCodec(), + safeguards=[ + dict(kind="eb", type="abs", eb=error_bound, equal_nan=True), + # guarantee that the global minimum and maximum are preserved, + # which simplifies the rescaling + dict(kind="sign", offset="$x_min"), + dict(kind="sign", offset="$x_max"), + dict( + kind="qoi_eb_pw", + qoi=""" + # we guarantee that + # min(data) = min(corrected) and + # max(data) = max(corrected) + # with the sign safeguards above + v["smin"] = c["$x_min"]; + v["smax"] = c["$x_max"]; + v["r"] = v["smax"] - v["smin"]; + + # re-scale to [0-1] and quantize to 256 bins + v["sc_a2"] = round_ties_even(((x - v["smin"]) / v["r"]) * 255) / 255; + + # force the quantized value to stay the same + return v["sc_a2"]; + """, + type="abs", + eb=0, + ), + ], + ) + + @staticmethod + def rel_bound_codec(error_bound, **kwargs): + return numcodecs_safeguards.SafeguardsCodec( + codec=numcodecs_zero.ZeroCodec(), + safeguards=[ + dict(kind="eb", type="rel", eb=error_bound, equal_nan=True), + # guarantee that the global minimum and maximum are preserved, + # which simplifies the rescaling + dict(kind="sign", offset="$x_min"), + dict(kind="sign", offset="$x_max"), + dict( + kind="qoi_eb_pw", + qoi=""" + # we guarantee that + # min(data) = min(corrected) and + # max(data) = max(corrected) + # with the sign safeguards above + v["smin"] = c["$x_min"]; + v["smax"] = c["$x_max"]; + v["r"] = v["smax"] - v["smin"]; + + # re-scale to [0-1] and quantize to 256 bins + v["sc_a2"] = round_ties_even(((x - v["smin"]) / v["r"]) * 255) / 255; + + # force the quantized value to stay the same + return v["sc_a2"]; + """, + type="abs", + eb=0, + ), + ], + ) diff --git a/src/climatebenchpress/compressor/compressors/safeguards/zfp_round.py b/src/climatebenchpress/compressor/compressors/safeguards/zfp_round.py new file mode 100644 index 0000000..81b8caa --- /dev/null +++ b/src/climatebenchpress/compressor/compressors/safeguards/zfp_round.py @@ -0,0 +1,34 @@ +__all__ = ["SafeguardsZfpRound"] + +import numcodecs_safeguards +import numcodecs_wasm_zfp + +from ..abc import Compressor + + +class SafeguardsZfpRound(Compressor): + """Safeguarded ZFP-ROUND compressor. + + This is an adjusted version of the ZFP compressor with an improved rounding mechanism + for the transform coefficients. + """ + + name = "safeguards-zfp-round" + description = "Safeguards(ZFP-ROUND)" + + # 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). + # 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_safeguards.SafeguardsCodec( + codec=numcodecs_wasm_zfp.Zfp(mode="fixed-accuracy", tolerance=error_bound), + safeguards=[ + dict(kind="eb", type="abs", eb=error_bound, equal_nan=True), + ], + )