Preserving a multi-variable pointwise quantity of interest (QoI) with safeguards¶

# Since numcodecs-safeguards only supports single-variable safeguarding, we
# stack the u and v variables into a combined variable.
ERA5_UV = np.stack([ERA5["u"].values, ERA5["v"].values], axis=0)
ERA5.u.dims, ERA5_UV.shape

from numcodecs_wasm_qpet_sperr import QpetSperr

qpet = QpetSperr(
    mode="qoi-symbolic", qoi="x^2", qoi_block_size=(2, 1, 1), qoi_pwe=ke_eb_abs
)

with observe.observe(qpet, observations):
    ERA5_UV_qpet_enc = qpet.encode(ERA5_UV)
    ERA5_UV_qpet = qpet.decode(ERA5_UV_qpet_enc)

ERA5_qpet = ERA5.copy(data=dict(u=ERA5_UV_qpet[0], v=ERA5_UV_qpet[1]))
ERA5_qpet_cr = ERA5_UV.nbytes / ERA5_UV_qpet_enc.nbytes

from numcodecs_wasm_pressio import Pressio

ERA5_optzconfig = dict()
ERA5_optzconfig_cr = dict()

for codec, parameter, lower_bound in [
    (zfp, "tolerance", 1e-3),  # initial guess
    (sz3, "eb_abs", 1e-3),  # initial guess
    (sperr, "pwe", 1e-3),  # initial guess
]:
    optzconfig = Pressio(
        compressor_id="opt",
        compressor_config={
            "opt:output": ["composite:score"],
            "opt:inputs": [f"numcodecs.rs:{parameter}"],
            "opt:lower_bound": np.log(lower_bound),
            "opt:upper_bound": np.log(eb_abs),
            "opt:max_iterations": 25,
            "opt:objective_mode_name": "max",
        },
        early_config={
            "opt:compressor": "pressio",
            "pressio:compressor": "numcodecs.rs",
            **{
                f"numcodecs.rs:{k}": f"e-{v}" if k == "id" else v
                for k, v in codec.get_config().items()
            },
            "opt:search": "fraz",
            "pressio:metric": "composite",
            "composite:plugins": ["size", "numcodecs.rs-metric"],
            "composite:scripts": [
                """
                violations = metrics["numcodecs.rs-metric:decompression"]
                if violations > 0 then
                    return "score", -violations
                else
                    return "score", metrics["size:compression_ratio"]
                end
                """
            ],
            "numcodecs.rs-metric:id": "safety-violations-metric",
        },
    )

    with observe.observe(optzconfig, observations):
        ERA5_UV_optzconfig_enc = optzconfig.encode(ERA5_UV)
        ERA5_UV_optzconfig = optzconfig.decode(ERA5_UV_optzconfig_enc)

    ERA5_optzconfig[codec.codec_id] = ERA5.copy(
        data=dict(u=ERA5_UV_optzconfig[0], v=ERA5_UV_optzconfig[1])
    )
    ERA5_optzconfig_cr[codec.codec_id] = ERA5_UV.nbytes / ERA5_UV_optzconfig_enc.nbytes

fig = earthkit.plots.Figure(
    size=(10, 23),
    rows=6,
    columns=2,
)

plot_kinetic_energy(
    ERA5, 1.0, fig.add_map(0, 0), "Original", span=50, ke_eb_abs=ke_eb_abs
)
plot_kinetic_energy(
    ERA5_zfp,
    ERA5_zfp_cr,
    fig.add_map(1, 0),
    r"ZFP($\epsilon_{{abs}}$)",
    span=0.25,
    ke_eb_abs=ke_eb_abs,
    error=True,
)
plot_kinetic_energy(
    ERA5_sz3,
    ERA5_sz3_cr,
    fig.add_map(2, 0),
    r"SZ3($\epsilon_{{abs}}$)",
    span=2.5,
    ke_eb_abs=ke_eb_abs,
    error=True,
)
plot_kinetic_energy(
    ERA5_sperr,
    ERA5_sperr_cr,
    fig.add_map(3, 0),
    r"SPERR($\epsilon_{{abs}}$)",
    span=2.5,
    ke_eb_abs=ke_eb_abs,
    error=True,
)

plot_kinetic_energy(
    ERA5_sg["zero"],
    ERA5_sg_cr["zero"],
    fig.add_map(0, 1),
    r"Safeguarded(0, $\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    corr=ERA5_zero,
    my_ERA5_it=ERA5_sg_it["zero"],
    cr_it=ERA5_sg_it_cr["zero"],
)
plot_kinetic_energy(
    ERA5_sg["zfp.rs"],
    ERA5_sg_cr["zfp.rs"],
    fig.add_map(1, 1),
    r"Safeguarded(ZFP, $\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    corr=ERA5_zfp,
    my_ERA5_it=ERA5_sg_it["zfp.rs"],
    cr_it=ERA5_sg_it_cr["zfp.rs"],
)
plot_kinetic_energy(
    ERA5_sg["sz3.rs"],
    ERA5_sg_cr["sz3.rs"],
    fig.add_map(2, 1),
    r"Safeguarded(SZ3, $\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    corr=ERA5_sz3,
    my_ERA5_it=ERA5_sg_it["sz3.rs"],
    cr_it=ERA5_sg_it_cr["sz3.rs"],
)
plot_kinetic_energy(
    ERA5_sg["sperr.rs"],
    ERA5_sg_cr["sperr.rs"],
    fig.add_map(3, 1),
    r"Safeguarded(SPERR, $\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    corr=ERA5_sperr,
    my_ERA5_it=ERA5_sg_it["sperr.rs"],
    cr_it=ERA5_sg_it_cr["sperr.rs"],
)

plot_kinetic_energy(
    ERA5_optzconfig["zfp.rs"],
    ERA5_optzconfig_cr["zfp.rs"],
    fig.add_map(4, 0),
    r"OptZConfig(ZFP, $\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    inset=False,
)
plot_kinetic_energy(
    ERA5_optzconfig["sz3.rs"],
    ERA5_optzconfig_cr["sz3.rs"],
    fig.add_map(4, 1),
    r"OptZConfig(SZ3, $\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    inset=False,
)
plot_kinetic_energy(
    ERA5_optzconfig["sperr.rs"],
    ERA5_optzconfig_cr["sperr.rs"],
    fig.add_map(5, 0),
    r"OptZConfig(SPERR, $\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    inset=False,
)

plot_kinetic_energy(
    ERA5_qpet,
    ERA5_qpet_cr,
    fig.add_map(5, 1),
    r"QPET-SPERR($\epsilon_{{QoI,abs}}$)",
    span=ke_eb_abs,
    ke_eb_abs=ke_eb_abs,
    error=True,
    inset=False,
)

fig.save(Path("plots") / "kinetic-energy.pdf")

ke_sg_table = pd.concat(
    [
        table_kinetic_energy(
            ERA5_sg_lossless["zero"],
            ERA5_sg_lossless_cr["zero"],
            ["0", r"$\epsilon_{QoI,abs}$", "lossless"],
            ke_eb_abs,
            ERA5_zero,
        ),
        table_kinetic_energy(
            ERA5_sg["zero"],
            ERA5_sg_cr["zero"],
            ["0", r"$\epsilon_{QoI,abs}$", "one-shot"],
            ke_eb_abs,
            ERA5_zero,
        ),
        table_kinetic_energy(
            ERA5_sg_it["zero"],
            ERA5_sg_it_cr["zero"],
            ["0", r"$\epsilon_{QoI,abs}$", "iterative"],
            ke_eb_abs,
            ERA5_zero,
        ),
        table_kinetic_energy(
            ERA5_zfp,
            ERA5_zfp_cr,
            [r"ZFP($\epsilon_{abs}$)", "-", ""],
            ke_eb_abs,
            None,
        ),
        table_kinetic_energy(
            ERA5_sg_lossless["zfp.rs"],
            ERA5_sg_lossless_cr["zfp.rs"],
            [r"ZFP($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "lossless"],
            ke_eb_abs,
            ERA5_zfp,
        ),
        table_kinetic_energy(
            ERA5_sg["zfp.rs"],
            ERA5_sg_cr["zfp.rs"],
            [r"ZFP($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "one-shot"],
            ke_eb_abs,
            ERA5_zfp,
        ),
        table_kinetic_energy(
            ERA5_sg_it["zfp.rs"],
            ERA5_sg_it_cr["zfp.rs"],
            [r"ZFP($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "iterative"],
            ke_eb_abs,
            ERA5_zfp,
        ),
        table_kinetic_energy(
            ERA5_optzconfig["zfp.rs"],
            ERA5_optzconfig_cr["zfp.rs"],
            [r"OptZConfig(ZFP)", r"$\epsilon_{QoI,abs}$", ""],
            ke_eb_abs,
            None,
        ),
        table_kinetic_energy(
            ERA5_sz3,
            ERA5_sz3_cr,
            [r"SZ3($\epsilon_{abs}$)", "-", ""],
            ke_eb_abs,
            None,
        ),
        table_kinetic_energy(
            ERA5_sg_lossless["sz3.rs"],
            ERA5_sg_lossless_cr["sz3.rs"],
            [r"SZ3($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "lossless"],
            ke_eb_abs,
            ERA5_sz3,
        ),
        table_kinetic_energy(
            ERA5_sg["sz3.rs"],
            ERA5_sg_cr["sz3.rs"],
            [r"SZ3($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "one-shot"],
            ke_eb_abs,
            ERA5_sz3,
        ),
        table_kinetic_energy(
            ERA5_sg_it["sz3.rs"],
            ERA5_sg_it_cr["sz3.rs"],
            [r"SZ3($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "iterative"],
            ke_eb_abs,
            ERA5_sz3,
        ),
        table_kinetic_energy(
            ERA5_optzconfig["sz3.rs"],
            ERA5_optzconfig_cr["sz3.rs"],
            [r"OptZConfig(SZ3)", r"$\epsilon_{QoI,abs}$", ""],
            ke_eb_abs,
            None,
        ),
        table_kinetic_energy(
            ERA5_sperr,
            ERA5_sperr_cr,
            [r"SPERR($\epsilon_{abs}$)", "-", ""],
            ke_eb_abs,
            None,
        ),
        table_kinetic_energy(
            ERA5_sg_lossless["sperr.rs"],
            ERA5_sg_lossless_cr["sperr.rs"],
            [r"SPERR($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "lossless"],
            ke_eb_abs,
            ERA5_sperr,
        ),
        table_kinetic_energy(
            ERA5_sg["sperr.rs"],
            ERA5_sg_cr["sperr.rs"],
            [r"SPERR($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "one-shot"],
            ke_eb_abs,
            ERA5_sperr,
        ),
        table_kinetic_energy(
            ERA5_sg_it["sperr.rs"],
            ERA5_sg_it_cr["sperr.rs"],
            [r"SPERR($\epsilon_{abs}$)", r"$\epsilon_{QoI,abs}$", "iterative"],
            ke_eb_abs,
            ERA5_sperr,
        ),
        table_kinetic_energy(
            ERA5_optzconfig["sperr.rs"],
            ERA5_optzconfig_cr["sperr.rs"],
            [r"OptZConfig(SPERR)", r"$\epsilon_{QoI,abs}$", ""],
            ke_eb_abs,
            None,
        ),
        table_kinetic_energy(
            ERA5_qpet,
            ERA5_qpet_cr,
            ["QPET-SPERR", r"$\epsilon_{QoI,abs}$", ""],
            ke_eb_abs,
            None,
        ),
        table_kinetic_energy(
            ERA5_zstd,
            ERA5_zstd_cr,
            ["ZSTD(22)", "-", ""],
            ke_eb_abs,
            None,
        ),
    ]
).set_index(["Compressor", "Safeguarded", "Corrections"])

Path("tables").joinpath("kinetic-energy.tex").write_text(
    ke_sg_table.to_latex(escape=False)
    .replace("%", r"\%")
    .replace("\\cline{1-10} \\cline{2-10}\n\\bottomrule", "\\bottomrule")
)

ke_sg_table

fig = earthkit.plots.Figure(
    size=(10, 5),
    rows=1,
    columns=2,
)

chart = fig.add_map(0, 0)

# Original
da = ERA5_KE.sel(longitude=slice(180, 300))
# plot the square root of kinetic energy to better capture scale
da = np.sqrt(da)
da.attrs.update(long_name=f"sqrt({da.long_name})", units="m**1 s**-1")

# compute the default style that earthkit.maps would apply
source_original = earthkit.plots.sources.XarraySource(da)
style_original = copy.deepcopy(
    earthkit.plots.styles.auto.guess_style(
        source_original,
        units=source_original.units,
    )
)
style_original._levels = earthkit.plots.styles.levels.Levels(np.linspace(0, 50, 21))
style_original._legend_kwargs["ticks"] = np.linspace(0, 50, 5)
style_original._colors = "viridis"
style_original._legend_kwargs["extend"] = "neither"

chart.pcolormesh(da, style=style_original, zorder=-11)

t = chart.ax.text(
    1 / 6,
    0.9,
    "Original",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

chart.legend()

# # SPERR
my_ERA5_KE = compute_kinetic_energy(ERA5_sperr)
with xr.set_options(keep_attrs=True):
    da = (my_ERA5_KE - ERA5_KE).compute()
da.attrs.update(long_name=f"Absolute error over {da.long_name}")

is_err = ~(np.abs(my_ERA5_KE - ERA5_KE) <= ke_eb_abs)

# compute the default style that earthkit.maps would apply
source_error = earthkit.plots.sources.XarraySource(da)
style_error = copy.deepcopy(
    earthkit.plots.styles.auto.guess_style(
        source_error,
        units=source_error.units,
    )
)
style_error._levels = earthkit.plots.styles.levels.Levels(
    np.linspace(-ke_eb_abs, ke_eb_abs, 21)
)
style_error._legend_kwargs["ticks"] = np.linspace(-ke_eb_abs, ke_eb_abs, 5)
style_error._colors = "coolwarm"
style_error._legend_kwargs["extend"] = "both"

chart.pcolormesh(da.sel(longitude=slice(300, 360)), style=style_error, zorder=-11)
chart.pcolormesh(da.sel(longitude=slice(0, 60)), style=style_error, zorder=-11)

t = chart.ax.text(
    0.5,
    0.9,
    "SPERR",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

t = chart.ax.text(
    0.5,
    0.5,
    rf"$\times$ {np.round(ERA5_sperr_cr, 2)}",
    ha="center",
    va="center",
    transform=chart.ax.transAxes,
    color="mistyrose",
    fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])

err_v_sperr = np.mean(~(np.abs(my_ERA5_KE - ERA5_KE) <= ke_eb_abs))
err_v_sperr = (
    0
    if err_v_sperr == 0
    else np.format_float_positional(100 * err_v_sperr, precision=1, min_digits=1) + "%"
)
if err_v_sperr == "0.0%":
    err_v_sperr = "<0.05%"
t = chart.ax.text(
    0.5,
    0.1,
    f"V={err_v_sperr}",
    ha="center",
    va="bottom",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

# QPET-SPERR
my_ERA5_KE = compute_kinetic_energy(ERA5_qpet)
with xr.set_options(keep_attrs=True):
    da = (my_ERA5_KE - ERA5_KE).compute()
da.attrs.update(long_name=f"Absolute error over {da.long_name}")
da = da.sel(longitude=slice(60, 180))
chart.pcolormesh(da, style=style_error, zorder=-11)

t = chart.ax.text(
    5 / 6,
    0.9,
    "QPET-SPERR",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

t = chart.ax.text(
    5 / 6,
    0.5,
    rf"$\times$ {np.round(ERA5_qpet_cr, 2)}",
    ha="center",
    va="center",
    transform=chart.ax.transAxes,
    color="lightgreen",
    fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])

err_v = np.mean(~(np.abs(my_ERA5_KE - ERA5_KE) <= ke_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(
    5 / 6,
    0.1,
    f"V={err_v}",
    ha="center",
    va="bottom",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

chart.ax.set_rasterization_zorder(-10)

chart.ax.axvline(-60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(-60, c="black", ls=(6, (4, 4)), lw=2)
chart.ax.axvline(+60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(+60, c="black", ls=(6, (4, 4)), lw=2)

chart.title("Without safeguards")

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 = fig.add_map(0, 1)

# Corrections and Errors
is_corr = (ERA5_sg["sperr.rs"]["u"] != ERA5_sperr["u"]) | (
    ERA5_sg["sperr.rs"]["v"] != ERA5_sperr["v"]
)
is_corr_it = (ERA5_sg_it["sperr.rs"]["u"] != ERA5_sperr["u"]) | (
    ERA5_sg_it["sperr.rs"]["v"] != ERA5_sperr["v"]
)

chart.pcolormesh(
    is_corr.sel(longitude=slice(300, 340)),
    no_style=True,
    cmap=mpl.colors.ListedColormap(["white", "green"]),
    zorder=-11,
    legend_style=None,
)
chart.pcolormesh(
    is_err.sel(longitude=slice(340, 360)),
    no_style=True,
    cmap=mpl.colors.ListedColormap(["white", "red"]),
    zorder=-11,
    legend_style=None,
)
chart.pcolormesh(
    is_err.sel(longitude=slice(0, 20)),
    no_style=True,
    cmap=mpl.colors.ListedColormap(["white", "red"]),
    zorder=-11,
    legend_style=None,
)
chart.pcolormesh(
    is_corr_it.sel(longitude=slice(20, 60)),
    no_style=True,
    cmap=mpl.colors.ListedColormap(["white", "limegreen"]),
    zorder=-11,
    legend_style=None,
)

t = chart.ax.text(
    0.5,
    0.9,
    "SPERR",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
    0.5,
    0.5,
    f"V={err_v_sperr}",
    ha="center",
    va="center",
    transform=chart.ax.transAxes,
    rotation=90,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

corr = compute_corrections_percentage(ERA5_sg["sperr.rs"], ERA5_sperr)
corr = (
    0
    if corr == 0
    else np.format_float_positional(100 * corr, precision=1, min_digits=1) + "%"
)
if corr == "0.0%":
    corr = "<0.05%"
t = chart.ax.text(
    7 / 18,
    0.9,
    "Sg",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
    7 / 18,
    0.5,
    f"C={corr}",
    ha="center",
    va="center",
    transform=chart.ax.transAxes,
    rotation=90,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

corr_it = compute_corrections_percentage(ERA5_sg_it["sperr.rs"], ERA5_sperr)
corr_it = (
    0
    if corr_it == 0
    else np.format_float_positional(100 * corr_it, precision=1, min_digits=1) + "%"
)
if corr_it == "0.0%":
    corr_it = "<0.05%"
t = chart.ax.text(
    11 / 18,
    0.9,
    "Sg[it]",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
    11 / 18,
    0.5,
    f"C={corr_it}",
    ha="center",
    va="center",
    transform=chart.ax.transAxes,
    rotation=90,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

# Safeguarded(SPERR)
my_ERA5_KE = compute_kinetic_energy(ERA5_sg["sperr.rs"])
with xr.set_options(keep_attrs=True):
    da = (my_ERA5_KE - ERA5_KE).compute()
da.attrs.update(long_name=f"Absolute error over {da.long_name}")
da = da.sel(longitude=slice(180, 300))
chart.pcolormesh(da, style=style_error, zorder=-11)

t = chart.ax.text(
    1 / 6,
    0.9,
    "Sg(SPERR)",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

t = chart.ax.text(
    1 / 6,
    0.5,
    rf"$\times$ {np.round(ERA5_sg_cr['sperr.rs'], 2)}",
    ha="center",
    va="center",
    transform=chart.ax.transAxes,
    color="lightgreen",
    fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])

err_v = np.mean(~(np.abs(my_ERA5_KE - ERA5_KE) <= ke_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(
    1 / 6,
    0.1,
    f"V={err_v}",
    ha="center",
    va="bottom",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

chart.legend(extend="neither")

# Safeguarded[it](SPERR)
my_ERA5_KE = compute_kinetic_energy(ERA5_sg_it["sperr.rs"])
with xr.set_options(keep_attrs=True):
    da = (my_ERA5_KE - ERA5_KE).compute()
da.attrs.update(long_name=f"Absolute error over {da.long_name}")
da = da.sel(longitude=slice(60, 180))
chart.pcolormesh(da, style=style_error, zorder=-11)

t = chart.ax.text(
    5 / 6,
    0.9,
    "Sg[it](SPERR)",
    ha="center",
    va="top",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

t = chart.ax.text(
    5 / 6,
    0.5,
    rf"$\times$ {np.round(ERA5_sg_it_cr['sperr.rs'], 2)}",
    ha="center",
    va="center",
    transform=chart.ax.transAxes,
    color="lightgreen",
    fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])

err_v = np.mean(~(np.abs(my_ERA5_KE - ERA5_KE) <= ke_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(
    5 / 6,
    0.1,
    f"V={err_v}",
    ha="center",
    va="bottom",
    transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

chart.ax.set_rasterization_zorder(-10)

chart.ax.axvline(-20, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(-20, c="black", ls=(6, (4, 4)), lw=2)
chart.ax.axvline(+20, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(+20, c="black", ls=(6, (4, 4)), lw=2)

chart.ax.axvline(-60, c="black", lw=1)
chart.ax.axvline(+60, c="black", lw=1)

chart.title("Safeguarded: preserve kinetic energy")

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)()

fig.save(Path("plots") / "kinetic-energy-summary.pdf")