Zarr JPEG-XL encoder squeezes chunks dimension

[JeppeKlitgaard]

Once again, thank you for putting together this very useful library!

I am running into a problem/inconsistency using the Jpegxl codec for Zarr where the number of dimensions of the array passed to the underlying JXL encoder changes based on the chunk shape.

Concretely, I have an array of shape (100, 2000, 2000) where the first axis is understood to be bands/channels, while the latter two make up a 2D image. For this, I would want to specify the planar=True argument. This works well, except when chunking with a shape (1, ..., ...), in which case the array gets squeezed via this line into 2D, which causes the dimensionality check in these lines to fail.

I am wondering if this is the intended behaviour, or whether this is inadvertent? The consequence is that the encoder breaks when changing from chunking (2, ..., ...) to (1, ..., ...), which seems like a bit of a sharp edge.

[cgohlke]

Just for completeness, can you provide a minimal reproducible example for this issue?

I agree that shape (1, h, w) should work with planar=True.

The bigger issue is that these Array-Bytes codecs should not (need to) squeeze input arrays and that zarr-python lacks a reshape codec.

[cgohlke]

This will be fixed in the next release.

[JeppeKlitgaard]

Just for completeness, can you provide a minimal reproducible example for this issue?

I agree that shape (1, h, w) should work with planar=True.

The bigger issue is that these Array-Bytes codecs should not (need to) squeeze input arrays and that zarr-python lacks a reshape codec.

Yea, I was also a bit puzzled as to how to do reshaping on the chunk if required by the codec — currently this seems to be the responsibility of the codec instance, which may not be desirable.

Apologies about the lack of a reproducer, I reckon you have already made one, but if not then:

import numpy as np
import zarr
import zarr.storage
from imagecodecs import zarr as codecs

codecs.register_codecs()

dtype = np.uint16
shape = (100, 2000, 2000)

arr = np.random.randint(0, 65536, size=shape, dtype=dtype)

z = zarr.create_array(
    store=zarr.storage.MemoryStore(),
    shape=arr.shape,
    dtype=arr.dtype,
    chunks=(1, 2000, 2000),
    serializer=(codecs.Jpegxl(planar=True, effort=2, lossless=True)),
)

z[:] = arr

Notably changing the chunking along the first axes does not result in an error.

[cgohlke]

Thanks. You might be interested in the Zarr Jpegxl codec specification draft, which explicitly excludes planar layout and instead recommends to compress the channels independently.

[JeppeKlitgaard]

Thanks. You might be interested in the Zarr Jpegxl codec specification draft, which explicitly excludes planar layout and instead recommends to compress the channels independently.

Thank you so much, I was not aware of this and it seems very relevant to my use-case.

I reckon this would be problematic for my use-case, however, as I want my data to be projection-major (P, X, Y) where P are monochrome channels due to convention (and optimal access pattern, potentially, though maybe order=F would work?). If you apply the fix suggested here to not squeeze the dimensions, I am not sure how I would be able to achieve this?

To keep (P, X, Y) wouldn't I need planar=True or alternatively some bidirectional reshape/squeeze transform to be part of Zarr (or the codec instance, though that seems a little hacky).

Am I missing something?

[fcollman]

no you are not missing something, and I think the idea is that indeed you should be using a reshape codec in order to get the data from the jpegxl encoded data into the shape of your data. You need to do this in general for general high dimensional data, and I think the idea is that this is more explicit and extensible a paradigm than implicitly having rules that map jpegxl chunks into dimensions of the zarr.