Add stochround-pco compressor

compressed-datasets/.gitignore

@@ -1,2 +1,3 @@
 /*/*/*/decompressed.zarr
 /*/*/*/measurements.json
+/*/*/*/error.out

plots/.gitignore

@@ -1,3 +1,3 @@
-/*.png
-/*/*.png
+/*.pdf
+/*/*.pdf
 /*/*/*.png

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "numcodecs-wasm-uniform-noise~=0.3.0",
     "numcodecs-wasm-zfp~=0.5.1",
     "numcodecs-wasm-zfp-classic~=0.3.1",
-    "numcodecs-wasm-zlib~=0.3.0",
+    "numcodecs-wasm-zstd~=0.3.0",
     "pandas~=2.2",
     "scipy~=1.14",
     "seaborn~=0.13.2",

src/climatebenchpress/compressor/compressors/__init__.py

@@ -3,6 +3,7 @@
     "BitRoundPco",
     "Jpeg2000",
     "StochRound",
+    "StochRoundPco",
     "Sz3",
     "Tthresh",
     "Zfp",
@@ -14,6 +15,7 @@
 from .bitround_pco import BitRoundPco
 from .jpeg2000 import Jpeg2000
 from .stochround import StochRound
+from .stochround_pco import StochRoundPco
 from .sz3 import Sz3
 from .tthresh import Tthresh
 from .zfp import Zfp

src/climatebenchpress/compressor/compressors/bitround.py

@@ -1,7 +1,7 @@
 __all__ = ["BitRound"]
 
 import numcodecs_wasm_bit_round
-import numcodecs_wasm_zlib
+import numcodecs_wasm_zstd
 from numcodecs_combinators.stack import CodecStack
 
 from .abc import Compressor
@@ -17,5 +17,5 @@ def rel_bound_codec(dtype, error_bound):
         keepbits = compute_keepbits(dtype, error_bound)
         return CodecStack(
             numcodecs_wasm_bit_round.BitRound(keepbits=keepbits),
-            numcodecs_wasm_zlib.Zlib(level=6),
+            numcodecs_wasm_zstd.Zstd(level=3),
         )

src/climatebenchpress/compressor/compressors/stochround.py

@@ -2,7 +2,7 @@
 
 import numcodecs_wasm_round
 import numcodecs_wasm_uniform_noise
-import numcodecs_wasm_zlib
+import numcodecs_wasm_zstd
 from numcodecs_combinators.stack import CodecStack
 
 from .abc import Compressor
@@ -18,5 +18,5 @@ def abs_bound_codec(dtype, error_bound):
         return CodecStack(
             numcodecs_wasm_uniform_noise.UniformNoise(scale=precision / 2, seed=42),
             numcodecs_wasm_round.Round(precision=precision),
-            numcodecs_wasm_zlib.Zlib(level=6),
+            numcodecs_wasm_zstd.Zstd(level=3),
         )

src/climatebenchpress/compressor/compressors/stochround_pco.py

@@ -0,0 +1,27 @@
+__all__ = ["StochRoundPco"]
+
+import numcodecs_wasm_pco
+import numcodecs_wasm_round
+import numcodecs_wasm_uniform_noise
+from numcodecs_combinators.stack import CodecStack
+
+from .abc import Compressor
+
+
+class StochRoundPco(Compressor):
+    name = "stochround-pco"
+    description = "Stochastic Rounding + PCodec"
+
+    @staticmethod
+    def abs_bound_codec(dtype, error_bound):
+        precision = error_bound
+        return CodecStack(
+            numcodecs_wasm_uniform_noise.UniformNoise(scale=precision / 2, seed=42),
+            numcodecs_wasm_round.Round(precision=precision),
+            numcodecs_wasm_pco.Pco(
+                level=8,
+                mode="auto",
+                delta="auto",
+                paging="equal-pages-up-to",
+            ),
+        )

src/climatebenchpress/compressor/plotting/plot_metrics.py

@@ -11,20 +11,24 @@
 COMPRESSOR2LINEINFO = {
     "jpeg2000": ("#EE7733", "-"),
     "zfp": ("#EE3377", "--"),
+    "zfp-round": ("#DDAA33", "--"),
     "sz3": ("#CC3311", "-."),
     "bitround-pco-conservative-rel": ("#0077BB", ":"),
     "bitround-conservative-rel": ("#33BBEE", "-"),
     "stochround": ("#009988", "--"),
-    "tthresh": ("#BBBBBB", "-."),
+    "stochround-pco": ("#BBBBBB", "--"),
+    "tthresh": ("#000000", "-."),
 }
 
 COMPRESSOR2LEGEND_NAME = {
     "jpeg2000": "JPEG2000",
     "zfp": "ZFP",
+    "zfp-round": "ZFP-ROUND",
     "sz3": "SZ3",
     "bitround-pco-conservative-rel": "BitRound + PCO",
     "bitround-conservative-rel": "BitRound + Zlib",
-    "stochround": "StochRound",
+    "stochround": "StochRound + Zlib",
+    "stochround-pco": "StochRound + PCO",
     "tthresh": "TTHRESH",
 }