zarr_memory_contiguity.py

# /// script
# requires-python = "<=3.13"
# dependencies = [
#   "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main",
#   "numpy",
# ]
# ///

"""
The decoded memory layout is a runtime setting, not a property of the array.

A `(t, z, y, x)` array is written and read back under different values of
the `order` config option. The stored bytes are constant. The
strides of the decoded array are variable, and therefore so is which named dimension varies
fastest in memory.
"""

import numpy as np

import zarr

DIMS = ("t", "z", "y", "x")

z = zarr.create_array(
    store={},
    shape=(2, 3, 4, 5),
    chunks=(2, 3, 4, 5),
    dtype="int32",
    dimension_names=DIMS,
    compressors=None,
)
data = np.arange(2 * 3 * 4 * 5, dtype="int32").reshape(z.shape)
z[:] = data


def describe(out: np.ndarray) -> str:
    fastest = DIMS[int(np.argmin(out.strides))]
    return (
        f"strides={out.strides}  C={out.flags.c_contiguous} "
        f"F={out.flags.f_contiguous}  fastest-varying={fastest!r}"
    )


# Per-array config: same stored bytes, decoded into a different memory layout.
for order in ("C", "F"):
    out = z.with_config({"order": order})[:]
    print(f"with_config order={order!r}   {describe(out)}")
    assert np.array_equal(out, data)