fix Safeguarded(0, dSSIM) for multiple 2D slices

Author: juntyr

src/climatebenchpress/compressor/compressors/abc.py

@@ -90,6 +90,8 @@ def abs_bound_codec(
         data_max: Optional[float] = None,
         data_abs_min: Optional[float] = None,
         data_abs_max: Optional[float] = None,
+        data_min_2d: Optional[np.ndarray] = None,
+        data_max_2d: Optional[np.ndarray] = None,
     ) -> Codec:
         """Create a codec with an absolute error bound."""
         pass
@@ -104,6 +106,8 @@ def rel_bound_codec(
         data_max: Optional[float] = None,
         data_abs_min: Optional[float] = None,
         data_abs_max: Optional[float] = None,
+        data_min_2d: Optional[np.ndarray] = None,
+        data_max_2d: Optional[np.ndarray] = None,
     ) -> Codec:
         """Create a codec with a relative error bound."""
         pass
@@ -116,6 +120,8 @@ def build(
         data_abs_max: dict[VariableName, float],
         data_min: dict[VariableName, float],
         data_max: dict[VariableName, float],
+        data_min_2d: dict[VariableName, np.ndarray],
+        data_max_2d: dict[VariableName, np.ndarray],
         error_bounds: list[dict[VariableName, ErrorBound]],
     ) -> dict[VariantName, list[NamedPerVariableCodec]]:
         """
@@ -139,6 +145,12 @@ def build(
             Dict mapping from variable name to minimum value for the variable.
         data_max : dict[VariableName, float]
             Dict mapping from variable name to maximum value for the variable.
+        data_min_2d : dict[VariableName, np.ndarray]
+            Dict mapping from variable name to per-lat-lon-slice minimum value for the
+            variable.
+        data_max_2d : dict[VariableName, np.ndarray]
+            Dict mapping from variable name to per-lat-lon-slice maximum value for the
+            variable.
         error_bounds: list[ErrorBound]
             List of error bounds to use for the compressor.
 
@@ -173,6 +185,8 @@ def build(
                         data_max=data_max[var],
                         data_abs_min=data_abs_min[var],
                         data_abs_max=data_abs_max[var],
+                        data_min_2d=data_min_2d[var],
+                        data_max_2d=data_max_2d[var],
                     )
                 elif eb.rel_error is not None and cls.has_rel_error_impl:
                     new_codecs[var] = partial(
@@ -183,6 +197,8 @@ def build(
                         data_max=data_max[var],
                         data_abs_min=data_abs_min[var],
                         data_abs_max=data_abs_max[var],
+                        data_min_2d=data_min_2d[var],
+                        data_max_2d=data_max_2d[var],
                     )
                 else:
                     # This should never happen as we have already transformed the error bounds.

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

@@ -13,9 +13,9 @@ class SafeguardedZeroDssim(Compressor):
     description = "Safeguarded(0, dSSIM)"
 
     @staticmethod
-    def abs_bound_codec(error_bound, data_min=None, data_max=None, **kwargs):
-        assert data_min is not None, "data_min must be provided"
-        assert data_max is not None, "data_max must be provided"
+    def abs_bound_codec(error_bound, data_min_2d=None, data_max_2d=None, **kwargs):
+        assert data_min_2d is not None, "data_min_2d must be provided"
+        assert data_max_2d is not None, "data_max_2d must be provided"
 
         return numcodecs_safeguards.SafeguardedCodec(
             codec=numcodecs_zero.ZeroCodec(),
@@ -46,14 +46,15 @@ def abs_bound_codec(error_bound, data_min=None, data_max=None, **kwargs):
                     eb=0,
                 ),
             ],
-            # use data_min instead of $x_min to allow for chunking
-            fixed_constants=dict(x_min=data_min, x_max=data_max),
+            # use data_min_2d instead of $x_min since we need the minimum per
+            #  2d latitude-longitude slice
+            fixed_constants=dict(x_min=data_min_2d, x_max=data_max_2d),
         )
 
     @staticmethod
-    def rel_bound_codec(error_bound, data_min=None, data_max=None, **kwargs):
-        assert data_min is not None, "data_min must be provided"
-        assert data_max is not None, "data_max must be provided"
+    def rel_bound_codec(error_bound, data_min_2d=None, data_max_2d=None, **kwargs):
+        assert data_min_2d is not None, "data_min_2d must be provided"
+        assert data_max_2d is not None, "data_max_2d must be provided"
 
         return numcodecs_safeguards.SafeguardedCodec(
             codec=numcodecs_zero.ZeroCodec(),
@@ -84,6 +85,7 @@ def rel_bound_codec(error_bound, data_min=None, data_max=None, **kwargs):
                     eb=0,
                 ),
             ],
-            # use data_min instead of $x_min to allow for chunking
-            fixed_constants=dict(x_min=data_min, x_max=data_max),
+            # use data_min_2d instead of $x_min since we need the minimum per
+            #  2d latitude-longitude slice
+            fixed_constants=dict(x_min=data_min_2d, x_max=data_max_2d),
         )

src/climatebenchpress/compressor/plotting/plot_metrics.py

@@ -207,8 +207,9 @@ def _normalize(data):
 
         # Normalize each variable by its mean and std
         normalized[new_col] = normalized.apply(
-            lambda x: (x[col] - mean_std[x["Variable"]][0])
-            / mean_std[x["Variable"]][1],
+            lambda x: (
+                (x[col] - mean_std[x["Variable"]][0]) / mean_std[x["Variable"]][1]
+            ),
             axis=1,
         )
 

src/climatebenchpress/compressor/scripts/compress.py

@@ -87,6 +87,8 @@ def compress(
         ds_abs_maxs: dict[str, float] = dict()
         ds_mins: dict[str, float] = dict()
         ds_maxs: dict[str, float] = dict()
+        ds_min_2ds: dict[str, np.ndarray] = dict()
+        ds_max_2ds: dict[str, np.ndarray] = dict()
         for v in ds:
             vs: str = str(v)
             abs_vals = xr.ufuncs.abs(ds[v])
@@ -96,6 +98,16 @@ def compress(
             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()
+            ds_min_2ds[vs] = (
+                ds[v]
+                .min(dim=[ds[v].cf["Y"].name, ds[v].cf["X"].name], keepdims=True)
+                .values
+            )
+            ds_max_2ds[vs] = (
+                ds[v]
+                .max(dim=[ds[v].cf["Y"].name, ds[v].cf["X"].name], keepdims=True)
+                .values
+            )
 
         if chunked:
             for v in ds:
@@ -115,7 +127,14 @@ def compress(
 
             compressor_variants: dict[str, list[NamedPerVariableCodec]] = (
                 compressor.build(
-                    ds_dtypes, ds_abs_mins, ds_abs_maxs, ds_mins, ds_maxs, error_bounds
+                    ds_dtypes,
+                    ds_abs_mins,
+                    ds_abs_maxs,
+                    ds_mins,
+                    ds_maxs,
+                    ds_min_2ds,
+                    ds_max_2ds,
+                    error_bounds,
                 )
             )
 
@@ -189,6 +208,15 @@ def compress_decompress(
         if not isinstance(codec, CodecStack):
             codec = CodecStack(codec)
 
+        # HACK: Safeguarded(0, dSSIM) requires the per-lat-lon-slice minimum
+        #  and maximum
+        # for potentially-chunked data we should really use xarray-safeguards,
+        #  but not using chunks also works (for now)
+        is_safeguarded_zero_dssim = (
+            "# === pointwise dSSIM quantity of interest === #"
+            in json.dumps(codec.get_config())
+        )
+
         with numcodecs_observers.observe(
             codec,
             observers=[
@@ -197,7 +225,9 @@ def compress_decompress(
                 timing,
             ],
         ) as codec_:
-            variables[v] = codec_.encode_decode_data_array(ds[v]).compute()
+            variables[v] = codec_.encode_decode_data_array(
+                ds[v].compute() if is_safeguarded_zero_dssim else ds[v]
+            ).compute()
 
             cs = [c._codec for c in codec_.__iter__()]