diff --git a/src/zarr/codecs/transpose.py b/src/zarr/codecs/transpose.py index 373a27cab9..9fcee4e66b 100644 --- a/src/zarr/codecs/transpose.py +++ b/src/zarr/codecs/transpose.py @@ -4,6 +4,8 @@ from dataclasses import dataclass, replace from typing import TYPE_CHECKING, cast +import numpy as np + from zarr.abc.codec import ArrayArrayCodec from zarr.buffer import NDBuffer from zarr.codecs.registry import register_codec @@ -76,9 +78,7 @@ async def _decode_single( chunk_array: NDBuffer, chunk_spec: ArraySpec, ) -> NDBuffer: - inverse_order = [0] * chunk_spec.ndim - for x, i in enumerate(self.order): - inverse_order[x] = i + inverse_order = np.argsort(self.order) chunk_array = chunk_array.transpose(inverse_order) return chunk_array diff --git a/tests/v3/test_codecs.py b/tests/v3/test_codecs.py index 251570f767..514294c4b0 100644 --- a/tests/v3/test_codecs.py +++ b/tests/v3/test_codecs.py @@ -406,6 +406,23 @@ async def test_transpose( assert await (store / "transpose/0.0").get() == await (store / "transpose_zarr/0.0").get() +@pytest.mark.parametrize("order", [[1, 2, 0], [1, 2, 3, 0], [3, 2, 4, 0, 1]]) +def test_transpose_non_self_inverse(store: Store, order): + shape = [i + 3 for i in range(len(order))] + data = np.arange(0, np.prod(shape), dtype="uint16").reshape(shape) + a = Array.create( + store / "transpose_non_self_inverse", + shape=data.shape, + chunk_shape=data.shape, + dtype=data.dtype, + fill_value=0, + codecs=[TransposeCodec(order=order), BytesCodec()], + ) + a[:, :] = data + read_data = a[:, :] + assert np.array_equal(data, read_data) + + def test_transpose_invalid( store: Store, ):