Preserving the structural similarity index for floating-point data (dSSIM) with safeguards¶

import ssl

ssl._create_default_https_context = ssl._create_stdlib_context

from collections import defaultdict
from pathlib import Path

import earthkit.plots
import humanize
import matplotlib as mpl
import numpy as np
import pandas as pd
import xarray as xr
from matplotlib import patheffects as PathEffects

# Retrieve the data
ERA5 = xr.open_dataset(Path() / "data" / "era5-uv" / "data.nc")
ERA5_U = (
    ERA5["u"]
    .sel(valid_time="2024-04-02T12:00:00", pressure_level=500)
    .astype(np.float64)
)

def dssim_mat(
    a1: np.ndarray,
    a2: np.ndarray,
    eps: float = 1e-8,
    kernel_size: int = 3,
    reproducible: bool = True,
) -> np.ndarray:
    """
    Implementation adapted from the official dSSIM implementation at
    https://github.com/NCAR/ldcpy/blob/6c5bcb8149ec7876a4f53b0e784e9c528f6f14cb/ldcpy/calcs.py#L2516

    The official implementation makes assumptions about the input data that are
    specific to models developed at NCAR which is why we cannot use the official
    implementation directly.

    The implementation has been further adapted to ensure bit-reproducible
    evaluation with the below dSSIM quantity of interest, with the following
    notable changes:

    - we take extra care to ensure all operations are computed in the data
      dtype, i.e. without accidentally upcasting to np.float64
    - the original uses np.round, which rounds halfway cases away from zero,
      we use np.rint which rounds halfway cases to the nearest even integer
    - the original uses astropy.convolution.convolve, which internally upcasts
      to float64 - to avoid rounding differences, we manually implement it here
    - astropy.convolution.convolve handles NaN values by replacing them during
      convolution with a Gaussian kernel by interpolating with the same
      Gaussian kernel, to avoid NaNs bleeding over, and by returning NaN for
      elements that were NaN in the input. We implement a simpler version by
      excluding NaN values from the convolution and returning NaN for elements
      that were NaN in the input
    - convolution requires summation but np.sum internally uses partial sums
      for better rounding, which the safeguards currently do not use, so we
      manually sum
    - we reduce the 11x11 Gaussian kernel to 3x3 for the sake of compute time

    Parameters
    ----------
    x : np.ndarray
        Shape: (latitude, longitude)
    y : np.ndarray
        Shape: (latitude, longitude)
    kernel_size : int
        The size of the Gaussian kernel for the convolution operation in SSIM. Has to be
        an odd number. The default is 3.

    Returns
    -------
    dssim : np.ndarray
        The pointwise data-SSIM values where the k x k kernel is valid.
    """
    assert kernel_size % 2 == 1, "kernel_size must be an odd number."
    assert a1.shape == a2.shape
    assert a1.dtype == a2.dtype

    # re-scale  to [0,1] - if not constant
    smin = min(np.nanmin(a1), np.nanmin(a2))
    smax = max(np.nanmax(a1), np.nanmax(a2))
    r = smax - smin
    if r == 0:  # scale by smax if field is a constant (and smax != 0)
        if smax == 0:
            sc_a1 = a1
            sc_a2 = a2
        else:
            sc_a1 = a1 / smax
            sc_a2 = a2 / smax
    else:
        sc_a1 = (a1 - smin) / r
        sc_a2 = (a2 - smin) / r

    # now quantize to 256 bins
    sc_a1 = np.rint(sc_a1 * 255) / 255
    sc_a2 = np.rint(sc_a2 * 255) / 255

    # 2D gaussian filter
    sigma = a1.dtype.type(1.5)
    pi = a1.dtype.type(np.pi)

    i = np.arange(kernel_size).astype(a1.dtype) - (kernel_size // 2)
    k = 1 / (np.sqrt(2 * pi) * sigma) * np.exp((-np.square(i)) / (2 * np.square(sigma)))
    kernel = np.array([k]).T @ np.array([k])

    def convolve(a, k):
        # pad with zeros
        ap = np.pad(a, kernel_size // 2, mode="constant", constant_values=0)
        # reshape a to (*a.shape, *k.shape)
        av = np.lib.stride_tricks.as_strided(
            ap, a.shape + k.shape, ap.strides + ap.strides
        )

        # broadcast k to (*a.shape, *k.shape)
        kv = np.copy(
            np.broadcast_to(k.reshape((1, 1, kernel_size, kernel_size)), av.shape)
        )

        # exclude weights in kv where av is NaN
        kv[np.isnan(av)] = 0
        # renormalize kv
        kv_flat = kv.reshape(a.shape + (-1,))
        if reproducible:
            # manual np.sum over the last dimension
            kv_acc = np.copy(kv_flat[..., 0])
            for i in range(1, kv_flat.shape[-1]):
                kv_acc += kv_flat[..., i]
        else:
            kv_acc = np.sum(kv_flat, axis=-1)
        kv /= kv_acc.reshape(a.shape + (1, 1))

        # pointwise multiply av with kv, then reshape to (*a.shape, -1)
        ag = (av * kv).reshape(a.shape + (-1,))
        if reproducible:
            # manual np.sum over the last dimension
            acc = np.copy(ag[..., 0])
            for i in range(1, ag.shape[-1]):
                acc += ag[..., i]
        else:
            acc = np.sum(ag, axis=-1)

        # keep NaN for elements that are NaN in a
        return np.where(np.isnan(a), np.nan, acc)

    a1_mu = convolve(sc_a1, kernel)
    a2_mu = convolve(sc_a2, kernel)

    a1a1 = convolve(np.square(sc_a1), kernel)
    a2a2 = convolve(np.square(sc_a2), kernel)

    a1a2 = convolve(sc_a1 * sc_a2, kernel)

    ###########
    var_a1 = a1a1 - np.square(a1_mu)
    var_a2 = a2a2 - np.square(a2_mu)
    cov_a1a2 = a1a2 - a1_mu * a2_mu

    # ssim constants
    C1 = a1.dtype.type(eps)
    C2 = a1.dtype.type(eps)

    ssim_t1 = 2 * a1_mu * a2_mu + C1
    ssim_t2 = 2 * cov_a1a2 + C2

    ssim_b1 = np.square(a1_mu) + np.square(a2_mu) + C1
    ssim_b2 = var_a1 + var_a2 + C2

    ssim_1 = ssim_t1 / ssim_b1
    ssim_2 = ssim_t2 / ssim_b2
    ssim_mat = ssim_1 * ssim_2

    # Cropping the border region of the 2D field where the convolution kernel is not
    # fully overlapping with the 2D input field.
    k = (kernel_size - 1) // 2
    return ssim_mat[k : ssim_mat.shape[0] - k, k : ssim_mat.shape[1] - k]


def dssim(
    a1: np.ndarray,
    a2: np.ndarray,
    eps: float = 1e-8,
    kernel_size: int = 3,
    reproducible: bool = True,
) -> float:
    """
    Implementation adapted from the official dSSIM implementation at
    https://github.com/NCAR/ldcpy/blob/6c5bcb8149ec7876a4f53b0e784e9c528f6f14cb/ldcpy/calcs.py#L2516

    The official implementation makes assumptions about the input data that are
    specific to models developed at NCAR which is why we cannot use the official
    implementation directly.

    Parameters
    ----------
    x : np.ndarray
        Shape: (latitude, longitude)
    y : np.ndarray
        Shape: (latitude, longitude)
    kernel_size : int
        The size of the Gaussian kernel for the convolution operation in SSIM. Has to be
        an odd number. The default is 3.

    Returns
    -------
    dssim : float
        The data-SSIM value between the two input arrays.
    """
    return np.nanmean(
        dssim_mat(a1, a2, eps=eps, kernel_size=kernel_size, reproducible=reproducible)
    )

old_cmap_and_norm = earthkit.plots.styles.colors.cmap_and_norm

def my_cmap_and_norm(colors, levels, normalize=True, extend=None, extend_levels=True):
    return old_cmap_and_norm(colors, levels, normalize, extend, True)


earthkit.plots.styles.colors.cmap_and_norm = my_cmap_and_norm

def compute_corrections_percentage(
    my_ERA5_U: xr.DataArray, orig_ERA5_U: xr.DataArray
) -> float:
    return np.mean(my_ERA5_U != orig_ERA5_U)

def plot_u_wind_dssim(
    my_ERA5_U: xr.DataArray,
    cr,
    chart,
    title,
    span,
    dssim_bound,
    error=False,
    corr=None,
    my_ERA5_U_it=None,
    cr_it=None,
    inset=True,
    corr_colours=["white", "green", "lawngreen"],
):
    import copy

    if error:
        with xr.set_options(keep_attrs=True):
            da = ERA5_U[1:-1, 1:-1].copy(
                data=dssim_mat(
                    ERA5_U.values,
                    my_ERA5_U.values,
                    reproducible=False,
                )
            )
            da.attrs.update(long_name=f"dSSIM({da.long_name.lower()})", units=None)
    else:
        da = my_ERA5_U

    # compute the default style that earthkit.maps would apply
    source = earthkit.plots.sources.XarraySource(da)
    style = copy.deepcopy(
        earthkit.plots.styles.auto.guess_style(
            source,
            units=source.units,
        )
    )

    style._levels = earthkit.plots.styles.levels.Levels(np.linspace(*span, 22))
    style._legend_kwargs["ticks"] = np.linspace(*span, 5)

    if error:
        style._colors = "cool_r"

    extend_left = np.nanmin(da) < span[0]
    extend_right = np.nanmax(da) > span[1]

    extend = {
        (False, False): "neither",
        (True, False): "min",
        (False, True): "max",
        (True, True): "both",
    }[(extend_left, extend_right)]

    if error:
        style._legend_kwargs["extend"] = extend
        chart.pcolormesh(da, style=style, zorder=-11)

        if corr is not None:
            da_hatch = my_ERA5_U == corr

            if my_ERA5_U_it is None:
                da_corr = da_hatch.astype(float)
            else:
                with xr.set_options(keep_attrs=True):
                    da_hatch_it = my_ERA5_U_it == corr
                da_corr = (~da_hatch).astype(float) + (~da_hatch_it).astype(float)

            old_process_projection_requirements = (
                chart.ax.get_figure()._process_projection_requirements
            )

            def _process_projection_requirements(
                *, axes_class=None, polar=False, projection=None, **kwargs
            ):
                if axes_class is not None and projection is not None:
                    return axes_class, dict(projection=projection, **kwargs)
                return old_process_projection_requirements(
                    axes_class=axes_class, polar=polar, projection=projection, **kwargs
                )

            chart.ax.get_figure()._process_projection_requirements = (
                _process_projection_requirements
            )

            axin = chart.ax.inset_axes(
                [0.025, 0.05, 1 / 3, 1 / 3],
                xticklabels=[],
                yticklabels=[],
                axes_class=type(chart.ax),
                projection=chart.ax.projection,
            )
            axin.pcolormesh(
                da_corr.longitude.values,
                da_corr.latitude.values,
                np.squeeze(da_corr.values),
                cmap=mpl.colors.ListedColormap(corr_colours),
                vmin=0,
                vmax=2,
                rasterized=True,
            )
            axin.coastlines(color="#555555")
            axin.spines["geo"].set_edgecolor("black")
            axin.set_title(
                "Corrections",
                path_effects=[PathEffects.withStroke(linewidth=3, foreground="white")],
            )
        elif inset:
            da_err = ~(da >= dssim_bound)

            old_process_projection_requirements = (
                chart.ax.get_figure()._process_projection_requirements
            )

            def _process_projection_requirements(
                *, axes_class=None, polar=False, projection=None, **kwargs
            ):
                if axes_class is not None and projection is not None:
                    return axes_class, dict(projection=projection, **kwargs)
                return old_process_projection_requirements(
                    axes_class=axes_class, polar=polar, projection=projection, **kwargs
                )

            chart.ax.get_figure()._process_projection_requirements = (
                _process_projection_requirements
            )

            axin = chart.ax.inset_axes(
                [0.025, 0.05, 1 / 3, 1 / 3],
                xticklabels=[],
                yticklabels=[],
                axes_class=type(chart.ax),
                projection=chart.ax.projection,
            )
            axin.pcolormesh(
                da_err.longitude.values,
                da_err.latitude.values,
                np.squeeze(da_err.values),
                cmap=mpl.colors.ListedColormap(["white", "darkorange"]),
                vmin=0,
                vmax=1,
                rasterized=True,
            )
            axin.coastlines(color="#555555")
            axin.spines["geo"].set_edgecolor("black")
            axin.set_title(
                f"dSSIM < {dssim_bound}",
                path_effects=[PathEffects.withStroke(linewidth=3, foreground="white")],
            )
    else:
        chart.quickplot(da, style=style, extend=extend, zorder=-11)

    chart.ax.set_rasterization_zorder(-10)

    chart.title(title)

    if error:
        dssim_mean = np.nanmean(da)
        err_s = ~(dssim_mean >= dssim_bound)
        err_s = r"$\times$" if err_s else r"$\checkmark$"

        t = chart.ax.text(
            0.95,
            0.1,
            f"E[dSSIM]={dssim_mean:.05} ({err_s})",
            ha="right",
            va="bottom",
            transform=chart.ax.transAxes,
        )
        t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

    t = chart.ax.text(
        0.95,
        0.9,
        (
            rf"$\times$ {np.round(cr, 2)}"
            + ("" if cr_it is None else rf" ($\times$ {np.round(cr_it, 2)})")
        )
        if error
        else humanize.naturalsize(ERA5_U.nbytes, binary=True),
        ha="right",
        va="top",
        transform=chart.ax.transAxes,
    )
    t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

    for m in earthkit.plots.schemas.schema.quickmap_subplot_workflow:
        if m != "title":
            getattr(chart, m)()

    for m in earthkit.plots.schemas.schema.quickmap_figure_workflow:
        getattr(chart, m)()

    counts, bins = np.histogram(da.values.flatten(), range=span, bins=21)
    midpoints = bins[:-1] + np.diff(bins) / 2
    cb = chart.ax.collections[0].colorbar
    if error:
        if extend_left:
            cb._extend_patches[0].set_hatch("xx")
            cb._extend_patches[0].set_ec("white")
        cb.ax.fill_between(
            [span[0], dssim_bound],
            *cb.ax.get_ylim(),
            hatch="xx",
            ec="w",
            fc="none",
            lw=0,
        )
    extend_width = (bins[-1] - bins[-2]) / (bins[-1] - bins[0])
    cax = cb.ax.inset_axes(
        [
            0.0 - extend_width * extend_left,
            1.25,
            1.0 + extend_width * (0 + extend_left + extend_right),
            1.0,
        ]
    )
    cax.bar(
        midpoints,
        height=counts,
        width=(bins[-1] - bins[0]) / len(counts),
        color=cb.cmap(cb.norm(midpoints)),
        **(
            dict(
                hatch=["xx" if np.abs(m) < dssim_bound else "" for m in midpoints],
                ec="white",
                lw=0,
            )
            if error
            else dict()
        ),
    )
    if extend_left:
        cax.bar(
            bins[0] - (bins[1] - bins[0]) / 2,
            height=np.sum(da < span[0]),
            width=(bins[-1] - bins[0]) / len(counts),
            color=cb.cmap(cb.norm(midpoints[0])),
            **(
                dict(
                    hatch="xx",
                    ec="white",
                    lw=0,
                )
                if error and span[0] < dssim_bound
                else dict()
            ),
        )
    if extend_right:
        cax.bar(
            bins[-1] + (bins[-1] - bins[-2]) / 2,
            height=np.sum(da > span[1]),
            width=(bins[-1] - bins[0]) / len(counts),
            color=cb.cmap(cb.norm(midpoints[-1])),
            **(
                dict(
                    hatch="xx",
                    ec="white",
                    lw=0,
                )
                if error and span[1] < dssim_bound
                else dict()
            ),
        )
    q1, q2, q3 = da.quantile([0.25, 0.5, 0.75]).values
    cax.axvline(da.mean().item(), ls=":", ymin=0.1, ymax=0.9, c="w", lw=2)
    cax.axvline(q1, ymin=0.25, ymax=0.75, c="w", lw=2)
    cax.axvline(q2, ymin=0.1, ymax=0.9, c="w", lw=2)
    cax.axvline(q3, ymin=0.25, ymax=0.75, c="w", lw=2)
    cax.axvline(da.mean().item(), ymin=0.1, ymax=0.9, ls=":", c="k", lw=1)
    cax.axvline(q1, ymin=0.25, ymax=0.75, c="k", lw=1)
    cax.axvline(q2, ymin=0.1, ymax=0.9, c="k", lw=1)
    cax.axvline(q3, ymin=0.25, ymax=0.75, c="k", lw=1)
    cax.set_xlim(
        span[0] - (bins[-1] - bins[-2]) * extend_left,
        span[1] + (bins[-1] - bins[-2]) * extend_right,
    )
    cax.set_xticks([])
    cax.set_yticks([])
    cax.spines[:].set_visible(False)

def table_u_wind_dssim(
    my_ERA5_U: xr.DataArray,
    cr,
    title,
    dssim_bound,
    corr,
) -> pd.DataFrame:
    err_inf_U = np.amax(np.abs(my_ERA5_U - ERA5_U))
    err_2_U = np.sqrt(np.mean(np.square(my_ERA5_U - ERA5_U)))

    with xr.set_options(keep_attrs=True):
        da = ERA5_U[1:-1, 1:-1].copy(
            data=dssim_mat(
                ERA5_U.values,
                my_ERA5_U.values,
                reproducible=False,
            )
        )

    dssim_mean = np.nanmean(da)
    dssim_min = np.nanmin(da)
    dssim_max = np.nanmax(da)

    dssim_pw_v = np.mean(~(da >= dssim_bound))
    dssim_pw_v = (
        0
        if dssim_pw_v == 0
        else np.format_float_positional(100 * dssim_pw_v, precision=1, min_digits=1)
        + "%"
    )
    if dssim_pw_v == "0.0%":
        dssim_pw_v = "<0.05%"

    err_s = ~(np.mean(da) >= dssim_bound)
    err_s = r"$\times$" if err_s else r"$\checkmark$"

    corr = None if corr is None else compute_corrections_percentage(my_ERA5_U, corr)
    corr = (
        ""
        if corr is None
        else (
            0
            if corr == 0
            else np.format_float_positional(100 * corr, precision=1, min_digits=1) + "%"
        )
    )
    if corr == "0.0%":
        corr = "<0.05%"

    return pd.DataFrame(
        {
            "Compressor": [title[0]],
            r"$\epsilon_{abs}$": [title[1]],
            "Safeguarded": [title[2]],
            "Corrections": [title[3]],
            r"$L_{\infty}(\hat{u})$": [
                f"{err_inf_U:.03}",
            ],
            r"$L_{2}(\hat{u})$": [
                f"{err_2_U:.03}",
            ],
            r"$E[\text{D}]$": [
                f"{dssim_mean:.05}",
            ],
            r"$\min(\text{D})$": [
                f"{dssim_min:.03}",
            ],
            r"$\max(\text{D})$": [
                f"{dssim_max:.03}",
            ],
            rf"$\text{{D}} < {dssim_bound}$": [dssim_pw_v],
            "S": [err_s],
            "C": [corr],
            "CR": [
                rf"$\times$ {np.round(cr, 2)}",
            ],
        }
    )

import observe

observations = []

from numcodecs_wasm_zstd import Zstd

zstd = Zstd(level=22)

with observe.observe(zstd, observations):
    ERA5_U_zstd_enc = zstd.encode(ERA5_U.values)
    ERA5_U_zstd = ERA5_U.copy(data=zstd.decode(ERA5_U_zstd_enc))
ERA5_U_zstd_cr = ERA5_U.nbytes / np.asarray(ERA5_U_zstd_enc).nbytes

from numcodecs_wasm_sperr import Sperr
from numcodecs_wasm_sz3 import Sz3
from numcodecs_wasm_zfp import Zfp
from numcodecs_zero import ZeroCodec

eb_abs = [1.0, 0.1]

ERA5_U_codec = defaultdict(dict)
ERA5_U_codec_cr = defaultdict(dict)

for ea in eb_abs:
    for codec in [
        Zfp(mode="fixed-accuracy", tolerance=ea),
        Sz3(eb_mode="abs", eb_abs=ea),
        Sperr(mode="pwe", pwe=ea),
        ZeroCodec(),
    ]:
        with observe.observe(codec, observations):
            ERA5_U_codec_enc = codec.encode(ERA5_U.values)
            ERA5_U_codec[codec.codec_id][ea] = ERA5_U.copy(
                data=codec.decode(ERA5_U_codec_enc)
            )
        ERA5_U_codec_cr[codec.codec_id][ea] = (
            ERA5_U.nbytes / np.asarray(ERA5_U_codec_enc).nbytes
        )

from compression_safeguards import SafeguardKind

# Baker et al. recommend one of [0.99, 0.995, 0.99919]
dssim_bound = 0.995

qoi_eb_stencil = SafeguardKind.qoi_eb_stencil.value(
    qoi="""
    # dssim constants
    V["sigma"] = 1.5;
    V["C1"] = 1e-8;
    V["C2"] = 1e-8;
    
    # we guarantee that
    #  min(data) = min(corrected) and
    #  max(data) = max(corrected)
    # with the sign safeguards above
    V["smin"] = c["$x_min"];
    V["smax"] = c["$x_max"];
    V["r"] = V["smax"] - V["smin"];
    
    # re-scale to [0-1] and quantize to 256 bins
    V["sc_a1"] = round_ties_even(((C["$X"] - V["smin"]) / V["r"]) * 255) / 255;
    V["sc_a2"] = round_ties_even(((X - V["smin"]) / V["r"]) * 255) / 255;

    # create a 2D 3x3 Gaussian kernel
    V["i"] = [-1, 0, 1];
    V["k"] = 1/(sqrt(2*pi)*V["sigma"]) * exp((-square(V["i"])) / (2*square(V["sigma"])));
    V["kern"] = matmul([V["k"]].T, [V["k"]]);

    # handle sparse NaN values by excluding them from the convolution,
    # then renormalize the kernel
    V["kern_no_NaN1"] = where(isnan(V["sc_a1"]), 0, V["kern"]);
    V["kernel1"] = V["kern_no_NaN1"] / sum(V["kern_no_NaN1"]);
    V["kern_no_NaN2"] = where(isnan(V["sc_a2"]), 0, V["kern"]);
    V["kernel2"] = V["kern_no_NaN2"] / sum(V["kern_no_NaN2"]);
    V["kern_no_NaN12"] = where(isnan(V["sc_a1"] * V["sc_a2"]), 0, V["kern"]);
    V["kernel12"] = V["kern_no_NaN12"] / sum(V["kern_no_NaN12"]);

    # apply the Gaussian filter via convolution
    V["a1_mu_noNaN"] = sum(V["sc_a1"] * V["kernel1"]);
    V["a2_mu_noNaN"] = sum(V["sc_a2"] * V["kernel2"]);
    V["a1a1_noNaN"] = sum(square(V["sc_a1"]) * V["kernel1"]);
    V["a2a2_noNaN"] = sum(square(V["sc_a2"]) * V["kernel2"]);
    V["a1a2_noNaN"] = sum(V["sc_a1"] * V["sc_a2"] * V["kernel12"]);

    # keep NaN for elements that are originally NaN
    V["a1_mu"] = where(isnan(V["sc_a1"])[I], NaN, V["a1_mu_noNaN"]);
    V["a2_mu"] = where(isnan(V["sc_a2"])[I], NaN, V["a2_mu_noNaN"]);
    V["a1a1"] = where(isnan(V["sc_a1"])[I], NaN, V["a1a1_noNaN"]);
    V["a2a2"] = where(isnan(V["sc_a2"])[I], NaN, V["a2a2_noNaN"]);
    V["a1a2"] = where(isnan(V["sc_a1"] * V["sc_a2"])[I], NaN, V["a1a2_noNaN"]);

    ###########
    V["var_a1"] = V["a1a1"] - square(V["a1_mu"]);
    V["var_a2"] = V["a2a2"] - square(V["a2_mu"]);
    V["cov_a1a2"] = V["a1a2"] - V["a1_mu"] * V["a2_mu"];
    
    # compute the SSIM components
    V["ssim_t1"] = 2 * V["a1_mu"] * V["a2_mu"] + V["C1"];
    V["ssim_t2"] = 2 * V["cov_a1a2"] + V["C2"];
    V["ssim_b1"] = square(V["a1_mu"]) + square(V["a2_mu"]) + V["C1"];
    V["ssim_b2"] = V["var_a1"] + V["var_a2"] + V["C2"];

    V["ssim_1"] = V["ssim_t1"] / V["ssim_b1"];
    V["ssim_2"] = V["ssim_t2"] / V["ssim_b2"];
    
    # compute the pointwise dSSIM
    return V["ssim_1"] * V["ssim_2"];
    """,
    type="abs",
    eb=1 - dssim_bound,
    # 3x3 neighbourhood
    neighbourhood=[
        # latitude
        dict(axis=0, before=1, after=1, boundary="valid"),
        # longitude
        dict(axis=1, before=1, after=1, boundary="valid"),
    ],
)
qoi_eb_stencil_expr = qoi_eb_stencil._qoi_expr._expr

import tqdm

from compression_safeguards.safeguards._qois.expr import map_expr
from compression_safeguards.safeguards._qois.expr.abc import AnyExpr
from compression_safeguards.safeguards._qois.expr.reporter import (
    Reporter,
    ReportingExpr,
)


def data_expr_size(expr: AnyExpr) -> int:
    if not expr.has_data:
        return 0

    args_size = sum(data_expr_size(a) for a in expr.args)

    if isinstance(expr, ReportingExpr):
        return args_size

    return args_size + 1


class TqdmReporter(Reporter):
    __slots__ = ("_tqdm", "_current", "_stack")
    _tqdm: None | tqdm.tqdm
    _current: int
    _stack: list[int]

    def __init__(self) -> None:
        self._tqdm = None
        self._current = 0
        self._stack = []

    def enter(self, expr: AnyExpr) -> None:
        size = data_expr_size(expr)
        if self._tqdm is None:
            self._tqdm = tqdm.tqdm(
                total=size, desc="QoI", postfix=dict(depth=len(self._stack))
            )
            self._current = 0
        self._stack.append(self._current + size)
        self._tqdm.set_postfix(refresh=False, depth=len(self._stack))

    def exit(self, expr: AnyExpr) -> None:
        after = self._stack.pop()
        self._tqdm.set_postfix(refresh=False, depth=len(self._stack))
        self._tqdm.update(after - self._current)
        self._current = after
        if len(self._stack) > 0:
            return
        self._tqdm.close()
        self._tqdm = None
        self._current = 0


reporter = TqdmReporter()
qoi_eb_stencil._qoi_expr._expr = map_expr(
    qoi_eb_stencil_expr, mapper=lambda e: ReportingExpr(e, reporter)
)

from numcodecs_safeguards import SafeguardedCodec

ERA5_U_sg_lossless = defaultdict(dict)
ERA5_U_sg_lossless_cr = defaultdict(dict)

for ea in eb_abs:
    for codec in [
        ZeroCodec(),
        Zfp(mode="fixed-accuracy", tolerance=ea),
        Sz3(eb_mode="abs", eb_abs=ea),
        Sperr(mode="pwe", pwe=ea),
    ]:
        codec_sg = SafeguardedCodec(
            codec=codec,
            safeguards=[
                # guarantee that the global minimum and maximum are preserved,
                #  which simplifies the rescaling
                dict(kind="sign", offset="$x_min"),
                dict(kind="sign", offset="$x_max"),
                qoi_eb_stencil,
            ],
            # produce lossless corrections and refine them with iteration
            compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
        )

        with observe.observe(codec_sg, observations):
            ERA5_U_sg_lossless_enc = codec_sg.encode(ERA5_U.values)
            ERA5_U_sg_lossless[codec.codec_id][ea] = ERA5_U.copy(
                data=codec_sg.decode(ERA5_U_sg_lossless_enc)
            )

        ERA5_U_sg_lossless_cr[codec.codec_id][ea] = (
            ERA5_U.nbytes / np.asarray(ERA5_U_sg_lossless_enc).nbytes
        )

noa = Sz3(
    eb_mode="abs", eb_abs=(np.nanmax(ERA5_U.values) - np.nanmin(ERA5_U.values)) / 255
)

with observe.observe(noa, observations):
    ERA5_U_noa_enc = noa.encode(ERA5_U.values)
    ERA5_U_noa = ERA5_U.copy(data=noa.decode(ERA5_U_noa_enc))

ERA5_U_noa_cr = ERA5_U.nbytes / ERA5_U_noa_enc.nbytes

codec_sg_rint = SafeguardedCodec(
    codec=ZeroCodec(),
    safeguards=[
        # guarantee that the global minimum and maximum are preserved,
        #  which simplifies the rescaling
        dict(kind="sign", offset="$x_min"),
        dict(kind="sign", offset="$x_max"),
        dict(
            kind="qoi_eb_pw",
            qoi="""
            # we guarantee that
            #  min(data) = min(corrected) and
            #  max(data) = max(corrected)
            # with the sign safeguards above
            v["smin"] = c["$x_min"];
            v["smax"] = c["$x_max"];
            v["r"] = v["smax"] - v["smin"];
            
            # re-scale to [0-1] and quantize to 256 bins
            v["sc_a2"] = round_ties_even(((x - v["smin"]) / v["r"]) * 255) / 255;
            
            # force the quantized value to stay the same
            return v["sc_a2"];
            """,
            type="abs",
            eb=0,
        ),
    ],
)

with observe.observe(codec_sg_rint, observations):
    ERA5_U_sg_rint_enc = codec_sg_rint.encode(ERA5_U.values)
    ERA5_U_sg_rint = ERA5_U.copy(data=codec_sg_rint.decode(ERA5_U_sg_rint_enc))

ERA5_U_sg_rint_cr = ERA5_U.nbytes / np.asarray(ERA5_U_sg_rint_enc).nbytes

codec_sg_it_rint = SafeguardedCodec(
    codec=ZeroCodec(),
    safeguards=codec_sg_rint.safeguards,
    # use iteration to refine the corrections
    compute=dict(unstable_iterative=True),
)

with observe.observe(codec_sg_it_rint, observations):
    ERA5_U_sg_it_rint_enc = codec_sg_it_rint.encode(ERA5_U.values)
    ERA5_U_sg_it_rint = ERA5_U.copy(data=codec_sg_it_rint.decode(ERA5_U_sg_it_rint_enc))

ERA5_U_sg_it_rint_cr = ERA5_U.nbytes / np.asarray(ERA5_U_sg_it_rint_enc).nbytes

codec_sg_lossless_rint = SafeguardedCodec(
    codec=ZeroCodec(),
    safeguards=codec_sg_rint.safeguards,
    # produce lossless corrections and refine them with iteration
    compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
)

with observe.observe(codec_sg_lossless_rint, observations):
    ERA5_U_sg_lossless_rint_enc = codec_sg_lossless_rint.encode(ERA5_U.values)
    ERA5_U_sg_lossless_rint = ERA5_U.copy(
        data=codec_sg_lossless_rint.decode(ERA5_U_sg_lossless_rint_enc)
    )

ERA5_U_sg_lossless_rint_cr = (
    ERA5_U.nbytes / np.asarray(ERA5_U_sg_lossless_rint_enc).nbytes
)

import numcodecs


class SafetyViolationsMetric(numcodecs.abc.Codec):
    codec_id = "safety-violations-metric"

    def __init__(self):
        self._data = None

    def encode(self, buf):
        # store the original data for later
        self._data = np.array(buf, copy=True)
        # return no metric
        return np.empty(0, dtype=np.float64)

    def decode(self, buf, out=None):
        # compute the violations
        buf_dssim = np.nanmean(
            dssim_mat(
                self._data,
                buf,
                reproducible=False,
            )
        )
        # use distance to the dSSIM threshold as the violations score
        violations = dssim_bound - buf_dssim
        self._data = None
        # return the violations score metric
        return numcodecs.compat.ndarray_copy(np.float64(violations), out)


numcodecs.registry.register_codec(SafetyViolationsMetric)

class ExponentialZfp(Zfp):
    codec_id = "e-zfp.rs"

    def __new__(cls, tolerance: float, **kwargs):
        codec = super().__new__(cls, tolerance=np.exp(tolerance), **kwargs)
        codec._tolerance = tolerance
        return codec

    def get_config(self):
        return {
            **super().get_config(),
            "id": type(self).codec_id,
            "tolerance": self._tolerance,
        }


class ExponentialSz3(Sz3):
    codec_id = "e-sz3.rs"

    def __new__(cls, eb_abs: float, **kwargs):
        codec = super().__new__(cls, eb_abs=np.exp(eb_abs), **kwargs)
        codec._eb_abs = eb_abs
        return codec

    def get_config(self):
        return {
            **super().get_config(),
            "id": type(self).codec_id,
            "eb_abs": self._eb_abs,
        }


class ExponentialSperr(Sperr):
    codec_id = "e-sperr.rs"

    def __new__(cls, pwe: float, **kwargs):
        codec = super().__new__(cls, pwe=np.exp(pwe), **kwargs)
        codec._pwe = pwe
        return codec

    def get_config(self):
        return {**super().get_config(), "id": type(self).codec_id, "pwe": self._pwe}


numcodecs.registry.register_codec(ExponentialZfp)
numcodecs.registry.register_codec(ExponentialSz3)
numcodecs.registry.register_codec(ExponentialSperr)