fixes bug in transpose (zarr-developers#1949)

normanrz · web-flow · commit 5c4dc2339d68 · 2024-06-04T12:57:55.000+02:00
diff --git 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
@@ -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,
 ):
