xarray_safeguards

xarray_safeguards

Fearless (chunked) lossy compression with xarray-safeguards

Lossy compression can be scary as valuable information or features of the data may be lost.

By using Safeguards to guarantee your safety requirements, lossy compression can be applied safely and without fear.

Overview

This package provides functionality to use safeguards with (chunked) xarray.DataArrays and cross-chunk boundary conditions. Since applying safeguards to chunked data by yourself is very difficult and error-prone, this package handles all of the complexity for you with a simple and safe API.

In particular, xarray-safeguards provides the produce_data_array_correction function to produce a correction such that certain properties of the original data are preserved after compression, which can be stored in the same or a different dataset (or file).

This correction can be applied to the decompressed data using the apply_data_array_correction function or the .safeguarded accessor on datasets.

This package also provides the .safeguards accessor on correction or corrected data arrays to inspect the safeguards that were applied.

By applying the corrections produced by produce_data_array_correction, data that was compressed with badly-behaving lossy compressors can be safely used, at the cost of potentially less efficient compression, and lossy compression can be applied without fear.

Example

You can produce and apply the corrections to uphold an absolute error bound of \(eb_{abs} = 0.1\) for an already-compressed data array as follows:

import numpy as np
import xarray as xr
from xarray_safeguards import apply_data_array_correction, produce_data_array_correction

# some (chunked) n-dimensional data array
da = xr.DataArray(np.linspace(-10, 10, 21), name="da").chunk(10)
# lossy-compressed approximation of the data, here all zeros
da_approximation = xr.DataArray(np.zeros_like(da.values), name="da").chunk(10)

da_correction = produce_data_array_correction(
    data=da,
    approximation=da_approximation,
    # guarantee an absolute error bound of 0.1:
    #   |x - x'| <= 0.1
    safeguards=[dict(kind="eb", type="abs", eb=0.1)],
)

## (a) manual correction ##

da_corrected = apply_data_array_correction(da_approximation, da_correction)
np.testing.assert_allclose(da_corrected.values, da.values, rtol=0, atol=0.1)

## (b) automatic correction with xarray accessors ##

# combine the lossy approximation and the correction into one dataset
#  e.g. by loading them from different files using `xarray.open_mfdataset`
ds = xr.Dataset({
    da_approximation.name: da_approximation,
    da_correction.name: da_correction,
})

# access the safeguarded dataset that applies all corrections
ds_safeguarded: xr.Dataset = ds.safeguarded
np.testing.assert_allclose(ds_safeguarded["da"].values, da.values, rtol=0, atol=0.1)

Please refer to the compression_safeguards.SafeguardKind for an enumeration of all supported safeguards.

Referenced by:

Classes:

Functions:

  • produce_data_array_correction –

    Produce the correction required to make the approximation data array satisfy the safeguards relative to the data array.

  • apply_data_array_correction –

    Apply the correction to the approximation array to satisfy the safeguards for which the correction was produced.

DataValue module-attribute 

DataValue: TypeAlias = int | float | number | DataArray

Parameter value type that includes scalar numbers and data arrays thereof.

Used by:

produce_data_array_correction 

produce_data_array_correction(
    data: DataArray,
    approximation: DataArray,
    *,
    safeguards: Collection[dict[str, JSON] | Safeguard],
    late_bound: Mapping[str, DataValue] = MappingProxyType(
        dict()
    ),
    check_chunks_first: bool = True,
    allow_unsafe_safeguards_override: bool = False,
) -> DataArray

Produce the correction required to make the approximation data array satisfy the safeguards relative to the data array.

The data array may be chunked1 and the approximation array must use the same chunking. Importantly, the data array must contain the complete data, i.e. not just a sub-chunk of the data, so that non-pointwise safeguards are correctly applied.

If the the data array is chunked, the correction is produced lazily, otherwise it is computed eagerly.

The data array must have a name and the approximation array must use the same name.

  1. At the moment, only chunking with dask is supported. ↩

Parameters:
  • data (DataArray) –

    The data array, relative to which the safeguards are enforced.

  • approximation (DataArray) –

    The approximation array for which the correction is produced.

  • safeguards (Collection[dict[str, JSON] | Safeguard]) –

    The safeguards that will be applied relative to the data array.

    They can either be passed as a safeguard configuration dict or an already initialized Safeguard.

    Please refer to the SafeguardKind for an enumeration of all supported safeguards.

  • late_bound (Mapping[str, DataValue], default: MappingProxyType(dict()) ) –

    The bindings for all late-bound parameters of the safeguards.

    The bindings must resolve all late-bound parameters and include no extraneous parameters.

    If a binding resolves to a xr.DataArray, its dimensions must be a subset of data.dims and the shape must be broadcastable to the data shape, i.e. each axis must have either size 1 or the same size as the matching dimension in the data.

    This method automatically provides the following built-in constants to the safeguards, which must not be included:

    • $x and $X: the original data as a constant
    • $x_min and $x_max: the global minimum/maximum of the data
    • $d_DIM for each dimension DIM of the data array
  • check_chunks_first (bool, default: True ) –

    If True, all chunks are first checked to determine if no corrections are required, which requires eagerly computing the check result across all chunks. If False, the check is skipped and corrections for all chunks are produced, even if no corrections are required.

  • allow_unsafe_safeguards_override (bool, default: False ) –

    WARNING: This option is unsafe and must only be passed if instructed. Do not pass this option otherwise.
Returns:
  • correction( DataArray ) –

    The correction array
Raises:
  • ValueError –

    if the data has already been corrected by safeguards, i.e. if it is not the original uncorrected data, which will create a printer problem.

  • ValueError –

    if the approximation has already been corrected, which may create a printer problem.

  • TypeSetError –

    if the data uses an unsupported data type.

  • ValueError –

    if the data's dimensions, shape, dtype, chunks, or name do not match the approximation's, or if the data name is None.

  • LateBoundParameterResolutionError –

    if late_bound does not resolve all late-bound parameters of the safeguards or includes any extraneous parameters.

  • TypeError –

    if a late-bound parameter is not a scalar or xarray.DataArray.

  • ValueError –

    if a late-bound parameter's dimensions are not a subset of the data dimensions, or it is not broadcastable to the data shape.

  • ... –

    if instantiating a safeguard, checking a safeguard, or computing the correction for a safeguard raises an exception.
Referenced by:

apply_data_array_correction 

apply_data_array_correction(
    approximation: DataArray, correction: DataArray
) -> DataArray

Apply the correction to the approximation array to satisfy the safeguards for which the correction was produced.

The approximation must be bitwise equivalent to the approximation that was used to produce the correction.

The approximation array may be chunked1 and the correction array must use the same chunking, though this chunking may differ from the one that was used to produce the correction.

If the the approximation array is chunked, the correction is applied lazily, otherwise it its application is computed eagerly.

  1. Any chunking supported by xarray is supported, including but not limited to dask, please see https://docs.xarray.dev/en/stable/internals/chunked-arrays.html. ↩

Parameters:
  • approximation (DataArray) –

    The approximation array for which the correction has been produced.

  • correction (DataArray) –

    The correction array.
Returns:
  • corrected( DataArray ) –

    The corrected array, which satisfies the safeguards.
Raises:
  • ValueError –

    if the correction's dimensions, shape, or chunks do not match the approximation's, or if the correction's dtype does not match the correction dtype for the approximation's dtype.

  • ValueError –

    if the correction does not contain metadata about the safeguards that it was produced with.

  • ... –

    if re-instantiating the safeguards from their configuration metadata raises an exception.
Referenced by:

DatasetSafeguardedAccessor

An extension for an xarray.Dataset that provides the .safeguarded property that applies safeguards corrections in the dataset to their respective variables.

The safeguarded variables must be bitwise equivalent to the variables for which their corrections were originally produced.

For instance, for a dataset ds that contains both a variable ds.foo and its correction, you can access the corrected variable using ds.safeguarded.foo.

Raises:
  • AttributeError –

    if the dataset does not contain not-yet-applied safeguards corrections.

  • RuntimeError –

    if the attributes of the safeguards correction are invalid.

  • ValueError –

    if applying a safeguards correction raises an exception.

  • ... –

    if re-instantiating the safeguards from their configuration metadata raises an exception.

DataArraySafeguardsAccessor

An extension for an xarray.DataArray that provides the .safeguards property that exposes the collection of Safeguards for a safeguards correction or corrected array.

Raises:
  • AttributeError –

    if the data array does not have associated safeguards.

  • ... –

    if re-instantiating the safeguards from their configuration metadata raises an exception.