From 89fd46c57de17a7adf2023b370cada61016c6cb9 Mon Sep 17 00:00:00 2001 From: David Ackerman Date: Tue, 3 Nov 2026 12:03:29 -0500 Subject: [PATCH] Present axes in another order on request (0.10.0) TensorSwitch writes OME-Zarr v3 with dimension_names x, y, z, and worldarray presents the store's C order, so a reader that indexes axes by position (mesh-n-bone, cellmap-analyze's read_onto, cellmap-flow's blockwise input) meshes or reads such an array transposed, or refuses it. open, open_pyramid and metadata take axes_order: "stored" (the default, for good: an x, y, z store read faithfully is x, y, z with its names), "zyx" (the axes named z, y and x last, the others first in stored order) or a sequence of names. Names move with the data, and shape, chunks, shards, units, voxel size and translation with them, through reads, writes, a pyramid's placements and level_meta, level choice, base_voxel_size, reopen and a pickle. Chunk keys take and give presented indices with the store's keys, and add_level takes its arguments in the pyramid's order. open_tensorstore stays in the store's order. Claude-Session: https://claude.ai/code/session_0141wgF8mkGbjGfBFKfKhp64 --- CHANGELOG.md | 18 ++++ README.md | 5 +- docs/spec.md | 55 ++++++++-- pyproject.toml | 2 +- src/worldarray/__init__.py | 2 +- src/worldarray/array.py | 133 ++++++++++++++++++++---- src/worldarray/placement.py | 84 ++++++++++++++- tests/test_axes_order.py | 197 ++++++++++++++++++++++++++++++++++++ uv.lock | 2 +- 9 files changed, 459 insertions(+), 39 deletions(-) create mode 100644 tests/test_axes_order.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d7dea75..c14c32b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,23 @@ # Changelog +## 0.10.0 (2026-11-03) + +Axes in another order, on request. No placement change, no unit table change, and no +default changes: without `axes_order`, every array opens as it did in 0.9. + +- `open`, `open_pyramid` and `metadata` take `axes_order`: `"stored"` (the default, for + good), `"zyx"` (the axes named z, y and x last, in that order, the others first in + stored order) or a sequence naming each axis once. An OME-Zarr written x, y, z + (TensorSwitch's) then reads as z, y, x, and a channel-last array channel first. The + names move with the data, and shape, chunks, shards, units, voxel size and + translation with them, in `.ts`, `meta`, `grid`, reads, writes, `blocks`, a pyramid's + placements and `level_meta`, level choice by voxel size, `base_voxel_size`, `reopen` + and a pickle. `open_tensorstore` stays in the store's order. +- On such a handle, `chunk_key`, `parse_chunk_key` and `list_chunks` take and give + indices in the presented order (a selected axis in its stored place); the keys are + the store's. `add_level` takes its per-axis arguments in the pyramid's order and + creates the array in the group's. + ## 0.9.1 (2026-10-26) A fix for exit. No placement change, no unit table change, and no default changes. diff --git a/README.md b/README.md index c0db081..632b71a 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,7 @@ vol = wa.open(path, units=None) vol.meta # shape, dtype, chunk_shape, axes, units, voxel_size, translation, attrs, ... vol.ts # the tensorstore itself, C order, origin 0 vol.read_index((0, 0, 0), (64, 64, 64), padding=0) +wa.open(path, axes_order="zyx") # an x, y, z store (TensorSwitch's) presented z, y, x, names and all pyr = wa.open_pyramid(path) # every level, opened on first use pyr.select((16, 16, 16), mode="floor") # the level to read at 16 nm @@ -133,7 +134,9 @@ Dialects that store the corner are moved half a voxel on reading: | precomputed `voxel_offset` | corner | `(voxel_offset + 1/2) * resolution` | | nothing | a plain array | voxels of 1, corner at the origin | -Axes are the store's C order; N5 and precomputed are transposed to it. Unlisted or +Axes are the store's C order; N5 and precomputed are transposed to it. `axes_order="zyx"` +presents the axes named z, y and x last, in that order, for a reader that indexes axes by +position; the store's order stays the default. Unlisted or conflicting attributes warn; a voxel size that disagrees with the OME scale raises, since that is how stale metadata is caught. Metadata the parser did not consume stays available raw. The full precedence per format, the half-voxel rules, and every read behaviour are diff --git a/docs/spec.md b/docs/spec.md index 83ce1a1..ddc9005 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -131,6 +131,9 @@ vol = wa.open( # part of a box outside the data method="nearest", # "nearest" | "linear": the resampling kernel; no dtype heuristic (Behaviour 5) select=None, # {"c": 0, "t": 0}: pin non-spatial axes to one index; else they are kept whole + axes_order="stored", # "stored": the store's order (Convention 2) | "zyx": the axes named z, y and x last, in + # that order, the others first in stored order | a sequence naming each axis left after + # `select` once (0.10); every per-axis field and the data move with their names units="nanometer", # what world coordinates are in on spatial axes; a dict per axis kind # ({"space": "micrometer", "time": "second"}); None keeps the store's own units context=None, # a tensorstore.Context; default shared_context(0) @@ -144,7 +147,7 @@ vol = wa.open( strict_units=True, # False: an unknown unit ("pixel") is taken as the target unit, with a warning ) -vol.ts # tensorstore.TensorStore: C-order axes, origin 0, selected axes dropped +vol.ts # tensorstore.TensorStore: the axes in `axes_order`, origin 0, selected axes dropped vol.meta # ArrayMeta in the store's own units vol.grid # ArrayMeta in the requested units, every axis vol.shape, vol.dtype, vol.axes, vol.spatial, vol.level_name, vol.levels, vol.remote @@ -168,7 +171,7 @@ vol.index_box(box) -> (begin, end) # the voxels a world box selects on this gri vol.chunks_covering(begin, end) # chunk indices a voxel box touches vol.reopen(context=None, **overrides) -> WorldArray # the same array on another context (a worker changing its concurrency) -pyr = wa.open_pyramid(path, context=None, units="nanometer", select=None) +pyr = wa.open_pyramid(path, context=None, units="nanometer", select=None, axes_order="stored") pyr.names; pyr.levels; pyr.by_name("s1"); pyr.select(voxel_size, mode="floor") pyr.metas() ; pyr.level_meta("s1") # ArrayMeta per level from the documents alone: shape, dtype, chunks (a sharded # array's inner chunk), fill and placement, with no tensorstore opened @@ -177,11 +180,13 @@ pyr.voxel_sizes() ; pyr.translations() # where it lists the levels (OME-NGFF # choosing or placing a level opens nothing, whether the pyramid was opened by its # group path or by one of its level paths; else by opening each level -wa.metadata(path, level=None, units=None, open=True, retries=None) -> ArrayMeta # open=False: from the documents alone +wa.metadata(path, level=None, units=None, open=True, retries=None, axes_order="stored") -> ArrayMeta + # open=False: from the documents alone wa.detect(path, missing_ok=False, retries=None) -> Node | None # format, kind (array|group), path, docs, scale_key; .attrs (raw user # attributes), .multiscales (OME 0.5's `ome` unwrapped) wa.open_tensorstore(path, level=None, context=None, retries=None) -> tensorstore.TensorStore # reads no world metadata, so an - # array with odd placement attributes (a coordinate array) still opens + # array with odd placement attributes (a coordinate array) still opens; + # always the store's order (Convention 2) wa.corner_to_centre(corner, voxel_size, axes=None) ; wa.default_axes(ndim) ; wa.index_range(begin, end, voxel_size, translation) wa.units.TABLE # the unit tables, public: {"length": {...metres per unit}, "time": {...seconds per unit}} wa.select_level(levels, voxel_size, mode) ; wa.Box ; wa.Retry(max_attempts=10, timeout=None, base_delay=1.0) @@ -196,7 +201,7 @@ wa.list_chunk_keys(path, separator=".", ndim=None) # chunk indices present und # metadata document needed, so a directory of chunk files with no .zarray lists too ``` -`ArrayMeta` is a frozen dataclass with one entry per axis, in the array's C order, and +`ArrayMeta` is a frozen dataclass with one entry per axis, in the array's order (Convention 2), and deliberately the same field names as chunkmirage's `ArrayInfo`: `shape`, `dtype`, `chunk_shape` (tensorstore's read chunk), `axes` (as stored; `("t","c","z","y","x")[-ndim:]` when the store names none), `axes_named`, `units` (as stored, any string; `""` when none), @@ -276,7 +281,8 @@ wa.ensure_group_entry(group_path, name, voxel_size, translation, axes, units="na # document and missing ancestors created, nothing written when the entry already says the same # (Behaviour 12). `create` and `add_level` call this wa.write.codecs_of(path_or_node) -> (codecs, shards) # an array's codecs in its format's dialect and its shard shape, for - # a mirror created like it (cellmap-flow's sync) + # a mirror created like it (cellmap-flow's sync); the shards in the store's order, so a + # mirror of a handle opened with `axes_order` takes vol.meta.shard_shape (Grid.from_meta does) vol = wa.create( path, # ".../vol.zarr/s0": the last component is the level, its parent the multiscale group, made if absent @@ -422,9 +428,21 @@ this document and that one differ, fix one of them. 2. **Axis order is the store's C order; N5 and precomputed are transposed to it** on read, so `(z, y, x)` for a volume, `(c, z, y, x)` with channels, `(t, c, z, y, x)` for OME - 0.5. Nothing is reordered beyond that transposition; `select` pins non-spatial axes. A + 0.5. Nothing is reordered beyond that transposition unless asked. `axes_order="zyx"` + (0.10, on `open`, `open_pyramid` and `metadata`) presents the axes named z, y and x last, + in that order, and the others first in their stored order, so an OME-Zarr written x, y, z + (TensorSwitch's) reads as z, y, x and a channel-last array channel first; an array naming + none of z, y and x is unchanged. A sequence of names presents exactly that order and must + name each axis left after `select` once. The names move with the data, and shape, chunks, + shards, units, voxel size and translation with them, in `.ts`, `meta`, `grid`, every read + and write, the pyramid's placements and `level_meta` (where a selected axis, still there, + keeps its stored place) and a pickle; `open_tensorstore` stays in the store's order. The + default is the store's order for good, not only until 1.0: an x, y, z store read + faithfully is x, y, z with its names, which a reader placing axes by name needs + (Neuroglancer, chunkmirage, a mirror), and a reader indexing axes by position asks for its + order or checks the names (David, 2026-11-03). `select` pins non-spatial axes. A non-zero stored origin (precomputed `voxel_offset`) is moved to 0 in `.ts` and carried in - `translation`. \1 Sharded stores (zarr v3 `sharding_indexed`, sharded precomputed) read like any + `translation`. Sharded stores (zarr v3 `sharding_indexed`, sharded precomputed) read like any other through tensorstore; `chunk_shape` is the inner read chunk, not the shard. Unnamed axes are called `("t", "c", "z", "y", "x")[-ndim:]` (chunkmirage's default) with `axes_named` False, so a consumer can apply its own default instead. @@ -910,6 +928,10 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3. name is replaced; the group's other content is untouched (Behaviour 12). It writes no voxel data: downsampling is the caller's, per block, through `write_index`, labels by mode and images by mean as each repo decides. `exists` means what it means for `create`. + On a pyramid opened with `axes_order` (0.10), the per-axis arguments are in the order it + presents (`factors`, `origin`, `voxel_size` and `translation` on its spatial axes, `shape`, + `chunks` and `shards` as `level_meta` lists them), the array is created in the group's own + order, and the level comes back in the pyramid's. 15. **Chunk keys.** `chunk_key(index)` is the key, relative to the array, of the write unit at `index`: zarr v2 indices joined by the array's `dimension_separator` (`0.1.2`); zarr @@ -921,7 +943,11 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3. (`zarr.json`, `.zarray`, `attributes.json`, `info`, a `c/` key of the wrong arity). `list_chunk_keys` and `list_chunks` use the same parser. The free functions take the format and the array's facts instead of an open array, for an array served from nowhere - (chunkmirage's frontends) or about to be created. + (chunkmirage's frontends) or about to be created. On a handle opened with `axes_order` + (0.10), `chunk_key`, `parse_chunk_key` and `list_chunks` take and give indices in the + order the handle presents, a selected axis in its stored place, as an N5 index is in C + order and its key x first; the key is the store's. The free functions stay in the + store's order. 16. **Write units.** `write_shape` is the shard shape of a sharded zarr v3 array, else the chunk shape; `ArrayMeta.shard_shape` is None when unsharded; `chunk_shape` stays the @@ -1098,7 +1124,16 @@ array under `exists="error"`. 0.9: `WORLDARRAY_THREADS` ahead of every other var blank as unset and a bad value refused; `file_io_floor` sizing file I/O under the budget, ignored with a thread count, and carried through a pickle; a handle pickled on the driver sized in a spawned worker by that worker's `WORLDARRAY_THREADS`. 0.9.1: a process whose -finalization takes a while, exiting while the watchdog kicks flat out, exits cleanly. +finalization takes a while, exiting while the watchdog kicks flat out, exits cleanly. 0.10: +an x, y, z array and a c, x, y, z one opened with `axes_order="zyx"` reading as the +transpose of the stored data with every per-axis field moved, a write through the +reordered view landing in the stored array's order, its chunk keys unchanged and its +indices presented; a sequence order, and a bad order refused; an N5 array and an +unnamed one unchanged under "zyx"; a selected channel keeping its place in a key index; +`metadata` with and without opening; a pyramid's placements and `level_meta` reordered, a +level chosen by an anisotropic voxel size, `base_voxel_size` applied to the right axes, and +`add_level` taking its factors in the presented order; the order carried through +`reopen` and a pickle. Opt-in (`WORLDARRAY_STRESS=1`, about a minute, `tests/test_lost_wakeup.py`): worker processes aim tasks at the moment a one-thread pool's worker exits idle, so a plain context strands some for good (13 of about 7,000 on diff --git a/pyproject.toml b/pyproject.toml index baeca93..68f21b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "worldarray" -version = "0.9.1" +version = "0.10.0" description = "Read zarr, N5 and Neuroglancer precomputed volumes in world coordinates through tensorstore: correct voxel placement, pyramid levels, padded region reads, resampling. numpy + tensorstore only." readme = "README.md" license = { file = "LICENSE" } diff --git a/src/worldarray/__init__.py b/src/worldarray/__init__.py index 8cfe397..a0ffb3b 100644 --- a/src/worldarray/__init__.py +++ b/src/worldarray/__init__.py @@ -28,7 +28,7 @@ parse_chunk_key, ) -__version__ = "0.9.1" +__version__ = "0.10.0" __all__ = [ "ArrayMeta", diff --git a/src/worldarray/array.py b/src/worldarray/array.py index 1f4c772..b78d3f4 100644 --- a/src/worldarray/array.py +++ b/src/worldarray/array.py @@ -29,11 +29,16 @@ from .levels import LevelRef, Mode, discover, discover_group, same_voxel_size, select_level from .placement import ( ArrayMeta, + AxesOrder, _dims_from_docs, + axes_permutation, + check_axes_order, default_axes, light_levels, read_meta, read_meta_light, + spatial_axes, + spatial_order, ) from .retry import Retry, read_with_retries, write_with_retries from .store import ( @@ -92,6 +97,7 @@ def __init__( strict_units: bool = True, mode: str = "r", store_fill: bool = False, + axes_order: AxesOrder = "stored", _ref: LevelRef | None = None, _node: Node | None = None, ): @@ -112,6 +118,7 @@ def __init__( self.context = context if context is not None else shared_context(0) self.retries = retries if retries is not None else Retry() self._select = dict(select or {}) + self._axes_order = check_axes_order(axes_order) self._opts = dict( level=level, voxel_size=voxel_size, @@ -124,6 +131,7 @@ def __init__( strict_units=strict_units, mode=mode, store_fill=store_fill, + axes_order=self._axes_order, ) self._anchor = tuple(float(a) for a in anchor) if anchor is not None else None @@ -147,7 +155,21 @@ def __init__( self._meta = read_meta( node_for_level, self._store, scale_index=ref.scale_index, conflicts=metadata_conflicts ) + stored_axes = self._meta.axes self._apply_select() + # the order the axes are presented in (Conventions 2), and the stored axis at each + # place of a chunk key's index: a selected axis keeps its stored place there + order = axes_permutation(self._meta.axes, self._axes_order) + if order != tuple(range(len(order))): + self._store = self._store[ts.d[order].transpose[:]] + self._meta = self._meta.permuted(order) + if node_for_level.format == "precomputed": # a precomputed key names the spatial chunk + sp = spatial_axes(stored_axes) + self._key_order = tuple( + sp.index(i) for i in spatial_order(stored_axes, self._meta.axes) + ) + else: + self._key_order = axes_permutation(stored_axes, self._meta.axes, fixed=self._select) if base_voxel_size is not None: self._meta = self._rescaled_to_base(base_voxel_size, base_units) space, time = _unit_targets(units) @@ -198,7 +220,12 @@ def _choose_level(self, node: Node, level, voxel_size, level_mode, on_mismatch) warnings.warn(f"{self.path} is an array; level={level!r} is ignored", stacklevel=3) return LevelRef(node.path.rsplit("/", 1)[-1], node.path) # the array itself, always self._pyramid = Pyramid( - self.path, context=self.context, units=self.units, select=self._select, _node=node + self.path, + context=self.context, + units=self.units, + select=self._select, + axes_order=self._axes_order, + _node=node, ) if level is not None: return self._pyramid._ref(level) @@ -227,7 +254,13 @@ def _rescaled_to_base(self, base: Sequence[float], base_units: str | None) -> Ar by the same per-axis factor, so stores whose metadata is in voxels of level 0 (as BigStitcher writes) or unitless (a BigDataViewer N5) can be placed correctly.""" base = tuple(float(v) for v in base) - native = Pyramid(self.path, context=self.context, units=None, select=self._select) + native = Pyramid( + self.path, + context=self.context, + units=None, + select=self._select, + axes_order=self._axes_order, + ) level0 = native.levels[0].meta if len(native) > 1 else self._meta spatial = self._meta.spatial if len(base) != len(spatial): @@ -394,7 +427,11 @@ def levels(self) -> list[tuple[str, tuple[float, ...]]]: def levels_pyramid(self) -> Pyramid: if self._pyramid is None: self._pyramid = Pyramid( - self._where, context=self.context, units=self.units, select=self._select + self._where, + context=self.context, + units=self.units, + select=self._select, + axes_order=self._axes_order, ) return self._pyramid @@ -507,19 +544,37 @@ def _key_params(self) -> dict: params["origin"] = tuple(int(o) for o in scale.get("voxel_offset", [0, 0, 0]))[::-1] return params + def _stored_index(self, index: Sequence[int]) -> tuple[int, ...]: + """A key index in the presented order, in the store's.""" + index, order = tuple(index), self._key_order + if len(index) != len(order): + return index # the key function says what is wrong + stored = [0] * len(order) + for place, axis in enumerate(order): + stored[axis] = index[place] + return tuple(stored) + + def _presented_index(self, index: tuple[int, ...] | None) -> tuple[int, ...] | None: + """A key index in the store's order, in the presented one.""" + if index is None or len(index) != len(self._key_order): + return index + return tuple(index[axis] for axis in self._key_order) + def chunk_key(self, index: Sequence[int]) -> str: """The storage key, relative to the array, of the write unit at ``index`` on the - storage grid (Behaviour 15). A precomputed array's key names the spatial chunk only.""" - return _write.chunk_key(index, **self._key_params()) + storage grid, in the order the axes are presented (Behaviour 15). A precomputed + array's key names the spatial chunk only.""" + return _write.chunk_key(self._stored_index(index), **self._key_params()) def parse_chunk_key(self, key: str) -> tuple[int, ...] | None: """``chunk_key`` inverted; None for a key that is not a chunk.""" - return _write.parse_chunk_key(key, **self._key_params()) + return self._presented_index(_write.parse_chunk_key(key, **self._key_params())) def list_chunks(self) -> list[tuple[int, ...]]: - """The chunk indices that exist in storage (C order, the storage grid), by listing the - array's keys with ``parse_chunk_key``: zarr v2 and v3 (a sharded array lists its - shards), N5, and unsharded precomputed. Works on any kvstore, URLs included.""" + """The chunk indices that exist in storage (the storage grid, in the order the axes + are presented), by listing the array's keys with ``parse_chunk_key``: zarr v2 and v3 + (a sharded array lists its shards), N5, and unsharded precomputed. Works on any + kvstore, URLs included.""" params = self._key_params() kv = ts.KvStore.open(self._store.kvstore.spec()).result() keys = [k.decode() if isinstance(k, bytes) else str(k) for k in kv.list().result()] @@ -529,7 +584,7 @@ def list_chunks(self) -> list[tuple[int, ...]]: k.rsplit("/", 1)[-1] if params["format"] == "precomputed" else k, **params ) if idx is not None: - found.add(idx) + found.add(self._presented_index(idx)) return sorted(found) # --- writing (Behaviour 13) ---------------------------------------------------- @@ -1013,6 +1068,7 @@ def __init__( units: Units = "nanometer", select: dict[str, int] | None = None, mode: str = "r", + axes_order: AxesOrder = "stored", _node: Node | None = None, **kwargs, ): @@ -1023,7 +1079,8 @@ def __init__( self.units = units self._select = select or {} self._mode = mode - self._kwargs = {**kwargs, "mode": mode} + self._axes_order = check_axes_order(axes_order) + self._kwargs = {**kwargs, "mode": mode, "axes_order": self._axes_order} self._node = _node if _node is not None else detect(path) self._group, self._refs = discover_group(self._node) # the node listing the levels if not self._refs: @@ -1079,7 +1136,11 @@ def placements(self) -> list[tuple[str, tuple[float, ...], tuple[float, ...]]]: space, _ = _unit_targets(self.units) # the group the levels were discovered from, whichever path opened the pyramid group = self._group - light = light_levels(group, self.names, space) if group is not None else None + light = ( + light_levels(group, self.names, space, self._axes_order) + if group is not None + else None + ) if light is None: light = [ (lvl.level_name, lvl.source_voxel_size, lvl.translation) for lvl in self.levels @@ -1090,12 +1151,15 @@ def placements(self) -> list[tuple[str, tuple[float, ...], tuple[float, ...]]]: def level_meta(self, level: str | int) -> ArrayMeta: """One level's ``ArrayMeta`` in native units from its documents alone, with no tensorstore opened (shape, dtype, chunks and placement; a sharded array's inner - chunk).""" + chunk), every axis in the order the pyramid presents them, a selected one in its + stored place.""" ref = self._ref(level) conflicts = self._kwargs.get("metadata_conflicts", "error") if ref.scale_index is not None: - return read_meta_light(self._node, scale_index=ref.scale_index, conflicts=conflicts) - return read_meta_light(detect(ref.path), conflicts=conflicts) + meta = read_meta_light(self._node, scale_index=ref.scale_index, conflicts=conflicts) + else: + meta = read_meta_light(detect(ref.path), conflicts=conflicts) + return meta.permuted(axes_permutation(meta.axes, self._axes_order, fixed=self._select)) def metas(self) -> list[ArrayMeta]: """``level_meta`` for every level, in order.""" @@ -1141,6 +1205,12 @@ def add_level( base_node = detect(base_ref.path) base = read_meta_light(base_node, conflicts=self._kwargs.get("metadata_conflicts", "error")) ndim, spatial = base.ndim, base.spatial + # the arguments are in the order the pyramid presents the axes; the array is the group's + order = axes_permutation(base.axes, self._axes_order, fixed=self._select) + factors, origin, voxel_size, translation, shape, chunks, shards = ( + _in_stored_order(v, order, spatial) + for v in (factors, origin, voxel_size, translation, shape, chunks, shards) + ) if voxel_size is None or translation is None: if factors is None: raise ValueError("pass factors (and origin), or voxel_size and translation") @@ -1184,7 +1254,7 @@ def add_level( ) self._group, self._refs = discover_group(detect(group.path)) self._light, self._levels = None, {} - return vol + return vol if self._axes_order == "stored" else vol.reopen(axes_order=self._axes_order) def _select_ref(self, voxel_size, mode: Mode = "floor") -> LevelRef: name, _ = select_level(self.voxel_sizes(), voxel_size, mode) @@ -1337,6 +1407,27 @@ def adopt(normalized: str, fmt: str) -> WorldArray: ) +def _in_stored_order(values, order: tuple[int, ...], spatial: tuple[int, ...]): + """Per-axis ``values`` given in a presented order (``order[k]`` is the stored axis at + place ``k``), in the store's: one per axis, or one per spatial axis. A scalar, None or a + list of another length passes through.""" + if values is None or np.isscalar(values): + return values + values = tuple(values) + if len(values) == len(order): + stored = [None] * len(order) + for place, axis in enumerate(order): + stored[axis] = values[place] + return tuple(stored) + presented = [axis for axis in order if axis in spatial] + if len(values) == len(presented): + stored = [None] * len(spatial) + for place, axis in enumerate(presented): + stored[spatial.index(axis)] = values[place] + return tuple(stored) + return values + + def _reopen(path: str, opts: dict, recipe=None) -> WorldArray: return WorldArray(path, context=from_recipe(recipe), **opts) @@ -1357,13 +1448,14 @@ def metadata( units: Units = None, open: bool = True, retries: Retry | None = None, + axes_order: AxesOrder = "stored", ) -> ArrayMeta: """The ``ArrayMeta`` of an array (or of ``level`` under a group), in native units unless - ``units`` says otherwise. ``open=False`` reads it from the documents alone, with no - tensorstore opened. ``retries`` times out and retries each document read and the open - (Behaviour 6).""" + ``units`` says otherwise, its axes in ``axes_order`` (Conventions 2). ``open=False`` + reads it from the documents alone, with no tensorstore opened. ``retries`` times out and + retries each document read and the open (Behaviour 6).""" if open: - vol = WorldArray(path, level=level, units=units, retries=retries) + vol = WorldArray(path, level=level, units=units, retries=retries, axes_order=axes_order) return vol.meta if units is None else vol.grid node = detect(path) if level == "first": @@ -1391,6 +1483,7 @@ def metadata( meta = read_meta_light(detect(ref.path)) else: meta = read_meta_light(node) + meta = meta.permuted(axes_permutation(meta.axes, axes_order)) if units is None: return meta space, time = _unit_targets(units) diff --git a/src/worldarray/placement.py b/src/worldarray/placement.py index e23100b..f626da8 100644 --- a/src/worldarray/placement.py +++ b/src/worldarray/placement.py @@ -17,6 +17,7 @@ import json import logging import re +from collections.abc import Sequence from dataclasses import dataclass, field, replace import numpy as np @@ -143,12 +144,82 @@ def drop_axis(self, axis: int) -> ArrayMeta: shard_shape=pick(self.shard_shape) if self.shard_shape is not None else None, ) + def permuted(self, order: Sequence[int]) -> ArrayMeta: + """The same metadata with axis ``order[k]`` at place ``k``: every per-axis field moves + with its name.""" + pick = lambda t: tuple(t[i] for i in order) # noqa: E731 + return replace( + self, + shape=pick(self.shape), + chunk_shape=pick(self.chunk_shape), + axes=pick(self.axes), + units=pick(self.units), + voxel_size=pick(self.voxel_size), + translation=pick(self.translation), + shard_shape=pick(self.shard_shape) if self.shard_shape is not None else None, + ) + @property def write_shape(self) -> tuple[int, ...]: """The unit one writer may own: the shard shape when sharded, else the chunk shape.""" return self.shard_shape if self.shard_shape is not None else self.chunk_shape +AxesOrder = str | Sequence[str] + + +def check_axes_order(order: AxesOrder) -> str | tuple[str, ...]: + """``"stored"``, ``"zyx"`` or a tuple of axis names (Conventions 2).""" + if isinstance(order, str): + if order not in ("stored", "zyx"): + raise ValueError( + f"axes_order is 'stored', 'zyx' or a sequence of axis names, not {order!r}" + ) + return order + names = tuple(str(a) for a in order) + if len(set(names)) != len(names): + raise ValueError(f"axes_order {names} names an axis twice") + return names + + +def axes_permutation( + axes: Sequence[str], order: AxesOrder = "stored", fixed: Sequence[str] = () +) -> tuple[int, ...]: + """The axis of ``axes`` at each place once ``order`` presents them (Conventions 2): + ``"stored"`` keeps them; ``"zyx"`` moves the axes named z, y and x last, in that order, + the others first in stored order; a sequence of names is that order. Axes named in + ``fixed`` (selected ones, still in a key index or a pyramid's ``level_meta``) keep their + places, and the others fill the rest in order.""" + axes = tuple(axes) + order = check_axes_order(order) + movable = [i for i, a in enumerate(axes) if a not in fixed] + if order == "stored": + ranked = movable + elif order == "zyx": + named = sorted( + (i for i in movable if axes[i] in SPATIAL_NAMES), key=lambda i: "zyx".index(axes[i]) + ) + ranked = [i for i in movable if axes[i] not in SPATIAL_NAMES] + named + else: + names = [axes[i] for i in movable] + if sorted(order) != sorted(names): + raise ValueError(f"axes_order {order} must name each of the axes {tuple(names)} once") + ranked = [movable[names.index(a)] for a in order] + places = iter(ranked) + return tuple(i if axes[i] in fixed else next(places) for i in range(len(axes))) + + +def spatial_order(axes: Sequence[str], order: AxesOrder = "stored") -> tuple[int, ...]: + """The spatial axes of ``axes`` in the order ``order`` presents them.""" + axes = tuple(axes) + sp = spatial_axes(axes) + others = tuple(a for i, a in enumerate(axes) if i not in sp) + order = check_axes_order(order) + if not isinstance(order, str): + order = tuple(a for a in order if a in {axes[i] for i in sp}) + return tuple(i for i in axes_permutation(axes, order, fixed=others) if i in sp) + + def spatial_axes(axes: tuple[str, ...]) -> tuple[int, ...]: named = tuple(i for i, a in enumerate(axes) if a in SPATIAL_NAMES) if named: @@ -730,29 +801,32 @@ def _precomputed(info: dict, scale_index: int, ndim: int) -> _Resolved: def light_voxel_sizes( - group: Node, names: list[str], space: str | None + group: Node, names: list[str], space: str | None, axes_order: AxesOrder = "stored" ) -> list[tuple[str, tuple[float, ...]]] | None: - levels = light_levels(group, names, space) + levels = light_levels(group, names, space, axes_order) return None if levels is None else [(name, vs) for name, vs, _ in levels] def light_levels( - group: Node, names: list[str], space: str | None + group: Node, names: list[str], space: str | None, axes_order: AxesOrder = "stored" ) -> list[tuple[str, tuple[float, ...], tuple[float, ...]]] | None: """``(name, spatial voxel size)`` for every level of a group, from the group's own document, in ``space`` units (native when None): OME-NGFF datasets, COSEM N5 ``multiscales[].datasets[].transform`` and precomputed ``scales``. None when the document does not say (then the levels have to be opened), or when a unit cannot be converted. ``light_levels`` adds each level's spatial translation (the centre of its - voxel 0, in the same units), so a pyramid can be placed without opening any level.""" + voxel 0, in the same units), so a pyramid can be placed without opening any level. + Both on the spatial axes in the order ``axes_order`` presents them.""" try: if group.format == "precomputed": f = U.factor("nm", space) if space else 1.0 + order = spatial_order(("z", "y", "x"), axes_order) out = [] for i, scale in enumerate(group.docs["info"]["scales"]): res = [float(v) * f for v in scale["resolution"]][::-1] vo = [float(v) for v in scale.get("voxel_offset", [0, 0, 0])][::-1] tr = [(o + 0.5) * r for o, r in zip(vo, res)] + res, tr = [res[k] for k in order], [tr[k] for k in order] out.append((str(scale.get("key", i)), tuple(res), tuple(tr))) return out multiscales = group.multiscales @@ -788,7 +862,7 @@ def light_levels( axes = found.axes or default_axes(ndim) units = found.units or ("",) * ndim sizes, trs = [], [] - for i in spatial_axes(axes): + for i in spatial_order(axes, axes_order): f = U.factor(units[i], space, missing_as=U.NANOMETER) if space else 1.0 sizes.append(found.voxel_size[i] * f) trs.append((found.translation[i] if found.translation is not None else 0.0) * f) diff --git a/tests/test_axes_order.py b/tests/test_axes_order.py new file mode 100644 index 0000000..de63719 --- /dev/null +++ b/tests/test_axes_order.py @@ -0,0 +1,197 @@ +"""Presenting the axes in another order: ``axes_order`` (spec, Conventions 2, Behaviours 14 +and 15).""" + +from __future__ import annotations + +import pickle + +import numpy as np +import pytest + +import worldarray as wa +from conftest import pattern, write_n5, write_zarr2 + +XYZ = ("x", "y", "z") + + +def _xyz(tmp_path, shape=(8, 6, 4), **kw) -> tuple[str, np.ndarray]: + """An x, y, z array (TensorSwitch's order) at 16, 16, 45 nm, filled with a pattern.""" + path = str(tmp_path / "xyz.zarr/s0") + vol = wa.create(path, shape, "uint16", (16, 16, 45), (4, 4, 0), axes=XYZ, **kw) + data = pattern(shape) + vol.write_index((0,) * len(shape), shape, data) + return path, data + + +def test_zyx_presents_an_xyz_array_as_its_transpose(tmp_path): + path, data = _xyz(tmp_path, chunks=(4, 3, 2)) + stored = wa.open(path) + assert stored.axes == XYZ and stored.shape == (8, 6, 4) # the default is the store's order + vol = wa.open(path, axes_order="zyx") + assert vol.axes == ("z", "y", "x") and vol.shape == (4, 6, 8) and vol.ts.shape == (4, 6, 8) + assert vol.voxel_size == (45, 16, 16) and vol.translation == (0, 4, 4) + assert vol.meta.chunk_shape == (2, 3, 4) and vol.write_shape == (2, 3, 4) + assert vol.meta.units == ("nanometer",) * 3 + np.testing.assert_array_equal(vol.read_index((0, 0, 0), (4, 6, 8)), data.transpose(2, 1, 0)) + np.testing.assert_array_equal(vol.read(vol.box), data.transpose(2, 1, 0)) + part = vol.read_index((1, 2, 3), (3, 5, 7)) + np.testing.assert_array_equal(part, data[3:7, 2:5, 1:3].transpose(2, 1, 0)) + assert wa.Grid.from_meta(vol.meta).axes == ("z", "y", "x") + + +def test_a_channel_axis_comes_first_and_select_still_pins_it(tmp_path): + path = str(tmp_path / "cxyz.zarr/s0") + data = pattern((2, 8, 6, 4)) + wa.create( + path, data.shape, "uint16", (16, 16, 45), (4, 4, 0), axes=("c", *XYZ), chunks=(1, 4, 3, 2) + ) + wa.open(path, mode="r+").write_index((0,) * 4, data.shape, data) + vol = wa.open(path, axes_order="zyx") + assert vol.axes == ("c", "z", "y", "x") and vol.meta.chunk_shape == (1, 2, 3, 4) + np.testing.assert_array_equal(vol.read_index((0,) * 4, vol.shape), data.transpose(0, 3, 2, 1)) + one = wa.open(path, axes_order="zyx", select={"c": 1}) + assert one.axes == ("z", "y", "x") + np.testing.assert_array_equal(one.read_index((0,) * 3, one.shape), data[1].transpose(2, 1, 0)) + # a channel-last array comes out channel first + last = str(tmp_path / "xyzc.zarr/s0") + wa.create( + last, (8, 6, 4, 2), "uint16", (16, 16, 45), (4, 4, 0), axes=(*XYZ, "c"), chunks=(4, 3, 2, 1) + ) + assert wa.open(last, axes_order="zyx").axes == ("c", "z", "y", "x") + + +def test_a_write_through_the_reordered_view_lands_in_the_stores_order(tmp_path): + path = str(tmp_path / "w.zarr/s0") + wa.create( + path, + (8, 6, 4), + "uint16", + (16, 16, 45), + (4, 4, 0), + axes=XYZ, + chunks=(2, 3, 2), + shards=(4, 6, 2), + ) + vol = wa.open(path, mode="r+", axes_order="zyx") + assert vol.meta.shard_shape == (2, 6, 4) and vol.write_shape == (2, 6, 4) + zyx = pattern((4, 6, 8)) + vol.write_index((0, 0, 0), (2, 6, 4), zyx[:2, :, :4]) # one shard: z 0-2, x 0-4 + vol.write_chunk((1, 0, 1), zyx[2:, :, 4:]) + stored = wa.open(path) + np.testing.assert_array_equal( + stored.read_index((0, 0, 0), (4, 6, 2)), zyx[:2, :, :4].transpose(2, 1, 0) + ) + np.testing.assert_array_equal( + stored.read_index((4, 0, 2), (8, 6, 4)), zyx[2:, :, 4:].transpose(2, 1, 0) + ) + with pytest.raises(ValueError, match="whole write units"): + vol.write_index((0, 0, 0), (1, 6, 4), zyx[:1, :, :4]) + # keys are the store's; the indices are in the presented order + assert vol.chunk_key((1, 0, 1)) == stored.chunk_key((1, 0, 1)[::-1]) == "c/1/0/1" + assert vol.chunk_key((0, 0, 1)) == "c/1/0/0" + assert vol.parse_chunk_key("c/1/0/0") == (0, 0, 1) + assert vol.list_chunks() == [(0, 0, 0), (1, 0, 1)] + assert stored.list_chunks() == [(0, 0, 0), (1, 0, 1)] + for b in vol.blocks(): + assert vol.parse_chunk_key(vol.chunk_key(b.index)) == b.index + + +def test_a_selected_axis_keeps_its_stored_place_in_a_key_index(tmp_path): + path = str(tmp_path / "k.zarr/s0") + wa.create( + path, (8, 6, 4, 2), "uint16", (16, 16, 45), (4, 4, 0), axes=(*XYZ, "c"), chunks=(4, 3, 2, 1) + ) + vol = wa.open(path, axes_order="zyx", select={"c": 1}) + assert vol.axes == ("z", "y", "x") + # stored (x, y, z, c); the key index is (z, y, x, c): c stays last, where it is stored + assert vol.chunk_key((1, 0, 1, 1)) == "c/1/0/1/1" + assert vol.chunk_key((0, 1, 1, 0)) == "c/1/1/0/0" + assert vol.parse_chunk_key("c/1/1/0/0") == (0, 1, 1, 0) + + +def test_a_sequence_names_the_order_and_a_bad_order_is_refused(tmp_path): + path, data = _xyz(tmp_path, chunks=(4, 3, 2)) + vol = wa.open(path, axes_order=("y", "z", "x")) + assert vol.axes == ("y", "z", "x") and vol.voxel_size == (16, 45, 16) + np.testing.assert_array_equal(vol.read_index((0, 0, 0), vol.shape), data.transpose(1, 2, 0)) + with pytest.raises(ValueError, match="axes_order is"): + wa.open(path, axes_order="xyz") + with pytest.raises(ValueError, match="must name each of the axes"): + wa.open(path, axes_order=("z", "y")) + with pytest.raises(ValueError, match="names an axis twice"): + wa.open(path, axes_order=("z", "z", "x")) + + +def test_n5_and_unnamed_arrays_are_unchanged_under_zyx(tmp_path): + data = pattern((4, 6, 8)) + write_n5(tmp_path / "a.n5", data, attrs={"resolution": [16, 16, 45], "axes": ["x", "y", "z"]}) + n5 = wa.open(str(tmp_path / "a.n5"), axes_order="zyx", units=None) + assert n5.axes == ("z", "y", "x") and n5.voxel_size == (45, 16, 16) + np.testing.assert_array_equal(n5.read_index((0, 0, 0), n5.shape), data) + write_zarr2(tmp_path / "u.zarr", data, chunks=(2, 3, 4)) + unnamed = wa.open(str(tmp_path / "u.zarr"), axes_order="zyx", units=None) + assert unnamed.axes == ("z", "y", "x") and not unnamed.meta.axes_named + np.testing.assert_array_equal(unnamed.read_index((0, 0, 0), unnamed.shape), data) + + +def test_metadata_with_and_without_opening(tmp_path): + path, _ = _xyz(tmp_path, chunks=(4, 3, 2)) + for opened in (True, False): + meta = wa.metadata(path, axes_order="zyx", open=opened) + assert meta.axes == ("z", "y", "x") and meta.shape == (4, 6, 8) + assert meta.chunk_shape == (2, 3, 4) and meta.voxel_size == (45, 16, 16) + assert meta.translation == (0, 4, 4) + assert wa.metadata(path, open=False).axes == XYZ + + +def _xyz_pyramid(tmp_path) -> str: + """s0 at 16, 16, 45 nm and s1 at 32, 32, 45 (x and y halved), stored x, y, z.""" + _xyz(tmp_path, chunks=(4, 3, 2)) + group = str(tmp_path / "xyz.zarr") + wa.open_pyramid(group, mode="r+").add_level("s1", factors=(2, 2, 1)) + return group + + +def test_a_pyramid_presents_its_placements_and_level_meta_in_the_order(tmp_path): + group = _xyz_pyramid(tmp_path) + pyr = wa.open_pyramid(group, axes_order="zyx") + assert pyr.placements() == [ + ("s0", (45.0, 16.0, 16.0), (0.0, 4.0, 4.0)), + ("s1", (45.0, 32.0, 32.0), (0.0, 12.0, 12.0)), + ] + assert pyr.level_meta("s1").shape == (4, 3, 4) and pyr.level_meta("s1").axes == ("z", "y", "x") + assert pyr.by_name("s1").shape == (4, 3, 4) and pyr.by_name("s1").axes == ("z", "y", "x") + # a level chosen by an anisotropic voxel size, given in the presented order + vol = wa.open(group, voxel_size=(45, 32, 32), axes_order="zyx") + assert vol.level_name == "s1" and vol.voxel_size == (45, 32, 32) + with pytest.raises(ValueError, match="no level at voxel size"): + wa.open(group, voxel_size=(45, 32, 32)) # read as x, y, z: no such level + + +def test_base_voxel_size_lands_on_the_presented_axes(tmp_path): + group = _xyz_pyramid(tmp_path) + vol = wa.open(group, level="s1", axes_order="zyx", base_voxel_size=(450, 160, 160)) + assert vol.voxel_size == (450, 320, 320) + + +def test_add_level_takes_its_arguments_in_the_presented_order(tmp_path): + group = _xyz_pyramid(tmp_path) + pyr = wa.open_pyramid(group, mode="r+", axes_order="zyx") + s2 = pyr.add_level("s2", factors=(2, 2, 2), from_level="s1") + s3 = pyr.add_level("s3", factors=(1, 2, 2), from_level="s1", chunks=(2, 3, 4)) + assert s2.axes == s3.axes == ("z", "y", "x") + assert s2.voxel_size == (90, 64, 64) and s3.voxel_size == (45, 64, 64) + stored = wa.open(f"{group}/s3") + assert stored.axes == XYZ and stored.voxel_size == (64, 64, 45) + assert stored.shape == (2, 2, 4) and stored.meta.chunk_shape == (4, 3, 2) + + +def test_the_order_survives_reopen_and_a_pickle(tmp_path): + path, data = _xyz(tmp_path, chunks=(4, 3, 2)) + vol = wa.open(path, axes_order="zyx") + for again in (vol.reopen(), pickle.loads(pickle.dumps(vol))): + assert again.axes == ("z", "y", "x") + np.testing.assert_array_equal( + again.read_index((0, 0, 0), again.shape), data.transpose(2, 1, 0) + ) + assert vol.reopen(axes_order="stored").axes == XYZ diff --git a/uv.lock b/uv.lock index 9403e16..293a063 100644 --- a/uv.lock +++ b/uv.lock @@ -1247,7 +1247,7 @@ wheels = [ [[package]] name = "worldarray" -version = "0.9.1" +version = "0.10.0" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },