read_tensors: allow an explicit z-window so multi-sensor tensors are co-registered

[espg]

🤖 from Claude

Problem

read_tensors derives its z-window per call, from the digests it is given. Two sensors' tensors of the same block therefore land on different axes and cannot be stacked, which is exactly what a multi-sensor consumer wants to do.

Measured on the public demo stores (englacial/zagg/demo/{atl03_tdigest_o9,gedi_flux_o9}.zarr), same block, identical call (n_bins=256, resolution=1.0, block_order=12, fit="degrade_resolution"):

block 4331422233444        atl03  z0 =  -71.0  gain = 1
                           gedi   z0 =  -59.0  gain = 1     -> 12 bins apart

block 4331422411132        atl03  z0 = -201.0  gain = 2
                           gedi   z0 =  -24.0  gain = 1     -> different origin AND bin height

The second case is the bad one: the gains differ, so the two tensors are not merely offset — they are not on the same scale, and no amount of shifting reconciles them. Nothing in the returned tuple flags this; the offsets come back in (offset, gain) and a caller who does not compare them across sensors gets silently unstackable arrays.

What a consumer has to do today

Rebuild the windowing by hand, outside the reader:

window = dict(n_bins=n_bins, resolution=resolution, bottom=0.05, top=0.95,
              fit="degrade_resolution")
z0, n_bins, dz = chunk_z_range([digests from BOTH sensors], **window)
cube = np.zeros((side, side, n_bins), dtype=np.float32)
for word, digest in cells:
    r, c = rank_to_rowcol(block_rank(word, block_order)[0], cell_order - block_order)
    cube[r, c] = rasterize_cell(digest, z0, dz, n_bins)

That is read_tensors' own internals, reimplemented in the consumer, and it means reading the ragged digests separately rather than using the tensor path at all. chunk_z_range and rasterize_cell are importable from moczarr.hhdc, so this works — but every multi-sensor consumer will write the same loop, and any that does not will produce quietly misregistered cubes.

Proposal

Let the caller supply the window instead of only deriving it:

read_tensors(store, field, ..., z_window=(z0, dz))   # or offset=/gain=

When given, skip chunk_z_range and rasterize onto the supplied axis; weight outside the window is dropped exactly as it is today for a window that does not fit. Then a co-registered read is:

z0, n_bins, dz = chunk_z_range(both_sensors_digests, **window)   # already public
a = next(read_tensors(atl03_store, afield, z_window=(z0, dz), n_bins=n_bins, ...))
g = next(read_tensors(gedi_store,  gfield, z_window=(z0, dz), n_bins=n_bins, ...))

which keeps the tensor path, the occupancy mask, and the subtree restriction, none of which the hand-rolled version gets for free.

Secondary, and cheaper if the above is unwanted: have fit="degrade_resolution" record that it degraded somewhere a caller can see without comparing gain to what it asked for. Right now a silently coarsened axis and an as-requested one are indistinguishable in the return value.

Context

Found while building the co-registered ATL03 + GEDI voxel export for the zagg reader demo (englacial/zagg#540). The consumer there wants two (128, 128, n_bins) cubes of identical shape to stack into a 2-channel model input. Worth noting the cost is real: on 4331422411132 the shared window degrades GEDI from the 0.5 m bins it would get alone to 4 m, because ATL03's tail sets the floor — that is inherent to sharing an axis, but a consumer can only discover it by computing both windows and comparing.