Update numcodecs-wasm, test NaNs in ZFP (#57)

* Update numcodecs-wasm, test NaNs in ZFP

* fix typing

* Fix minor bugs

---------

Co-authored-by: Tim Reichelt <[email protected]>

Author: juntyr

pyproject.toml

@@ -13,21 +13,21 @@ dependencies = [
     "matplotlib~=3.8",
     "netcdf4~=1.7.2",
     "numcodecs>=0.13.0,<0.17",
-    "numcodecs-combinators[xarray]~=0.2.4",
-    "numcodecs-observers~=0.1.1",
-    "numcodecs-wasm~=0.1.7",
-    "numcodecs-wasm-bit-round~=0.3.0",
-    "numcodecs-wasm-fixed-offset-scale~=0.3.0",
-    "numcodecs-wasm-jpeg2000~=0.2.0",
-    "numcodecs-wasm-pco~=0.2.1",
-    "numcodecs-wasm-round~=0.4.0",
-    "numcodecs-wasm-sperr~=0.1.0",
-    "numcodecs-wasm-stochastic-rounding~=0.1.1",
-    "numcodecs-wasm-sz3~=0.6.0",
-    "numcodecs-wasm-tthresh~=0.2.0",
-    "numcodecs-wasm-zfp~=0.5.3",
-    "numcodecs-wasm-zfp-classic~=0.3.3",
-    "numcodecs-wasm-zstd~=0.3.0",
+    "numcodecs-combinators[xarray]~=0.2.10",
+    "numcodecs-observers~=0.1.2",
+    "numcodecs-wasm~=0.2.1",
+    "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",
     "pandas~=2.2",
     "scipy~=1.14",
     "seaborn~=0.13.2",

src/climatebenchpress/compressor/compressors/abc.py

@@ -196,7 +196,7 @@ def build(
 
     # Class interface
     @classproperty
-    def registry(cls) -> Mapping:
+    def registry(cls) -> Mapping[str, type["Compressor"]]:
         return MappingProxyType(Compressor._registry)
 
     # Implementation details
@@ -247,11 +247,13 @@ def _get_variant_bounds(
         converted_bounds: dict[VariableName, dict[VariantName, ErrorBound]] = dict()
         variant_names = {cls.name}
         for var, error_bound in error_bounds.items():
+            cls_has_abs_error_impl: bool = cls.has_abs_error_impl  # type: ignore
             abs_bound_codec = (
-                error_bound.abs_error is not None and cls.has_abs_error_impl
+                error_bound.abs_error is not None and cls_has_abs_error_impl
             )
+            cls_has_rel_error_impl: bool = cls.has_rel_error_impl  # type: ignore
             rel_bound_codec = (
-                error_bound.rel_error is not None and cls.has_rel_error_impl
+                error_bound.rel_error is not None and cls_has_rel_error_impl
             )
             if abs_bound_codec or rel_bound_codec:
                 # If codec is compatible with the error bound no transformation

src/climatebenchpress/compressor/compressors/zfp.py

@@ -22,5 +22,5 @@ class Zfp(Compressor):
     @staticmethod
     def abs_bound_codec(error_bound, **kwargs):
         return numcodecs_wasm_zfp_classic.ZfpClassic(
-            mode="fixed-accuracy", tolerance=error_bound
+            mode="fixed-accuracy", tolerance=error_bound, non_finite="allow-unsafe"
         )

src/climatebenchpress/compressor/compressors/zfp_round.py

@@ -25,4 +25,6 @@ class ZfpRound(Compressor):
 
     @staticmethod
     def abs_bound_codec(error_bound, **kwargs):
-        return numcodecs_wasm_zfp.Zfp(mode="fixed-accuracy", tolerance=error_bound)
+        return numcodecs_wasm_zfp.Zfp(
+            mode="fixed-accuracy", tolerance=error_bound, non_finite="allow-unsafe"
+        )

src/climatebenchpress/compressor/scripts/compress.py

@@ -3,11 +3,12 @@
 import argparse
 import json
 import traceback
-from collections.abc import Container
+from collections.abc import Container, Mapping
 from pathlib import Path
 from typing import Callable
 
 import numcodecs_observers
+import numpy as np
 import xarray as xr
 from numcodecs.abc import Codec
 from numcodecs_combinators.stack import CodecStack
@@ -68,24 +69,24 @@ def compress(
             continue
 
         ds = xr.open_dataset(dataset, chunks=dict(), engine="zarr")
-        ds_dtypes, ds_abs_mins, ds_abs_maxs, ds_mins, ds_maxs = (
-            dict(),
-            dict(),
-            dict(),
-            dict(),
-            dict(),
-        )
+        ds_dtypes: dict[str, np.dtype] = dict()
+        ds_abs_mins: dict[str, float] = dict()
+        ds_abs_maxs: dict[str, float] = dict()
+        ds_mins: dict[str, float] = dict()
+        ds_maxs: dict[str, float] = dict()
         for v in ds:
+            vs: str = str(v)
             abs_vals = xr.ufuncs.abs(ds[v])
+            ds_dtypes[vs] = ds[v].dtype
             # Take minimum of non-zero absolute values to avoid division by zero.
-            ds_abs_mins[v] = abs_vals.where(abs_vals > 0).min().values.item()
-            ds_abs_maxs[v] = abs_vals.max().values.item()
-            ds_mins[v] = ds[v].min().values.item()
-            ds_maxs[v] = ds[v].max().values.item()
-            ds_dtypes[v] = ds[v].dtype
+            ds_abs_mins[vs] = abs_vals.where(abs_vals > 0).min().values.item()
+            ds_abs_maxs[vs] = abs_vals.max().values.item()
+            ds_mins[vs] = ds[v].min().values.item()
+            ds_maxs[vs] = ds[v].max().values.item()
 
         error_bounds = get_error_bounds(datasets_error_bounds, dataset.parent.name)
-        for compressor in Compressor.registry.values():
+        registry: Mapping[str, type[Compressor]] = Compressor.registry  # type: ignore
+        for compressor in registry.values():
             if compressor.name in exclude_compressor:
                 continue
             if include_compressor and compressor.name not in include_compressor:
@@ -216,7 +217,7 @@ def get_error_bounds(
     args = parser.parse_args()
 
     compress(
-        basepath=Path(),
+        basepath=args.basepath,
         exclude_dataset=args.exclude_dataset,
         include_dataset=args.include_dataset,
         exclude_compressor=args.exclude_compressor,

src/climatebenchpress/compressor/scripts/concatenate_metrics.py

@@ -189,4 +189,3 @@ def get_error_bound_name(
     args = parser.parse_args()
 
     concatenate_metrics(basepath=args.basepath)
-    concatenate_metrics(basepath=args.basepath)