Preview: Incrementally applying safeguards for tighter and less conservative corrections¶

import ssl

ssl._create_default_https_context = ssl._create_stdlib_context

import re
from collections.abc import Collection, Mapping, Set
from itertools import product
from pathlib import Path
from typing import ClassVar, Literal

import earthkit.plots
import humanize
import matplotlib as mpl
import numpy as np
import pandas as pd
import xarray as xr
from matplotlib import pyplot as plt
from numcodecs_safeguards.lossless import Lossless
from typing_extensions import override

from compression_safeguards import SafeguardKind, Safeguards
from compression_safeguards.safeguards._qois import StencilQuantityOfInterest
from compression_safeguards.safeguards._qois.typing import F, Ns, Ps
from compression_safeguards.safeguards.abc import Safeguard
from compression_safeguards.safeguards.pointwise.abc import PointwiseSafeguard
from compression_safeguards.safeguards.stencil.qoi.eb import (
    StencilQuantityOfInterestErrorBoundSafeguard,
)
from compression_safeguards.utils.bindings import Bindings, Parameter
from compression_safeguards.utils.intervals import Interval, IntervalUnion, Lower, Upper
from compression_safeguards.utils.typing import JSON, C, S, T

# 0.1 m/s absolute error bound on the box mean
u_mean_eb_abs = 0.1

def compute_mean_wind_u(ERA5_U: xr.DataArray) -> xr.DataArray:
    ERA5_U_mean = ERA5_U.copy(
        data=qoi.evaluate_qoi(ERA5_U.values, late_bound=Bindings.EMPTY)
    )
    ERA5_U_mean.attrs.update(long_name=f"Mean {ERA5_U.long_name.lower()}")
    return ERA5_U_mean

ERA5_U_mean = compute_mean_wind_u(ERA5_U)

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_mean_wind_u(
    my_ERA5_U: xr.DataArray,
    chart,
    title,
    span,
    u_eb_abs,
    mean=True,
    error=False,
    qoi_stencil=None,
    incremental_masks=None,
    correction=None,
    extra_bytes=0,
):
    import copy

    ERA5_U_pick = ERA5_U_mean if mean else ERA5_U
    my_ERA5_U_pick = compute_mean_wind_u(my_ERA5_U) if mean else my_ERA5_U

    if error:
        with xr.set_options(keep_attrs=True):
            da = (my_ERA5_U_pick - ERA5_U_pick).compute()

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

    # 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 = ("PiYG_r" if mean else "coolwarm") if error else style._colors

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

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

    chart.ax.set_rasterization_zorder(-10)

    if title:
        chart.title(title)

    if error and mean:
        err_v = np.mean(~(np.abs(my_ERA5_U_pick - ERA5_U_pick) <= u_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"))

    if correction is not None:
        encoded = corrections_codec.encode(correction)
        cr = ERA5_U.nbytes / (np.array(encoded).nbytes + extra_bytes)

        if mean:
            corr = np.mean(correction != 0)
            corr = (
                0
                if corr == 0
                else np.format_float_positional(100 * corr, precision=2, min_digits=2)
                + "%"
            )
            if corr == "0.00%":
                corr = "<0.005%"

            t = chart.ax.text(
                0.05,
                0.1,
                f"C={corr}",
                ha="left",
                va="bottom",
                transform=chart.ax.transAxes,
            )
            t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))

    if mean:
        t = chart.ax.text(
            0.95,
            0.9,
            rf"$\times$ {np.round(cr, 2)}"
            if error
            else humanize.naturalsize(ERA5_U.nbytes, binary=True),
            ha="right",
            va="top",
            transform=chart.ax.transAxes,
            zorder=99,
        )
        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)()

    legend_layer = chart.distinct_legend_layers[-1]
    cb = legend_layer.style.legend(legend_layer, location="bottom")

    stackcolors = np.array(mpl.color_sequences["Pastel1"])[[2, 3, 4, 0, 1]]
    stackcmap = mpl.colors.LinearSegmentedColormap.from_list(
        "stackcolors", stackcolors, N=5
    )
    boldstackcolors = np.array(mpl.color_sequences["Set1"])[[2, 3, 4, 0, 1]]

    if not mean:
        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.75, 0.5, 0.225, 0.45],
            xlim=(-0.5 * 7.5, 0.5 * 7.5),
            ylim=(50 - 0.5 * 7.5, 50 + 0.5 * 7.5),
            xticklabels=[],
            yticklabels=[],
            axes_class=type(chart.ax),
            projection=chart.ax.projection,
        )
        axin.pcolormesh(
            da.longitude.values,
            da.latitude.values,
            np.squeeze(da.values),
            cmap=cb.cmap,
            norm=cb.norm,
            rasterized=True,
        )
        axin.coastlines(color="#555555")
        axin.spines["geo"].set_edgecolor("black")
        inset_indicator = chart.ax.indicate_inset_zoom(axin, edgecolor="black", alpha=1)
        inset_indicator.connectors[1].set_visible(True)
        inset_indicator.connectors[3].set_visible(False)

        qoi_stencil_indices = (
            [
                tuple(
                    slice(None) if s == 1 else slice(j, None, s)
                    for j, s in zip(i, qoi_stencil)
                )
                for i in product(*[range(s) for s in qoi_stencil])
            ]
            if qoi_stencil is not None
            else incremental_masks
        )
        da_indices = np.zeros(da.shape, dtype=int)
        for i, s in enumerate(qoi_stencil_indices):
            da_indices[s] = i

        axin = chart.ax.inset_axes(
            [0.025, 0.5, 0.225, 0.45],
            xlim=(-0.5 * 7.5, 0.5 * 7.5),
            ylim=(50 - 0.5 * 7.5, 50 + 0.5 * 7.5),
            xticklabels=[],
            yticklabels=[],
            axes_class=type(chart.ax),
            projection=chart.ax.projection,
        )
        axin.pcolormesh(
            da.longitude.values,
            da.latitude.values,
            np.squeeze(da_indices),
            cmap=stackcmap,
            rasterized=True,
        )
        axin.coastlines(color="#555555")
        axin.spines["geo"].set_edgecolor("black")
        inset_indicator = chart.ax.indicate_inset_zoom(axin, edgecolor="black", alpha=1)
        inset_indicator.connectors[1].set_visible(False)
        inset_indicator.connectors[3].set_visible(True)

    counts, bins = np.histogram(da.values.flatten(), range=(-span, span), bins=21)
    midpoints = bins[:-1] + np.diff(bins) / 2
    extend_width = (bins[-1] - bins[-2]) / (bins[-1] - bins[0])
    cax = cb.ax.inset_axes(
        [
            0.0 - extend_width * extend_left,
            1.25 if (qoi_stencil is None and incremental_masks is None) else 3.5,
            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)),
    )
    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])),
        )
    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])),
        )
    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)

    if (qoi_stencil is not None) or (incremental_masks is not None):
        assert (qoi_stencil is None) or (incremental_masks is None)
        cax2 = cb.ax.inset_axes(
            [
                0.0 - extend_width * extend_left,
                1.25,
                1.0 + extend_width * (0 + extend_left + extend_right),
                2.0,
            ]
        )
        left_stacked = []
        counts_stacked = []
        right_stacked = []
        if qoi_stencil is not None:
            for i in product(*[range(s) for s in qoi_stencil]):
                incrindex = tuple(
                    slice(None) if s == 1 else slice(j, None, s)
                    for j, s in zip(i, qoi_stencil)
                )
                counts2, _bins2 = np.histogram(
                    da.values[incrindex].flatten(), range=(-span, span), bins=21
                )
                counts_stacked.append(counts2)
                left_stacked.append(np.sum(da.values[incrindex] < -span))
                right_stacked.append(np.sum(da.values[incrindex] > span))
        else:
            for next_incremental_mask in incremental_masks:
                counts2, _bins2 = np.histogram(
                    da.values.flatten()[next_incremental_mask.flatten()],
                    range=(-span, span),
                    bins=21,
                )
                counts_stacked.append(counts2)
                left_stacked.append(
                    np.sum(da.values.flatten()[next_incremental_mask.flatten()] < -span)
                )
                right_stacked.append(
                    np.sum(da.values.flatten()[next_incremental_mask.flatten()] > span)
                )
        counts_stacked = np.stack(counts_stacked, axis=0)
        counts_stacked = counts_stacked / np.sum(counts_stacked, axis=0)
        for i in range(len(counts_stacked)):
            cax2.bar(
                midpoints,
                counts_stacked[i],
                bottom=counts_stacked[:i].sum(0),
                width=(bins[-1] - bins[0]) / len(counts),
                color=stackcolors[i],
            )
        if extend_left:
            left_stacked = np.stack(left_stacked, axis=0)
            left_stacked = left_stacked / np.sum(left_stacked, axis=0)
            for i in range(len(left_stacked)):
                cax2.bar(
                    bins[0] - (bins[1] - bins[0]) / 2,
                    left_stacked[i],
                    bottom=left_stacked[:i].sum(0),
                    width=(bins[-1] - bins[0]) / len(counts),
                    color=stackcolors[i],
                )
        if extend_right:
            right_stacked = np.stack(right_stacked, axis=0)
            right_stacked = right_stacked / np.sum(right_stacked, axis=0)
            for i in range(len(right_stacked)):
                cax2.bar(
                    bins[-1] + (bins[-1] - bins[-2]) / 2,
                    right_stacked[i],
                    bottom=right_stacked[:i].sum(0),
                    width=(bins[-1] - bins[0]) / len(counts),
                    color=stackcolors[i],
                )
        cax2.set_xlim(
            -span - (bins[-1] - bins[-2]) * extend_left,
            span + (bins[-1] - bins[-2]) * extend_right,
        )
        cax2.set_xticks([])
        cax2.set_yticks([])
        cax2.spines[:].set_visible(False)

        if error and not mean:
            qoi_stencil_indices = (
                [
                    tuple(
                        slice(None) if s == 1 else slice(j, None, s)
                        for j, s in zip(i, qoi_stencil)
                    )
                    for i in product(*[range(s) for s in qoi_stencil])
                ]
                if qoi_stencil is not None
                else incremental_masks
            )
            if len(qoi_stencil_indices) > 1:
                k_cr = []
                k_c = []
                for i in qoi_stencil_indices:
                    correction_encoded = corrections_codec.encode(
                        np.ascontiguousarray(correction[i])
                    )
                    if qoi_stencil is not None:
                        correction_cr = ERA5_U[i].nbytes / (
                            np.array(correction_encoded).nbytes
                            + extra_bytes * ERA5_U[i].size / ERA5_U.size
                        )
                        corrections = np.mean(correction[i] != 0)
                    else:
                        correction_cr = ERA5_U.data[i].nbytes / (
                            np.array(correction_encoded).nbytes
                            + extra_bytes * ERA5_U.data[i].size / ERA5_U.size
                        )
                        corrections = np.mean(correction.flatten()[i.flatten()] != 0)
                    k_cr.append(correction_cr)
                    k_c.append(corrections * 100)

                axin = chart.ax.inset_axes([0.5 - 0.225 / 2, 0.05, 0.225, 0.45])
                axin2 = axin.twinx()

                axin.plot(k_c, c="gray", lw=1)
                axin2.plot(k_cr, c="gray", lw=1)

                for i, a, b, c in zip(range(5), k_c, k_cr, boldstackcolors):
                    axin.scatter([i], [a], color=c, s=15, marker="x")
                    axin2.scatter([i], [b], color=c, s=15, marker="x")

                axin.set_xlim([-0.25, 4.25])
                ylim = axin.get_ylim()
                axin.set_ylim(
                    ylim[0] - (ylim[1] - ylim[0]) * 0.25 / 4,
                    ylim[1] + (ylim[1] - ylim[0]) * 0.25 / 4,
                )
                # axin.set_yticks([y for y in axin.get_yticks() if y > (ylim[0]+(ylim[1]-ylim[0])*0.125) and y < ylim[1]-(ylim[1]-ylim[0])*0.125])
                # axin.set_yticklabels([f"{y.get_text()}%" for y in axin.get_yticklabels()])
                axin.set_yticks([])
                ylim2 = axin2.get_ylim()
                axin2.set_ylim(
                    ylim2[0] - (ylim2[1] - ylim2[0]) * 0.25 / 4,
                    ylim2[1] + (ylim2[1] - ylim2[0]) * 0.25 / 4,
                )
                # axin2.set_yticks([y for y in axin2.get_yticks() if y > (ylim2[0]+(ylim2[1]-ylim2[0])*0.125) and y < ylim2[1]-(ylim2[1]-ylim2[0])*0.125])
                # axin2.set_yticklabels([rf"$\times${y.get_text()}" for y in axin2.get_yticklabels()])
                axin2.set_yticks([])

                axin.set_ylabel("C", rotation="horizontal", ha="left", va="top")
                axin2.set_ylabel("CR", rotation="horizontal", ha="right", va="top")

                axin.set_xticks([])

                axin.tick_params(axis="x", pad=-1, labelsize=6)
                axin.tick_params(axis="y", direction="in", pad=-5)
                axin2.tick_params(axis="y", direction="in", pad=-5)

                axin.set_xticks([0, 1, 2, 3, 4])
                axin.set_xticklabels(["#1", "#2", "#3", "#4", "#5"])

                plt.setp(
                    axin.get_xticklabels(),
                    bbox=dict(
                        boxstyle="square,pad=0.1", ec="white", fc="white", alpha=0.5
                    ),
                )
                plt.setp(axin.get_yticklabels(), ha="left")
                plt.setp(axin2.get_yticklabels(), ha="right")
                axin.yaxis.set_label_coords(0.15, 0.95)
                axin2.yaxis.set_label_coords(0.85, 0.95)

                for t, c in zip(axin.xaxis.get_ticklabels(), boldstackcolors):
                    t.set_color(c)

                axin.set_facecolor((1.0, 1.0, 1.0, 0.5))

                axin.grid(False)
                axin2.grid(False)

    if not mean:
        chart.layers.pop()

def table_mean_wind_u(
    my_ERA5_U: xr.DataArray,
    title,
    u_mean_eb_abs,
    correction,
    extra_bytes,
) -> pd.DataFrame:
    my_ERA5_U_mean = compute_mean_wind_u(my_ERA5_U)

    err_inf_U = np.amax(np.abs(my_ERA5_U - ERA5_U))
    err_2_U = np.sqrt(np.mean(np.square(my_ERA5_U - ERA5_U)))
    err_inf_U_mean = np.amax(np.abs(my_ERA5_U_mean - ERA5_U_mean))
    err_2_U_mean = np.sqrt(np.mean(np.square(my_ERA5_U_mean - ERA5_U_mean)))

    err_v = np.mean(~(np.abs(my_ERA5_U_mean - ERA5_U_mean) <= u_mean_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 correction is None else np.mean(correction != 0)
    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%"

    if correction is None:
        cr = ERA5_U.nbytes / extra_bytes
    else:
        encoded = corrections_codec.encode(correction)
        cr = ERA5_U.nbytes / (np.array(encoded).nbytes + extra_bytes)

    return pd.DataFrame(
        {
            "Compressor": [title[0]],
            "Safeguarded": [title[1]],
            "Corrections": [title[2]],
            r"$L_{\infty}(\hat{u})$": [
                f"{err_inf_U:.02}",
            ],
            r"$L_{2}(\hat{u})$": [
                f"{err_2_U:.02}",
            ],
            r"$L_{\infty}(\overline{\hat{u}})$": [
                f"{err_inf_U_mean:.02}",
            ],
            r"$L_{2}(\overline{\hat{u}})$": [
                f"{err_2_U_mean:.02}",
            ],
            "V": [err_v],
            "C": [corr],
            "CR": [
                rf"$\times$ {np.round(cr, 2)}",
            ],
        }
    )

from numcodecs_zero import ZeroCodec

zero = ZeroCodec()

ERA5_U_zero_encoded = zero.encode(ERA5_U.values)
ERA5_U_zero_approximation = zero.decode(ERA5_U_zero_encoded)

sg = Safeguards(safeguards=[qoi])

ERA5_U_sg_zero_correction = sg.compute_correction(
    ERA5_U.values, ERA5_U_zero_approximation
)
ERA5_U_sg_zero = sg.apply_correction(
    ERA5_U_zero_approximation, ERA5_U_sg_zero_correction
)

assert sg.check(ERA5_U.values, ERA5_U_sg_zero)

ERA5_U_sg_sperr_correction = sg.compute_correction(
    ERA5_U.values, ERA5_U_sperr_approximation
)
ERA5_U_sg_sperr = sg.apply_correction(
    ERA5_U_sperr_approximation, ERA5_U_sg_sperr_correction
)

assert sg.check(ERA5_U.values, ERA5_U_sg_sperr)

INCREMENTAL_MASK: Parameter = Parameter("$incrmask")
INCREMENTAL_APPROXIMATION: Parameter = Parameter("$incrapprox")


class IncrementalSafeguard(PointwiseSafeguard):
    __slots__: tuple[str, ...] = ()

    kind: ClassVar[str] = "incremental"

    def __init__(self) -> None:
        pass

    @property
    @override
    def late_bound(self) -> Set[Parameter]:
        return frozenset([INCREMENTAL_MASK, INCREMENTAL_APPROXIMATION])

    @override
    def check_pointwise(
        self,
        data: np.ndarray[S, np.dtype[T]],
        approximation: np.ndarray[S, np.dtype[T]],
        *,
        late_bound: Bindings,
        where: Literal[True] | np.ndarray[S, np.dtype[np.bool]] = True,
    ) -> np.ndarray[S, np.dtype[np.bool]]:
        return np.ones_like(data, dtype=np.bool)  # type: ignore

    @override
    def compute_safe_intervals(
        self,
        data: np.ndarray[S, np.dtype[T]],
        *,
        late_bound: Bindings,
        where: Literal[True] | np.ndarray[S, np.dtype[np.bool]] = True,
    ) -> IntervalUnion[T, int, int]:
        mask = late_bound.resolve_ndarray_with_lossless_cast(
            INCREMENTAL_MASK, data.shape, np.dtype(np.bool)
        )
        approximation = late_bound.resolve_ndarray_with_lossless_cast(
            INCREMENTAL_APPROXIMATION, data.shape, data.dtype
        )

        valid = Interval.full_like(data)
        Lower(approximation.flatten()) <= valid[mask.flatten()] <= Upper(
            approximation.flatten()
        )

        validu = valid.into_union()

        # override the class of the valid interval union with one that accepts
        #  non-overlap with the data for elements that have already been
        #  corrected,
        # which is required to create a safeguard that forces the existing
        #  corrections to be recreated
        class IncrementalIntervalUnion(IntervalUnion):
            __slots__ = ()

            def contains(
                self, other: np.ndarray[S, np.dtype[T]]
            ) -> np.ndarray[S, np.dtype[np.bool]]:
                is_contained: np.ndarray[S, np.dtype[np.bool]] = super().contains(other)

                if np.array_equal(other, data, equal_nan=True):
                    is_contained |= mask

                return is_contained

        validu.__class__ = IncrementalIntervalUnion

        return validu

    @override
    def compute_footprint(
        self,
        foot: np.ndarray[S, np.dtype[np.bool]],
        *,
        late_bound: Bindings,
        where: Literal[True] | np.ndarray[S, np.dtype[np.bool]] = True,
    ) -> np.ndarray[S, np.dtype[np.bool]]:
        # the incremental safeguard does not affect the footprint
        return np.zeros_like(foot)

    @override
    def compute_inverse_footprint(
        self,
        foot: np.ndarray[S, np.dtype[np.bool]],
        *,
        late_bound: Bindings,
        where: Literal[True] | np.ndarray[S, np.dtype[np.bool]] = True,
    ) -> np.ndarray[S, np.dtype[np.bool]]:
        # the incremental safeguard does not affect the footprint
        return np.zeros_like(foot)

    @override
    def get_config(self) -> dict[str, JSON]:
        return dict(kind=type(self).kind)

def rewrite_stencil_qoi_eb_safeguard_into_incremental(
    qoi: StencilQuantityOfInterestErrorBoundSafeguard,
) -> Collection[Safeguard]:
    # extract the fully-resolved QoI expression, in text representation
    qoi_expr_repr = repr(qoi._qoi_expr)

    # replace mentions of the data stencil X with where(mask, approximation, X)
    #  to fill in the already-corrected approximation as a constant for elements
    #   where the correction has already been produced
    qoi_expr_incremental_repr = re.sub(
        r"X\[([^]]+)\]",
        rf'where(C["{INCREMENTAL_MASK}"][\1], C["{INCREMENTAL_APPROXIMATION}"][\1], X[\1])',
        qoi_expr_repr,
    )

    # create a new QoI expression with the incremental QoI
    qoi_incremental_expr = StencilQuantityOfInterestErrorBoundSafeguard(
        qoi=qoi_expr_incremental_repr,
        **{k: v for k, v in qoi.get_config().items() if k not in ("kind", "qoi")},
    )._qoi_expr

    # override (a copy of) the QoI's class with one that uses the incremental
    #  QoI, but only for deriving corrections, and thus requires the additional
    #  late-bound constants for incremental corrections
    class IncrementalStencilQuantityOfInterest(StencilQuantityOfInterest):
        __slots__ = ()

        @property
        def late_bound_constants(self) -> Set[Parameter]:
            return qoi_incremental_expr.late_bound_constants

        @np.errstate(divide="ignore", over="ignore", under="ignore", invalid="ignore")
        def compute_data_bounds(
            self,
            qoi_lower: np.ndarray[Ps, np.dtype[F]],
            qoi_upper: np.ndarray[Ps, np.dtype[F]],
            Xs: np.ndarray[Ns, np.dtype[F]],
            late_bound: Mapping[Parameter, np.ndarray[Ns, np.dtype[F]]],
        ) -> tuple[np.ndarray[Ns, np.dtype[F]], np.ndarray[Ns, np.dtype[F]]]:
            return qoi_incremental_expr.compute_data_bounds(
                qoi_lower, qoi_upper, Xs, late_bound
            )

    qoi_incremental = StencilQuantityOfInterestErrorBoundSafeguard(
        **{k: v for k, v in qoi.get_config().items() if k != "kind"}
    )
    qoi_incremental._qoi_expr.__class__ = IncrementalStencilQuantityOfInterest

    # combine the incremental safeguard, which takes care of recreating already
    #  computed corrections, with the now-incremental QoI, which takes these
    #  already computed corrections into account to produce tighter corrections
    return (IncrementalSafeguard(), qoi_incremental)


sg_incremental = Safeguards(
    safeguards=[*rewrite_stencil_qoi_eb_safeguard_into_incremental(qoi)]
)

# create stepby masks using array slices with steps
# instead of applying the safeguards incrementally one element at a time,
# apply them to the i'th element of non-overlapping windows of the same size as
#  the stencil
# this gives (asymptotically) the same extra knowledge when the stencil is
#  symmetric, but one element per axis less knowledge when it is asymmetric
stepby_incremental_masks = [
    tuple(slice(None) if s == 1 else slice(j, None, s) for j, s in zip(i, qoi_stencil))
    for i in product(*[range(s) for s in qoi_stencil])
]

rng = np.random.Generator(np.random.PCG64(seed=42))

# create pseudo-randomly shuffled masks that determine how the safeguards are
#  applied incrementally
# the number of masks is arbitrary (any number works), though setting it to the
#  size of the stencil makes sense as a default
shuffled_incremental_masks = [
    np.zeros(ERA5_U.size, dtype=np.bool) for _ in range(np.prod(qoi_stencil))
]

indices = rng.permutation(ERA5_U.size)

for i in range(len(shuffled_incremental_masks)):
    shuffled_incremental_masks[i][indices[i :: len(shuffled_incremental_masks)]] = True
    shuffled_incremental_masks[i] = shuffled_incremental_masks[i].reshape(ERA5_U.shape)

def compute_incremental_correction(
    data: np.ndarray[S, np.dtype[T]],
    approximation: np.ndarray[S, np.dtype[T]],
    incremental_masks: list[tuple[slice] | tuple[np.ndarray[S, np.dtype[np.bool]]]],
) -> np.ndarray[S, np.dtype[C]]:
    incremental_mask = np.zeros(data.shape, dtype=np.bool)
    incremental_correction = np.empty(
        data.shape, sg.correction_dtype_for_data(data.dtype)
    )
    incremental_approximation = np.empty_like(data)

    # instead of applying the safeguards incrementally one element at a time,
    # apply them to the several non-overlapping masks that partition the data
    for next_incremental_mask in incremental_masks:
        # compute and apply the correction
        incremental_correction_partial = sg_incremental.compute_correction(
            data,
            approximation,
            late_bound={
                # all boundary conditions work for the mask
                # - valid, edge, reflect, symmetric, and wrap work as expected
                # - for a constant zero boundary, the mask is False for
                #   boundary points even though the constant value there is
                #   known - while unfortunate this is more conservative and
                #   thus still produces the correct results
                # - for a constant non-zero boundary, the mask is True for
                #   boundary points and, since the constant boundary points
                #   are always known, this is correct
                INCREMENTAL_MASK: incremental_mask,
                # boundary conditions are applied as expected to the already-
                #  corrected approximation
                INCREMENTAL_APPROXIMATION: incremental_approximation,
            },
        )
        incremental_corrected_partial = sg_incremental.apply_correction(
            approximation, incremental_correction_partial
        )

        # save the correction and corrected elements for the part of the stencil
        #  that we are (re-)computing, then update the mask to reflect that these
        #  elements have now already been corrected
        incremental_correction[next_incremental_mask] = incremental_correction_partial[
            next_incremental_mask
        ]
        incremental_approximation[next_incremental_mask] = (
            incremental_corrected_partial[next_incremental_mask]
        )
        incremental_mask[next_incremental_mask] = True

    return incremental_correction

ERA5_U_sg_incremental_correction = dict()
ERA5_U_sg_incremental = dict()

for name, ERA5_U_approximation in dict(
    zero=ERA5_U_zero_approximation, sperr=ERA5_U_sperr_approximation
).items():
    # apply the combined correction using the non-incremental safeguard
    ERA5_U_sg_incremental_correction[f"{name}-stepby"] = compute_incremental_correction(
        ERA5_U.values,
        ERA5_U_approximation,
        stepby_incremental_masks,
    )
    ERA5_U_sg_incremental[f"{name}-stepby"] = sg.apply_correction(
        ERA5_U_approximation, ERA5_U_sg_incremental_correction[f"{name}-stepby"]
    )

    assert sg.check(ERA5_U.values, ERA5_U_sg_incremental[f"{name}-stepby"])

    # apply the combined correction using the non-incremental safeguard
    ERA5_U_sg_incremental_correction[f"{name}-shuffled"] = (
        compute_incremental_correction(
            ERA5_U.values, ERA5_U_approximation, shuffled_incremental_masks
        )
    )
    ERA5_U_sg_incremental[f"{name}-shuffled"] = sg.apply_correction(
        ERA5_U_approximation, ERA5_U_sg_incremental_correction[f"{name}-shuffled"]
    )

    assert sg.check(ERA5_U.values, ERA5_U_sg_incremental[f"{name}-shuffled"])