zarr lazy indexing demo

lazy.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "zarr @ git+https://github.com/d-v-b/zarr-python.git@feat/lazy-indexing",
#     "dask[array]",
#     "numpy",
# ]
# ///

import zarr
import numpy as np
from zarr.experimental.lazy_indexing import Array, merge
import dask.array as da

store = {}
np_data = np.arange(100)
zarr.create_array(store, data=np_data, chunks=(10,), fill_value=0, write_data=True)

# Use the lazy array
lazy_array = Array.open(store)
print(lazy_array)
# <Array memory://129773024766528 domain=IndexDomain([0, 100)) dtype=int64>

slice_a = slice(0, 10)
slice_b = slice(10, None)

# Declare the lower 10% of the array
subregion_a = lazy_array[slice_a]
print(subregion_a)
# <Array memory://129773024766528 domain=IndexDomain([0, 10)) dtype=int64>
assert np.array_equal(np.array(subregion_a), np_data[slice_a])

# Declare the upper 90% of the array
subregion_b = lazy_array[slice_b]
print(subregion_b)
# <Array memory://129773024766528 domain=IndexDomain([10, 100)) dtype=int64>
assert np.array_equal(np.array(subregion_a), np_data[slice_a])

# Test that merging inverts slicing
merged = merge([subregion_a, subregion_b])
assert  merged == lazy_array
assert np.array_equal(np.array(merged), np_data)

# Test with dask
assert np.array_equal(da.from_array(lazy_array).compute(), np_data)