Preserving (non-NaN) missing values with safeguards¶

import ssl

ssl._create_default_https_context = ssl._create_stdlib_context

import copy
from pathlib import Path

import earthkit.plots
import eccodes
import humanize
import matplotlib as mpl
import numpy as np
import pandas as pd
from earthkit.plots.resample import Interpolate
from matplotlib import patheffects as PathEffects
from matplotlib import pyplot as plt
from numcodecs_safeguards import SafeguardedCodec
from numcodecs_wasm_pressio import Pressio
from numcodecs_wasm_sperr import Sperr
from numcodecs_wasm_sz3 import Sz3
from numcodecs_wasm_zfp import Zfp
from numcodecs_wasm_zstd import Zstd
from numcodecs_zero import ZeroCodec

data = Path("data") / "cems-ercnfdr" / "data.grib"

with data.open("rb") as f:
    gid = eccodes.codes_grib_new_from_file(f)
    grib_keys = []
    it = eccodes.codes_keys_iterator_new(gid)
    while eccodes.codes_keys_iterator_next(it):
        grib_keys.append(eccodes.codes_keys_iterator_get_name(it))
    eccodes.codes_keys_iterator_delete(it)
    eccodes.codes_release(gid)

with data.open("rb") as f:
    gid = eccodes.codes_grib_new_from_file(f)
    (long_name,) = eccodes.codes_get_array(gid, "name", ktype=str)
    (missing_value,) = eccodes.codes_get_array(gid, "missingValue")
    latitudes = eccodes.codes_get_array(gid, "latitudes")
    longitudes = eccodes.codes_get_array(gid, "longitudes")
    ercnfdr = eccodes.codes_get_array(gid, "values")
    eccodes.codes_release(gid)

def compute_corrections_percentage(
    my_ercnfdr: np.ndarray, orig_ercnfdr: np.ndarray
) -> float:
    return np.mean(
        ~(
            (my_ercnfdr == orig_ercnfdr)
            | (np.isnan(my_ercnfdr) & np.isnan(orig_ercnfdr))
        )
    )

import observe

observations = []

zstd = Zstd(level=22)

with observe.observe(zstd, observations):
    ercnfdr_zstd_enc = zstd.encode(ercnfdr)
    ercnfdr_zstd = zstd.decode(ercnfdr_zstd_enc)
ercnfdr_zstd_cr = ercnfdr.nbytes / ercnfdr_zstd_enc.nbytes

eb_abs = 1

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

with observe.observe(zfp, observations):
    ercnfdr_zfp_enc = zfp.encode(ercnfdr)
    ercnfdr_zfp = zfp.decode(ercnfdr_zfp_enc)
ercnfdr_zfp_cr = ercnfdr.nbytes / ercnfdr_zfp_enc.nbytes

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

with observe.observe(sz3, observations):
    ercnfdr_sz3_enc = sz3.encode(ercnfdr)
    ercnfdr_sz3 = sz3.decode(ercnfdr_sz3_enc)
ercnfdr_sz3_cr = ercnfdr.nbytes / ercnfdr_sz3_enc.nbytes

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

with observe.observe(sperr, observations):
    ercnfdr_sperr_enc = sperr.encode(ercnfdr)
    ercnfdr_sperr = sperr.decode(ercnfdr_sperr_enc)
ercnfdr_sperr_cr = ercnfdr.nbytes / ercnfdr_sperr_enc.nbytes

zero = ZeroCodec()

with observe.observe(zero, observations):
    ercnfdr_zero_enc = zero.encode(ercnfdr)
    ercnfdr_zero = zero.decode(ercnfdr_zero_enc)

ercnfdr_sg = dict()
ercnfdr_sg_cr = dict()

for codec in [
    zero,
    zfp,
    sz3,
    sperr,
]:
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[
            dict(kind="eb", type="abs", eb=eb_abs),
            dict(kind="same", value="missing_value", exclusive=True),
        ],
        fixed_constants=dict(missing_value=missing_value),
    )

    with observe.observe(sg, observations):
        ercnfdr_sg_enc = sg.encode(ercnfdr)
        ercnfdr_sg[codec.codec_id] = sg.decode(ercnfdr_sg_enc)
    ercnfdr_sg_cr[codec.codec_id] = ercnfdr.nbytes / np.asarray(ercnfdr_sg_enc).nbytes

ercnfdr_sg_lossless = dict()
ercnfdr_sg_lossless_cr = dict()

for codec in [
    zero,
    zfp,
    sz3,
    sperr,
]:
    sg = SafeguardedCodec(
        codec=codec,
        safeguards=[
            dict(kind="eb", type="abs", eb=eb_abs),
            dict(kind="same", value="missing_value", exclusive=True),
        ],
        fixed_constants=dict(missing_value=missing_value),
        # produce lossless corrections and refine them with iteration
        compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
    )

    with observe.observe(sg, observations):
        ercnfdr_sg_lossless_enc = sg.encode(ercnfdr)
        ercnfdr_sg_lossless[codec.codec_id] = sg.decode(ercnfdr_sg_lossless_enc)
    ercnfdr_sg_lossless_cr[codec.codec_id] = (
        ercnfdr.nbytes / np.asarray(ercnfdr_sg_lossless_enc).nbytes
    )

from numcodecs.packbits import PackBits
from numcodecs_combinators.stack import CodecStack
from numcodecs_mask import MaskMetaCodec

ercnfdr_mask = dict()
ercnfdr_mask_cr = dict()

for codec in [
    zfp,
    sz3,
    sperr,
]:
    mask = MaskMetaCodec(
        mask=int(missing_value),
        codec=codec,
        # bitpack the mask and compress it with default Zstd
        bitmap_codec=CodecStack(PackBits(), Zstd(level=3)),
    )

    with observe.observe(mask, observations):
        ercnfdr_mask_enc = mask.encode(ercnfdr)
        ercnfdr_mask[codec.codec_id] = mask.decode(ercnfdr_mask_enc)
    ercnfdr_mask_cr[codec.codec_id] = (
        ercnfdr.nbytes / np.asarray(ercnfdr_mask_enc).nbytes
    )

ercnfdr_mask_sg = dict()
ercnfdr_mask_sg_cr = dict()

for codec in [
    zfp,
    sz3,
    sperr,
]:
    mask_sg = SafeguardedCodec(
        codec=MaskMetaCodec(
            mask=int(missing_value),
            codec=codec,
            # bitpack the mask and compress it with default Zstd
            bitmap_codec=CodecStack(PackBits(), Zstd(level=3)),
        ),
        safeguards=[
            dict(kind="eb", type="abs", eb=eb_abs),
            dict(kind="same", value="missing_value", exclusive=True),
        ],
        fixed_constants=dict(missing_value=missing_value),
    )

    with observe.observe(mask_sg, observations):
        ercnfdr_mask_sg_enc = mask_sg.encode(ercnfdr)
        ercnfdr_mask_sg[codec.codec_id] = mask_sg.decode(ercnfdr_mask_sg_enc)
    ercnfdr_mask_sg_cr[codec.codec_id] = (
        ercnfdr.nbytes / np.asarray(ercnfdr_mask_sg_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_nan = np.where(self._data == missing_value, np.nan, self._data)
        buf_nan = np.where(buf == missing_value, np.nan, buf)
        violations = np.mean(
            ~(
                (np.abs(buf_nan - data_nan) <= eb_abs)
                | (np.isnan(buf_nan) & np.isnan(data_nan))
            )
        )
        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)

old_scale_dashes = mpl.lines._scale_dashes


def scale_dashes(offset, dashes, lw):
    if lw == 0:
        return offset, dashes
    return old_scale_dashes(offset, dashes, lw)


mpl.lines._scale_dashes = scale_dashes

def plot_ercnfdr(
    my_ercnfdr: np.ndarray,
    cr,
    chart,
    title,
    eb_abs,
    error=False,
    corr=None,
    my_ercnfdr_mask=None,
    cr_mask=None,
    inset=True,
):
    x = longitudes
    y = latitudes
    z = np.where(my_ercnfdr == missing_value, np.nan, my_ercnfdr)

    source = earthkit.plots.sources.get_source(x=x, y=y, z=z)
    style = copy.deepcopy(earthkit.plots.styles.auto.guess_style(source))
    style._levels = earthkit.plots.styles.levels.Levels(np.linspace(-1, 81, 22))
    style._legend_kwargs["ticks"] = np.linspace(-1, 81, 5)

    extend_left = np.nanmin(z) < -1
    extend_right = np.nanmax(z) > 81

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

    style._legend_kwargs["extend"] = extend

    interpolate = Interpolate()

    x, y, z = interpolate.apply(
        x=x,
        y=y,
        z=z,
        source_crs=source.crs,
        target_crs=chart.crs,
    )

    ercnfdr_x = longitudes
    ercnfdr_y = latitudes
    ercnfdr_z = np.where(ercnfdr == missing_value, np.nan, ercnfdr)

    ercnfdr_x, ercnfdr_y, ercnfdr_z = interpolate.apply(
        x=ercnfdr_x,
        y=ercnfdr_y,
        z=ercnfdr_z,
        source_crs=source.crs,
        target_crs=chart.crs,
    )

    chart.ax.set_global()
    chart.ax.fill_between(
        [0, 1],
        [1, 1],
        hatch="X",
        edgecolor="white",
        facecolor="lightgrey",
        transform=chart.ax.transAxes,
        zorder=-12,
    )

    chart.pcolormesh(
        x=x,
        y=y,
        z=z,
        style=style,
        norm=mpl.colors.BoundaryNorm(np.linspace(-1, 81, 22), ncolors=21, clip=True),
        zorder=-11,
        rasterized=True,
    )

    with plt.rc_context(
        {
            "hatch.color": "black",
            "hatch.linewidth": 2,
        }
    ):
        chart.contourf(
            x=ercnfdr_x,
            y=ercnfdr_y,
            z=np.isnan(ercnfdr_z) & ~np.isnan(z),
            colors=["none"],
            levels=[-0.5, 0.9, 1.5],
            hatches=[None, "X"],
            legend_style=None,
            zorder=-10,
        )
    with plt.rc_context(
        {
            "hatch.color": "white",
            "hatch.linewidth": 1,
        }
    ):
        chart.contourf(
            x=ercnfdr_x,
            y=ercnfdr_y,
            z=np.isnan(ercnfdr_z) & ~np.isnan(z),
            colors=["none"],
            levels=[-0.5, 0.9, 1.5],
            hatches=[None, "X"],
            legend_style=None,
            zorder=-10,
        )

    if error:
        if corr is not None:
            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.set_global()
            axin.scatter(
                longitudes,
                latitudes,
                s=1,
                c=~(my_ercnfdr == corr),
                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:
            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.set_global()
            axin.scatter(
                longitudes,
                latitudes,
                s=1,
                c=~(
                    (np.abs(my_ercnfdr - ercnfdr) <= eb_abs)
                    & ((my_ercnfdr == missing_value) == (ercnfdr == missing_value))
                ),
                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")],
            )

    chart.title(title)

    if error:
        my_ercnfdr_nan = np.where(my_ercnfdr == missing_value, np.nan, my_ercnfdr)
        ercnfdr_nan = np.where(ercnfdr == missing_value, np.nan, ercnfdr)

        err_v = np.mean(
            ~(
                (np.abs(my_ercnfdr_nan - ercnfdr_nan) <= eb_abs)
                | (np.isnan(my_ercnfdr_nan) & np.isnan(ercnfdr_nan))
            )
        )
        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_ercnfdr_mask is not None:
            my_ercnfdr_mask_nan = np.where(
                my_ercnfdr_mask == missing_value, np.nan, my_ercnfdr_mask
            )

            err_mask_v = np.mean(
                ~(
                    (np.abs(my_ercnfdr_mask_nan - ercnfdr_nan) <= eb_abs)
                    | (np.isnan(my_ercnfdr_mask_nan) & np.isnan(ercnfdr_nan))
                )
            )
            err_mask_v = (
                0
                if err_mask_v == 0
                else np.format_float_positional(
                    100 * err_mask_v, precision=1, min_digits=1
                )
                + "%"
            )
            if err_mask_v == "0.0%":
                err_mask_v = "<0.05%"

        t = chart.ax.text(
            0.95,
            0.1,
            f"V={err_v}" + ("" if my_ercnfdr_mask is None else f" ({err_mask_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_mask is None else rf" ($\times$ {np.round(cr_mask, 2)})")
        if error
        else humanize.naturalsize(ercnfdr.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:
        if m != "legend":
            getattr(chart, m)()

    chart.legend(
        label=long_name.split(" (")[0] if error else long_name.replace(" (", "\n(")
    )

    counts, bins = np.histogram(z.flatten(), range=(-1, 81), 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 + 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(z < -1),
            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(z > 81),
            width=(bins[-1] - bins[0]) / len(counts),
            color=cb.cmap(cb.norm(midpoints[-1])),
        )
    cax.bar(
        bins[-1] + (bins[-1] - bins[-2]) * (extend_right * 2 + 2 + 1) / 2,
        height=np.sum(np.isnan(z)),
        width=(bins[-1] - bins[0]) / len(counts),
        color="lightgrey",
        edgecolor="white",
        lw=0,
        hatch="XXXX",
    )
    q1, q2, q3 = np.quantile(z.flatten(), [0.25, 0.5, 0.75])
    cax.axvline(z.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(z.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(
        -1 - (bins[-1] - bins[-2]) * extend_left,
        81 + (bins[-1] - bins[-2]) * (extend_right + 2),
    )
    cax.set_xticks([])
    cax.set_yticks([])
    cax.spines[:].set_visible(False)

def table_ercnfdr(
    my_ercnfdr: np.ndarray,
    cr,
    title,
    eb_abs,
    corr=None,
):
    err_inf = np.amax(np.abs(my_ercnfdr - ercnfdr))
    err_2 = np.sqrt(np.mean(np.square(my_ercnfdr - ercnfdr)))

    my_ercnfdr_nan = np.where(my_ercnfdr == missing_value, np.nan, my_ercnfdr)
    ercnfdr_nan = np.where(ercnfdr == missing_value, np.nan, ercnfdr)

    err_v = np.mean(
        ~(
            (np.abs(my_ercnfdr_nan - ercnfdr_nan) <= eb_abs)
            | (np.isnan(my_ercnfdr_nan) & np.isnan(ercnfdr_nan))
        )
    )
    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_ercnfdr, 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]],
            "Meta": [title[1]],
            "Safeguarded": [title[2]],
            "Corrections": [title[3]],
            r"$L_{\infty}$": [
                f"{err_inf:.02}",
            ],
            r"$L_{2}$": [
                f"{err_2:.02}",
            ],
            "V": [err_v],
            "C": [corr],
            "CR": [
                rf"$\times$ {np.round(cr, 2)}",
            ],
        }
    )

import json

with Path("observations").joinpath("missing.json").open("w") as f:
    json.dump(observations, f)