rechunk dataset

[tackandr]

Add a manual 2-pass rechunk module with a temporary store

Use a temporary chunking that merges only one leading axis while already splitting the spatial axes. The best default here is:

temp: (24, 1, 100, 100)
final: (24, 66, 100, 100)

Why this works well:

Pass 1 reads each source chunk once.
Pass 1 only needs a buffer of shape (24, 1, 1069, 949).
Pass 2 only needs a buffer of shape (24, 66, 100, 100).
You avoid repeatedly touching the same final chunk with tiny partial writes.

For your shape, that gives:

temp chunks: 3 * 66 * 11 * 10 = 21,780
final chunks: 3 * 1 * 11 * 10 = 330

Peak RAM is roughly:

pass 1 buffer: 24 * 1 * 1069 * 949 elements
about 97 MiB for float32, 195 MiB for float64
pass 2 buffer: 24 * 66 * 100 * 100 elements
about 60.4 MiB for float32, 120.8 MiB for float64

A sketch that suggested by chatgpt:
`import numpy as np
import zarr

source array

src = zarr.open("src.zarr", mode="r") # shape (67, 66, 1069, 949)

temp store: uncompressed or lightly compressed is usually best for scratch

tmp = zarr.open(
"tmp.zarr",
mode="w",
shape=src.shape,
chunks=(24, 1, 100, 100),
dtype=src.dtype,
compressor=None, # scratch store
)

final store

dst = zarr.open(
"dst.zarr",
mode="w",
shape=src.shape,
chunks=(24, 66, 100, 100),
dtype=src.dtype,
compressor=src.compressor, # or your desired final compressor
)

T, C, Y, X = src.shape

Pass 1: merge axis 0 into 24, keep axis 1 = 1, and split Y/X to 100x100

for t0 in range(0, T, 24):
t1 = min(t0 + 24, T)
tlen = t1 - t0
for c in range(C):
buf = np.empty((tlen, 1, Y, X), dtype=src.dtype)
for k, t in enumerate(range(t0, t1)):
buf[k, 0, :, :] = src[t, c, :, :]
tmp[t0:t1, c:c+1, :, :] = buf

Pass 2: merge axis 1 into 66

for t0 in range(0, T, 24):
t1 = min(t0 + 24, T)
for y0 in range(0, Y, 100):
y1 = min(y0 + 100, Y)
for x0 in range(0, X, 100):
x1 = min(x0 + 100, X)
buf = tmp[t0:t1, :, y0:y1, x0:x1]
dst[t0:t1, :, y0:y1, x0:x1] = buf`