[DO NOT MERGE] Test simplest safeguards compressor wrappers

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",

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

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

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),
+            ],
+        )

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),
+            ],
+        )

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),
+            ],
+        )

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,
+                ),
+            ],
+        )

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),
+            ],
+        )