safeguard (conservative) relative error bounds

Author: juntyr

pyproject.toml

@@ -13,9 +13,9 @@ dependencies = [
     "matplotlib~=3.8",
     "netcdf4==1.7.3",
     "numcodecs>=0.13.0,<0.17",
-    "numcodecs-combinators[xarray]~=0.2.10",
+    "numcodecs-combinators[xarray]~=0.2.13",
     "numcodecs-observers~=0.1.2",
-    "numcodecs-safeguards==0.1.0a1",
+    "numcodecs-safeguards==0.1.0b1",
     "numcodecs-wasm~=0.2.2",
     "numcodecs-wasm-bit-round~=0.4.0",
     "numcodecs-wasm-fixed-offset-scale~=0.4.0",
@@ -24,12 +24,12 @@ dependencies = [
     "numcodecs-wasm-round~=0.5.0",
     "numcodecs-wasm-sperr~=0.2.0",
     "numcodecs-wasm-stochastic-rounding~=0.2.0",
-    "numcodecs-wasm-sz3~=0.7.0",
+    "numcodecs-wasm-sz3~=0.8.0",
     "numcodecs-wasm-tthresh~=0.3.0",
     "numcodecs-wasm-zfp~=0.6.0",
     "numcodecs-wasm-zfp-classic~=0.4.0",
     "numcodecs-wasm-zstd~=0.4.0",
-    "numcodecs-zero~=0.1.0",
+    "numcodecs-zero~=0.1.2",
     "pandas~=2.2",
     "scipy~=1.14",
     "seaborn~=0.13.2",

src/climatebenchpress/compressor/compressors/abc.py

@@ -167,6 +167,8 @@ def build(
                         dtype=dtypes[var],
                         data_min=data_min[var],
                         data_max=data_max[var],
+                        data_abs_min=data_abs_min[var],
+                        data_abs_max=data_abs_max[var],
                     )
                 elif eb.rel_error is not None and cls.has_rel_error_impl:
                     new_codecs[var] = partial(
@@ -175,6 +177,8 @@ def build(
                         dtype=dtypes[var],
                         data_min=data_min[var],
                         data_max=data_max[var],
+                        data_abs_min=data_abs_min[var],
+                        data_abs_max=data_abs_max[var],
                     )
                 else:
                     # This should never happen as we have already transformed the error bounds.

src/climatebenchpress/compressor/compressors/safeguards/sperr.py

@@ -1,7 +1,10 @@
 __all__ = ["SafeguardsSperr"]
 
+import numcodecs
 import numcodecs_safeguards
 import numcodecs_wasm_sperr
+import numpy as np
+from numcodecs_combinators.stack import CodecStack
 
 from ..abc import Compressor
 
@@ -15,8 +18,38 @@ class SafeguardsSperr(Compressor):
     @staticmethod
     def abs_bound_codec(error_bound, **kwargs):
         return numcodecs_safeguards.SafeguardsCodec(
-            codec=numcodecs_wasm_sperr.Sperr(mode="pwe", pwe=error_bound),
+            codec=CodecStack(
+                NaNToZero(),
+                numcodecs_wasm_sperr.Sperr(mode="pwe", pwe=error_bound),
+            ),
             safeguards=[
                 dict(kind="eb", type="abs", eb=error_bound, equal_nan=True),
             ],
         )
+
+    @staticmethod
+    def rel_bound_codec(error_bound, *, data_abs_min=None, **kwargs):
+        assert data_abs_min is not None, "data_abs_min must be provided"
+
+        return numcodecs_safeguards.SafeguardsCodec(
+            codec=CodecStack(
+                NaNToZero(),
+                # conservative rel->abs error bound transformation,
+                #  same as convert_rel_error_to_abs_error
+                # so that we can inform the safeguards of the rel bound
+                numcodecs_wasm_sperr.Sperr(mode="pwe", pwe=error_bound * data_abs_min),
+            ),
+            safeguards=[
+                dict(kind="eb", type="rel", eb=error_bound, equal_nan=True),
+            ],
+        )
+
+
+class NaNToZero(numcodecs.abc.Codec):
+    codec_id = "nan-to-zero"
+
+    def encode(self, buf):
+        return np.nan_to_num(buf, nan=0, posinf=np.inf, neginf=-np.inf)
+
+    def decode(self, buf, out=None):
+        return numcodecs.compat.ndarray_copy(buf, out)

src/climatebenchpress/compressor/compressors/safeguards/zfp_round.py

@@ -32,3 +32,19 @@ def abs_bound_codec(error_bound, **kwargs):
                 dict(kind="eb", type="abs", eb=error_bound, equal_nan=True),
             ],
         )
+
+    @staticmethod
+    def rel_bound_codec(error_bound, *, data_abs_min=None, **kwargs):
+        assert data_abs_min is not None, "data_abs_min must be provided"
+
+        return numcodecs_safeguards.SafeguardsCodec(
+            # conservative rel->abs error bound transformation,
+            #  same as convert_rel_error_to_abs_error
+            # so that we can inform the safeguards of the rel bound
+            codec=numcodecs_wasm_zfp.Zfp(
+                mode="fixed-accuracy", tolerance=error_bound * data_abs_min
+            ),
+            safeguards=[
+                dict(kind="eb", type="rel", eb=error_bound, equal_nan=True),
+            ],
+        )