Integration with zarr

[Marston]

Hi,

Would like to test this with zarr, but this doesn't seem to be a part of the numcodecs used by zarr. How could I use this lossy compression to encode and decode geospatial float32 data?

[Marston]

I've gotten some help from the xarray discussion group but we are stuck on this point:

The snag is on this part:

bitinfo = xb.get_bitinformation(ds, dim="lon")  # calling bitinformation.jl.bitinformation

Is the compression applied here via a specified dim? So if I have an array with dims: time,xx,lev,lat,lon, how is this handled? All the data is covered by the lat and lon but what about the other dims? I'm writing this to a zarr file where xx=1 but I will append many NWP variables along this dimension. Time has also dimension 1.
Basically, I'm asking if this makes sense for a python environment:

bitinfo = xb.get_bitinformation(ds, dim=['lev', 'lat', 'lon'])  

The document is not very clear on this.

[observingClouds]

Hi @Marston,

get_bitinformation() is used to retrieve the information content within the dataset only. Depending on the dimension, each bit in e.g. a float32 can have a different information content.

Have a look at the individual steps shown in the HowTo.

The compression is done only in the last step and would be ds_bitrounded.to_compressed_zarr(outpath) in case of zarr. But you could also just use ds_bitrounded.to_zarr() and use either the default compression (https://numcodecs.readthedocs.io/en/stable/lz4.html) or define a different one (e.g. another one from the numcodes package).

The rounding to the requested keepbits is done with ds_bitrounded = xb.xr_bitround(ds, keepbits). There numcodecs implementation of BitRound is used. This is not the compression itself but the necessary step to achieve a good compression afterwards.

I hope this clarifies it.

[Marston]

Hi @observingClouds

I'm sorry, this is more confusing. Could you take the example in HowTo and apply it to a use case where one has a 5D array in a dataset, ds:

ds.shape
>>> (1,2,31,2000,2000) # ( time, features, lev, lat, lon)

How do I both write and read back such an array in python?

import xarray as xr
import xbitinfo as xb
import zarr as za

ds = xr.tutorial.load_dataset(inpath)

# psuedo code
bitinfo = xb.get_bitinformation(ds, dim=["lev", "lat", "lon"])  # calling bitinformation.jl.bitinformation
keepbits = xb.get_keepbits(bitinfo, inflevel=0.99)  # get number of mantissa bits to keep for 99% real information
ds_bitrounded = xb.xr_bitround(ds, keepbits)  # bitrounding keeping only keepbits mantissa bits
compressor = za.Blosc(cname="zstd", clevel=6, shuffle=2)
chunks3d = {'time':1, 'feature': 1, 'lev': 31, 'lat': 200, 'lon': 200}
ds_bitrounded.chunk(chunks3d).to_compressed_zarr('test.zarr', group='g2d', consolidated=True, encoding={'varb3d': {"compressor": compressor}})

rb_ds = xr.open_dataset(infile, engine='zarr', group='f3d')

My aim here is to find a good lossy compression for zarr.

[observingClouds]

We can create a 5D array from a tutorial dataset:

import xarray as xr
ds = xr.tutorial.load_dataset('air_temperature')
ds = ds.expand_dims({'lev':1,'ensemble':1})

you could write this array to disk as zarr with ds.to_zarr(output_path) and read it back with ds = xr.open_dataset(output_path, engine='zarr')

These steps work for any dataset. to_compressed_zarr is just an accessor that xbitinfo provides. Have a look at the source code of to_compressed_zarr. We just use to_zarr there ourselves and specify a compression algorithm that we think works better than the default of to_zarr when combined with bit rounding.

I don't see a particular problem with the code you are giving above, expect that you write to group g2d and read from group f3d. Personally I would try to avoid groups as much as possible as they are not necessarily supported by all software packages. Because you want to experiment with your own compressor, just use to_zarr with your compressor of choice instead of to_compressed_zarr:

compressor = za.Blosc(cname="zstd", clevel=6, shuffle=2)
chunks3d = {'time':1, 'ensemble': 1, 'lev': 31, 'lat': 200, 'lon': 200}

# Writing to disk
ds.chunk(chunks3d).to_zarr('test.zarr', group='g2d', consolidated=True, encoding={'air': {"compressor": compressor}})

# Reading from disk
ds_in = xr.open_dataset("test.zarr", group="g2d", engine="zarr")