Preserve a pointwise quantity of interest (QoI) with safeguards¶

import ssl

ssl._create_default_https_context = ssl._create_stdlib_context

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-q" / "data.nc")
ERA5_Q = ERA5["q"].sel(valid_time="2024-04-02T12:00:00", pressure_level=850)

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

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_specific_humidity_log10(
    my_ERA5_Q: xr.DataArray,
    cr,
    chart,
    title,
    span,
    qlog10_eb_abs,
    error=False,
    corr=None,
    my_ERA5_Q_ratio=None,
    cr_ratio=None,
    corr_ratio=None,
    inset=True,
):
    import copy

    with np.errstate(divide="ignore", invalid="ignore"):
        ERA5_Q_log10 = np.log10(ERA5_Q)
        my_ERA5_Q_log10 = np.log10(my_ERA5_Q)

        if my_ERA5_Q_ratio is not None:
            my_ERA5_Q_ratio_log10 = np.log10(my_ERA5_Q_ratio)

    if error:
        err_v = np.mean(~(np.abs(my_ERA5_Q_log10 - ERA5_Q_log10) <= qlog10_eb_abs))

        if my_ERA5_Q_ratio is not None:
            err_ratio_v = np.mean(
                ~(np.abs(my_ERA5_Q_ratio_log10 - ERA5_Q_log10) <= qlog10_eb_abs)
            )

        with xr.set_options(keep_attrs=True):
            da = (my_ERA5_Q_log10 - ERA5_Q_log10).compute()

        da.attrs.update(
            long_name=f"Absolute error over log10({da.long_name.lower()})",
            units=ERA5_Q.units,
        )
    else:
        # plot the decimal logarithm of specific humidity to better capture scale
        da = np.log10(my_ERA5_Q)
        da.attrs.update(
            long_name=f"log10({ERA5_Q.long_name.lower()})", units=ERA5_Q.units
        )

    # 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,
        )
    )

    if error:
        style._levels = earthkit.plots.styles.levels.Levels(
            np.linspace(-span, span, 22)
        )
        style._legend_kwargs["ticks"] = np.linspace(-span, span, 5)
        style._colors = "coolwarm"
    else:
        style._levels = earthkit.plots.styles.levels.Levels(np.linspace(*span, 22))
        style._legend_kwargs["ticks"] = np.linspace(*span, 5)
        style._colors = "viridis"

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

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

    chart.ax.fill_between(
        [0, 1],
        [1, 1],
        hatch="XX",
        edgecolor="magenta",
        facecolor="lavenderblush",
        transform=chart.ax.transAxes,
        zorder=-12,
    )

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

        if corr is not None:
            da_hatch = my_ERA5_Q == corr
            da_corr = (~da_hatch).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) <= qlog10_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.quickplot(da, style=style, extend=extend, zorder=-11)

    chart.ax.set_rasterization_zorder(-10)

    chart.title(title)

    if error:
        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%"

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

        t = chart.ax.text(
            0.95,
            0.1,
            f"V={err_v}" + ("" if my_ERA5_Q_ratio is None else f" ({err_ratio_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_ratio is None else rf" ($\times$ {np.round(cr_ratio, 2)})")
        if error
        else humanize.naturalsize(ERA5_Q.nbytes, binary=True),
        ha="right",
        va="top",
        transform=chart.ax.transAxes,
    )
    t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

    chart.contour(
        x=np.broadcast_to(
            da.longitude.values.reshape(1, -1),
            da.shape,
        ),
        y=np.broadcast_to(
            da.latitude.values.reshape(-1, 1),
            da.shape,
        ),
        z=~(np.isfinite(da).values),
        colors=["none"],
        linecolors=["magenta"],
        levels=[-0.5, 0.9, 1.5],
        legend_style=None,
        labels=False,
        zorder=-10,
    )

    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 if error else span[0], span if error else span[1]),
        bins=21,
    )
    midpoints = bins[:-1] + np.diff(bins) / 2
    cb = chart.ax.collections[1].colorbar
    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 + np.any(~np.isfinite(da)) * 2),
            1.0,
        ]
    )
    cax.bar(
        midpoints,
        height=counts,
        width=(bins[-1] - bins[0]) / len(counts),
        color=cb.cmap(cb.norm(midpoints)),
    )
    if extend_left:
        cax.bar(
            bins[0] - (bins[1] - bins[0]) / 2,
            height=np.sum(da < (-span if error else span[0])),
            width=(bins[-1] - bins[0]) / len(counts),
            color=cb.cmap(cb.norm(midpoints[0])),
        )
    if extend_right:
        cax.bar(
            bins[-1] + (bins[-1] - bins[-2]) / 2,
            height=np.sum(da > (span if error else span[1])),
            width=(bins[-1] - bins[0]) / len(counts),
            color=cb.cmap(cb.norm(midpoints[-1])),
        )
    if np.any(~np.isfinite(da)):
        cax.bar(
            bins[-1] + (bins[-1] - bins[-2]) * (extend_right * 2 + 2 + 1) / 2,
            height=np.sum(~np.isfinite(da)),
            width=(bins[-1] - bins[0]) / len(counts),
            color="lavenderblush",
            edgecolor="magenta",
            lw=0,
            hatch="XXXX",
        )
    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 if error else span[0]) - (bins[-1] - bins[-2]) * extend_left,
        (span if error else span[1])
        + (bins[-1] - bins[-2]) * (0 + extend_right + np.any(~np.isfinite(da)) * 2),
    )
    cax.set_xticks([])
    cax.set_yticks([])
    cax.spines[:].set_visible(False)

def table_specific_humidity_log10(
    my_ERA5_Q: xr.DataArray,
    cr,
    title,
    qlog10_eb_abs,
    corr=None,
):
    with np.errstate(divide="ignore", invalid="ignore"):
        ERA5_Q_log10 = np.log10(ERA5_Q)
        my_ERA5_Q_log10 = np.log10(my_ERA5_Q)

    err_Q_inf = np.amax(np.abs(my_ERA5_Q - ERA5_Q).values)
    err_Q_log10_inf = np.amax(np.abs(my_ERA5_Q_log10 - ERA5_Q_log10).values)
    err_Q_log10_fin_inf = np.nanmax(
        np.nan_to_num(
            np.abs(my_ERA5_Q_log10 - ERA5_Q_log10),
            nan=np.nan,
            posinf=np.nan,
            neginf=np.nan,
        )
    )
    err_Q_log10_2 = np.sqrt(np.mean(np.square(my_ERA5_Q_log10 - ERA5_Q_log10).values))
    err_Q_log10_fin_2 = np.sqrt(
        np.nanmean(
            np.nan_to_num(
                np.square(my_ERA5_Q_log10 - ERA5_Q_log10),
                nan=np.nan,
                posinf=np.nan,
                neginf=np.nan,
            )
        )
    )

    err_v = np.mean(~(np.abs(my_ERA5_Q_log10 - ERA5_Q_log10) <= qlog10_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_Q, 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]],
            "(Bound)": [title[1]],
            "Safeguarded": [title[2]],
            "Corrections": [title[3]],
            r"$L_{\infty}(\hat{q})$": [f"{err_Q_inf:.02}"],
            r"$L_{\infty}(\log_{10}(\hat{q}))$": [
                f"{err_Q_log10_inf:.03}".replace("nan", "NaN")
                + (
                    ""
                    if np.isfinite(err_Q_log10_inf)
                    else f" [{err_Q_log10_fin_inf:.03}]"
                )
            ],
            r"$L_{2}(\log_{10}(\hat{q}))$": [
                f"{err_Q_log10_2:.03}".replace("nan", "NaN")
                + ("" if np.isfinite(err_Q_log10_2) else f" [{err_Q_log10_fin_2:.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_Q_zstd_enc = zstd.encode(ERA5_Q.values)
    ERA5_Q_zstd = ERA5_Q.copy(data=zstd.decode(ERA5_Q_zstd_enc))

ERA5_Q_zstd_cr = ERA5_Q.nbytes / ERA5_Q_zstd_enc.nbytes

eb_abs = 0.0005

from numcodecs_wasm_zfp import Zfp

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

with observe.observe(zfp, observations):
    ERA5_Q_zfp_enc = zfp.encode(ERA5_Q.values)
    ERA5_Q_zfp = ERA5_Q.copy(data=zfp.decode(ERA5_Q_zfp_enc))

ERA5_Q_zfp_cr = ERA5_Q.nbytes / ERA5_Q_zfp_enc.nbytes

from numcodecs_wasm_sz3 import Sz3

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

with observe.observe(sz3, observations):
    ERA5_Q_sz3_enc = sz3.encode(ERA5_Q.values)
    ERA5_Q_sz3 = ERA5_Q.copy(data=sz3.decode(ERA5_Q_sz3_enc))

ERA5_Q_sz3_cr = ERA5_Q.nbytes / ERA5_Q_sz3_enc.nbytes

from numcodecs_wasm_sperr import Sperr

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

with observe.observe(sperr, observations):
    ERA5_Q_sperr_enc = sperr.encode(ERA5_Q.values)
    ERA5_Q_sperr = ERA5_Q.copy(data=sperr.decode(ERA5_Q_sperr_enc))

ERA5_Q_sperr_cr = ERA5_Q.nbytes / ERA5_Q_sperr_enc.nbytes

from numcodecs_zero import ZeroCodec

zero = ZeroCodec()

with observe.observe(zero, observations):
    ERA5_Q_zero_enc = zero.encode(ERA5_Q.values)
    ERA5_Q_zero = ERA5_Q.copy(data=zero.decode(ERA5_Q_zero_enc))

qlog10_eb_abs = 0.25

from numcodecs_safeguards import SafeguardedCodec

ERA5_Q_sg = dict()
ERA5_Q_sg_cr = dict()

for codec_id, codec in {
    "zero": zero,
    "zfp.rs": zfp,
    "sz3.rs": sz3,
    "sperr.rs": sperr,
}.items():
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[
            dict(
                kind="qoi_eb_pw",
                qoi="log10(x)",
                type="abs",
                eb=qlog10_eb_abs,
            )
        ],
    )

    with observe.observe(sg, observations):
        ERA5_Q_sg_enc = sg.encode(ERA5_Q.values)
        ERA5_Q_sg[codec_id] = ERA5_Q.copy(data=sg.decode(ERA5_Q_sg_enc))

    ERA5_Q_sg_cr[codec_id] = ERA5_Q.nbytes / np.asarray(ERA5_Q_sg_enc).nbytes

from numcodecs_safeguards import SafeguardedCodec

ERA5_Q_sg_lossless = dict()
ERA5_Q_sg_lossless_cr = dict()

for codec_id, codec in {
    "zero": zero,
    "zfp.rs": zfp,
    "sz3.rs": sz3,
    "sperr.rs": sperr,
}.items():
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[
            dict(
                kind="qoi_eb_pw",
                qoi="log10(x)",
                type="abs",
                eb=qlog10_eb_abs,
            )
        ],
        # produce lossless corrections and refine them with iteration
        compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
    )

    with observe.observe(sg, observations):
        ERA5_Q_sg_lossless_enc = sg.encode(ERA5_Q.values)
        ERA5_Q_sg_lossless[codec_id] = ERA5_Q.copy(
            data=sg.decode(ERA5_Q_sg_lossless_enc)
        )

    ERA5_Q_sg_lossless_cr[codec_id] = (
        ERA5_Q.nbytes / np.asarray(ERA5_Q_sg_lossless_enc).nbytes
    )

from numcodecs_pw_ratio import PointwiseRatioErrorBoundedCodec

zfp_ratio = PointwiseRatioErrorBoundedCodec(
    eb_ratio=np.power(10, 0.25),
    eb_abs_marker="$eb_abs",
    log_codec={**zfp.get_config(), "tolerance": "$eb_abs"},
    sign_codec=dict(id="zlib", level=9),
)

with observe.observe(zfp_ratio, observations):
    ERA5_Q_zfp_ratio_enc = zfp_ratio.encode(ERA5_Q.values)
    ERA5_Q_zfp_ratio = ERA5_Q.copy(data=zfp_ratio.decode(ERA5_Q_zfp_ratio_enc))

ERA5_Q_zfp_ratio_cr = ERA5_Q.nbytes / np.asarray(ERA5_Q_zfp_ratio_enc).nbytes

sz3_ratio = PointwiseRatioErrorBoundedCodec(
    eb_ratio=np.power(10, 0.25),
    eb_abs_marker="$eb_abs",
    log_codec={**sz3.get_config(), "eb_abs": "$eb_abs"},
    sign_codec=dict(id="zlib", level=9),
)

with observe.observe(sz3_ratio, observations):
    ERA5_Q_sz3_ratio_enc = sz3_ratio.encode(ERA5_Q.values)
    ERA5_Q_sz3_ratio = ERA5_Q.copy(data=sz3_ratio.decode(ERA5_Q_sz3_ratio_enc))

ERA5_Q_sz3_ratio_cr = ERA5_Q.nbytes / np.asarray(ERA5_Q_sz3_ratio_enc).nbytes

sperr_ratio = PointwiseRatioErrorBoundedCodec(
    eb_ratio=np.power(10, 0.25),
    eb_abs_marker="$eb_abs",
    log_codec={**sperr.get_config(), "pwe": "$eb_abs"},
    sign_codec=dict(id="zlib", level=9),
)

with observe.observe(sperr_ratio, observations):
    ERA5_Q_sperr_ratio_enc = sperr_ratio.encode(ERA5_Q.values)
    ERA5_Q_sperr_ratio = ERA5_Q.copy(data=sperr_ratio.decode(ERA5_Q_sperr_ratio_enc))

ERA5_Q_sperr_ratio_cr = ERA5_Q.nbytes / np.asarray(ERA5_Q_sperr_ratio_enc).nbytes

ERA5_Q_ratio_sg = dict()
ERA5_Q_ratio_sg_cr = dict()

for codec_id, codec in {
    "zfp.rs": zfp_ratio,
    "sz3.rs": sz3_ratio,
    "sperr.rs": sperr_ratio,
}.items():
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[
            dict(
                kind="qoi_eb_pw",
                qoi="log10(x)",
                type="abs",
                eb=qlog10_eb_abs,
            )
        ],
    )

    with observe.observe(sg, observations):
        ERA5_Q_ratio_sg_enc = sg.encode(ERA5_Q.values)
        ERA5_Q_ratio_sg[codec_id] = ERA5_Q.copy(data=sg.decode(ERA5_Q_ratio_sg_enc))

    ERA5_Q_ratio_sg_cr[codec_id] = (
        ERA5_Q.nbytes / np.asarray(ERA5_Q_ratio_sg_enc).nbytes
    )

sg_ratio = SafeguardedCodec(
    codec=ZeroCodec(),
    safeguards=[
        dict(
            kind="eb",
            type="ratio",
            eb=np.power(10.0, qlog10_eb_abs),
        )
    ],
)

with observe.observe(sg_ratio, observations):
    ERA5_Q_sg_ratio_enc = sg_ratio.encode(ERA5_Q.values)
    ERA5_Q_sg_ratio = ERA5_Q.copy(data=sg_ratio.decode(ERA5_Q_sg_ratio_enc))

ERA5_Q_sg_ratio_cr = ERA5_Q.nbytes / np.asarray(ERA5_Q_sg_ratio_enc).nbytes

sg_ratio_lossless = SafeguardedCodec(
    codec=ZeroCodec(),
    safeguards=[
        dict(
            kind="eb",
            type="ratio",
            eb=np.power(10.0, qlog10_eb_abs),
        )
    ],
    # produce lossless corrections and refine them with iteration
    compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
)

with observe.observe(sg_ratio_lossless, observations):
    ERA5_Q_sg_ratio_lossless_enc = sg_ratio_lossless.encode(ERA5_Q.values)
    ERA5_Q_sg_ratio_lossless = ERA5_Q.copy(
        data=sg_ratio_lossless.decode(ERA5_Q_sg_ratio_lossless_enc)
    )

ERA5_Q_sg_ratio_lossless_cr = (
    ERA5_Q.nbytes / np.asarray(ERA5_Q_sg_ratio_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
        with np.errstate(divide="ignore", invalid="ignore"):
            data_log10 = np.log10(self._data)
            buf_log10 = np.log10(buf)
        violations = np.mean(~(np.abs(buf_log10 - data_log10) <= qlog10_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)

import json

with Path("observations").joinpath("specific-humidity-log10.json").open("w") as f:
    json.dump(observations, f)