Cleanup with numcodecs-replace (#67)

Author: juntyr

pyproject.toml

@@ -15,21 +15,22 @@ dependencies = [
     "numcodecs>=0.13.0,<0.17",
     "numcodecs-combinators[xarray]~=0.2.13",
     "numcodecs-observers~=0.1.2",
+    "numcodecs-replace==0.1.0",
     "numcodecs-safeguards==0.1.0b2",
-    "numcodecs-wasm~=0.2.2",
-    "numcodecs-wasm-bit-round~=0.4.0",
-    "numcodecs-wasm-fixed-offset-scale~=0.4.0",
-    "numcodecs-wasm-jpeg2000~=0.3.0",
-    "numcodecs-wasm-pco~=0.3.0",
-    "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-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.2",
+    "numcodecs-wasm==0.2.2",
+    "numcodecs-wasm-bit-round==0.4.0",
+    "numcodecs-wasm-fixed-offset-scale==0.4.0",
+    "numcodecs-wasm-jpeg2000==0.3.0",
+    "numcodecs-wasm-pco==0.3.0",
+    "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-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.2",
     "pandas~=2.2",
     "scipy~=1.14",
     "seaborn~=0.13.2",

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

@@ -1,8 +1,6 @@
 __all__ = ["SafeguardedSperr"]
 
-import numcodecs
-import numcodecs.abc
-import numcodecs.compat
+import numcodecs_replace
 import numcodecs_safeguards
 import numcodecs_wasm_sperr
 import numpy as np
@@ -20,8 +18,11 @@ class SafeguardedSperr(Compressor):
     @staticmethod
     def abs_bound_codec(error_bound, **kwargs):
         return numcodecs_safeguards.SafeguardedCodec(
+            # inspired by H5Z-SPERR's treatment of NaN values:
+            # https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5z-sperr.c#L464-L473
+            # https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5zsperr_helper.cpp#L179-L212
             codec=CodecStack(
-                NaNToMean(),
+                numcodecs_replace.ReplaceFilterCodec(replacements={np.nan: "nan_mean"}),
                 numcodecs_wasm_sperr.Sperr(mode="pwe", pwe=error_bound),
             ),
             safeguards=[
@@ -35,7 +36,7 @@ def rel_bound_codec(error_bound, *, data_abs_min=None, **kwargs):
 
         return numcodecs_safeguards.SafeguardedCodec(
             codec=CodecStack(
-                NaNToMean(),
+                numcodecs_replace.ReplaceFilterCodec(replacements={np.nan: "nan_mean"}),
                 # 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
@@ -45,16 +46,3 @@ def rel_bound_codec(error_bound, *, data_abs_min=None, **kwargs):
                 dict(kind="eb", type="rel", eb=error_bound, equal_nan=True),
             ],
         )
-
-
-# inspired by H5Z-SPERR's treatment of NaN values:
-# https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5z-sperr.c#L464-L473
-# https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5zsperr_helper.cpp#L179-L212
-class NaNToMean(numcodecs.abc.Codec):
-    codec_id = "nan-to-mean"  # type: ignore
-
-    def encode(self, buf):
-        return np.nan_to_num(buf, nan=np.nanmean(buf), posinf=np.inf, neginf=-np.inf)
-
-    def decode(self, buf, out=None):
-        return numcodecs.compat.ndarray_copy(buf, out)