Safely applying safeguards to chunked datasets with xarray-safeguards¶

ERA5_U_chunked = ERA5_U.chunk(72)
print(ERA5_U_chunked.sizes)
ERA5_U_chunked.data

# When checking for safety violations, we need to remove a 2 grid-cell boundary
#  since our stencil safeguard above uses the valid boundary condition and thus
#  does not apply to the outer data boundary
print(
    f"Safeguards have {np.sum((np.abs(ERA5_U_sg - ERA5_U) >= 1.0)[:, 2:-2, 2:-2]).item()} safety requirement violations"
)

print(
    f"Unsafely chunked safeguards have {np.sum(ERA5_U_sg_unsafe_chunked != ERA5_U_sg).compute().item()} differently corrected values"
)

print(
    f"Unsafely chunked safeguards have {np.sum((np.abs(ERA5_U_sg_unsafe_chunked - ERA5_U) >= 1.0)[:, 2:-2, 2:-2]).compute().item()} safety requirement violations"
)

chart = earthkit.plots.Map()

with xr.set_options(keep_attrs=True):
    ERA5_U_sg_unsafe_chunked_error = ERA5_U_sg_unsafe_chunked - ERA5_U
ERA5_U_sg_unsafe_chunked_error.attrs.update(
    long_name=f"Absolute error over {ERA5_U_sg_unsafe_chunked_error.long_name.lower()}"
)

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

extend_left = np.nanmin(ERA5_U_sg_unsafe_chunked_error) < -2
extend_right = np.nanmax(ERA5_U_sg_unsafe_chunked_error) > 2

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

style._legend_kwargs["extend"] = extend
chart.pcolormesh(ERA5_U_sg_unsafe_chunked_error, style=style, zorder=-11)
chart.ax.set_rasterization_zorder(-10)

chart.title("Applying non-pointwise safeguards independently per-chunk is unsafe")

err_v = np.mean((np.abs(ERA5_U_sg_unsafe_chunked - ERA5_U) >= 1.0)[:, 2:-2, 2:-2])
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.975,
    0.05,
    f"V={err_v}",
    ha="right",
    va="bottom",
    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:
    getattr(chart, m)()

counts, bins = np.histogram(
    ERA5_U_sg_unsafe_chunked_error.values.flatten(), range=(-2, 2), bins=20
)
midpoints = bins[:-1] + np.diff(bins) / 2
cb = chart.ax.collections[0].colorbar
if extend_left:
    cb._extend_patches[0].set_hatch("xx")
    cb._extend_patches[0].set_ec("white")
cb.ax.fill_between([-2, -1], *cb.ax.get_ylim(), hatch="xx", ec="w", fc="none", lw=0)
cb.ax.fill_between([1, 2], *cb.ax.get_ylim(), hatch="xx", ec="w", fc="none", lw=0)
if extend_right:
    cb._extend_patches[-1].set_hatch("xx")
    cb._extend_patches[-1].set_ec("white")
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),
        1.0,
    ]
)
cax.bar(
    midpoints,
    height=counts,
    width=(bins[-1] - bins[0]) / len(counts),
    color=cb.cmap(cb.norm(midpoints)),
    hatch=["xx" if np.abs(m) > 1.0 else "" for m in midpoints],
    ec="white",
    lw=0,
)
if extend_left:
    cax.bar(
        bins[0] - (bins[1] - bins[0]) / 2,
        height=np.sum(ERA5_U_sg_unsafe_chunked_error < -2),
        width=(bins[-1] - bins[0]) / len(counts),
        color=cb.cmap(cb.norm(midpoints[0])),
        hatch="xx",
        ec="white",
        lw=0,
    )
if extend_right:
    cax.bar(
        bins[-1] + (bins[-1] - bins[-2]) / 2,
        height=np.sum(ERA5_U_sg_unsafe_chunked_error > 2),
        width=(bins[-1] - bins[0]) / len(counts),
        color=cb.cmap(cb.norm(midpoints[-1])),
        hatch="xx",
        ec="white",
        lw=0,
    )
q1, q2, q3 = ERA5_U_sg_unsafe_chunked_error.quantile([0.25, 0.5, 0.75]).values
cax.axvline(
    ERA5_U_sg_unsafe_chunked_error.mean().values.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(
    ERA5_U_sg_unsafe_chunked_error.mean().values.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(
    -2 - (bins[-1] - bins[-2]) * extend_left, 2 + (bins[-1] - bins[-2]) * extend_right
)
cax.set_xticks([])
cax.set_yticks([])
cax.spines[:].set_visible(False)

chart.save(Path("plots") / "chunked-unsafe.pdf")

# now we combine both variables when loading the dataset from both files
# and apply the safeguards using the `.safeguarded` accessor on the dataset
# we can now even use completely different chunks
ERA5_sg_safe_chunked = xr.open_mfdataset(
    str(output / "chunked-zero*.nc"), chunks=36, compat="identical"
).safeguarded
list(ERA5_sg_safe_chunked.variables)

# we can also inspect the applied safeguards with the `.safeguards` property on the variable
ERA5_U_sg_safe_chunked = ERA5_sg_safe_chunked["u"]
print(ERA5_U_sg_safe_chunked.sizes)
display(ERA5_U_sg_safe_chunked.data)
print(ERA5_U_sg_safe_chunked.safeguards)

print(
    f"Safely chunked safeguards have {np.sum(ERA5_U_sg_safe_chunked != ERA5_U_sg).compute().item()} differently corrected values"
)

print(
    f"Safely chunked safeguards have {np.sum((np.abs(ERA5_U_sg_safe_chunked - ERA5_U) >= 1.0)[:, 2:-2, 2:-2]).compute().item()} safety requirement violations"
)