Preserving a spatial quantity of interest (QoI) with safeguards¶

import ssl

ssl._create_default_https_context = ssl._create_stdlib_context

import copy
from pathlib import Path

import earthkit.plots
import humanize
import matplotlib as mpl
import matplotlib.pyplot as plt
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 = ERA5.sel(valid_time="2024-04-02T12:00:00", pressure_level=500)

def compute_relative_vorticity(ERA5: xr.Dataset) -> xr.DataArray:
    # reimplementation of np.deg2rad that matches the safeguards QoI
    def deg2rad(a: np.ndarray) -> np.ndarray:
        # np.deg2rad(a) = a * a.dtype.type(np.pi / 180)
        return a * a.dtype.type(np.pi) / a.dtype.type(180)

    earth_radius = 6371000  # [m], globally averaged

    # computing the derivative with a finite difference requires extending the
    #  data domain and tricking xarray for the coordinates
    # e.g. the data needs to be wrapped along the longitude axis,
    #  but the longitude coordinate needs to be extended with odd reflection
    #  ([0, 0.25, ..., 359.75, 360] -> [-0.25, 0, 0.25, ..., 359.75, 360, 360.25])
    #  since xarray cannot handle differentiating along a proper periodic axis
    # along the latitude axis, xr.differentiate -> np.gradient falls back to
    #  forward/backwards differences at the boundaries (poles)
    ERA5_wrapped = ERA5.pad(longitude=1, mode="wrap").assign_coords(
        longitude=ERA5.longitude.pad(longitude=1, mode="reflect", reflect_type="odd"),
    )

    # compute the relative vorticity
    ERA5_dUdTheta = (
        ERA5_wrapped["u"]
        * np.cos(deg2rad(ERA5_wrapped["latitude"].astype(ERA5_wrapped["u"].dtype)))
    ).differentiate("latitude")
    ERA5_dVdPhi = ERA5_wrapped["v"].differentiate("longitude")

    ERA5_VOR = (ERA5_dVdPhi - ERA5_dUdTheta) / (
        earth_radius
        * np.cos(deg2rad(ERA5_wrapped["latitude"].astype(ERA5_wrapped["u"].dtype)))
    )

    # remove the padding to extract just the valid values
    ERA5_VOR = ERA5_VOR.sel(longitude=slice(0, 359.9))
    ERA5_VOR.attrs.update(long_name="Relative vorticity", units="s**-1")

    return ERA5_VOR

ERA5_VOR = compute_relative_vorticity(ERA5)

def compute_corrections_percentage(my_ERA5: xr.Dataset, orig_ERA5: xr.Dataset) -> float:
    neq = np.sum(my_ERA5 != orig_ERA5)
    return int(neq.u + neq.v) / int(orig_ERA5.u.size + orig_ERA5.v.size)

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 plot_relative_vorticity(
    my_ERA5: xr.Dataset,
    cr,
    chart,
    title,
    span,
    vor_eb_abs,
    error=False,
    corr=None,
    my_ERA5_it=None,
    cr_it=None,
    inset=True,
):
    my_ERA5_VOR = compute_relative_vorticity(my_ERA5)

    if error:
        with xr.set_options(keep_attrs=True):
            da = (my_ERA5_VOR - ERA5_VOR).compute()

        da.attrs.update(long_name=f"Absolute error over {da.long_name.lower()}")
    else:
        da = my_ERA5_VOR

    # 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, span, 22))
    style._legend_kwargs["ticks"] = np.linspace(-span, span, 5)
    style._colors = "coolwarm" if error else "BrBG"

    extend_left = np.nanmin(da) < -span
    extend_right = np.nanmax(da) > span

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

    if error:
        chart.pcolormesh(da, style=style, zorder=-12)

        if corr is not None:
            da_hatch = (my_ERA5["u"] == corr["u"]) & (my_ERA5["v"] == corr["v"])

            if my_ERA5_it is None:
                da_corr = da_hatch.astype(float)
            else:
                with xr.set_options(keep_attrs=True):
                    da_hatch_it = (my_ERA5_it["u"] == corr["u"]) & (
                        my_ERA5_it["v"] == corr["v"]
                    )
                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(["white", "green", "lawngreen"]),
                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 = ~(np.abs(da) <= vor_eb_abs)

            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", "red"]),
                vmin=0,
                vmax=1,
                rasterized=True,
            )
            axin.coastlines(color="#555555")
            axin.spines["geo"].set_edgecolor("black")
            axin.set_title(
                "Violations",
                path_effects=[PathEffects.withStroke(linewidth=3, foreground="white")],
            )
    else:
        chart.pcolormesh(da, style=style, zorder=-11)

    chart.ax.set_rasterization_zorder(-10)

    chart.title(title)

    if error:
        err_v = np.mean(~(np.abs(my_ERA5_VOR - ERA5_VOR) <= vor_eb_abs))
        err_v = (
            0
            if err_v == 0
            else np.format_float_positional(100 * err_v, precision=1, min_digits=1)
            + "%"
        )
        if err_v == "0.0%":
            err_v = "<0.05%"

        t = chart.ax.text(
            0.95,
            0.1,
            f"V={err_v}",
            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 + ERA5["v"].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, 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, -vor_eb_abs], *cb.ax.get_ylim(), hatch="xx", ec="w", fc="none", lw=0
        )
        cb.ax.fill_between(
            [vor_eb_abs, span], *cb.ax.get_ylim(), hatch="xx", ec="w", fc="none", lw=0
        )
        if extend_right:
            cb._extend_patches[-1].set_hatch("xx")
            cb._extend_patches[-1].set_ec("white")
    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) > vor_eb_abs 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),
            width=(bins[-1] - bins[0]) / len(counts),
            color=cb.cmap(cb.norm(midpoints[0])),
            **(
                dict(
                    hatch="xx",
                    ec="white",
                    lw=0,
                )
                if error
                else dict()
            ),
        )
    if extend_right:
        cax.bar(
            bins[-1] + (bins[-1] - bins[-2]) / 2,
            height=np.sum(da > span),
            width=(bins[-1] - bins[0]) / len(counts),
            color=cb.cmap(cb.norm(midpoints[-1])),
            **(
                dict(
                    hatch="xx",
                    ec="white",
                    lw=0,
                )
                if error
                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 - (bins[-1] - bins[-2]) * extend_left,
        span + (bins[-1] - bins[-2]) * extend_right,
    )
    cax.set_xticks([])
    cax.set_yticks([])
    cax.spines[:].set_visible(False)

def table_relative_vorticity(
    my_ERA5: xr.Dataset,
    cr,
    title,
    vor_eb_abs,
    corr,
) -> pd.DataFrame:
    my_ERA5_VOR = compute_relative_vorticity(my_ERA5)

    err_inf_U = np.amax(np.abs(my_ERA5["u"] - ERA5["u"]))
    err_inf_V = np.amax(np.abs(my_ERA5["v"] - ERA5["v"]))
    err_inf_VOR = np.amax(np.abs(my_ERA5_VOR - ERA5_VOR))
    err_2_VOR = np.sqrt(np.mean(np.square(my_ERA5_VOR - ERA5_VOR)))

    err_v = np.mean(~(np.abs(my_ERA5_VOR - ERA5_VOR) <= vor_eb_abs))
    err_v = (
        0
        if err_v == 0
        else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
    )
    if err_v == "0.0%":
        err_v = "<0.05%"

    corr = None if corr is None else compute_corrections_percentage(my_ERA5, 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]],
            "Safeguarded": [title[1]],
            "Corrections": [title[2]],
            r"$L_{\infty}(\hat{u})$": [
                f"{err_inf_U:.03}",
            ],
            r"$L_{\infty}(\hat{v})$": [
                f"{err_inf_V:.03}",
            ],
            r"$L_{\infty}(\hat{\zeta})$": [
                f"{err_inf_VOR:.03}",
            ],
            r"$L_{2}(\hat{\zeta})$": [
                f"{err_2_VOR:.03}",
            ],
            "V": [err_v],
            "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_UV_zstd_enc = zstd.encode(ERA5_UV)
    ERA5_UV_zstd = zstd.decode(ERA5_UV_zstd_enc)

ERA5_zstd = ERA5.copy(data=dict(u=ERA5_UV_zstd[0], v=ERA5_UV_zstd[1]))
ERA5_zstd_cr = ERA5_UV.nbytes / ERA5_UV_zstd_enc.nbytes

eb_abs = 0.01

from numcodecs_wasm_zfp import Zfp

zfp = Zfp(mode="fixed-accuracy", tolerance=eb_abs)

with observe.observe(zfp, observations):
    ERA5_UV_zfp_enc = zfp.encode(ERA5_UV)
    ERA5_UV_zfp = zfp.decode(ERA5_UV_zfp_enc)

ERA5_zfp = ERA5.copy(data=dict(u=ERA5_UV_zfp[0], v=ERA5_UV_zfp[1]))
ERA5_zfp_cr = ERA5_UV.nbytes / ERA5_UV_zfp_enc.nbytes

from numcodecs_wasm_sz3 import Sz3

sz3 = Sz3(eb_mode="abs", eb_abs=eb_abs)

with observe.observe(sz3, observations):
    ERA5_UV_sz3_enc = sz3.encode(ERA5_UV)
    ERA5_UV_sz3 = sz3.decode(ERA5_UV_sz3_enc)

ERA5_sz3 = ERA5.copy(data=dict(u=ERA5_UV_sz3[0], v=ERA5_UV_sz3[1]))
ERA5_sz3_cr = ERA5_UV.nbytes / ERA5_UV_sz3_enc.nbytes

from numcodecs_wasm_sperr import Sperr

sperr = Sperr(mode="pwe", pwe=eb_abs)

with observe.observe(sperr, observations):
    ERA5_UV_sperr_enc = sperr.encode(ERA5_UV)
    ERA5_UV_sperr = sperr.decode(ERA5_UV_sperr_enc)

ERA5_sperr = ERA5.copy(data=dict(u=ERA5_UV_sperr[0], v=ERA5_UV_sperr[1]))
ERA5_sperr_cr = ERA5_UV.nbytes / ERA5_UV_sperr_enc.nbytes

from numcodecs_zero import ZeroCodec

zero = ZeroCodec()

with observe.observe(zero, observations):
    ERA5_UV_zero_enc = zero.encode(ERA5_UV)
    ERA5_UV_zero = zero.decode(ERA5_UV_zero_enc)

ERA5_zero = ERA5.copy(data=dict(u=ERA5_UV_zero[0], v=ERA5_UV_zero[1]))

vor_eb_abs = 1e-8

from compression_safeguards import SafeguardKind

qoi_eb_stencil = SafeguardKind.qoi_eb_stencil.value(
    qoi="""
    V["earth_radius"] = 6371000;  # [m], globally averaged

    # extract u and v from the first stacked dimension
    V["u"] = X[0, I[1], I[2]];
    V["v"] = X[1, I[1], I[2]];
    # convert latitude in degrees to radians
    V["latRad"] = c["lat"] * pi / 180;

    # approximate the spatial derivatives with finite differences
    # (1) along the latitude axis, we need to handle the polar boundary
    V["dUdTheta"] = where(
        abs(c["lat"]) < 90,
        # use central difference where possible
        finite_difference(
            V["u"] * cos(V["latRad"]),
            order=1, accuracy=2, type=0, axis=1,
            grid_centre=c["lat"],
        ),
        # at the poles, use forward/backwards difference
        where(
            c["lat"] == -90,
            finite_difference(
                V["u"] * cos(V["latRad"]),
                order=1, accuracy=1, type=-1, axis=1,
                grid_centre=c["lat"],
            ),
            finite_difference(
                V["u"] * cos(V["latRad"]),
                order=1, accuracy=1, type=+1, axis=1,
                grid_centre=c["lat"],
            ),
        ),
    );
    # (2) along the longitude axis, we handle the 360 degree periodic boundary
    V["dVdPhi"] = finite_difference(
        V["v"],
        order=1, accuracy=2, type=0, axis=2,
        grid_centre=c["lon"], grid_period=360,
    );

    # compute the relative vorticity
    return (V["dVdPhi"] - V["dUdTheta"]) / (
        V["earth_radius"] * cos(V["latRad"])
    );
    """,
    type="abs",
    eb=vor_eb_abs,
    neighbourhood=[
        # [u, v]: stacked variables
        dict(axis=0, before=0, after=1, boundary="valid"),
        # latitude: edge boundary is good enough (since we ensure in the QoI
        #           that the edge-extended values are not used)
        dict(axis=1, before=1, after=1, boundary="edge"),
        # longitude: wrapping boundary
        dict(axis=2, before=1, after=1, boundary="wrap"),
    ],
)

from compression_safeguards.utils.bindings import Bindings

vor_py = compute_relative_vorticity(ERA5)
vor_qoi = qoi_eb_stencil.evaluate_qoi(
    ERA5_UV,
    late_bound=Bindings(
        lat=ERA5.latitude.values.reshape(1, -1, 1),
        lon=ERA5.longitude.values.reshape(1, 1, -1),
    ),
)

assert np.all(vor_py == vor_qoi.squeeze())

from numcodecs_safeguards import SafeguardedCodec

ERA5_sg = dict()
ERA5_sg_cr = dict()

for codec in [
    zero,
    zfp,
    sz3,
    sperr,
]:
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[qoi_eb_stencil],
        fixed_constants=dict(
            lat=ERA5.latitude.values.reshape(1, -1, 1),
            lon=ERA5.longitude.values.reshape(1, 1, -1),
        ),
    )

    with observe.observe(sg, observations):
        ERA5_UV_sg_enc = sg.encode(ERA5_UV)
        ERA5_UV_sg = sg.decode(ERA5_UV_sg_enc)

    ERA5_sg[codec.codec_id] = ERA5.copy(data=dict(u=ERA5_UV_sg[0], v=ERA5_UV_sg[1]))
    ERA5_sg_cr[codec.codec_id] = ERA5_UV.nbytes / np.asarray(ERA5_UV_sg_enc).nbytes

ERA5_sg_it = dict()
ERA5_sg_it_cr = dict()

for codec in [
    zero,
    zfp,
    sz3,
    sperr,
]:
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[qoi_eb_stencil],
        fixed_constants=dict(
            lat=ERA5.latitude.values.reshape(1, -1, 1),
            lon=ERA5.longitude.values.reshape(1, 1, -1),
        ),
        # use iteration to refine the corrections
        compute=dict(unstable_iterative=True),
    )

    with observe.observe(sg, observations):
        ERA5_UV_sg_it_enc = sg.encode(ERA5_UV)
        ERA5_UV_sg_it = sg.decode(ERA5_UV_sg_it_enc)

    ERA5_sg_it[codec.codec_id] = ERA5.copy(
        data=dict(u=ERA5_UV_sg_it[0], v=ERA5_UV_sg_it[1])
    )
    ERA5_sg_it_cr[codec.codec_id] = (
        ERA5_UV.nbytes / np.asarray(ERA5_UV_sg_it_enc).nbytes
    )

ERA5_sg_lossless = dict()
ERA5_sg_lossless_cr = dict()

for codec in [
    zero,
    zfp,
    sz3,
    sperr,
]:
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[qoi_eb_stencil],
        fixed_constants=dict(
            lat=ERA5.latitude.values.reshape(1, -1, 1),
            lon=ERA5.longitude.values.reshape(1, 1, -1),
        ),
        # produce lossless corrections and refine them with iteration
        compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
    )

    with observe.observe(sg, observations):
        ERA5_UV_sg_lossless_enc = sg.encode(ERA5_UV)
        ERA5_UV_sg_lossless = sg.decode(ERA5_UV_sg_lossless_enc)

    ERA5_sg_lossless[codec.codec_id] = ERA5.copy(
        data=dict(u=ERA5_UV_sg_lossless[0], v=ERA5_UV_sg_lossless[1])
    )
    ERA5_sg_lossless_cr[codec.codec_id] = (
        ERA5_UV.nbytes / np.asarray(ERA5_UV_sg_lossless_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
        data_VOR = compute_relative_vorticity(
            ERA5.copy(data=dict(u=self._data[0], v=self._data[1]))
        )
        buf_VOR = compute_relative_vorticity(ERA5.copy(data=dict(u=buf[0], v=buf[1])))
        violations = np.mean(~(np.abs(buf_VOR - data_VOR) <= vor_eb_abs))
        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)