diff --git a/CHANGELOG.md b/CHANGELOG.md index 46efac2..5ccaf54 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,28 @@ # Changelog +## 0.3.0 (2026-10-22) + +Grids and blocks, built to spec draft 3 (Behaviours 17 to 19). Nothing on the read or +write side changes. + +- `blocks(shape, block_shape, context=0, box=None, clip=True)`: the array (or a voxel + box, which may reach past the array) cut into write boxes that tile it with no overlap + and no gap, each with a read box grown by `context` (one int, one per axis, or + `(before, after)` pairs) and clipped to the array when `clip`; C order, `len`, + `blocks[id]` in O(1), `index_of(begin)`; a `Block` carries `index`, `id`, `begin`, + `end`, `read_begin`, `read_end`, `at_start` and `at_end` per axis and `inner`, the write + box's slices inside the read box. Plain picklable records; nothing is scheduled. + `vol.blocks()` defaults the block to `write_shape`, so default blocks are aligned writes. + `vol.write_chunk(index, data)` writes one whole unit. +- `Grid(shape, voxel_size, translation, axes, units, chunks, shards)`: a `WorldArray` + without a store, with `box`, `corner`, `index_box`, `chunks_covering`, `write_shape`, + `blocks` and `meta(dtype)`; `Grid.from_meta(meta, voxel_size=)` is an existing array's + grid or the resample target grid over it. `create(path, grid=g, dtype=)` creates from it. +- `level_factors(voxel_size, ratio=1.5, ...)`: the anisotropy-aware pyramid schedule + mesh-n-bone and cellmap-analyze share, per level. +- `units.abbreviation`, and `store.join`, `store.parent`, `store.basename` as public + URL-aware path helpers. + ## 0.2.0 (2026-10-22) Metadata and placed creation, built to spec draft 3 (Behaviours 10 to 16). Nothing on the diff --git a/README.md b/README.md index c3bc0dc..088de67 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,11 @@ s1 = wa.open_pyramid("out.zarr", mode="r+").add_level("s1", factors=2) # place and only two answers to "it exists": raise, or adopt an identical array after checking; whole-unit writes from many processes with no coordination; `add_level` as metadata plus an empty array; chunk keys for every format. Nothing is ever deleted or overwritten. +- **Plans outputs and blocks** (0.3): a store-less `Grid` with the same geometry as an + array, so an output is planned, cut into blocks and only then created; `blocks` tiling an + array or a box into write boxes with context-grown read boxes, in C order, as plain + picklable records a cluster job takes by id; the anisotropy-aware pyramid schedule. + Nothing is scheduled or run. ## The placement rule @@ -143,6 +148,14 @@ wa.ome_multiscales(axes, units, levels, version="0.5", type=None) ; wa.level_pla wa.level_shape(shape, factors, origin=0) ; wa.chunk_key(index, format, ...) ; wa.parse_chunk_key(key, format, ...) wa.ensure_group_entry(group_path, name, voxel_size, translation, axes, units="nanometer") # list an array you made yourself wa.units.ome_name("nm") -> "nanometer" ; wa.write.codecs_of(path) -> (codecs, shards) + +# grids and blocks (0.3) +g = wa.Grid(shape, voxel_size, translation, axes=None, units="nanometer", chunks=None, shards=None) # or Grid.from_meta(meta, voxel_size=) +wa.create(path, grid=g, dtype="uint64") +for b in wa.blocks(shape, block_shape, context=0, box=None, clip=True): # also vol.blocks() and g.blocks() + b.begin, b.end, b.read_begin, b.read_end, b.inner, b.at_end, b.id # plain records; blocks[id] is O(1) +wa.level_factors(voxel_size, ratio=1.5, max_levels=None, stop_shape=None, shape=None) +wa.units.abbreviation("nanometer") ; wa.store.join(path, "s0") ; wa.store.parent(path) ; wa.store.basename(path) ``` `WorldArray`: `ts`, `meta`, `grid`, `shape`, `dtype`, `axes`, `voxel_size`, `translation`, diff --git a/docs/spec.md b/docs/spec.md index bdea8fe..d1bdf12 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -765,24 +765,38 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3. 17. **Blocks tile, context grows, nothing runs.** `blocks(shape, block_shape, context, box, clip)` yields, in C order, blocks whose write boxes `[begin, end)` partition `[0, shape)`, or `box` (a voxel box, the grid anchored at its begin) when given, with no - overlap and no gap, the last block on an axis clipped. A block's read box is its write - box grown by `context` on each side, where `context` is one int, one per axis, or a - `(before, after)` pair per axis, clipped to the array when `clip=True` (the default; - cellmap-analyze's `create_block` edge rule) and left unclipped when `clip=False` for a - padded `read_index`. `len(blocks)`, `blocks[b.id]`, `blocks.index_of(b.begin)` and + overlap and no gap, the last block on an axis clipped to the box. `box` may reach past + the array on any side (a ROI rounded outward to a factor multiple, mesh-n-bone's halo + mode): the write boxes tile it exactly, outside the array too. A block's read box is + its write box grown by `context` on each side, where `context` is one int, one per + axis, or a `(before, after)` pair per axis, clipped to the **array** (never to `box`, + never a write box) when `clip=True` (the default; cellmap-analyze's `create_block` edge + rule) and left unclipped when `clip=False` for a padded `read_index` (its + `neg_with_edge_pos` mode, which reads past the array for fill). A block also carries + `at_start` and `at_end` per axis, whether it is the first or last block on that axis, + so a caller adds context at the array's far faces itself, and `inner`, the slices of + the write box inside the read box, correct also when `clip=False` lets `read_begin` + go below 0. `len(blocks)`, `blocks[b.id]` (O(1)), `blocks.index_of(b.begin)` and `b.index` agree. `vol.blocks()` defaults `block_shape` to `write_shape`, so default blocks are aligned writes. Nothing is scheduled: the iterable and its blocks are plain, - picklable records. `write_chunk(index, data)` writes the unit at `index` whole; - `data.shape` must equal that unit's clipped shape. + picklable records, and the iterable pickles without its blocks, so a worker takes an + id. One grid per call: a consumer that writes several levels from one block derives + the per-level boxes itself (mesh-n-bone's direct mode, two lines). `write_chunk(index, + data)` writes the unit at `index` whole; `data.shape` must equal that unit's clipped + shape. 18. **A `Grid` is a `WorldArray` without a store.** `Grid(shape, voxel_size, translation, axes, units, chunks, shards)` answers `box`, `corner`, `index_box`, `chunks_covering`, `write_shape` and `blocks` exactly as a `WorldArray` of those facts does, and - `g.meta(dtype)` equals `wa.create(path, grid=g, dtype=dtype).meta`, so a pipeline plans + `g.meta(dtype)` agrees with `wa.create(path, grid=g, dtype=dtype).meta` in every + field a grid has: shape, dtype, chunk and shard shape, axes, units, voxel size and + translation (a grid has no path, format, documents or fill value), so a pipeline plans its output before anything exists. `Grid.from_meta(meta, voxel_size=t)` is Behaviour 4's target grid over that array: anchored at its corner with `ceil(n * s / t)` voxels, - so a `read(voxel_size=t)` of the source fills it exactly; this is the output grid - cellmap-analyze's commands compute today in the integer frame. + so a `read(voxel_size=t)` of the source fills it exactly, and its `box` is the target + grid's extent, which may reach past the source's own `box` by less than one target + voxel; this is the output grid cellmap-analyze's commands compute today in the integer + frame. 19. **Schedules and helpers.** `level_factors(voxel_size, ratio, max_levels, stop_shape, shape)` returns per level the per-axis factors, 2 on every axis whose current voxel size @@ -1017,11 +1031,12 @@ over funlib geometry, by this spec's own assignment. attributes? If so, should the package offer `funlib_attrs(voxel_size, translation, axes, units)`, writing `offset` as the corner, funlib's meaning, for `create(attrs=)`, or do the repos keep writing their own? -6. **The block iterator's shape** (Behaviour 17) is written from cellmap-analyze's - `create_block` (directional context, the edge rule, a fourth axis) and mesh-n-bone's - super-chunk grids (factor-aligned strides per level) without either session confirming - it. Confirm with both before 0.3 is built. mesh-n-bone's direct mode, where one task - writes several levels at once, may need per-level grids that `blocks` does not give. +6. **The block iterator's shape** (Behaviour 17). Resolved 2026-10-22: both sessions + confirmed it. mesh-n-bone asked that `box` may reach past the array and that `clip` + clips read boxes only, and derives its per-level boxes from one grid itself; + cellmap-analyze asked for `at_start`/`at_end` per axis and `inner`, and confirmed that + `(before, after)` pairs cover its context modes and that `clip=False` with a padded + `read_index` is its edge-reading mode. All are in Behaviour 17. 7. **Sidecar attributes.** Resolved 2026-10-22: registration-station keeps `fs` and `fix_origin_um` on the group through `group_attrs` and the part attributes (`y0`, `y1`, `y0_core`, `y1_core`, `feather_vox`) on the array through `attrs`, all known at creation; @@ -1056,8 +1071,8 @@ over funlib geometry, by this spec's own assignment. 13. **World boxes on blocks.** cellmap-flow reads its input at one voxel size by a world box grown by a context in nanometres and writes at another; `blocks` (0.3) gives index boxes on one grid. `Grid.blocks` could carry world boxes too (`b.box`, `b.read_box` as - `Box`, a context in world units). Two lines on cellmap-flow's side today; a candidate - for 0.4. + `Box`, a context in world units). Two lines on cellmap-flow's side today. Not planned + (2026-10-22): daisy does cellmap-flow's tiling and stays; reopen if cellmap-flow drops it. ## 8. Constraints for whoever builds on this diff --git a/pyproject.toml b/pyproject.toml index e3c1575..5ddc66f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "worldarray" -version = "0.2.0" +version = "0.3.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 7288e27..dada170 100644 --- a/src/worldarray/__init__.py +++ b/src/worldarray/__init__.py @@ -2,16 +2,19 @@ Open zarr v2, zarr v3, N5 and Neuroglancer precomputed arrays and pyramids by path, know where each sits (voxel size, the centre of voxel 0, units, axes), and read voxel boxes or -world boxes as numpy arrays, padded where the box leaves the data. ``docs/spec.md`` is the -contract; the README is its first five sections. +world boxes as numpy arrays, padded where the box leaves the data; create placed OME-NGFF +levels and write them from any number of jobs; plan an output grid and cut it into blocks +before anything exists. ``docs/spec.md`` is the contract; the README is its first five +sections. """ from __future__ import annotations -from . import store, units, write +from . import grid, store, units, write from .array import Pyramid, WorldArray, create, metadata, open, open_pyramid, open_tensorstore from .context import capped_context, capped_context_spec, shared_context from .geometry import Box, index_range +from .grid import Block, Grid, blocks, level_factors from .levels import select_level from .placement import ArrayMeta, corner_to_centre, default_axes from .retry import Retry @@ -25,19 +28,24 @@ parse_chunk_key, ) -__version__ = "0.2.0" +__version__ = "0.3.0" __all__ = [ "ArrayMeta", + "Block", "Box", + "Grid", "Node", + "blocks", "chunk_key", "corner_to_centre", "create", "default_axes", "detect", "ensure_group_entry", + "grid", "index_range", + "level_factors", "level_placement", "level_shape", "list_chunk_keys", diff --git a/src/worldarray/array.py b/src/worldarray/array.py index e5e3987..8828432 100644 --- a/src/worldarray/array.py +++ b/src/worldarray/array.py @@ -551,6 +551,34 @@ def write(self, box, data: np.ndarray, partial: bool = False) -> None: ) self.write_index(begin, end, data, partial=partial) + def blocks(self, block_shape=None, context=0, box=None, clip: bool = True): + """The array cut into blocks (Behaviour 17): write boxes tiling the stored grid (or a + voxel ``box``) with ``block_shape`` (default ``write_shape``, so the default blocks are + aligned writes), read boxes grown by ``context`` and clipped to the array when + ``clip``. Plain records, nothing scheduled.""" + from .grid import blocks as _blocks + + return _blocks( + self._grid.shape, + self.write_shape if block_shape is None else block_shape, + context=context, + box=box, + clip=clip, + ) + + def write_chunk(self, index: Sequence[int], data: np.ndarray) -> None: + """Write the whole write unit at ``index`` on the storage grid (the shard of a sharded + array); ``data`` has that unit's shape, clipped at the array's end.""" + unit, shape = self.write_shape, self._grid.shape + index = tuple(int(i) for i in index) + if len(index) != len(shape): + raise ValueError(f"{self.path}: index needs one entry per axis of {shape}") + begin = tuple(i * u for i, u in zip(index, unit)) + if any(i < 0 or b >= n for i, b, n in zip(index, begin, shape)): + raise ValueError(f"{self.path}: no write unit at {index}") + end = tuple(min(b + u, n) for b, u, n in zip(begin, unit, shape)) + self.write_index(begin, end, data) + def read_index( self, begin: Sequence[int], @@ -1051,11 +1079,12 @@ def __len__(self) -> int: def create( path: str, - shape: Sequence[int], - dtype, - voxel_size, - translation, + shape: Sequence[int] | None = None, + dtype=None, + voxel_size=None, + translation=None, *, + grid=None, axes: Sequence[str] | None = None, units: Units = "nanometer", chunks: Sequence[int] | None = None, @@ -1079,6 +1108,19 @@ def create( raises ``FileExistsError``; with ``"adopt"`` an identical array (shape, dtype, chunks, shards, placement) is opened for writing and a different one raises ``ValueError`` naming every difference. Returns a writable ``WorldArray``.""" + if grid is not None: + # a Grid planned the output (Behaviour 18): its facts are the array's, chunks and + # shards overridable; the placement arguments must not also be given + if any(v is not None for v in (shape, voxel_size, translation, axes)): + raise ValueError("pass grid= or shape, voxel_size, translation and axes, not both") + shape, voxel_size, translation = grid.shape, grid.voxel_size, grid.translation + axes, units = grid.axes, grid.units + chunks = grid.chunks if chunks is None else chunks + shards = grid.shards if shards is None else shards + if shape is None or dtype is None or voxel_size is None or translation is None: + raise ValueError( + "create needs shape, dtype, voxel_size and translation, or grid= and dtype" + ) if chunks is None: raise ValueError("chunks is required: the package never picks a chunk shape") ndim = len(shape) diff --git a/src/worldarray/grid.py b/src/worldarray/grid.py new file mode 100644 index 0000000..89c14ec --- /dev/null +++ b/src/worldarray/grid.py @@ -0,0 +1,492 @@ +"""Grids and blocks: an output planned, and cut into jobs, before anything is stored. + +``blocks(shape, block_shape, context, box, clip)`` tiles ``[0, shape)``, or a voxel +``box`` with the block grid anchored at its begin, into write boxes in C order with no +overlap and no gap, the last block on an axis clipped. Each block's read box is its write +box grown by ``context``. ``Grid`` is a ``WorldArray`` without a store: the same ``box``, +``corner``, ``index_box``, ``chunks_covering``, ``write_shape`` and ``blocks`` from the +facts alone. ``level_factors`` is the pyramid schedule mesh-n-bone and cellmap-analyze +share. Nothing here schedules or runs anything (spec, Behaviours 17 to 19). + +Choices the spec leaves open, made here: + +- ``box`` may reach past the array on any side (a ROI rounded outward to a factor + multiple). The write boxes tile ``box`` exactly, outside the array too; ``clip=True`` + clips the **read** boxes to the array, ``[0, shape)``, never to ``box`` and never a write + box. A write box wholly outside the array then has an empty read box at the array's + nearest face. ``clip=False`` leaves read boxes unclipped, negative begins included, for a + padded ``read_index``. +- ``context`` is one int, one per axis, or a ``(before, after)`` pair per axis; the per-axis + entries may mix ints and pairs. Negative context raises. +- ``Blocks`` is a plain frozen dataclass of the five arguments, normalised: ``box`` to a + pair of int tuples (``[0, shape)`` when none is given), ``context`` to one + ``(before, after)`` pair per axis. Blocks are computed from their id on demand, so + ``blocks[id]`` is O(1) and the record pickles to a worker without the blocks. Ids run from + 0 to ``len - 1``; any other id is an ``IndexError``. ``index_of(begin)`` takes a begin on + the box-anchored grid. +- ``Block.at_start`` and ``Block.at_end`` say, per axis, whether the block is the first or + last on the block grid (the box's edge, or the array's). ``Block.inner`` is the write box + inside the read box, ``slice(begin - read_begin, end - read_begin)`` per axis, which stays + a correct offset into a padded read when ``clip=False`` lets ``read_begin`` go below 0. +- ``Grid`` keeps ``voxel_size`` and ``translation`` with one entry per axis, as + ``ArrayMeta`` does; given over the spatial axes only, the other axes get 1.0 and 0.0 (the + short-list rule of Convention 4). ``units`` is kept per axis too: a string is the spatial + axes' unit and every other axis gets ``""``; a dict by kind gives ``"space"`` to the + spatial axes and ``"time"`` to an axis named ``t``, ``""`` elsewhere; a sequence is per + axis. ``shards`` need ``chunks`` and must be a whole number of them on every axis. +- ``Grid.chunks_covering(begin, end, write=False)`` takes ``write`` as ``WorldArray``'s + does in 0.2: ``write=True`` runs on ``write_shape``. A grid without that shape raises. +- ``Grid.meta(dtype)`` has ``format`` None, ``path`` ``""``, ``parent`` None, empty + ``attrs`` and ``group_attrs``, ``chunk_shape`` the chunks (the shape when none), + ``axes_named`` True (a created array names its axes), ``placed_by`` ``"ome"`` and + ``fill_value`` None: a grid has no store, no documents and no fill. +- ``Grid.from_meta(meta)`` keeps the array's chunks (and its ``shard_shape`` when the meta + carries one); the target grid at ``voxel_size`` keeps neither, since its chunks are the + new array's to choose. ``voxel_size`` is in the meta's units. +- ``level_factors`` compares voxel sizes with ``rtol 1e-6``, as Behaviour 3 does. The stop + rule is checked before each level, so a ``shape`` already within ``stop_shape`` gives no + levels and the last level returned is the first that fits. ``ratio`` below 1 raises, + since no axis would ever qualify. + +The ``WorldArray`` side of 0.3 (``vol.blocks``, ``vol.write_chunk`` and +``wa.create(path, grid=g)``) is wired in with the write module of 0.2; nothing here opens a +store. +""" + +from __future__ import annotations + +import itertools +import math +import operator +from collections.abc import Iterable, Iterator, Sequence +from dataclasses import dataclass + +from .geometry import Box, index_range +from .placement import ArrayMeta, default_axes, spatial_axes +from .resample import target_shape + +RTOL = 1e-6 + + +def _whole(values, what: str) -> tuple[int, ...]: + values = tuple(values) + out = tuple(int(v) for v in values) + if any(o != v for o, v in zip(out, values)): + raise ValueError(f"{what} must be whole numbers, got {values}") + return out + + +def _sized(values, what: str, ndim: int, minimum: int) -> tuple[int, ...]: + out = _whole(values, what) + if len(out) != ndim: + raise ValueError(f"{what} has {len(out)} entries; expected {ndim}") + if any(v < minimum for v in out): + raise ValueError(f"{what} entries must be at least {minimum}, got {out}") + return out + + +def _context(context, ndim: int) -> tuple[tuple[int, int], ...]: + """One ``(before, after)`` pair per axis from an int, one int per axis, or a pair per axis.""" + if not isinstance(context, Iterable): + entries = [context] * ndim + else: + entries = list(context) + if len(entries) != ndim: + raise ValueError(f"context has {len(entries)} entries; the blocks have {ndim} axes") + pairs = [] + for e in entries: + pair = _whole((e, e) if not isinstance(e, Iterable) else e, "context") + if len(pair) != 2: + raise ValueError(f"a context entry is an int or a (before, after) pair, got {e!r}") + if min(pair) < 0: + raise ValueError(f"context cannot be negative, got {e!r}") + pairs.append(pair) + return tuple(pairs) + + +# --------------------------------------------------------------------------- +# Blocks + + +@dataclass(frozen=True) +class Block: + """One block: a plain, picklable record. ``begin``/``end`` is the write box, + ``read_begin``/``read_end`` the write box grown by the context (clipped to the array + when the blocks clip).""" + + index: tuple[int, ...] # position on the block grid + id: int # linear C-order id + begin: tuple[int, ...] + end: tuple[int, ...] + read_begin: tuple[int, ...] + read_end: tuple[int, ...] + at_start: tuple[bool, ...] # per axis: the first block on the grid + at_end: tuple[bool, ...] # per axis: the last block on the grid + + @property + def inner(self) -> tuple[slice, ...]: + """The write box inside the read box: ``data[block.inner]`` of a read of + ``[read_begin, read_end)`` is the block's own voxels.""" + return tuple( + slice(b - rb, e - rb) for b, e, rb in zip(self.begin, self.end, self.read_begin) + ) + + +@dataclass(frozen=True) +class Blocks: + """The blocks tiling ``[0, shape)``, or ``box``, in C order. A frozen record of its + arguments that computes each block on demand: ``len``, iteration, ``blocks[id]`` (O(1)) + and ``index_of(begin)``.""" + + shape: tuple[int, ...] + block_shape: tuple[int, ...] + context: tuple[tuple[int, int], ...] = 0 # normalised to (before, after) per axis + box: tuple[tuple[int, ...], tuple[int, ...]] | None = None # normalised; [0, shape) if None + clip: bool = True + + def __post_init__(self) -> None: + shape = _whole(self.shape, "shape") + if any(s < 0 for s in shape): + raise ValueError(f"shape cannot be negative, got {shape}") + n = len(shape) + block_shape = _sized(self.block_shape, "block_shape", n, 1) + if self.box is None: + box = ((0,) * n, shape) + else: + b = Box.coerce(self.box) + if b.ndim != n: + raise ValueError(f"box has {b.ndim} axes; the blocks have {n}") + box = (_whole(b.begin, "box"), _whole(b.end, "box")) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "block_shape", block_shape) + object.__setattr__(self, "context", _context(self.context, n)) + object.__setattr__(self, "box", box) + object.__setattr__(self, "clip", bool(self.clip)) + + @property + def _counts(self) -> tuple[int, ...]: + """Blocks per axis.""" + return tuple(-(-max(e - b, 0) // s) for b, e, s in zip(*self.box, self.block_shape)) + + @property + def count(self) -> int: + return math.prod(self._counts) + + def __len__(self) -> int: + return self.count + + def __iter__(self) -> Iterator[Block]: + counts = self._counts + for i, index in enumerate(itertools.product(*(range(c) for c in counts))): + yield self._block(index, i, counts) + + def __getitem__(self, block_id: int) -> Block: + i = operator.index(block_id) + counts = self._counts + total = math.prod(counts) + if not 0 <= i < total: + raise IndexError(f"block id {i} is outside [0, {total})") + index, rest = [], i + for c in reversed(counts): + rest, k = divmod(rest, c) + index.append(k) + return self._block(tuple(reversed(index)), i, counts) + + def index_of(self, begin: Sequence[int]) -> tuple[int, ...]: + """The block index of the block whose write box starts at ``begin``.""" + begin = tuple(begin) + if len(begin) != len(self.shape): + raise ValueError(f"begin has {len(begin)} entries; the blocks have {len(self.shape)}") + index = [] + for b, b0, s, c in zip(begin, self.box[0], self.block_shape, self._counts): + k, r = divmod(b - b0, s) + if r != 0 or not 0 <= k < c: + raise ValueError( + f"{begin} is not a block's first voxel: the blocks start at {self.box[0]} " + f"in steps of {self.block_shape}, {self._counts} per axis" + ) + index.append(int(k)) + return tuple(index) + + def _block(self, index: tuple[int, ...], block_id: int, counts: tuple[int, ...]) -> Block: + origin, stop = self.box + begin = tuple(o + k * s for o, k, s in zip(origin, index, self.block_shape)) + end = tuple(min(b + s, e) for b, s, e in zip(begin, self.block_shape, stop)) + read_begin = tuple(b - c[0] for b, c in zip(begin, self.context)) + read_end = tuple(e + c[1] for e, c in zip(end, self.context)) + if self.clip: + read_begin = tuple(min(max(r, 0), n) for r, n in zip(read_begin, self.shape)) + read_end = tuple( + max(min(r, n), lo) for r, n, lo in zip(read_end, self.shape, read_begin) + ) + return Block( + index=tuple(index), + id=block_id, + begin=begin, + end=end, + read_begin=read_begin, + read_end=read_end, + at_start=tuple(k == 0 for k in index), + at_end=tuple(k == c - 1 for k, c in zip(index, counts)), + ) + + +def blocks(shape, block_shape, context=0, box=None, clip: bool = True) -> Blocks: + """The blocks tiling ``[0, shape)``, or the voxel ``box`` (anchored at its begin), by + ``block_shape``, each read box grown by ``context`` and clipped to the array when + ``clip`` (Behaviour 17).""" + return Blocks(shape, block_shape, context, box, clip) + + +# --------------------------------------------------------------------------- +# Grids + + +def _per_axis(values, what: str, ndim: int, spatial: tuple[int, ...], other: float): + values = tuple(float(v) for v in values) + if len(values) == ndim: + return values + if len(values) == len(spatial): + out = [other] * ndim + for k, axis in enumerate(spatial): + out[axis] = values[k] + return tuple(out) + raise ValueError( + f"{what} has {len(values)} entries; expected {len(spatial)} (spatial) or {ndim}" + ) + + +def _axis_units(units, axes: tuple[str, ...], spatial: tuple[int, ...]) -> tuple[str, ...]: + sp = set(spatial) + if isinstance(units, str): + return tuple(units if i in sp else "" for i in range(len(axes))) + if isinstance(units, dict): + unknown = set(units) - {"space", "time"} + if unknown: + raise ValueError(f"units by kind take 'space' and 'time', got {sorted(unknown)}") + return tuple( + units.get("space", "") if i in sp else units.get("time", "") if a == "t" else "" + for i, a in enumerate(axes) + ) + units = tuple(str(u) for u in units) + if len(units) != len(axes): + raise ValueError(f"units has {len(units)} entries; expected {len(axes)}") + return units + + +@dataclass(frozen=True) +class Grid: + """A store-less grid: ``shape`` voxels of ``voxel_size`` with voxel 0 centred at + ``translation`` (in ``units``), optionally chunked and sharded. Answers what a + ``WorldArray`` of the same facts answers (Behaviour 18).""" + + shape: tuple[int, ...] + voxel_size: tuple[float, ...] # one entry per axis once constructed + translation: tuple[float, ...] # the CENTRE of voxel 0, one entry per axis + axes: tuple[str, ...] | None = None # default_axes(ndim) when None + units: str | dict | tuple[str, ...] = "nanometer" # one string per axis once constructed + chunks: tuple[int, ...] | None = None + shards: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + shape = _whole(self.shape, "shape") + if any(s < 0 for s in shape): + raise ValueError(f"shape cannot be negative, got {shape}") + n = len(shape) + axes = default_axes(n) if self.axes is None else tuple(str(a) for a in self.axes) + if len(axes) != n: + raise ValueError(f"axes has {len(axes)} entries; the shape has {n}") + spatial = spatial_axes(axes) + voxel_size = _per_axis(self.voxel_size, "voxel_size", n, spatial, 1.0) + if any(v <= 0 for v in voxel_size): + raise ValueError(f"voxel_size must be positive, got {voxel_size}") + translation = _per_axis(self.translation, "translation", n, spatial, 0.0) + chunks = None if self.chunks is None else _sized(self.chunks, "chunks", n, 1) + shards = None if self.shards is None else _sized(self.shards, "shards", n, 1) + if shards is not None: + if chunks is None: + raise ValueError("shards need chunks: a shard is a whole number of chunks") + if any(s % c for s, c in zip(shards, chunks)): + raise ValueError(f"shards {shards} are not a whole number of chunks {chunks}") + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "axes", axes) + object.__setattr__(self, "voxel_size", voxel_size) + object.__setattr__(self, "translation", translation) + object.__setattr__(self, "units", _axis_units(self.units, axes, spatial)) + object.__setattr__(self, "chunks", chunks) + object.__setattr__(self, "shards", shards) + + @classmethod + def from_meta(cls, meta: ArrayMeta, voxel_size: Sequence[float] | None = None) -> Grid: + """The grid of an existing array; with ``voxel_size`` (per spatial axis, in the meta's + units), Behaviour 4's target grid over it: anchored at its corner, ``ceil(n * s / t)`` + voxels per spatial axis, voxel 0 centred at ``corner + t / 2``, so a + ``read(voxel_size=t)`` of the source fills it exactly.""" + if voxel_size is None: + return cls( + meta.shape, + meta.voxel_size, + meta.translation, + meta.axes, + meta.units, + meta.chunk_shape, + getattr(meta, "shard_shape", None), + ) + target = tuple(float(v) for v in voxel_size) + spatial = meta.spatial + if len(target) != len(spatial): + raise ValueError( + f"voxel_size has {len(target)} entries; the array has {len(spatial)} spatial axes" + ) + shape, vs, tr = list(meta.shape), list(meta.voxel_size), list(meta.translation) + corner = meta.corner + for k, axis in enumerate(spatial): + shape[axis] = target_shape(meta.shape[axis], meta.voxel_size[axis], target[k]) + vs[axis] = target[k] + tr[axis] = corner[axis] + target[k] / 2 + return cls(tuple(shape), tuple(vs), tuple(tr), meta.axes, meta.units) + + @property + def ndim(self) -> int: + return len(self.shape) + + @property + def spatial(self) -> tuple[int, ...]: + """Indices of the axes named z, y or x when any are; else the last three.""" + return spatial_axes(self.axes) + + @property + def corner(self) -> tuple[float, ...]: + """``translation - voxel_size / 2`` on the spatial axes.""" + return tuple(self.translation[i] - self.voxel_size[i] / 2 for i in self.spatial) + + @property + def box(self) -> Box: + """The grid's extent on the spatial axes, ``[corner, corner + shape * voxel_size)``.""" + sp = self.spatial + return Box.from_index( + (0,) * len(sp), + tuple(self.shape[i] for i in sp), + tuple(self.voxel_size[i] for i in sp), + tuple(self.translation[i] for i in sp), + ) + + def index_box(self, box) -> tuple[tuple[int, ...], tuple[int, ...]]: + """The voxel index range a world box selects (Behaviour 2's centre rule), over the + spatial axes or every axis; the other axes come back whole.""" + box = Box.coerce(box) + spatial, ndim = self.spatial, self.ndim + if box.ndim == len(spatial): + axes = list(spatial) + elif box.ndim == ndim: + axes = list(range(ndim)) + else: + raise ValueError( + f"box has {box.ndim} axes; the grid has {ndim} ({len(spatial)} spatial)" + ) + lo, hi = index_range( + box.begin, + box.end, + [self.voxel_size[a] for a in axes], + [self.translation[a] for a in axes], + ) + begin, end = [0] * ndim, list(self.shape) + for a, b, e in zip(axes, lo, hi): + begin[a], end[a] = b, e + return tuple(begin), tuple(end) + + @property + def write_shape(self) -> tuple[int, ...] | None: + """The unit one writer may own: the shards when given, else the chunks; None when + the grid has neither.""" + return self.shards if self.shards is not None else self.chunks + + def chunks_covering( + self, begin: Sequence[int], end: Sequence[int], write: bool = False + ) -> Iterator[tuple[int, ...]]: + """The chunk indices a voxel box touches, clipped to the grid; ``write=True`` on the + write-unit grid.""" + cs = self.write_shape if write else self.chunks + if cs is None: + raise ValueError(f"this grid has no {'write_shape' if write else 'chunks'}") + first = [max(0, int(b)) // c for b, c in zip(begin, cs)] + last = [(min(int(e), n) - 1) // c for e, n, c in zip(end, self.shape, cs)] + if any(hi < lo for lo, hi in zip(first, last)): + return iter(()) + return ( + tuple(int(i + lo) for i, lo in zip(idx, first)) + for idx in itertools.product(*(range(hi - lo + 1) for lo, hi in zip(first, last))) + ) + + def blocks(self, block_shape=None, context=0, box=None, clip: bool = True) -> Blocks: + """``blocks`` over this grid's shape; ``block_shape`` defaults to ``write_shape``, so + default blocks are aligned writes.""" + if block_shape is None: + block_shape = self.write_shape + if block_shape is None: + raise ValueError("this grid has no chunks or shards; pass block_shape") + return Blocks(self.shape, block_shape, context, box, clip) + + def meta(self, dtype) -> ArrayMeta: + """The ``ArrayMeta`` an array created from this grid reports, placed by OME.""" + return ArrayMeta( + shape=self.shape, + dtype=dtype, + chunk_shape=self.chunks if self.chunks is not None else self.shape, + axes=self.axes, + axes_named=True, + units=self.units, + voxel_size=self.voxel_size, + translation=self.translation, + format=None, + path="", + placed_by="ome", + shard_shape=self.shards, + ) + + +# --------------------------------------------------------------------------- +# Schedules + + +def level_factors( + voxel_size: Sequence[float], + ratio: float = 1.5, + max_levels: int | None = None, + stop_shape: Sequence[int] | None = None, + shape: Sequence[int] | None = None, +) -> list[tuple[int, ...]]: + """Per level, the per-axis downsampling factors from the level before: 2 on every axis + whose current voxel size is within ``ratio`` of the smallest current one, 1 elsewhere, + compounding level over level, until ``max_levels`` levels or until a level's shape + (``shape`` ceil-divided by the cumulative factors) fits within ``stop_shape`` on every + axis. Voxel size (30, 8, 8) gives (1, 2, 2), (1, 2, 2), then (2, 2, 2) on: z waits until + y and x reach 32. Isotropic data gives all 2s (mesh-n-bone's schedule, Behaviour 19).""" + current = [float(v) for v in voxel_size] + n = len(current) + if n == 0 or any(v <= 0 for v in current): + raise ValueError(f"voxel_size must be positive, got {tuple(voxel_size)}") + if ratio < 1: + raise ValueError(f"ratio must be at least 1, got {ratio}") + if max_levels is None and stop_shape is None: + raise ValueError("pass max_levels, or stop_shape with shape") + if max_levels is not None and max_levels < 0: + raise ValueError(f"max_levels cannot be negative, got {max_levels}") + if stop_shape is not None: + if shape is None: + raise ValueError("stop_shape needs shape") + stop_shape = _sized(stop_shape, "stop_shape", n, 1) + shape = _sized(shape, "shape", n, 0) + cumulative = [1] * n + out: list[tuple[int, ...]] = [] + while max_levels is None or len(out) < max_levels: + if stop_shape is not None and all( + -(-s // c) <= t for s, c, t in zip(shape, cumulative, stop_shape) + ): + break + limit = ratio * min(current) * (1 + RTOL) + step = tuple(2 if v <= limit else 1 for v in current) + out.append(step) + current = [v * f for v, f in zip(current, step)] + cumulative = [c * f for c, f in zip(cumulative, step)] + return out diff --git a/src/worldarray/store.py b/src/worldarray/store.py index 3a4cc6a..ea0acf9 100644 --- a/src/worldarray/store.py +++ b/src/worldarray/store.py @@ -77,10 +77,15 @@ def is_root(path: str | None) -> bool: def basename(path: str) -> str: + """The last component of a path or URL (``s3://b/a/s0`` gives ``s0``); a trailing slash + is dropped first.""" return normalize_path(path).rsplit("/", 1)[-1] def join(base: str, *parts: str) -> str: + """``base`` and ``parts`` joined by single slashes, keeping a URL's scheme + (``join("s3://b/a/", "s0")`` gives ``s3://b/a/s0``). Public, with ``parent`` and + ``basename``; like them it normalises ``base`` first (a local path made absolute).""" return "/".join([normalize_path(base).rstrip("/"), *(p.strip("/") for p in parts)]) diff --git a/src/worldarray/units.py b/src/worldarray/units.py index d85ac02..14944ae 100644 --- a/src/worldarray/units.py +++ b/src/worldarray/units.py @@ -77,6 +77,27 @@ def factor( return a[1] / b[1] +# The OME short form of each unit, by its full name; every spelling of a unit reaches its +# entry through the table above, so "nanometre", "nm" and "Nanometers" agree. +_SHORT = { + "picometer": "pm", "angstrom": "Å", "nanometer": "nm", "micrometer": "µm", + "millimeter": "mm", "centimeter": "cm", "meter": "m", + "nanosecond": "ns", "microsecond": "µs", "millisecond": "ms", "second": "s", + "minute": "min", "hour": "h", "day": "d", +} # fmt: skip +_SHORT_BY_SIZE = {_lookup(name): short for name, short in _SHORT.items()} + + +def abbreviation(unit: str) -> str: + """The OME short form of a unit: ``"nm"`` for any spelling of nanometre the table knows, + ``"µm"`` (U+00B5) for any spelling of micrometre, ``"Å"``, ``"s"``, ``"min"`` and so + on. A unit the table does not know (``""``, ``"pixel"``) is a ``ValueError``.""" + found = _lookup(unit) if unit else None + if found is None or found not in _SHORT_BY_SIZE: + raise ValueError(f"unknown unit {unit!r}") + return _SHORT_BY_SIZE[found] + + # OME-NGFF's name for each unit the table knows, by (kind, size): every spelling of a unit # reaches its entry through the table above, so "nanometre", "nm" and "Nanometers" agree. # The American spelling is OME's. diff --git a/tests/test_grid.py b/tests/test_grid.py new file mode 100644 index 0000000..e145361 --- /dev/null +++ b/tests/test_grid.py @@ -0,0 +1,385 @@ +"""Grids and blocks: tiling, context, store-less grids, the factor schedule, the helpers +(spec, Behaviours 17 to 19).""" + +from __future__ import annotations + +import pickle + +import numpy as np +import pytest + +import worldarray as wa +from conftest import ome_multiscales, pattern, write_zarr2, write_zarr2_group + + +def _cover(blocks, shape, origin=None) -> np.ndarray: + """How many write boxes hold each voxel of ``shape`` voxels starting at ``origin``.""" + origin = origin or (0,) * len(shape) + count = np.zeros(shape, dtype=int) + for b in blocks: + count[tuple(slice(s - o, e - o) for s, e, o in zip(b.begin, b.end, origin))] += 1 + return count + + +@pytest.fixture +def placed(tmp_path): + """One OME zarr v2 level with voxels (3, 4, 5), translation (1.5, 10, -2.5), chunks (4, 5, 6).""" + data = pattern((10, 12, 14)) + group = tmp_path / "placed.zarr" + write_zarr2_group(group, ome_multiscales("zyx", [("s0", (3, 4, 5), (1.5, 10, -2.5))])) + write_zarr2(group / "s0", data, chunks=(4, 5, 6)) + return str(group / "s0"), data + + +# --- blocks ------------------------------------------------------------------- + + +def test_blocks_partition_the_array_in_c_order(): + bl = wa.blocks((10, 7, 5), (4, 3, 5)) + assert len(bl) == bl.count == 9 + assert np.all(_cover(bl, (10, 7, 5)) == 1) # no overlap, no gap + assert [b.index for b in bl] == list(np.ndindex(3, 3, 1)) + assert [b.id for b in bl] == list(range(9)) + assert (bl[8].begin, bl[8].end) == ((8, 6, 0), (10, 7, 5)) # the last block clipped + assert all((b.read_begin, b.read_end) == (b.begin, b.end) for b in bl) + + +def test_blocks_partition_a_box_anchored_at_its_begin(): + bl = wa.blocks((10, 7, 5), (4, 3, 5), box=((2, 3, 1), (9, 7, 5))) + expected = np.zeros((10, 7, 5), dtype=int) + expected[2:9, 3:7, 1:5] = 1 + assert np.array_equal(_cover(bl, (10, 7, 5)), expected) + assert [(b.begin, b.end) for b in bl] == [ + ((2, 3, 1), (6, 6, 5)), + ((2, 6, 1), (6, 7, 5)), + ((6, 3, 1), (9, 6, 5)), + ((6, 6, 1), (9, 7, 5)), + ] + + +def test_a_box_past_the_array_tiles_exactly_and_only_reads_are_clipped(): + box = ((0, 0, 0), (40, 40, 40)) + bl = wa.blocks((32, 32, 32), (16, 16, 16), context=2, box=box) + assert len(bl) == 27 + assert np.all( + _cover(bl, (40, 40, 40)) == 1 + ) # the write boxes are the box, outside the array too + for b in bl: + assert all(0 <= lo <= hi <= 32 for lo, hi in zip(b.read_begin, b.read_end)) + last = bl[26] + assert (last.begin, last.end) == ((32, 32, 32), (40, 40, 40)) + assert (last.read_begin, last.read_end) == ((30, 30, 30), (32, 32, 32)) + unclipped = wa.blocks((32, 32, 32), (16, 16, 16), context=2, box=box, clip=False)[26] + assert (unclipped.read_begin, unclipped.read_end) == ((30, 30, 30), (42, 42, 42)) + + below = wa.blocks((32, 32, 32), (16, 16, 16), context=2, box=((-8, 0, 0), (24, 32, 32))) + assert np.all(_cover(below, (32, 32, 32), origin=(-8, 0, 0)) == 1) + assert (below[0].begin, below[0].end) == ((-8, 0, 0), (8, 16, 16)) + assert (below[0].read_begin, below[0].read_end) == ((0, 0, 0), (10, 18, 18)) + assert below.index_of((-8, 0, 0)) == (0, 0, 0) and below.index_of((8, 16, 16)) == (1, 1, 1) + with pytest.raises(ValueError): + below.index_of((0, 0, 0)) # not a start on the box-anchored grid + + far = wa.blocks((32,), (16,), context=2, box=((0,), (80,)))[4] # wholly outside the array + assert (far.begin, far.end, far.read_begin, far.read_end) == ((64,), (80,), (32,), (32,)) + + +def test_context_is_an_int_one_per_axis_or_before_after_pairs(): + def grown(context): + b = wa.blocks((12, 12, 12), (4, 4, 4), context=context, clip=False)[13] # index (1, 1, 1) + assert (b.begin, b.end) == ((4, 4, 4), (8, 8, 8)) + return b.read_begin, b.read_end + + assert grown(1) == ((3, 3, 3), (9, 9, 9)) + assert grown((0, 1, 2)) == ((4, 3, 2), (8, 9, 10)) + assert grown(((1, 0), (0, 2), (3, 3))) == ((3, 4, 1), (8, 10, 11)) + assert grown((1, (0, 2), 0)) == ((3, 4, 4), (9, 10, 8)) + for bad in (-1, (1, 2), ((1, 2, 3), 0, 0)): + with pytest.raises(ValueError): + wa.blocks((12, 12, 12), (4, 4, 4), context=bad) + + +def test_clip_clips_read_boxes_to_the_array_and_never_write_boxes(): + on = wa.blocks((10,), (4,), context=3) + off = wa.blocks((10,), (4,), context=3, clip=False) + assert [(b.read_begin, b.read_end) for b in on] == [((0,), (7,)), ((1,), (10,)), ((5,), (10,))] + assert [(b.read_begin, b.read_end) for b in off] == [ + ((-3,), (7,)), + ((1,), (11,)), + ((5,), (13,)), + ] + assert [(b.begin, b.end) for b in on] == [(b.begin, b.end) for b in off] + assert [(b.begin, b.end) for b in on] == [((0,), (4,)), ((4,), (8,)), ((8,), (10,))] + + +def test_edges_and_inner_on_a_middle_and_an_edge_block(): + clipped = wa.blocks((12, 12), (4, 4), context=2) + padded = wa.blocks((12, 12), (4, 4), context=2, clip=False) + middle = clipped[4] + assert middle.index == (1, 1) + assert middle.at_start == (False, False) and middle.at_end == (False, False) + assert (middle.read_begin, middle.read_end) == ((2, 2), (10, 10)) + assert middle.inner == (slice(2, 6), slice(2, 6)) + assert padded[4] == middle + + edge = clipped[2] + assert edge.index == (0, 2) # first along the first axis, last along the second + assert edge.at_start == (True, False) and edge.at_end == (False, True) + assert (edge.read_begin, edge.read_end) == ((0, 6), (6, 12)) + assert edge.inner == (slice(0, 4), slice(2, 6)) + edge = padded[2] + assert (edge.read_begin, edge.read_end) == ((-2, 6), (6, 14)) + assert edge.inner == (slice(2, 6), slice(2, 6)) + + +def test_inner_selects_the_write_box_from_a_read_clipped_or_padded(cube): + path, data = cube + vol = wa.open(path) + for clip in (True, False): + for b in wa.blocks(vol.shape, (4, 4, 4), context=((1, 2), 2, (0, 3)), clip=clip): + read = vol.read_index(b.read_begin, b.read_end, padding=0) + own = data[tuple(slice(s, e) for s, e in zip(b.begin, b.end))] + assert np.array_equal(read[b.inner], own) + + +def test_len_ids_and_index_of_agree(): + bl = wa.blocks((9, 10, 11), (2, 3, 4), context=1, box=((1, 0, 2), (9, 10, 11))) + seen = list(bl) + assert len(seen) == len(bl) == bl.count == 4 * 4 * 3 + for b in seen: + assert bl[b.id] == b + assert bl.index_of(b.begin) == b.index + for bad in (len(bl), -1): + with pytest.raises(IndexError): + bl[bad] + for begin in ((2, 0, 2), (1, 0, 14)): # off the grid; past the last block + with pytest.raises(ValueError): + bl.index_of(begin) + + +def test_blocks_and_a_block_pickle_as_plain_records(): + bl = wa.blocks( + (100, 100, 100), + (16, 16, 16), + context=((1, 2), 3, 0), + box=((-4, 0, 0), (100, 100, 100)), + clip=False, + ) + shipped = pickle.loads(pickle.dumps(bl)) + assert shipped == bl and len(shipped) == len(bl) == 343 + assert shipped[200] == bl[200] + assert pickle.loads(pickle.dumps(bl[200])) == bl[200] + assert len(pickle.dumps(bl)) < 1000 # the arguments, not the blocks + + +# --- grids -------------------------------------------------------------------- + + +def test_a_grid_answers_as_a_worldarray_of_the_same_facts(placed): + path, _ = placed + vol = wa.open(path) + g = wa.Grid((10, 12, 14), (3, 4, 5), (1.5, 10, -2.5), chunks=(4, 5, 6)) + assert g.box == vol.box + assert g.corner == vol.corner + for box in ( + vol.box, + ((0, 8, -5), (9, 30, 40)), + ((1.5, 10, -2.5), (7.5, 18, 7.5)), # centre-aligned + ((-10, -10, -10), (100, 100, 100)), + ((3, 12, 0), (3.1, 13, 1)), + ): + assert g.index_box(box) == vol.index_box(box) + for begin, end in ( + ((0, 0, 0), (10, 12, 14)), + ((3, 4, 5), (9, 11, 13)), + ((-5, -5, -5), (2, 2, 2)), + ((9, 11, 13), (40, 40, 40)), + ((5, 5, 5), (5, 6, 7)), + ): + assert list(g.chunks_covering(begin, end)) == list(vol.chunks_covering(begin, end)) + meta = g.meta(vol.dtype) + for name in ("shape", "dtype", "chunk_shape", "axes", "axes_named", "units", "voxel_size"): + assert getattr(meta, name) == getattr(vol.grid, name) + assert meta.translation == vol.grid.translation and meta.placed_by == "ome" + assert wa.Grid.from_meta(vol.grid) == g + + +def test_a_grid_with_a_channel_axis(tmp_path): + group = tmp_path / "czyx.zarr" + write_zarr2_group(group, ome_multiscales("czyx", [("s0", (1, 4, 4, 4), (0, 2, 2, 2))])) + write_zarr2(group / "s0", pattern((2, 6, 8, 10)), chunks=(1, 3, 4, 5)) + vol = wa.open(str(group / "s0")) + g = wa.Grid((2, 6, 8, 10), (4, 4, 4), (2, 2, 2), axes="czyx", chunks=(1, 3, 4, 5)) + assert g.spatial == vol.spatial == (1, 2, 3) + assert g.box == vol.box and g.corner == vol.corner + for box in (((4, 4, 4), (20, 20, 20)), ((1, 4, 4, 4), (2, 20, 20, 20))): + assert g.index_box(box) == vol.index_box(box) + assert g.index_box(((4, 4, 4), (20, 20, 20)))[0][0] == 0 # the channel axis whole + meta = g.meta(vol.dtype) + for name in ("shape", "chunk_shape", "axes", "units", "voxel_size", "translation"): + assert getattr(meta, name) == getattr(vol.grid, name) + assert wa.Grid.from_meta(vol.grid) == g + + +def test_write_shape_and_default_blocks(): + plain = wa.Grid((64, 64, 64), (8, 8, 8), (4, 4, 4)) + assert plain.write_shape is None + with pytest.raises(ValueError): + plain.blocks() + with pytest.raises(ValueError): + plain.chunks_covering((0, 0, 0), (1, 1, 1)) + assert plain.blocks((32, 32, 32)) == wa.blocks((64, 64, 64), (32, 32, 32)) + chunked = wa.Grid((64, 64, 64), (8, 8, 8), (4, 4, 4), chunks=(16, 16, 16)) + assert chunked.write_shape == (16, 16, 16) + assert chunked.blocks() == wa.blocks((64, 64, 64), (16, 16, 16)) + sharded = wa.Grid((64, 64, 64), (8, 8, 8), (4, 4, 4), chunks=(16, 16, 16), shards=(32, 32, 32)) + assert sharded.write_shape == (32, 32, 32) + assert sharded.blocks(context=2) == wa.blocks((64, 64, 64), (32, 32, 32), context=2) + assert list(sharded.chunks_covering((0, 0, 0), (40, 1, 1), write=True)) == [ + (0, 0, 0), + (1, 0, 0), + ] + assert len(list(sharded.chunks_covering((0, 0, 0), (40, 1, 1)))) == 3 + with pytest.raises(ValueError): + wa.Grid((64, 64, 64), (8, 8, 8), (4, 4, 4), chunks=(16, 16, 16), shards=(24, 24, 24)) + with pytest.raises(ValueError): + wa.Grid((64, 64, 64), (8, 8, 8), (4, 4, 4), shards=(32, 32, 32)) + + +@pytest.mark.parametrize("target", [(8, 8, 8), (6, 4, 2.5), (10.48, 7, 5)]) +def test_from_meta_is_behaviour_4s_target_grid(placed, target): + path, _ = placed + vol = wa.open(path) + resampled = wa.open(path, voxel_size=target, on_mismatch="resample") + g = wa.Grid.from_meta(vol.grid, voxel_size=target) + assert g.shape == resampled.shape + assert g.voxel_size == tuple(float(t) for t in target) + assert g.corner == pytest.approx(resampled.corner) + assert [g.translation[i] for i in g.spatial] == pytest.approx(resampled.translation) + assert g.chunks is None + assert vol.read(g.box, voxel_size=target).shape == g.shape # the read fills it exactly + + +# --- schedules and helpers ---------------------------------------------------- + + +def test_level_factors_follow_mesh_n_bones_anisotropic_table(): + factors = wa.level_factors((30, 8, 8), ratio=1.5, max_levels=5) + assert factors == [(1, 2, 2), (1, 2, 2), (2, 2, 2), (2, 2, 2), (2, 2, 2)] + sizes, vs = [], np.array([30.0, 8.0, 8.0]) + for f in factors: + vs = vs * f + sizes.append(tuple(vs)) + assert sizes[:3] == [(30, 16, 16), (30, 32, 32), (60, 64, 64)] # z waits for y and x at 32 + assert wa.level_factors((8, 8, 30), max_levels=3) == [(2, 2, 1), (2, 2, 1), (2, 2, 2)] + assert wa.level_factors((8, 8, 20), max_levels=3) == [(2, 2, 1), (2, 2, 2), (2, 2, 2)] + for iso in ((4, 4, 4), (10.48, 10.48, 10.48)): + assert wa.level_factors(iso, max_levels=4) == [(2, 2, 2)] * 4 + + +def test_level_factors_stop_when_a_level_fits(): + stop = (64, 64, 64) + assert wa.level_factors((8, 8, 8), stop_shape=stop, shape=(100, 1000, 1000)) == [(2, 2, 2)] * 4 + assert wa.level_factors((30, 8, 8), stop_shape=stop, shape=(64, 256, 256)) == [(1, 2, 2)] * 2 + assert ( + wa.level_factors((8, 8, 8), max_levels=2, stop_shape=stop, shape=(100, 1000, 1000)) + == [(2, 2, 2)] * 2 + ) + assert wa.level_factors((8, 8, 8), stop_shape=stop, shape=(64, 64, 64)) == [] + for kwargs in ({}, {"stop_shape": stop}, {"max_levels": 2, "ratio": 0.5}): + with pytest.raises(ValueError): + wa.level_factors((8, 8, 8), **kwargs) + + +SHORT_FORMS = { + "pm": ("picometer", "picometre", "pm"), + "Å": ("angstrom", "ångström", "å", "Å"), + "nm": ("nanometer", "nanometre", "nm"), + "µm": ("micrometer", "micrometre", "micron", "um", "µm", "μm"), + "mm": ("millimeter", "millimetre", "mm"), + "cm": ("centimeter", "centimetre", "cm"), + "m": ("meter", "metre", "m"), + "ns": ("nanosecond", "ns"), + "µs": ("microsecond", "us", "µs", "μs"), + "ms": ("millisecond", "ms"), + "s": ("second", "sec", "s"), + "min": ("minute", "min"), + "h": ("hour", "h"), + "d": ("day", "d"), +} + + +def test_abbreviation_for_every_spelling_in_the_table(): + spellings = {s for group in SHORT_FORMS.values() for s in group} + assert spellings == set(wa.units.TABLE["length"]) | set(wa.units.TABLE["time"]) + for short, group in SHORT_FORMS.items(): + for spelling in group: + assert wa.units.abbreviation(spelling) == short + assert wa.units.abbreviation("micrometer") == "µm" # the micro sign, not Greek mu + assert wa.units.abbreviation("Nanometers") == "nm" + for unknown in ("", "pixel", "degrees_north"): + with pytest.raises(ValueError): + wa.units.abbreviation(unknown) + + +def test_path_helpers_are_url_aware(tmp_path): + s = wa.store + assert s.parent("s3://b/a/b") == s.parent("s3://b/a/b/") == "s3://b/a" + assert s.basename("s3://b/a/b") == s.basename("s3://b/a/b/") == "b" + assert s.join("s3://b/a", "b") == s.join("s3://b/a/", "/b/") == "s3://b/a/b" + assert s.join("gs://bucket/vol.zarr", "seg", "s0") == "gs://bucket/vol.zarr/seg/s0" + assert s.parent("https://host/data/vol.zarr/s0") == "https://host/data/vol.zarr" + assert s.parent("precomputed://gs://b/vol/s0") == "gs://b/vol" + assert s.parent("s3://b") is None # a bucket root has no parent + local = s.join(str(tmp_path), "vol.zarr") + assert local == f"{tmp_path}/vol.zarr" + assert s.parent(local) == str(tmp_path) and s.basename(local) == "vol.zarr" + + +# --- the WorldArray side of 0.3 (wired with the 0.2 write module) ------------------------ + + +def test_array_blocks_default_to_write_units_and_write_chunk_fills_them(tmp_path): + vol = wa.create( + str(tmp_path / "s.zarr/s0"), + (20, 16, 16), + "uint8", + (1, 1, 1), + (0.5,) * 3, + chunks=(4, 4, 4), + shards=(8, 8, 8), + ) + bl = vol.blocks() + assert bl.count == 3 * 2 * 2 and bl[0].end == (8, 8, 8) and bl[len(bl) - 1].end == (20, 16, 16) + vol.write_chunk((1, 1, 1), np.full((8, 8, 8), 2, np.uint8)) # a middle unit + vol.write_chunk((2, 0, 0), np.full((4, 8, 8), 3, np.uint8)) # the clipped last unit on z + assert sorted(vol.list_chunks()) == [(1, 1, 1), (2, 0, 0)] + assert int(vol.read_index((16, 0, 0), (20, 8, 8)).sum()) == 3 * 4 * 64 + with pytest.raises(ValueError, match="data has shape"): + vol.write_chunk((2, 0, 0), np.zeros((8, 8, 8), np.uint8)) + with pytest.raises(ValueError, match="no write unit"): + vol.write_chunk((3, 0, 0), np.zeros((4, 8, 8), np.uint8)) + grown = vol.blocks(context=2, clip=False)[0] + assert grown.read_begin == (-2, -2, -2) and grown.inner == (slice(2, 10),) * 3 + + +def test_create_from_a_grid_reports_the_grid(tmp_path): + g = wa.Grid((16, 16, 16), (8, 8, 8), (4, 4, 4), chunks=(4, 4, 4), shards=(8, 8, 8)) + vol = wa.create(str(tmp_path / "g.zarr/s0"), grid=g, dtype="uint16") + planned, made = g.meta("uint16"), vol.meta + for field in ( + "shape", + "dtype", + "chunk_shape", + "shard_shape", + "axes", + "units", + "voxel_size", + "translation", + ): + assert getattr(planned, field) == getattr(made, field), field + assert vol.box == g.box and list(vol.blocks()) == list(g.blocks()) + with pytest.raises(ValueError, match="not both"): + wa.create(str(tmp_path / "h.zarr/s0"), (16, 16, 16), grid=g, dtype="uint16") + with pytest.raises(ValueError, match="needs shape"): + wa.create(str(tmp_path / "i.zarr/s0"), dtype="uint16", chunks=(4, 4, 4)) diff --git a/uv.lock b/uv.lock index c73739b..9b6fbf0 100644 --- a/uv.lock +++ b/uv.lock @@ -1247,7 +1247,7 @@ wheels = [ [[package]] name = "worldarray" -version = "0.2.0" +version = "0.3.0" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },