Compression Safeguards

from pathlib import Path

import netCDF4
import numpy as np
import xarray as xr

data = Path("data")

import earthkit.plots

from quickplot import quickplot

from numcodecs_safeguards import SafeguardedCodec
from numcodecs_zero import ZeroCodec

?SafeguardedCodec

eb_abs = 0.1

SafeguardedCodec(
    codec=ZeroCodec(),  # for example, any numcodecs codec works
    safeguards=[
        {"kind": "eb", "type": "abs", "eb": eb_abs},
    ],
)

eb_rel = 0.01  # 1 %

SafeguardedCodec(
    codec=ZeroCodec(),  # for example, any numcodecs codec works
    safeguards=[
        {"kind": "eb", "type": "rel", "eb": eb_rel},
    ],
)

eb_ratio = 1.01  # between x / 1.01 and x * 1.01

SafeguardedCodec(
    codec=ZeroCodec(),  # for example, any numcodecs codec works
    safeguards=[
        {"kind": "eb", "type": "ratio", "eb": eb_ratio},
    ],
)

from numcodecs_combinators.stack import CodecStack
from numcodecs_replace import Replacement, ReplaceFilterCodec

eb_abs = 0.1

SafeguardedCodec(
    # the replace filter codec must be used *inside* the safeguarded codec
    #  so that the safeguards see the NaN values in the original data and
    #  can reconstruct them during decompression
    codec=CodecStack(
        ReplaceFilterCodec(replacements={np.nan: "finite_mean"}),
        ZeroCodec(),  # for example, any numcodecs codec works
    ),
    safeguards=[
        {"kind": "eb", "type": "abs", "eb": eb_abs},
    ],
)

eb_abs = 0.1

sentinel = 9999

SafeguardedCodec(
    codec=ZeroCodec(),  # for example, any numcodecs codec works
    safeguards=[
        {"kind": "eb", "type": "abs", "eb": eb_abs},
        {"kind": "same", "value": sentinel, "exclusive": True},
    ],
)

eb_abs_qoi = 0.1

SafeguardedCodec(
    codec=ZeroCodec(),  # for example, any numcodecs codec works
    safeguards=[
        {"kind": "qoi_eb_pw", "qoi": "square(x)", "type": "abs", "eb": eb_abs_qoi},
    ],
)

# Load the data
ds = xr.open_dataset(
    data / "hplp" / "hplp_sfc_regridded_t_025deg_levels_steps_204_216_228_240.nc",
    engine="netcdf4",
    decode_timedelta=True,
)
da = ds["2t"]

eb_abs = 0.1  # 0.1 K

codec = SafeguardedCodec(
    codec=ZeroCodec(),  # safeguards need to correct constant zero prediction
    safeguards=[
        # bound absolute error over the data
        {"kind": "eb", "type": "abs", "eb": eb_abs},
        # preserve global minimum and maximum
        {"kind": "sign", "offset": "$x_min"},
        {"kind": "sign", "offset": "$x_max"},
    ],
)

# encode and decode the data
da_enc = codec.encode(da.values)
da_dec = da.copy(data=codec.decode(da_enc))

# plot a comparison figure
fig = earthkit.plots.Figure(
    size=(15, 4),
    rows=1,
    columns=3,
)

quickplot(da, fig.add_map(0, 0), title="Original {default_title}")
quickplot(
    da_dec,
    fig.add_map(0, 1),
    title="Safeguarded(0)",
    cr=da.nbytes / np.array(da_enc).nbytes,
)
quickplot(da_dec - da, fig.add_map(0, 2), error=True, title="Compression Error")

fig.show()

da_chunked = da.chunk(lat=70, lon=100)
display(da_chunked.data)
da_chunked.dims

import xarray_safeguards

eb_rel_qoi = 0.01  # 1 % over the quantity of interest

safeguards = [
    {
        "kind": "qoi_eb_stencil",
        "qoi": """
            return finite_difference(
                x,                        # over the data
                order=1,                  # 1st order derivative
                accuracy=2,               # 2nd order accurate
                type=0,                   # central
                axis=-1,                  # longitude axis
                grid_centre=c["$d_lon"],  # coordinate, provided by xarray-safeguards
                grid_period=360,          # periodic over 360 degrees
            );
        """,
        "neighbourhood": [
            # periodic longitude axis, -1,+1 stencil for finite difference
            {"axis": -1, "before": 1, "after": 1, "boundary": "wrap"},
        ],
        "type": "rel",
        "eb": eb_rel_qoi,
    }
]

from numcodecs_combinators.stack import CodecStack

# lazily produce a prediction, i.e. the lossy decompressed data
da_prediction = CodecStack(
    ZeroCodec(),  # for example, any numcodecs codec works
).encode_decode_data_array(da_chunked)

# lazily produce the correction, later persist the result
da_correction = xarray_safeguards.produce_data_array_correction(
    data=da_chunked,
    prediction=da_prediction,
    safeguards=safeguards,
).persist()

# lazily apply the correction, later persist the result
da_corrected = xarray_safeguards.apply_data_array_correction(
    prediction=da_prediction,
    correction=da_correction,
).persist()

# xarray-safeguards, like compression-safeguards, performs no lossless
# compression of the corrections
# here we use the lossless corrections compression from numcodecs-safeguards
import numcodecs_safeguards

da_correction_nbytes = np.array(
    numcodecs_safeguards.lossless._default_lossless_for_corrections().encode(
        da_correction.values
    )
).nbytes

def differentiate_along_longitude(da: xr.DataArray) -> xr.DataArray:
    da_wrapped = da.pad(lon=1, mode="wrap").assign_coords(
        lon=da.lon.pad(lon=1, mode="reflect", reflect_type="odd"),
    )

    da_dXdLon = da_wrapped.differentiate("lon")
    da_dXdLon.attrs.update(
        long_name=f"{da.long_name} derivative along longitude",
        units=f"{da.units} degree**-1",
    )

    return da_dXdLon.sel(lon=slice(da.lon.min(), da.lon.max()))

# plot a comparison figure
fig = earthkit.plots.Figure(
    size=(15, 4),
    rows=1,
    columns=3,
)

da_deriv = differentiate_along_longitude(da_chunked.compute())
da_corrected_deriv = differentiate_along_longitude(da_corrected.compute())

quickplot(
    da_deriv,
    fig.add_map(0, 0),
    title="Original {default_title}",
    error=True,
    vrange=(-10, 10),
)
quickplot(
    da_corrected_deriv,
    fig.add_map(0, 1),
    title="Safeguarded(0), chunked",
    cr=da.nbytes / da_correction_nbytes,
    error=True,
    vrange=(-10, 10),
)
quickplot(
    ((da_corrected_deriv - da_deriv) / da_deriv).assign_attrs(
        long_name="relative error", units="%"
    ),
    fig.add_map(0, 2),
    error=True,
    title="Relative Compression Error",
)

fig.show()