Lazy look up table in a custom backend

[mraspaud]

Hi,

I have data for which I am writing a backend that I would like to load lazily. One of the variable is a lookup table that I would like to apply to the data, lazily. I'm trying to use the same approach as the encoder/decoders present in xarray, but for some reason the function application doesn't like the lazily indexed arrays...

For clarity, applying a look up table to the data means in my world that data is a n-dimensional array of indices (unsigned shorts in the present case), and the lookup table (lut) has floating point values in it, and if using numpy I would do: new_data = lut[data].

from xarray.backends import BackendEntrypoint, BackendArray
from xarray.core import indexing
from xarray.coding.variables import VariableCoder, lazy_elemwise_func, unpack_for_decoding
import xarray as xr
import h5netcdf
import numpy as np
from functools import partial
import h5py

def apply_lut(data, lut):
    return lut[data]

def decode_lut(var, lut):
    dtype = lut.dtype
    dims, data, attrs, encoding = unpack_for_decoding(var)
    transform = partial(decode_lut, lut=lut)
    return xr.Variable(dims, lazy_elemwise_func(data, transform, dtype), attrs, encoding)


class H5Array(BackendArray):
    def __init__(self, array):
        self.shape = array.shape
        self.dtype = array.dtype
        self.array = array

    def __getitem__(self, key):
        return indexing.explicit_indexing_adapter(
            key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
        )

    def _getitem(self, key):
        return self.array[key]


class I3DBackend(BackendEntrypoint):
    """Backend for Insat 3D in hdf5 format."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.file_obj = None

    def open_dataset(self, filename, *, drop_variables=None):
        h5f = h5py.File(filename, mode="r")
        ds = xr.Dataset()
        h5_arr = h5f["IMG_VIS"]
        vis = xr.Variable(("time", "GeoY2", "GeoX2"),  # h5_arr.dimensions,
                          indexing.LazilyIndexedArray(H5Array(h5_arr)),
                          encoding={"preferred_chunks": h5_arr.chunks or {}})
        ds["IMG_VIS"] = vis

        lut = np.array(h5f["IMG_VIS_CALIBRATION_LUT"])
        decoded = decode_lut(vis, lut)
        ds["IMG_VIS_RADIANCE"] = decoded

        self.file_obj = h5f
        return ds

vis_shape = (1, 112, 112)

def insat_filename():
    dimensions = {"GeoX2": vis_shape[2],
                  "GeoY2": vis_shape[1],
                  "time": vis_shape[0],
                  "GreyCount": 1024}
    filename = "3DIMG_25OCT2022_0400_L1B_STD_V01R00.h5"
    with h5netcdf.File(filename, mode="w") as h5f:
        h5f.dimensions = dimensions
        var = h5f.create_variable("IMG_VIS", ("time", "GeoY2", "GeoX2"), np.uint16)
        var[:] = np.random.randint(0, 1000, vis_shape, dtype=np.uint16)
        var = h5f.create_variable("IMG_VIS_CALIBRATION_LUT", ("GreyCount",), float)
        var[:] = np.arange(0, 1024 * 2, 2)
    return filename

filename = insat_filename()
ds = xr.open_dataset(filename, engine=I3DBackend)
ds["IMG_VIS"].compute()
ds["IMG_VIS_RADIANCE"].compute()

The above code crashes with a AttributeError: 'LazilyIndexedArray' object has no attribute 'dims'

Any idea about what I'm doing wrong?

[keewis]

my initial guess would be that this is due to a typo in your code:

def decode_lut(var, lut):
    dtype = lut.dtype
    dims, data, attrs, encoding = unpack_for_decoding(var)
    transform = partial(decode_lut, lut=lut)
    return xr.Variable(dims, lazy_elemwise_func(data, transform, dtype), attrs, encoding)

where the partial should wrap apply_lut instead of decode_lut.

[mraspaud]

Man, I've been looking at this code for hours... sorry about the noise...

[mraspaud]

So fixing the typo make the code execute without errors. However ds["IMG_VIS_RADIANCE"] doesn't seem to be lazy...

[mraspaud]

oh, yes, of course, because I didn't ask it to be.
opening the dataset should read ds = xr.open_dataset(filename, engine=I3DBackend, chunks={})

[keewis]

I think the lazy indexing classes were built for the inverse: if lut was lazy, we could feed the in-memory data to VectorizedIndexer and lazily index lut with that. That being the case, those classes might not support what you're trying to do?

I'm no expert on the backend machinery, though, so there might still be a way to get it to work. @pydata/xarray, any ideas?

Edit: oh, great, then never mind if it already works