diff --git a/CHANGELOG.md b/CHANGELOG.md index a0d0fd6..46efac2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,14 +1,50 @@ # Changelog -## Unreleased +## 0.2.0 (2026-10-22) -- Specification only, no code change: `docs/spec.md` draft 3 adds the contract for 0.2 - (metadata and placed creation: the OME 0.4/0.5 document builder, the level rule, chunk - keys, `create`, opening for write, aligned writes, `add_level`; Behaviours 10 to 16) and - 0.3 (grids and blocks: `Grid`, `blocks`, `write_chunk`, `level_factors`, public path - helpers; Behaviours 17 to 19), written from the five writer sessions' answers and a - census of the I/O code the five repos still carry. The package stays read-only until - 0.2 ships. +Metadata and placed creation, built to spec draft 3 (Behaviours 10 to 16). Nothing on the +read side changes. + +- `ome_multiscales(axes, units, levels, version)` builds the complete OME-NGFF 0.4 or 0.5 + group document with one explicit placement per level; `level_placement` and + `level_shape` derive a level from factors and an origin (Convention 1's rule inverted, + so a level built from a ROI origin is placed from it; `level_shape` absorbs float-noise + factors such as `0.3 / 0.1`). Units are written in OME's names (`"nm"` becomes + `"nanometer"`; `units.ome_name`), an unknown unit as given; the multiscale's `type` is + written when given. Expected documents are pinned in `tests/fixtures/write.json`, shared + with chunkmirage's TypeScript builder, each case with its resolved per-axis units and + types, fifteen cases. +- `create(path, shape, dtype, voxel_size, translation, chunks=..., ...)` writes the group + document (made when absent, else its one dataset entry added or replaced, every other + attribute kept) and then the array, as zarr v2 with OME 0.4 or zarr v3 with OME 0.5 + following the group's format. `chunks` is required, codecs are explicit with a zstd + level 3 default, shards are a zarr v3 option. `exists` is `"error"` or `"adopt"` (an + identical array is opened for writing, a different one raises naming every difference, + a creation lost to another process is adopted after a short backoff). Nothing is ever + deleted, truncated or overwritten. Write units equal to the fill value are not stored + unless `store_fill=True`, and a stored unit written so loses its key, as zarr-python's + `write_empty_chunks=False` does. `create` also writes a bare group document for every + missing ancestor down from the `.zarr` root and follows the nearest existing ancestor's + zarr version, so a level made in a zarr v2 container is zarr v2 with OME 0.4 and the + hierarchy browses (cellmap-analyze, cellmap-flow). Adopting an identical array writes no + group document; after a write the document is read again and a lost entry is written + again, then raises (registration-station's concurrent block jobs). +- `ensure_group_entry` is public: list an array the package did not create, such as a + symlinked `s0` (mesh-n-bone), or a level filled before it is listed. +- `open(path, mode="r+")` opens an existing array for writing; a `WorldArray` pickles by + its arguments and its context's spec and reopens in the receiving process, a capped + context staying capped (mesh-n-bone's workers). `write.codecs_of` takes a path. +- `write_index(begin, end, data)` writes whole write units only (the shard when sharded, + else the chunk; the clipped last unit allowed), so writers in different processes need + no coordination; `partial=True` does tensorstore's read-modify-write. `write(box, data)` + maps a world box by the centre rule first. No cast, no padding, no resampling on write. +- `open_pyramid(path, mode="r+").add_level(name, factors=, origin=)` creates a level's + entry and empty array, placed and sized from the level below with its dtype, chunks, + shards, codecs and fill value copied; the downsampling itself stays with the caller. +- `chunk_key` and `parse_chunk_key` for zarr v2, zarr v3 (both encodings, sharded arrays + keyed by shard), N5 and precomputed, as free functions and on `WorldArray`; + `list_chunks` uses the same parser. `write_shape`, `ArrayMeta.shard_shape`, + `chunks_covering(write=True)`. ## 0.1.1 (2026-10-22) - funlib `voxel_size`/`offset` on the parent group place an array that carries none, for diff --git a/README.md b/README.md index 4bedf7f..c3bc0dc 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,10 @@ resampled when no level matches, fast enough for a model server or a thousand cl workers, with retries for flaky object stores and thread caps for shared hosts. No one package did all of that, so every tool wrote its own reader, and the readers drifted apart. worldarray is that reader, written once, with the rules stated and tested. -Its dependencies are `tensorstore` and `numpy`. It is read-only in this version. The +Its dependencies are `tensorstore` and `numpy`. Since 0.2 it also writes the one thing +every tool wrote for itself and got wrong in its own way: the OME-NGFF metadata of a placed +array or pyramid level, and the array itself, created once and filled by aligned writes +from any number of jobs. It never deletes, overwrites or downsamples. The [specification](https://janeliascicomp.github.io/worldarray/spec/) is the contract it is built to; every rule below is a test. @@ -56,6 +59,13 @@ vol.read_index((0, 0, 0), (64, 64, 64), padding=0) 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 + +# writing: a placed level, its OME-NGFF group document written for it (zarr v3 + OME 0.5 here) +out = wa.create("out.zarr/s0", shape=(512, 512, 512), dtype="uint64", + voxel_size=(8, 8, 8), translation=(4, 4, 4), chunks=(64, 64, 64)) +out.write_index((0, 0, 0), (64, 64, 64), block) # whole chunks only: many jobs, no locks +out.write(box, block) # or by world box +s1 = wa.open_pyramid("out.zarr", mode="r+").add_level("s1", factors=2) # placed and sized; you fill it ``` ## What it does @@ -80,6 +90,13 @@ pyr.select((16, 16, 16), mode="floor") # the level to read at 16 nm cancelled future on timeout. - **Carries the geometry**: `Box` with `intersect`, `union`, `contains`, `grow`, `shift`, `snap` to a grid, and the voxel indices a world box selects. +- **Writes placed arrays** (0.2): the OME-NGFF 0.4 or 0.5 document from voxel size, centre + translation, axes and units, with each level's placement explicit or derived from + factors and an origin; `create` following the container's format and writing the group + documents an output needs down from its `.zarr` root, with explicit chunks and codecs, + 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. ## The placement rule @@ -93,7 +110,7 @@ moved half a voxel on reading: | OME-NGFF `translation` | centre | as is | | COSEM `transform.translate` | centre | as is | | `offset` beside `resolution` / `voxel_size` (`offset` missing: 0) | corner | `offset + voxel_size / 2` | -| N5 `pixelResolution` or `resolution` with `downsamplingFactors f`, no offset | BigDataViewer | `(f - 1) / 2 * base` | +| N5 `pixelResolution` or `resolution` with `downsamplingFactors f`, no offset | corner at the origin at every level, as Neuroglancer draws N5 | `f / 2 * base` | | precomputed `voxel_offset` | corner | `(voxel_offset + 1/2) * resolution` | | nothing | a plain array | voxels of 1, corner at the origin | @@ -117,13 +134,24 @@ wa.open_tensorstore(path, level=None, context=None) wa.shared_context(cache_bytes=0) ; wa.capped_context(cache_bytes=0, concurrency=1, network=True) ; wa.Retry(...) wa.list_chunk_keys(path, separator=".") # chunk indices from keys alone, no metadata needed wa.Box ; wa.units.TABLE + +# writing (0.2) +wa.create(path, shape, dtype, voxel_size, translation, *, axes=None, units="nanometer", chunks=..., shards=None, + fill_value=0, codecs="zstd", format=None, attrs=None, group_attrs=None, exists="error", store_fill=False) -> WorldArray +wa.open(path, mode="r+") ; wa.open_pyramid(path, mode="r+").add_level(name, factors=, origin=0, ...) +wa.ome_multiscales(axes, units, levels, version="0.5", type=None) ; wa.level_placement(voxel_size, translation, factors, origin=0) +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) ``` `WorldArray`: `ts`, `meta`, `grid`, `shape`, `dtype`, `axes`, `voxel_size`, `translation`, -`corner`, `box`, `levels`, `level_name`, `level_path`, `source_voxel_size`, `resampled`; +`corner`, `box`, `levels`, `level_name`, `level_path`, `source_voxel_size`, `resampled`, +`writable`, `write_shape`; `read(box, voxel_size=, padding=, method=, mask=, anchor=, transform=, post_transform=, select=, strict=)`, `read_index(begin, end, padding=, mask=, select=, transform=)`, `index_box(box)`, -`chunks_covering(begin, end)`, `list_chunks()`, `reopen(context=)`. +`chunks_covering(begin, end, write=)`, `list_chunks()`, `chunk_key(index)`, `parse_chunk_key(key)`, +`write_index(begin, end, data, partial=)`, `write(box, data, partial=)`, `reopen(context=)`. A box is a `(begin, end)` pair of floats in your units, half-open per axis, over the spatial axes or every axis; any object with `.begin`/`.end` or `.offset`/`.shape` also diff --git a/docs/spec.md b/docs/spec.md index 25a088b..bdea8fe 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -242,20 +242,30 @@ read side, so a read-only consumer sees no change. ```python doc = wa.ome_multiscales( axes, # axis names in the array's C order: ("z","y","x"), ("c","z","y","x"), ("t","c","z","y","x") - units, # per axis ("nanometer"; "" on a channel axis) or by kind {"space": "nanometer", "time": "second"} + units, # per axis ("nanometer"; "" on a channel axis) or by kind {"space": "nanometer", "time": "second"}; + # written in OME's names ("nm" -> "nanometer", wa.units.ome_name), an unknown unit as given levels, # [("s0", voxel_size, translation), ...]: one explicit placement per level over the spatial axes # (or every axis). Never derived here: levels in the wild carry their own translations and # non-whole factors (chunkmirage's pyramids) version="0.5", # "0.5" (zarr v3: attributes["ome"]) | "0.4" (zarr v2: attributes["multiscales"]) axis_types=None, # per axis "space" | "time" | "channel"; default from the name (z y x space, t time, c channel) - name="", metadata=None, + name="", metadata=None, type=None, # the multiscale's own fields; `type` (the downsampling method) written when given ) -> dict # the complete group attributes to write; `wa.metadata` reads them back exactly (Behaviour 10) wa.level_placement(voxel_size, translation, factors, origin=0) -> (voxel_size, translation) # the level whose voxel j covers source voxels origin + j f .. origin + (j + 1) f - 1 on each axis: # (voxel_size * f, translation + voxel_size * (origin + (f - 1) / 2)). Convention 1's BigDataViewer # rule inverted, and the centre Behaviour 4's block mean produces -wa.level_shape(shape, factors, origin=0) -> shape # ceil((n - origin) / f) per axis +wa.level_shape(shape, factors, origin=0) -> shape # ceil((n - origin) / f) per axis, the quotient snapped to a whole + # number within 1e-6 first (0.3 / 0.1 is a factor of 3) +wa.units.ome_name(unit) -> str # OME's name for any spelling the unit table knows ("nanometer" for "nm"); unknown raises +wa.ensure_group_entry(group_path, name, voxel_size, translation, axes, units="nanometer", axis_types=None, + group_attrs=None, format=None) # list an array the package did not create (a symlinked s0, a level + # filled before it is listed): the group document's entry for `name` added or replaced, the + # 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) vol = wa.create( path, # ".../vol.zarr/s0": the last component is the level, its parent the multiscale group, made if absent @@ -267,18 +277,24 @@ vol = wa.create( fill_value=0, codecs="zstd", # "zstd" (level 3, which is what zarr-python's default `zstd(level=0)` means) | "none" | the format's # own JSON: a zarr v2 `compressor` dict or a zarr v3 `codecs` list - format=None, # "zarr3" | "zarr2"; default the group's format when a group document exists, else "zarr3". + format=None, # "zarr3" | "zarr2"; default the nearest existing ancestor group's format, else "zarr3"; missing + # ancestors up to the ".zarr" root get bare group documents (Behaviour 11). # OME 0.5 goes with v3 and 0.4 with v2, always attrs=None, # extra array attributes, merged into the array's own group_attrs=None, # extra group attributes, merged; no one else's key is dropped or rewritten (Behaviour 12) exists="error", # "error": anything already at `path` raises | "adopt": an identical array (shape, dtype, chunks, # shards, placement within rtol 1e-6) is opened for writing; a different one raises naming each - # difference. There is no third value + # difference. There is no third value. An adopt writes no group document when the entry + # already says the same (Behaviour 12) store_fill=False, # True stores write units equal to fill_value instead of leaving them absent + separator=None, # zarr v2 dimension_separator, zarr v3 chunk-key separator; the format's default when None + axis_types=None, # as ome_multiscales context=None, retries=None, ) -> WorldArray # writable; every read-side property and method as before -vol = wa.open(path, mode="r+", ...) # an existing array for writing; vol.writable; the metadata rules of mode "r" unchanged +vol = wa.open(path, mode="r+", store_fill=False, ...) # an existing array for writing; vol.writable; the metadata + # rules of mode "r" unchanged; a handle pickles by its arguments + # and its context's spec (Behaviour 11) vol.write_index(begin, end, data, partial=False) # whole write units only (Behaviour 13) unless partial vol.write(box, data, partial=False) # a world box by Behaviour 2, then write_index; no padding, no resampling vol.write_shape # the unit one writer may own: the shard shape when sharded, else the chunk shape @@ -625,7 +641,16 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3. for OME 0.4 and 0.5, for `(z, y, x)`, `(c, z, y, x)` and `(t, c, z, y, x)`, and for levels with non-whole factors and their own translations (chunkmirage's case). The documents equal the expected documents in `tests/fixtures/write.json` after JSON - normalisation, so chunkmirage's TypeScript builder can be held to the same output. + normalisation, so chunkmirage's TypeScript builder can be held to the same output; each + case also carries its resolved per-axis units and types, and the call with those gives + the same document, so the TypeScript side need not reimplement the broadcasting. Units + are written in OME's names: any spelling the unit table knows (`"nm"`, `"µm"`, + `"nanometre"`) becomes `"nanometer"`, `"micrometer"` and so on (`wa.units.ome_name`), + and a unit the table does not know is written as given, Convention 3's welcome kept. + The multiscale's `name`, `metadata` and `type` are written when given. `level_shape` + snaps the quotient `(n - origin) / f` to a whole number within `1e-6` before the + ceiling, so a factor that is a ratio of float voxel sizes (`0.3 / 0.1 = + 2.9999999999999996`, chunkmirage's case) gives 3 voxels for 9, not 4. `level_placement` and `level_shape` invert Convention 1's BigDataViewer rule and agree with Behaviour 4: the level a block mean by factor `f` from source voxel `origin` produces has voxel 0 centred at `translation + voxel_size * (origin + (f - 1) / 2)`, @@ -638,10 +663,16 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3. this order, the group document (created when absent; when present, read, the dataset entry for this level added or replaced with the datasets kept in order of voxel size, everything else kept as it was) and then the array with its chunks, shards, codecs, - fill value, dimension names and `attrs`. The format follows the group: a `.zgroup` - means zarr v2 and OME 0.4, a `zarr.json` means zarr v3 and OME 0.5, `format` decides - only when no group document exists, and a `format` that contradicts an existing group - raises `ValueError` rather than mixing versions. `chunks` is required (alignment is the + fill value, dimension names and `attrs`. The format follows the container: walking up + from the level's parent, the nearest existing node decides, a `.zgroup` group meaning + zarr v2 and OME 0.4, a `zarr.json` group zarr v3 and OME 0.5, an array, an N5 or a + precomputed node a `ValueError`; a `format` that contradicts it raises rather than + mixing versions. When no ancestor exists, the walk ends at the nearest ancestor named + `*.zarr`, the store root, `format` or zarr v3 decides, and every missing group from + that root down to the level's parent gets a document, bare above the parent, so + zarr-python browses the hierarchy and a level made in a zarr v2 container is zarr v2 + (cellmap-analyze's `prepare_ds`, cellmap-flow's `box_n` groups); with no `.zarr` + ancestor either, only the parent gets a document. `chunks` is required (alignment is the whole safety argument, so the package never picks a chunk shape); `shards` that are not a whole number of chunks on every axis raise. `codecs` defaults to zstd level 3 in the format's own dialect, what zarr-python's default `zstd(level=0)` means, so @@ -657,21 +688,35 @@ Behaviours 10 to 16 are 0.2; 17 to 19 are 0.3. absent, is re-read with backoff, 5 attempts at `0.5 s * attempt`, and then treated as an adopt (registration-station's `shared_array`, written because metadata appears late on NRS). The package deletes, truncates and follows no symlink, ever; a caller that wants a - fresh array removes the old one with its own code first. Write units whose values all + fresh array removes the old one with its own code first. Write units whose values all equal `fill_value` are not stored unless `store_fill=True`, so chunk presence stays the - signal cellmap-flow's trainer and manifest read. `create` returns a writable - `WorldArray`; `wa.open(path, mode="r+")` opens an existing one with every rule of - Conventions 3 to 5 unchanged, and a handle pickles by its arguments, so each worker - process opens its own cheaply (mesh-n-bone's per-process cache). + signal cellmap-flow's trainer and manifest read; a stored unit written so loses its + key, as zarr-python's `write_empty_chunks=False` does, the one key a write may remove + being the one it replaces. `create` returns a writable `WorldArray`; + `wa.open(path, mode="r+")` opens an existing one with every rule of Conventions 3 to 5 + unchanged, and a handle pickles by its arguments and its context's spec, so each worker + process opens its own cheaply, on the shared context when the handle had one, else on + a context rebuilt once per process from the same spec, so a capped context stays + capped (mesh-n-bone's per-process cache and thread caps). 12. **Group attributes belong to everyone.** The package rewrites only the `multiscales` list (0.4) or `ome.multiscales` (0.5), and within it only the dataset entry whose `path` is the level being created or added; other multiscales, other datasets and every other key of the group and of the array (`fs`, a registry record, `y0`/`y1`, anything) come - back unchanged. The group's attributes before and after a `create` differ in that one - entry. Rewriting a group document is a read-modify-write with no lock, so `create` and - `add_level` on one group are called from one process at a time, the coordinator, never - from block jobs; the docstrings say so. + back unchanged. The group's attributes before and after a `create` differ in that one + entry, and not at all when the entry already says the same: adopting an identical array + (placement within rtol 1e-6) reads the document, compares and writes nothing, so + registration-station's twenty block jobs calling `create(exists="adopt")` on one output + write the document in the first creation only. Rewriting a group document is a + read-modify-write with no lock, so `create` and `add_level` of different levels in one + group are called from one process at a time, the coordinator, never from block jobs; + the docstrings say so. After a write the document is read again and, when this level's + entry is missing because another process's write landed in between, merged and written + again, up to 5 times, then `RuntimeError`: a lost update becomes a repair or an error + where the store shows it, never a misplaced level. `ensure_group_entry` is this listing + step on its own, public, for an array the package did not create (mesh-n-bone's + symlinked `s0`) or one listed only after it is filled; it follows the same rules, + ancestors included. 13. **Aligned writes need no coordination.** `write_index(begin, end, data)` requires `data.shape == end - begin`, `data.dtype == vol.dtype` (no silent cast; `ValueError`), @@ -831,7 +876,14 @@ sizing from factors and origin, inserting in order, replacing a same-named entry dtype, chunks and codecs, writing no data; `chunk_key` and `parse_chunk_key` round trips for zarr v2 with `.` and `/`, zarr v3 `default` with both separators and `v2`, N5, precomputed with a `voxel_offset` and a clipped last chunk, and sharded zarr v3, with -`list_chunk_keys` agreeing; `write_shape` on sharded and unsharded arrays. +`list_chunk_keys` agreeing; `write_shape` on sharded and unsharded arrays; units written in +OME's names and an unknown unit as given, `type` when given, `level_shape` on a float-noise +factor; missing ancestor documents down from the `.zarr` root and none above it, the nearest +ancestor's format followed and a contradiction refused, an array in the way refused; an +adopt writing no document and a changed placement writing one; a lost entry written again +and the error after the attempts; `ensure_group_entry` on a foreign array, idempotent; a +pickled handle keeping a capped context's spec and sharing the shared one; `codecs_of` by +path; an all-fill write removing its own key. Grids (0.3): `blocks` partitioning with and without a box, `context` as an int, per axis and as before/after pairs, clipping on and off, `len`, `blocks[id]` and `index_of`; @@ -970,21 +1022,42 @@ over funlib geometry, by this spec's own assignment. 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. -7. **Sidecar attributes.** registration-station's `fs` and `fix_origin_um` duplicate the OME - scale and translation, and its `y0`, `y1`, `y0_core`, `y1_core`, `feather_vox` part - attributes encode placement. With `ArrayMeta` round-tripping, are they still wanted, or - does Behaviour 12's merge rule simply carry them? +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; + merge-on-create carries them and no attribute-update call is wanted. 8. **Creation on NRS.** The adopt-with-backoff numbers in Behaviour 11 (5 attempts, 0.5 s times the attempt) are registration-station's and are tested here only against a fake store. Validate from two LSF jobs on the cluster before anyone relies on them. -9. **The codec default.** zstd level 3 reproduces three writer repos' outputs; cellmap-flow's - blockwise outputs are blosc lz4 through funlib and would change compression unless its - subclass passes `codecs`. Decide whether it pins blosc or accepts zstd. +9. **The codec default.** Resolved 2026-10-22: cellmap-flow pins blosc explicitly (lz4 + level 5 for blockwise outputs, zstd level 3 for annotation volumes) as zarr v2 compressor + dicts; cellmap-analyze, mesh-n-bone and registration-station take the zstd default, which + cellmap-analyze checked byte for byte against zarr-python's. No shortcut name is wanted. 10. **Sharding.** Nobody writes shards today. registration-station wants them because a 2.4 GB chunk makes Neuroglancer decompress a whole band for one slice; mesh-n-bone would take them as an opt-in for v3 sources; cellmap-analyze would size blocks to shards. 0.2 - offers `shards=` at creation and `write_shape` for planners and nothing more; whether - that is enough is learned from the first repo that uses it. + offers `shards=` at creation and `write_shape` for planners and nothing more; whether + that is enough is learned from the first repo that uses it. registration-station checked + that a shard larger than the array is accepted and that band writes on shard edges, + the clipped last band included, go through. +11. **Listing after filling.** cellmap-analyze's #90 creates a level's array, fills it, and + lists it last, so a crashed job never leaves a half-filled level listed. `create` and + `add_level` list first. `ensure_group_entry` is the listing step, but a level created + unlisted has its placement nowhere until it is listed, so a `create(listed=False)` + handle, and its pickles, would carry the placement from its arguments alone. Decide + whether that is worth having, or whether an unlisted level should describe itself on + the array (open question 5's funlib attributes, or an OME-shaped array attribute). +12. **Per-open object-store credentials.** cellmap-flow's MinIO writers open with a + per-session key, secret and endpoint while the same process reads public S3 + anonymously. tensorstore 0.1.85's `aws_credentials` takes no inline key and secret, and + its endpoint and profile settings are process-wide through the environment, so those + three writers stay on zarr-python and s3fs. Revisit when tensorstore takes static + credentials per open. +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. ## 8. Constraints for whoever builds on this diff --git a/pyproject.toml b/pyproject.toml index 044797a..e3c1575 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "worldarray" -version = "0.1.1" +version = "0.2.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 d66f1b0..7288e27 100644 --- a/src/worldarray/__init__.py +++ b/src/worldarray/__init__.py @@ -8,26 +8,41 @@ from __future__ import annotations -from . import units -from .array import Pyramid, WorldArray, metadata, open, open_pyramid, open_tensorstore +from . import 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 .levels import select_level from .placement import ArrayMeta, corner_to_centre, default_axes from .retry import Retry from .store import Node, detect, list_chunk_keys +from .write import ( + chunk_key, + ensure_group_entry, + level_placement, + level_shape, + ome_multiscales, + parse_chunk_key, +) -__version__ = "0.1.1" +__version__ = "0.2.0" __all__ = [ "ArrayMeta", "Box", "Node", + "chunk_key", "corner_to_centre", + "create", "default_axes", "detect", + "ensure_group_entry", "index_range", + "level_placement", + "level_shape", "list_chunk_keys", + "ome_multiscales", + "parse_chunk_key", "Pyramid", "Retry", "WorldArray", @@ -39,5 +54,7 @@ "open_tensorstore", "select_level", "shared_context", + "store", "units", + "write", ] diff --git a/src/worldarray/array.py b/src/worldarray/array.py index 888159b..e5e3987 100644 --- a/src/worldarray/array.py +++ b/src/worldarray/array.py @@ -13,7 +13,6 @@ from __future__ import annotations import logging -import re import warnings from collections.abc import Callable, Iterator, Sequence from dataclasses import replace @@ -24,12 +23,20 @@ from . import resample as _resample from . import units as U -from .context import shared_context +from . import write as _write +from .context import from_recipe, recipe_of, shared_context from .geometry import Box, index_range, is_aligned from .levels import LevelRef, Mode, discover, discover_group, same_voxel_size, select_level -from .placement import ArrayMeta, light_levels, read_meta, read_meta_light +from .placement import ( + ArrayMeta, + _dims_from_docs, + default_axes, + light_levels, + read_meta, + read_meta_light, +) from .retry import Retry, read_with_retries -from .store import Node, detect, is_remote, open_array, precomputed_scale_index +from .store import Node, detect, is_remote, join, open_array, precomputed_scale_index logger = logging.getLogger(__name__) @@ -73,12 +80,18 @@ def __init__( anchor: Sequence[float] | None = None, transform: Transform | None = None, strict_units: bool = True, + mode: str = "r", + store_fill: bool = False, _ref: LevelRef | None = None, _node: Node | None = None, ): if on_mismatch not in ("error", "resample", "relabel", "keep"): raise ValueError("on_mismatch is 'error', 'resample', 'relabel' or 'keep'") + if mode not in ("r", "r+"): + raise ValueError("mode is 'r' (read) or 'r+' (read and write)") self.path = path + self._mode = mode + self._store_fill = bool(store_fill) self.padding = padding self.method = method self.units = units @@ -96,6 +109,8 @@ def __init__( base_units=base_units, anchor=anchor, strict_units=strict_units, + mode=mode, + store_fill=store_fill, ) self._anchor = tuple(float(a) for a in anchor) if anchor is not None else None @@ -109,7 +124,13 @@ def __init__( node if (ref.scale_index is not None or ref.path == node.path) else detect(ref.path) ) self._node = node_for_level - self._store = open_array(node_for_level, context=self.context, scale_index=ref.scale_index) + self._store = open_array( + node_for_level, + context=self.context, + scale_index=ref.scale_index, + write=mode == "r+", + store_fill=store_fill, + ) self._meta = read_meta( node_for_level, self._store, scale_index=ref.scale_index, conflicts=metadata_conflicts ) @@ -354,6 +375,17 @@ def level_path(self) -> str: def remote(self) -> bool: return is_remote(self._node.path) + @property + def writable(self) -> bool: + """Opened with ``mode="r+"`` (or by ``create``): ``write_index`` and ``write`` work.""" + return self._mode == "r+" + + @property + def write_shape(self) -> tuple[int, ...]: + """The unit one writer may own: the shard shape when sharded, else the chunk shape + (Behaviour 16).""" + return self._meta.write_shape + def __repr__(self) -> str: return ( f"WorldArray({self.path!r}, level={self.level_name!r}, shape={self.shape}, dtype={self.dtype}, " @@ -376,13 +408,30 @@ def reopen(self, *, context: ts.Context | None = None, **overrides) -> WorldArra self.path, context=context if context is not None else self.context, **opts ) + def __reduce__(self): + """A handle pickles by its arguments and its context's recipe, and reopens in the + receiving process (Behaviour 11): a worker opens its own store, cheaply, on a shared + context when the handle had one, else on a context rebuilt from the same spec, so a + capped context stays capped.""" + opts = { + **self._opts, + "padding": self.padding, + "method": self.method, + "select": self._select, + "units": self.units, + "retries": self.retries, + "transform": self.transform, + } + return (_reopen, (self.path, opts, recipe_of(self.context))) + # --- index space ------------------------------------------------------------ def chunks_covering( - self, begin: Sequence[int], end: Sequence[int] + self, begin: Sequence[int], end: Sequence[int], write: bool = False ) -> Iterator[tuple[int, ...]]: - """The chunk indices a voxel box touches, clipped to the array.""" - cs = self._meta.chunk_shape + """The chunk indices a voxel box touches, clipped to the array; ``write=True`` on the + write-unit grid (the shard grid of a sharded array).""" + cs = self.write_shape if write else self._meta.chunk_shape 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._grid.shape, cs)] if any(hi < lo for lo, hi in zip(first, last)): @@ -390,49 +439,117 @@ def chunks_covering( for idx in np.ndindex(*[hi - lo + 1 for lo, hi in zip(first, last)]): yield tuple(int(i + lo) for i, lo in zip(idx, first)) - 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: zarr v2 and v3, N5, and unsharded precomputed. For a sharded zarr v3 - array the storage grid is the shard grid. Works on any kvstore, URLs included.""" + def _key_params(self) -> dict: + """The facts ``chunk_key`` and ``parse_chunk_key`` need, from the array's documents: + the storage grid of the whole array (a sharded array's shards; a selected axis is + still a key component), the separator and encoding, a precomputed scale's origin.""" node = self._node - 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()] + shape, _dtype, chunks, _fill, shard = _dims_from_docs( + node, self._ref.scale_index if self._ref.scale_index is not None else None + ) + params: dict[str, Any] = { + "format": node.format, + "chunk_shape": shard or chunks, + "shape": shape, + } if node.format == "zarr2": - sep = node.docs.get(".zarray", {}).get("dimension_separator", ".") - pattern = re.compile(r"^\d+(?:" + re.escape(sep) + r"\d+)*$") - found = [tuple(int(p) for p in k.split(sep)) for k in keys if pattern.match(k)] + params["separator"] = node.docs.get(".zarray", {}).get("dimension_separator", ".") elif node.format == "zarr3": enc = node.docs.get("zarr.json", {}).get("chunk_key_encoding", {}) or {} - sep = (enc.get("configuration") or {}).get("separator", "/") - prefix = "c" if enc.get("name", "default") == "default" else "" - found = [] - for k in keys: - parts = k.split(sep) - if prefix and (not parts or parts[0] != prefix): - continue - parts = parts[1:] if prefix else parts - if parts and all(p.isdigit() for p in parts): - found.append(tuple(int(p) for p in parts)) - elif node.format == "n5": - found = [ - tuple(int(p) for p in k.split("/"))[::-1] - for k in keys - if k and all(p.isdigit() for p in k.split("/")) - ] - else: # precomputed, unsharded: "x0-x1_y0-y1_z0-z1" + params["encoding"] = enc.get("name", "default") + params["separator"] = (enc.get("configuration") or {}).get("separator") + elif node.format == "precomputed": scale = node.docs["info"]["scales"][self._ref.scale_index or 0] if "sharding" in scale: - raise NotImplementedError( - "list_chunks: sharded precomputed stores list shards, not chunks" + raise NotImplementedError("chunk keys of a sharded precomputed store name shards") + spatial = tuple(scale["size"])[::-1] + params["chunk_shape"] = tuple(scale["chunk_sizes"][0])[::-1] + params["shape"] = spatial + params["origin"] = tuple(int(o) for o in scale.get("voxel_offset", [0, 0, 0]))[::-1] + return params + + 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()) + + 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()) + + 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.""" + 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()] + found = set() + for k in keys: + idx = _write.parse_chunk_key( + k.rsplit("/", 1)[-1] if params["format"] == "precomputed" else k, **params + ) + if idx is not None: + found.add(idx) + return sorted(found) + + # --- writing (Behaviour 13) ---------------------------------------------------- + + def write_index( + self, begin: Sequence[int], end: Sequence[int], data: np.ndarray, partial: bool = False + ) -> None: + """Write ``data`` over the voxel box ``[begin, end)`` of the stored grid. The box must + be a union of whole write units (every bound a multiple of ``write_shape`` on its axis, + or the array's end), so writers in different processes need no coordination; a + unit written whole is one put. ``partial=True`` lifts the check and does tensorstore's + read-modify-write of the touched units, unsafe while anything else writes the same + unit. No cast, no padding, no resampling; the write commits before returning.""" + if not self.writable: + raise ValueError( + f"{self.path} was opened read-only; open it with mode='r+' or create it" + ) + if self._target is not None: + raise ValueError( + f"{self.path}: no resampling on write; open the array without voxel_size=" + ) + begin, end = tuple(int(b) for b in begin), tuple(int(e) for e in end) + data = np.asarray(data) + shape = self._grid.shape + if len(begin) != len(shape) or len(end) != len(shape): + raise ValueError(f"{self.path}: begin and end need one entry per axis of {shape}") + wanted = tuple(e - b for b, e in zip(begin, end)) + if data.shape != wanted: + raise ValueError(f"{self.path}: data has shape {data.shape}; the box is {wanted}") + if data.dtype != self.dtype: + raise ValueError( + f"{self.path}: data is {data.dtype}, the array {self.dtype}; nothing is cast on write" + ) + for axis, (b, e, n) in enumerate(zip(begin, end, shape)): + if b < 0 or e > n or e <= b: + raise ValueError( + f"{self.path}: write [{b}, {e}) on axis {axis} ({self.axes[axis]}) lies outside the array of {n}" ) - chunk = scale["chunk_sizes"][0] - found = [] - for k in keys: - m = re.fullmatch(r"(\d+)-\d+_(\d+)-\d+_(\d+)-\d+", k.rsplit("/", 1)[-1]) - if m: - starts = [int(m.group(i)) for i in (1, 2, 3)] - found.append(tuple(s // c for s, c in zip(starts, chunk))[::-1]) - return sorted(set(found)) + if not partial: + for axis, (b, e, u, n) in enumerate(zip(begin, end, self.write_shape, shape)): + if b % u != 0 or (e % u != 0 and e != n): + raise ValueError( + f"{self.path}: write [{b}, {e}) on axis {axis} ({self.axes[axis]}) is not whole write " + f"units of {u}; align the box to write_shape {self.write_shape} or pass partial=True" + ) + self._store[tuple(slice(b, e) for b, e in zip(begin, end))].write(data).result() + + def write(self, box, data: np.ndarray, partial: bool = False) -> None: + """Write ``data`` over a world box: the voxels Behaviour 2's centre rule selects, then + ``write_index``. A box reaching outside the array raises; nothing is padded or + resampled on write.""" + begin, end = self.index_box(box) + for axis, (b, e, n) in enumerate(zip(begin, end, self._grid.shape)): + if b < 0 or e > n: + raise ValueError( + f"{self.path}: the box reaches outside the array on axis {axis} ({self.axes[axis]}): " + f"voxels [{b}, {e}) of {n}; nothing is padded on write" + ) + self.write_index(begin, end, data, partial=partial) def read_index( self, @@ -740,14 +857,18 @@ def __init__( context: ts.Context | None = None, units: Units = "nanometer", select: dict[str, int] | None = None, + mode: str = "r", _node: Node | None = None, **kwargs, ): + if mode not in ("r", "r+"): + raise ValueError("mode is 'r' (read) or 'r+' (read and write)") self.path = path self.context = context if context is not None else shared_context(0) self.units = units self._select = select or {} - self._kwargs = kwargs + self._mode = mode + self._kwargs = {**kwargs, "mode": mode} 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: @@ -833,6 +954,83 @@ def translations(self) -> list[tuple[str, tuple[float, ...]]]: """``(name, spatial translation)`` per level; see ``placements``.""" return [(name, tr) for name, _, tr in self.placements()] + def add_level( + self, + name: str, + factors=None, + origin=0, + *, + voxel_size=None, + translation=None, + shape=None, + dtype=None, + chunks=None, + shards=None, + codecs=None, + fill_value=None, + attrs: dict | None = None, + exists: str = "error", + from_level: str | int | None = None, + ) -> WorldArray: + """A new level's dataset entry and its empty array (Behaviour 14). Placement by + ``level_placement`` over ``from_level`` (the finest listed when None) from ``factors`` + and ``origin``, unless ``voxel_size`` and ``translation`` are given; shape by + ``level_shape`` unless given; dtype, chunks, shards, codecs and fill value copied from + that level unless given. Writes no voxel data. One process at a time.""" + if self._mode != "r+": + raise ValueError(f"{self.path}: open the pyramid with mode='r+' to add a level") + group = self._group + if group is None or group.kind != "group" or group.format not in ("zarr2", "zarr3"): + raise ValueError(f"{self.path}: add_level needs an OME zarr group (v2 or v3)") + base_ref = self._ref(from_level if from_level is not None else 0) + 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 + if voxel_size is None or translation is None: + if factors is None: + raise ValueError("pass factors (and origin), or voxel_size and translation") + f = _write.expand_spatial( + (factors,) * len(spatial) if np.isscalar(factors) else factors, spatial, ndim, 1 + ) + o = _write.expand_spatial( + (origin,) * len(spatial) if np.isscalar(origin) else origin, spatial, ndim, 0 + ) + vs, tr = _write.level_placement(base.voxel_size, base.translation, f, o) + if shape is None: + shape = _write.level_shape(base.shape, f, o) + else: + vs = _write.expand_spatial(voxel_size, spatial, ndim, 1.0) + tr = _write.expand_spatial(translation, spatial, ndim, 0.0) + if shape is None: + raise ValueError("pass shape when placing a level by voxel_size and translation") + base_codecs, base_shards = _write.codecs_of(base_node) + if codecs is None: + codecs = base_codecs if base_codecs is not None else "none" + if shards is None: + shards = base_shards + vol = create( + join(group.path, name), + shape, + base.dtype if dtype is None else dtype, + vs, + tr, + axes=base.axes, + units=base.units, + chunks=base.chunk_shape if chunks is None else chunks, + shards=shards, + fill_value=(base.fill_value if base.fill_value is not None else 0) + if fill_value is None + else fill_value, + codecs=codecs, + format=group.format, + attrs=attrs, + exists=exists, + context=self.context, + ) + self._group, self._refs = discover_group(detect(group.path)) + self._light, self._levels = None, {} + return vol + def _select_ref(self, voxel_size, mode: Mode = "floor") -> LevelRef: name, _ = select_level(self.voxel_sizes(), voxel_size, mode) return self._ref(name) @@ -851,6 +1049,127 @@ def __len__(self) -> int: # Functions +def create( + path: str, + shape: Sequence[int], + dtype, + voxel_size, + translation, + *, + axes: Sequence[str] | None = None, + units: Units = "nanometer", + chunks: Sequence[int] | None = None, + shards: Sequence[int] | None = None, + fill_value=0, + codecs="zstd", + format: str | None = None, + attrs: dict | None = None, + group_attrs: dict | None = None, + exists: str = "error", + store_fill: bool = False, + separator: str | None = None, + axis_types: Sequence[str | None] | None = None, + context: ts.Context | None = None, + retries: Retry | None = None, +) -> WorldArray: + """Create a placed level: the group document (made when absent, else its one dataset + entry added or replaced) and then the array, as OME-NGFF zarr v2 (OME 0.4) or v3 (OME + 0.5) following the group's format (Behaviours 11 and 12). ``chunks`` is required. + Nothing is ever deleted or overwritten: with ``exists="error"`` anything at ``path`` + 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 chunks is None: + raise ValueError("chunks is required: the package never picks a chunk shape") + ndim = len(shape) + axes = tuple(str(a) for a in axes) if axes is not None else default_axes(ndim) + context = context if context is not None else shared_context(0) + kinds = _write.axis_kinds(axes, axis_types) + vs = _write.per_axis(voxel_size, kinds, 1.0, f"{path}: voxel_size") + tr = _write.per_axis(translation, kinds, 0.0, f"{path}: translation") + vol_units: Units = units if isinstance(units, (str, dict)) or units is None else None + + def adopt(normalized: str, fmt: str) -> WorldArray: + vol = WorldArray( + normalized, + mode="r+", + units=None, + context=context, + retries=retries, + store_fill=store_fill, + ) + diffs = _write.differences( + vol.meta, + shape=shape, + dtype=dtype, + chunks=chunks, + shards=shards, + voxel_size=vs, + translation=tr, + ) + if diffs: + raise ValueError( + f"{normalized} exists and differs from what was asked for; the package never " + "overwrites:\n " + "\n ".join(diffs) + ) + group_path = normalized.rsplit("/", 1)[0] + _write.ensure_group_entry( + group_path, + normalized.rsplit("/", 1)[-1], + vs, + tr, + axes, + units, + axis_types, + group_attrs, + format=fmt, + ) + return WorldArray( + normalized, + mode="r+", + units=vol_units, + context=context, + retries=retries, + store_fill=store_fill, + ) + + normalized, fmt, adopted = _write.create_array( + path, + shape, + dtype, + vs, + tr, + axes=axes, + units=units, + chunks=chunks, + shards=shards, + fill_value=fill_value, + codecs=codecs, + format=format, + attrs=attrs, + group_attrs=group_attrs, + exists=exists, + store_fill=store_fill, + separator=separator, + axis_types=axis_types, + context=context, + ) + if adopted: + return adopt(normalized, fmt) + return WorldArray( + normalized, + mode="r+", + units=vol_units, + context=context, + retries=retries, + store_fill=store_fill, + ) + + +def _reopen(path: str, opts: dict, recipe=None) -> WorldArray: + return WorldArray(path, context=from_recipe(recipe), **opts) + + def open(path: str, **kwargs) -> WorldArray: # noqa: A001 return WorldArray(path, **kwargs) @@ -939,4 +1258,13 @@ def open_tensorstore( return open_array(node, context=context) -__all__ = ["Pyramid", "WorldArray", "metadata", "open", "open_pyramid", "open_tensorstore", "U"] +__all__ = [ + "Pyramid", + "WorldArray", + "create", + "metadata", + "open", + "open_pyramid", + "open_tensorstore", + "U", +] diff --git a/src/worldarray/context.py b/src/worldarray/context.py index f1cd4ae..8d283f2 100644 --- a/src/worldarray/context.py +++ b/src/worldarray/context.py @@ -13,11 +13,14 @@ from __future__ import annotations +import json import threading +from typing import Any import tensorstore as ts _shared: dict[int, ts.Context] = {} +_rebuilt: dict[str, ts.Context] = {} _lock = threading.Lock() @@ -51,3 +54,31 @@ def capped_context_spec(cache_bytes: int = 0, concurrency: int = 1, network: boo def capped_context(cache_bytes: int = 0, concurrency: int = 1, network: bool = True) -> ts.Context: return ts.Context(capped_context_spec(cache_bytes, concurrency, network)) + + +def recipe_of(context: ts.Context) -> tuple[str, Any]: + """How another process rebuilds ``context``: ``("shared", cache_bytes)`` for a shared + context, else ``("spec", json)`` from the context's own spec, so a capped context stays + capped across a pickle (Behaviour 11, mesh-n-bone's workers).""" + with _lock: + for cache_bytes, ctx in _shared.items(): + if ctx is context: + return ("shared", cache_bytes) + spec = context.spec + return ("spec", json.dumps(spec.to_json() if spec is not None else {}, sort_keys=True)) + + +def from_recipe(recipe: tuple[str, Any] | None) -> ts.Context | None: + """``recipe_of`` inverted; one rebuilt context per spec per process, so handles that + shared a context keep sharing one.""" + if recipe is None: + return None + kind, value = recipe + if kind == "shared": + return shared_context(int(value)) + with _lock: + ctx = _rebuilt.get(value) + if ctx is None: + ctx = ts.Context(ts.Context.Spec(json.loads(value))) + _rebuilt[value] = ctx + return ctx diff --git a/src/worldarray/placement.py b/src/worldarray/placement.py index a4f77fb..e23100b 100644 --- a/src/worldarray/placement.py +++ b/src/worldarray/placement.py @@ -55,6 +55,7 @@ class ArrayMeta: group_attrs: dict = field(default_factory=dict) placed_by: str | None = None # which rule gave voxel_size and translation; None: nothing did fill_value: object = None # the store's fill value, as tensorstore reports it + shard_shape: tuple[int, ...] | None = None # a sharded zarr v3 array's shard; None otherwise def __post_init__(self) -> None: object.__setattr__(self, "dtype", np.dtype(self.dtype)) @@ -62,6 +63,10 @@ def __post_init__(self) -> None: for name in ("chunk_shape", "axes", "units", "voxel_size", "translation"): if len(getattr(self, name)) != n: raise ValueError(f"{name} must have length {n}, got {getattr(self, name)}") + if self.shard_shape is not None: + if len(self.shard_shape) != n: + raise ValueError(f"shard_shape must have length {n}, got {self.shard_shape}") + object.__setattr__(self, "shard_shape", tuple(int(s) for s in self.shard_shape)) object.__setattr__(self, "voxel_size", tuple(float(v) for v in self.voxel_size)) object.__setattr__(self, "translation", tuple(float(v) for v in self.translation)) @@ -135,8 +140,14 @@ def drop_axis(self, axis: int) -> ArrayMeta: 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 + def spatial_axes(axes: tuple[str, ...]) -> tuple[int, ...]: named = tuple(i for i, a in enumerate(axes) if a in SPATIAL_NAMES) @@ -189,14 +200,17 @@ def read_meta( ``conflicts``: what a per-array ``voxel_size`` that disagrees with the OME scale does, ``"error"`` (cellmap-analyze's choice, catches stale attributes) or ``"warn"`` (a server that would rather show the image).""" + read_chunk = tuple(int(c) for c in store.chunk_layout.read_chunk.shape) + write_chunk = tuple(int(c) for c in store.chunk_layout.write_chunk.shape) return _assemble( node, tuple(int(s) for s in store.shape), store.dtype.numpy_dtype, - tuple(int(c) for c in store.chunk_layout.read_chunk.shape), + read_chunk, _fill_value(store), scale_index=scale_index, conflicts=conflicts, + shard_shape=write_chunk if write_chunk != read_chunk else None, ) @@ -207,23 +221,40 @@ def read_meta_light( dtype, chunks and fill value from ``.zarray``, ``zarr.json`` (the inner chunk of a sharded array), ``attributes.json`` or ``info``. For listing a pyramid's levels cheaply; a precomputed single channel is dropped as ``open_array`` drops it.""" - shape, dtype, chunk_shape, fill = _dims_from_docs(node, scale_index) + shape, dtype, chunk_shape, fill, shard = _dims_from_docs(node, scale_index) return _assemble( - node, shape, dtype, chunk_shape, fill, scale_index=scale_index, conflicts=conflicts + node, + shape, + dtype, + chunk_shape, + fill, + scale_index=scale_index, + conflicts=conflicts, + shard_shape=shard, ) def _dims_from_docs(node: Node, scale_index: int | None): + """``(shape, dtype, chunk_shape, fill_value, shard_shape)`` from the documents: the inner + chunk and the shard of a sharded zarr v3 array, N5 and precomputed in C order, a + precomputed single channel dropped.""" if node.format == "zarr2": d = node.docs[".zarray"] - return tuple(d["shape"]), np.dtype(d["dtype"]), tuple(d["chunks"]), d.get("fill_value") + return ( + tuple(d["shape"]), + np.dtype(d["dtype"]), + tuple(d["chunks"]), + d.get("fill_value"), + None, + ) if node.format == "zarr3": d = node.docs["zarr.json"] - chunks = tuple(d["chunk_grid"]["configuration"]["chunk_shape"]) + grid = tuple(d["chunk_grid"]["configuration"]["chunk_shape"]) + chunks, shard = grid, None for codec in d.get("codecs", []) or []: if codec.get("name") == "sharding_indexed": - chunks = tuple(codec["configuration"]["chunk_shape"]) - return tuple(d["shape"]), np.dtype(d["data_type"]), chunks, d.get("fill_value") + chunks, shard = tuple(codec["configuration"]["chunk_shape"]), grid + return tuple(d["shape"]), np.dtype(d["data_type"]), chunks, d.get("fill_value"), shard if node.format == "n5": d = node.docs["attributes.json"] return ( @@ -231,14 +262,15 @@ def _dims_from_docs(node: Node, scale_index: int | None): np.dtype(d["dataType"]), tuple(d["blockSize"])[::-1], None, + None, ) info = node.docs["info"] scale = info["scales"][scale_index or 0] size, chunk = tuple(scale["size"])[::-1], tuple(scale["chunk_sizes"][0])[::-1] channels = int(info.get("num_channels", 1)) if channels > 1: - return (channels, *size), np.dtype(info["data_type"]), (channels, *chunk), None - return size, np.dtype(info["data_type"]), chunk, None + return (channels, *size), np.dtype(info["data_type"]), (channels, *chunk), None, None + return size, np.dtype(info["data_type"]), chunk, None, None def _assemble( @@ -250,6 +282,7 @@ def _assemble( *, scale_index: int | None, conflicts: str, + shard_shape: tuple[int, ...] | None = None, ) -> ArrayMeta: if conflicts not in ("error", "warn"): raise ValueError("conflicts is 'error' or 'warn'") @@ -289,6 +322,7 @@ def _assemble( group_attrs=group_attrs, placed_by=found.placed_by if found.voxel_size is not None else None, fill_value=fill_value, + shard_shape=shard_shape, ) diff --git a/src/worldarray/store.py b/src/worldarray/store.py index b4e4f8e..3a4cc6a 100644 --- a/src/worldarray/store.py +++ b/src/worldarray/store.py @@ -205,6 +205,13 @@ def _read_key(spec: dict, key: str) -> bytes | None: return None if result.state == "missing" else bytes(result.value) +def _write_key(spec: dict, key: str, data: bytes) -> None: + kv = _metadata_kvstore( + json.dumps(spec, sort_keys=True), os.environ.get("TENSORSTORE_GCS_HTTP_URL") + ) + kv.write(key, data).result() + + class Location: """A directory (local or a URL) and the ways in to it.""" @@ -219,6 +226,14 @@ def read_json(self, key: str) -> dict | None: data = self.read(key) return None if data is None else json.loads(data) + def write(self, key: str, data: bytes) -> None: + """One document under this location (the write side's group documents and + ``.zattrs``); the same ways in as a read, an anonymous S3 way refused and skipped.""" + with_access(self.specs, lambda spec: _write_key(spec, key, data)) + + def write_json(self, key: str, doc: dict) -> None: + self.write(key, (json.dumps(doc, indent=2) + "\n").encode()) + def read_many(self, keys: Sequence[str]) -> dict[str, bytes | None]: """All of ``keys`` under one way in. A bucket with listing off answers a refusal for a missing key, so a refusal on one key counts as missing when another key was read @@ -387,16 +402,22 @@ def open_array( context: ts.Context, scale_index: int | None = None, squeeze_channel: bool = True, + write: bool = False, + store_fill: bool = False, ) -> ts.TensorStore: """The tensorstore of an array node (or one scale of a precomputed node), in C order with its origin at 0: N5 and precomputed are transposed, a precomputed ``voxel_offset`` is moved to 0 (and carried in the metadata as the translation), and a precomputed - single channel is dropped.""" + single channel is dropped. ``write`` opens it for writing as well (``mode="r+"``); + ``store_fill`` then stores write units equal to the fill value instead of leaving them + absent (a per-open setting of tensorstore's, not metadata).""" if node.kind != "array" and node.format != "precomputed": raise ValueError(f"{node.path} is a {node.format} group, not an array") def attempt(kvstore: dict) -> ts.TensorStore: spec: dict[str, Any] = {"driver": _DRIVER[node.format], "kvstore": kvstore} + if write: + spec["store_data_equal_to_fill_value"] = bool(store_fill) if node.format == "zarr2": spec["metadata"] = cleaned_zarray(node.docs[".zarray"]) spec["assume_metadata"] = True @@ -412,7 +433,7 @@ def attempt(kvstore: dict) -> ts.TensorStore: spec["assume_metadata"] = True if node.format == "precomputed" and scale_index is not None: spec["scale_index"] = scale_index - return ts.open(spec, read=True, context=context).result() + return ts.open(spec, read=True, write=write, context=context).result() store = with_access(kvstore_specs(node.path), attempt) if node.format in ("n5", "precomputed"): diff --git a/src/worldarray/units.py b/src/worldarray/units.py index edfafe5..d85ac02 100644 --- a/src/worldarray/units.py +++ b/src/worldarray/units.py @@ -75,3 +75,25 @@ def factor( if a[0] != b[0]: raise ValueError(f"cannot convert {unit_from!r} ({a[0]}) to {unit_to!r} ({b[0]})") return a[1] / b[1] + + +# 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. +_OME_NAMES = ( + "picometer", "angstrom", "nanometer", "micrometer", "millimeter", "centimeter", "meter", + "nanosecond", "microsecond", "millisecond", "second", "minute", "hour", "day", +) # fmt: skip +_OME_BY_SIZE = {_lookup(name): name for name in _OME_NAMES} + + +def ome_name(unit: str) -> str: + """The OME-NGFF name of a unit: ``"nanometer"`` for any spelling of nanometre the table + knows (``"nm"``, ``"nanometre"``, ``"Nanometers"``), ``"micrometer"`` for ``"um"``, + ``"µm"`` and ``"micron"``, ``"second"`` for ``"s"`` and so on. A unit the table does + not know (``""``, ``"pixel"``) is a ``ValueError``; `ome_multiscales` writes such a + unit as given.""" + found = _lookup(unit) if unit else None + if found is None or found not in _OME_BY_SIZE: + raise ValueError(f"unknown unit {unit!r}") + return _OME_BY_SIZE[found] diff --git a/src/worldarray/write.py b/src/worldarray/write.py new file mode 100644 index 0000000..d3e5884 --- /dev/null +++ b/src/worldarray/write.py @@ -0,0 +1,829 @@ +"""Writing: OME-NGFF documents and placed creation (spec: Convention 6, Behaviours 10 to 16). + +The package writes OME-NGFF zarr only: zarr v2 with OME 0.4 (``.zgroup`` and ``.zattrs`` +holding ``multiscales``) or zarr v3 with OME 0.5 (``zarr.json`` with ``attributes.ome``), +never 0.4 inside v3, never N5 or precomputed. A written ``translation`` is voxel 0's centre, +written always, zeros included; ``scale`` is the voxel size; every number is the float it +was given. Nothing here deletes, truncates or overwrites: ``create_array`` raises on an +existing path unless told to adopt, and a group document is read, its one dataset entry +replaced or added, and written back with every other key as it was. + +The pure functions (``ome_multiscales``, ``level_placement``, ``level_shape``, +``chunk_key``, ``parse_chunk_key``) need no store. ``create_array`` writes the group +document and then the array through tensorstore; ``worldarray.array.create`` wraps it in a +writable ``WorldArray`` and does the adopt check. +""" + +from __future__ import annotations + +import json +import logging +import math +import re +import time +from collections.abc import Sequence +from typing import Any + +import numpy as np +import tensorstore as ts + +from .geometry import snap_integral +from .placement import spatial_axes +from .store import ( + _DRIVER, + Location, + Node, + basename, + detect, + is_refused, + is_root, + kvstore_specs, + normalize_path, + parent, + with_access, +) +from .units import ome_name + +logger = logging.getLogger(__name__) + +SPACE, TIME, CHANNEL = "space", "time", "channel" +ZSTD_LEVEL = 3 # what zarr-python's default zstd(level=0) means +RACE_ATTEMPTS = 5 # a creation that lost the race to another process is re-read this often +RACE_DELAY = 0.5 # seconds, times the attempt (registration-station's shared_array) +RTOL = 1e-6 + +Levels = Sequence[tuple[str, Sequence[float], Sequence[float]]] + + +# --------------------------------------------------------------------------- +# Axes, units, per-axis values + + +def axis_kinds(axes: Sequence[str], types: Sequence[str | None] | None = None) -> tuple: + """The OME ``type`` of each axis: as given, else from the name (``z``, ``y``, ``x`` + space, ``t`` time, ``c`` channel; an unnamed array's last three axes space).""" + axes = tuple(str(a) for a in axes) + if types is not None: + types = tuple(types) + if len(types) != len(axes): + raise ValueError(f"axis_types has {len(types)} entries for {len(axes)} axes") + for t in types: + if t not in (SPACE, TIME, CHANNEL, None): + raise ValueError(f"axis type {t!r}; use 'space', 'time', 'channel' or None") + return types + sp = set(spatial_axes(axes)) + out: list[str | None] = [] + for i, a in enumerate(axes): + if a == "t": + out.append(TIME) + elif a == "c": + out.append(CHANNEL) + elif a in ("z", "y", "x") or i in sp: + out.append(SPACE) + else: + out.append(None) + return tuple(out) + + +def axis_units(units, axes: Sequence[str], kinds: Sequence[str | None]) -> tuple[str, ...]: + """Per-axis unit strings: a str is the unit of the space axes, a dict by kind gives + ``space`` and ``time``, a sequence is one per axis; ``""`` where there is none.""" + n = len(axes) + if units is None: + return ("",) * n + if isinstance(units, str): + return tuple(units if k == SPACE else "" for k in kinds) + if isinstance(units, dict): + return tuple(str(units.get(k) or "") if k in (SPACE, TIME) else "" for k in kinds) + seq = tuple(str(u) if u else "" for u in units) + if len(seq) != n: + raise ValueError(f"units has {len(seq)} entries for {n} axes") + return seq + + +def per_axis(values, kinds: Sequence[str | None], fill: float, what: str) -> tuple[float, ...]: + """``values`` as one float per axis: given for every axis, or for the space axes only + (the others take ``fill``).""" + values = tuple(float(v) for v in values) + n = len(kinds) + space = [i for i, k in enumerate(kinds) if k == SPACE] + if len(values) == n: + return values + if len(values) == len(space): + out = [float(fill)] * n + for i, v in zip(space, values): + out[i] = v + return tuple(out) + raise ValueError( + f"{what} has {len(values)} entries; the array has {n} axes, {len(space)} spatial" + ) + + +def _broadcast(value, n: int, what: str) -> tuple[float, ...]: + if np.isscalar(value): + return (float(value),) * n + out = tuple(float(v) for v in value) + if len(out) != n: + raise ValueError(f"{what} has {len(out)} entries for {n} axes") + return out + + +# --------------------------------------------------------------------------- +# Documents (Behaviour 10) + + +def ome_multiscales( + axes: Sequence[str], + units="nanometer", + levels: Levels = (), + version: str = "0.5", + axis_types: Sequence[str | None] | None = None, + name: str = "", + metadata: dict | None = None, + type: str | None = None, +) -> dict: + """The complete group attributes of an OME-NGFF multiscale group: ``{"multiscales": + [...]}`` for 0.4 (zarr v2), ``{"ome": {"version": "0.5", "multiscales": [...]}}`` for + 0.5 (zarr v3). ``levels`` are ``(path, voxel_size, translation)`` with the placement + explicit per level, over the space axes or every axis; ``translation`` is voxel 0's + centre and is written always. Units are written in OME's names: any spelling the unit + table knows (``"nm"``, ``"µm"``, ``"nanometre"``) becomes ``"nanometer"``, + ``"micrometer"``, ...; a unit it does not know is written as given. ``type`` (the + downsampling method, OME's optional field), ``name`` and ``metadata`` are the + multiscale's own.""" + axes = tuple(str(a) for a in axes) + kinds = axis_kinds(axes, axis_types) + unit_list = tuple(ome_unit(u) for u in axis_units(units, axes, kinds)) + axis_docs = [] + for a, k, u in zip(axes, kinds, unit_list): + doc: dict[str, Any] = {"name": a} + if k is not None: + doc["type"] = k + if u and k in (SPACE, TIME, None): + doc["unit"] = u + axis_docs.append(doc) + datasets = [dataset_entry(path, vs, tr, kinds) for path, vs, tr in levels] + multiscale: dict[str, Any] = {"name": name, "axes": axis_docs, "datasets": datasets} + if type is not None: + multiscale["type"] = str(type) + if metadata: + multiscale["metadata"] = dict(metadata) + if version == "0.4": + return {"multiscales": [{"version": "0.4", **multiscale}]} + if version == "0.5": + return {"ome": {"version": "0.5", "multiscales": [multiscale]}} + raise ValueError(f"OME version {version!r}; the package writes 0.4 (zarr v2) or 0.5 (zarr v3)") + + +def ome_unit(unit: str) -> str: + """``unit`` in OME's name when the unit table knows it, else as given (``""`` stays + ``""``, ``"pixel"`` stays ``"pixel"``: Convention 3 welcomes unknown units).""" + try: + return ome_name(unit) if unit else "" + except ValueError: + return unit + + +def dataset_entry(path: str, voxel_size, translation, kinds: Sequence[str | None]) -> dict: + vs = per_axis(voxel_size, kinds, 1.0, f"{path}: voxel_size") + tr = per_axis(translation, kinds, 0.0, f"{path}: translation") + return { + "path": str(path).strip("/"), + "coordinateTransformations": [ + {"type": "scale", "scale": list(vs)}, + {"type": "translation", "translation": list(tr)}, + ], + } + + +def level_placement( + voxel_size, translation, factors, origin=0 +) -> tuple[tuple[float, ...], tuple[float, ...]]: + """The placement of the level whose voxel ``j`` covers source voxels ``origin + j f .. + origin + (j + 1) f - 1`` on each axis: ``(voxel_size f, translation + voxel_size (origin + + (f - 1) / 2))``. Convention 1's rule inverted, and the centre a block mean produces.""" + vs = tuple(float(v) for v in voxel_size) + tr = tuple(float(t) for t in translation) + if len(tr) != len(vs): + raise ValueError("voxel_size and translation must have the same length") + f = _broadcast(factors, len(vs), "factors") + o = _broadcast(origin, len(vs), "origin") + return ( + tuple(v * fi for v, fi in zip(vs, f)), + tuple(t + v * (oi + (fi - 1) / 2) for t, v, fi, oi in zip(tr, vs, f, o)), + ) + + +def level_shape(shape, factors, origin=0) -> tuple[int, ...]: + """``ceil((n - origin) / f)`` per axis: the voxels of a level that covers source voxels + from ``origin`` on. The quotient is snapped to a whole number within ``1e-6`` first + (Convention 3), so a factor that is a ratio of float voxel sizes, ``0.3 / 0.1 = + 2.9999999999999996``, gives 3 voxels for 9, not 4.""" + shape = tuple(int(n) for n in shape) + f = _broadcast(factors, len(shape), "factors") + o = _broadcast(origin, len(shape), "origin") + return tuple( + max(0, int(math.ceil(float(snap_integral((n - oi) / fi))))) + for n, fi, oi in zip(shape, f, o) + ) + + +def expand_spatial(values, spatial: Sequence[int], ndim: int, fill) -> tuple: + """Spatial-length ``values`` spread onto the ``spatial`` axes of an ``ndim`` array, the + other axes taking ``fill``; values already of length ``ndim`` pass through.""" + values = tuple(values) + if len(values) == ndim: + return values + if len(values) != len(spatial): + raise ValueError(f"expected {ndim} or {len(spatial)} values, got {len(values)}") + out = [fill] * ndim + for i, v in zip(spatial, values): + out[i] = v + return tuple(out) + + +# --------------------------------------------------------------------------- +# Chunk keys (Behaviour 15) + +_PRECOMPUTED_KEY = re.compile(r"^(\d+)-(\d+)(?:_(\d+)-(\d+))*$") + + +def chunk_key( + index: Sequence[int], + format: str, + separator: str | None = None, + encoding: str = "default", + chunk_shape: Sequence[int] | None = None, + shape: Sequence[int] | None = None, + origin: Sequence[int] | None = None, +) -> str: + """The storage key, relative to the array, of the write unit at ``index`` (C order, the + storage grid: shards for a sharded zarr v3 array). zarr v2 joins with + ``dimension_separator`` (default ``.``); zarr v3 ``default`` encoding gives ``c/0/1/2`` + with its separator (default ``/``) and ``v2`` encoding ``0.1.2``; N5 is x first; precomputed + is ``x0-x1_y0-y1_z0-z1`` in the scale's voxel coordinates, ``origin`` (its + ``voxel_offset``) included and the last range clipped to ``shape``.""" + index = tuple(int(i) for i in index) + if format == "zarr2": + return (separator or ".").join(str(i) for i in index) + if format == "zarr3": + if encoding == "default": + sep = separator or "/" + return sep.join(["c", *(str(i) for i in index)]) + if encoding == "v2": + return (separator or ".").join(str(i) for i in index) + raise ValueError(f"zarr v3 chunk_key_encoding {encoding!r}; use 'default' or 'v2'") + if format == "n5": + return "/".join(str(i) for i in reversed(index)) + if format == "precomputed": + if chunk_shape is None or shape is None: + raise ValueError("precomputed chunk keys need chunk_shape and shape") + origin = tuple(int(o) for o in origin) if origin is not None else (0,) * len(index) + parts = [] + for i, c, n, o in zip( + reversed(index), reversed(chunk_shape), reversed(shape), reversed(origin) + ): + start = o + i * int(c) + stop = min(start + int(c), o + int(n)) + parts.append(f"{start}-{stop}") + return "_".join(parts) + raise ValueError(f"unknown format {format!r}") + + +def parse_chunk_key( + key: str, + format: str, + separator: str | None = None, + encoding: str = "default", + chunk_shape: Sequence[int] | None = None, + shape: Sequence[int] | None = None, + origin: Sequence[int] | None = None, + ndim: int | None = None, +) -> tuple[int, ...] | None: + """``chunk_key`` inverted: the index of a write unit, or None for a key that is not a + chunk (``zarr.json``, ``.zarray``, ``attributes.json``, ``info``, a ``c/`` key of the + wrong arity).""" + key = str(key) + if format == "zarr2" or (format == "zarr3" and encoding == "v2"): + parts = key.split(separator or ".") + found = _digits(parts) + elif format == "zarr3": + parts = key.split(separator or "/") + found = _digits(parts[1:]) if parts and parts[0] == "c" and len(parts) > 1 else None + elif format == "n5": + found = _digits(key.split("/")) + found = found[::-1] if found is not None else None + elif format == "precomputed": + if chunk_shape is None or shape is None or not _PRECOMPUTED_KEY.match(key): + return None + ranges = [tuple(int(v) for v in r.split("-")) for r in key.split("_")] + if len(ranges) != len(chunk_shape): + return None + origin = tuple(int(o) for o in origin) if origin is not None else (0,) * len(ranges) + idx = [] + for (start, _stop), c, o in zip(ranges, reversed(chunk_shape), reversed(origin)): + if (start - o) % int(c) != 0 or start < o: + return None + idx.append((start - o) // int(c)) + found = tuple(idx[::-1]) + else: + raise ValueError(f"unknown format {format!r}") + if found is None or (ndim is not None and len(found) != ndim): + return None + return found + + +def _digits(parts: Sequence[str]) -> tuple[int, ...] | None: + if not parts or not all(p.isdigit() for p in parts): + return None + return tuple(int(p) for p in parts) + + +# --------------------------------------------------------------------------- +# Codecs + + +def codec_metadata(codecs, format: str) -> dict: + """The array-metadata members for ``codecs`` in the format's dialect: ``{"codecs": + [...]}`` for zarr v3, ``{"compressor": ...}`` for zarr v2. ``"zstd"`` is level 3, + ``"none"`` is uncompressed; a list is a zarr v3 codecs list and a dict a zarr v2 + compressor, taken as given.""" + if codecs in ("zstd", None): + if format == "zarr3": + return { + "codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "zstd", "configuration": {"level": ZSTD_LEVEL}}, + ] + } + return {"compressor": {"id": "zstd", "level": ZSTD_LEVEL}} + if codecs == "none": + if format == "zarr3": + return {"codecs": [{"name": "bytes", "configuration": {"endian": "little"}}]} + return {"compressor": None} + if isinstance(codecs, list): + if format != "zarr3": + raise ValueError("a codecs list is zarr v3's dialect; zarr v2 takes a compressor dict") + return {"codecs": list(codecs)} + if isinstance(codecs, dict): + if format != "zarr2": + raise ValueError("a compressor dict is zarr v2's dialect; zarr v3 takes a codecs list") + return {"compressor": dict(codecs)} + raise ValueError( + f"codecs {codecs!r}; use 'zstd', 'none', a zarr v3 codecs list or a zarr v2 compressor dict" + ) + + +def codecs_of(node: Node | str) -> tuple[Any, tuple[int, ...] | None]: + """An existing array's codecs in its format's dialect, and its shard shape (None when + unsharded): what ``add_level`` copies from the level below, and what a mirror of an + array passes to ``create(codecs=, shards=)`` (cellmap-flow's sync). A path is detected.""" + if isinstance(node, str): + node = detect(node) + if node.format == "zarr2": + return node.docs[".zarray"].get("compressor"), None + if node.format == "zarr3": + doc = node.docs["zarr.json"] + codecs = list(doc.get("codecs", []) or []) + for codec in codecs: + if codec.get("name") == "sharding_indexed": + conf = codec["configuration"] + return list(conf.get("codecs", [])), tuple( + doc["chunk_grid"]["configuration"]["chunk_shape"] + ) + return codecs, None + raise ValueError(f"{node.path}: the package writes zarr v2 and v3 only, not {node.format}") + + +# --------------------------------------------------------------------------- +# Group documents (Behaviours 11 and 12) + + +def read_group_attrs(loc: Location, format: str) -> tuple[dict, dict | None]: + """``(attributes, raw zarr.json)`` of a group: the raw document is kept for zarr v3 so a + rewrite changes nothing but the attributes.""" + if format == "zarr3": + doc = loc.read_json("zarr.json") or {} + return dict(doc.get("attributes", {}) or {}), doc + return dict(loc.read_json(".zattrs") or {}), None + + +def write_group_attrs(loc: Location, format: str, attrs: dict, raw: dict | None) -> None: + if format == "zarr3": + doc = {"zarr_format": 3, "node_type": "group", **(raw or {}), "attributes": attrs} + loc.write_json("zarr.json", doc) + return + if loc.read(".zgroup") is None: + loc.write_json(".zgroup", {"zarr_format": 2}) + loc.write_json(".zattrs", attrs) + + +def _multiscales_slot(attrs: dict) -> str | None: + """Where the multiscales list lives: ``"ome"`` (0.5 nesting) or ``"top"`` (0.4), None + when the group has none.""" + ome = attrs.get("ome") + if isinstance(ome, dict) and ome.get("multiscales"): + return "ome" + if attrs.get("multiscales"): + return "top" + return None + + +def merged_group_attrs( + attrs: dict, + format: str, + name: str, + voxel_size, + translation, + axes: Sequence[str], + units, + axis_types: Sequence[str | None] | None = None, + group_attrs: dict | None = None, +) -> dict: + """``attrs`` with the dataset entry for ``name`` added or replaced in the group's first + multiscale, the datasets kept in order of voxel size, and nothing else touched. A group + without multiscales gets a fresh document in the format's OME version. ``group_attrs`` + are merged on top.""" + axes = tuple(str(a) for a in axes) + kinds = axis_kinds(axes, axis_types) + merged = dict(attrs) + slot = _multiscales_slot(merged) + if slot is None: + version = "0.4" if format == "zarr2" else "0.5" + fresh = ome_multiscales(axes, units, [(name, voxel_size, translation)], version, axis_types) + if version == "0.4": + merged["multiscales"] = fresh["multiscales"] + else: + merged["ome"] = { + **(merged.get("ome") if isinstance(merged.get("ome"), dict) else {}), + **fresh["ome"], + } + else: + container = merged["ome"] if slot == "ome" else merged + multiscales = list(container["multiscales"]) + ms = dict(multiscales[0]) + existing_axes = tuple( + str(a["name"] if isinstance(a, dict) else a) for a in ms.get("axes", []) + ) + if existing_axes and existing_axes != axes: + raise ValueError( + f"the group's multiscales axes are {existing_axes}; the level's are {axes}" + ) + existing_units = tuple( + str((a.get("unit") if isinstance(a, dict) else None) or "") for a in ms.get("axes", []) + ) + wanted_units = axis_units(units, axes, kinds) + if existing_units and any( + e != w for e, w, k in zip(existing_units, wanted_units, kinds) if k in (SPACE, TIME) + ): + raise ValueError( + f"the group's axis units are {existing_units}; the level's are {wanted_units}" + ) + wanted = dataset_entry(name, voxel_size, translation, kinds) + current = [d for d in ms.get("datasets", []) if str(d.get("path", "")).strip("/") == name] + if not (len(current) == 1 and same_entry(current[0], wanted)): + others = [d for d in ms.get("datasets", []) if d not in current] + ms["datasets"] = sorted([*others, wanted], key=_dataset_fineness) + multiscales[0] = ms + if slot == "ome": + merged["ome"] = {**container, "multiscales": multiscales} + else: + merged["multiscales"] = multiscales + if group_attrs: + merged.update(group_attrs) + return merged + + +def _entry_placement(dataset: dict) -> tuple[list[float], list[float] | None]: + scale: list[float] = [] + translation: list[float] | None = None + for t in dataset.get("coordinateTransformations", []) or []: + if t.get("type") == "scale": + scale = [float(v) for v in t["scale"]] + elif t.get("type") == "translation": + translation = [float(v) for v in t["translation"]] + return scale, translation + + +def same_entry(a: dict, b: dict, rtol: float = 1e-6) -> bool: + """Two dataset entries that say the same: same path, same scale and translation within + ``rtol`` (a missing translation is zeros). What lets an adopt write nothing.""" + if str(a.get("path", "")).strip("/") != str(b.get("path", "")).strip("/"): + return False + sa, ta = _entry_placement(a) + sb, tb = _entry_placement(b) + if len(sa) != len(sb) or not sa: + return False + ta = ta if ta is not None else [0.0] * len(sa) + tb = tb if tb is not None else [0.0] * len(sb) + return bool( + np.allclose(sa, sb, rtol=rtol, atol=0.0) and np.allclose(ta, tb, rtol=rtol, atol=1e-9) + ) + + +def _dataset_fineness(dataset: dict) -> float: + for t in dataset.get("coordinateTransformations", []) or []: + if t.get("type") == "scale": + return float(np.prod([float(v) for v in t["scale"]])) + return float("inf") + + +# --------------------------------------------------------------------------- +# Creation (Behaviour 11) + + +def group_format(group_path: str, format: str | None) -> tuple[str, list[str]]: + """The format to write at ``group_path`` and the ancestors lacking a group document, + outermost first, ``group_path`` itself included when it has none. The walk goes up from + ``group_path`` until a node exists: a zarr v2 or v3 group decides the format (a + contradicting ``format`` raises; an array, an N5 or a precomputed node raise). When no + ancestor exists the walk ends at the nearest ancestor named ``*.zarr``, the store root, + and ``format`` or zarr v3 decides; with no such ancestor either, only ``group_path`` + gets a document (Behaviour 11; cellmap-analyze's and cellmap-flow's containers).""" + group_path = normalize_path(group_path) + missing: list[str] = [] + p: str | None = group_path + while p is not None and not is_root(p): + gnode = detect(p, missing_ok=True) + if gnode is not None: + if gnode.kind != "group": + raise ValueError(f"{p} is an array; a level's parent must be a group") + if gnode.format not in ("zarr2", "zarr3"): + raise ValueError( + f"{p} is {gnode.format}; the package writes zarr v2 and v3 groups only" + ) + if format is not None and format != gnode.format: + raise ValueError( + f"{p} is a {gnode.format} group; format={format!r} would mix zarr versions" + ) + return gnode.format, missing[::-1] + missing.append(p) + if basename(p).endswith(".zarr"): + return format or "zarr3", missing[::-1] + p = parent(p) + return format or "zarr3", [group_path] + + +def write_bare_group(path: str, format: str) -> None: + """A group document with no attributes: what a missing ancestor gets.""" + loc = Location(path) + if format == "zarr3": + loc.write_json("zarr.json", {"zarr_format": 3, "node_type": "group"}) + else: + loc.write_json(".zgroup", {"zarr_format": 2}) + + +def ensure_group_entry( + group_path: str, + name: str, + voxel_size, + translation, + axes: Sequence[str], + units="nanometer", + axis_types: Sequence[str | None] | None = None, + group_attrs: dict | None = None, + format: str | None = None, +) -> None: + """List the level ``name`` in the group document at ``group_path``: its dataset entry + added, or replaced when it says something else, nothing else touched; the document, and + any missing ancestor's up to the ``.zarr`` root, created when absent (``format`` as in + ``create``). An entry that already says the same, placement within rtol 1e-6, leaves + the document unwritten, so adopting an existing array writes nothing. The write is a + read-modify-write with no lock, re-read afterwards: an entry lost to another process's + concurrent write is written again, ``RACE_ATTEMPTS`` times, then ``RuntimeError`` + (Behaviour 12). Public for an array the package did not create: mesh-n-bone's symlinked + ``s0``, a level filled before it is listed.""" + group_path = normalize_path(group_path) + fmt, missing = group_format(group_path, format) + for p in missing: + if p != group_path: + write_bare_group(p, fmt) + loc = Location(group_path) + for _ in range(RACE_ATTEMPTS): + attrs, raw = read_group_attrs(loc, fmt) + merged = merged_group_attrs( + attrs, fmt, name, voxel_size, translation, axes, units, axis_types, group_attrs + ) + if json.loads(json.dumps(merged)) == attrs: + return # the group already says this + write_group_attrs(loc, fmt, merged, raw) + raise RuntimeError( + f"{group_path}: the dataset entry for {name!r} was lost {RACE_ATTEMPTS} times to another " + "process writing the group document; create and add_level on one group belong to one " + "process (Behaviour 12)" + ) + + +def array_spec( + kvstore: dict, + format: str, + shape: Sequence[int], + dtype, + chunks: Sequence[int], + shards: Sequence[int] | None, + fill_value, + codecs, + axes: Sequence[str], + attrs: dict | None, + separator: str | None, + store_fill: bool, +) -> dict: + dtype = np.dtype(dtype) + spec: dict[str, Any] = { + "driver": _DRIVER[format], + "kvstore": kvstore, + "store_data_equal_to_fill_value": bool(store_fill), + } + members = codec_metadata(codecs, format) + if format == "zarr3": + grid = list(shards) if shards is not None else list(chunks) + codec_list = members["codecs"] + if shards is not None: + codec_list = [ + { + "name": "sharding_indexed", + "configuration": { + "chunk_shape": list(chunks), + "codecs": codec_list, + "index_codecs": [ + {"name": "bytes", "configuration": {"endian": "little"}}, + {"name": "crc32c"}, + ], + }, + } + ] + metadata: dict[str, Any] = { + "shape": list(shape), + "data_type": dtype.name, + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": grid}}, + "codecs": codec_list, + "fill_value": _json_fill(fill_value, dtype), + "dimension_names": [str(a) for a in axes], + } + if separator is not None: + metadata["chunk_key_encoding"] = { + "name": "default", + "configuration": {"separator": separator}, + } + if attrs: + metadata["attributes"] = dict(attrs) + else: + metadata = { + "shape": list(shape), + "chunks": list(chunks), + "dtype": dtype.str + if dtype.byteorder != "=" or dtype.itemsize == 1 + else dtype.newbyteorder("<").str, + "compressor": members["compressor"], + "fill_value": _json_fill(fill_value, dtype), + "dimension_separator": separator or ".", + } + spec["metadata"] = metadata + return spec + + +def _json_fill(fill_value, dtype: np.dtype): + if fill_value is None: + return 0 if dtype.kind != "b" else False + if isinstance(fill_value, float) and math.isnan(fill_value): + return "NaN" + if dtype.kind == "b": + return bool(fill_value) + return float(fill_value) if dtype.kind == "f" else int(fill_value) + + +def create_array( + path: str, + shape: Sequence[int], + dtype, + voxel_size, + translation, + *, + axes: Sequence[str], + units, + chunks: Sequence[int], + shards: Sequence[int] | None, + fill_value, + codecs, + format: str | None, + attrs: dict | None, + group_attrs: dict | None, + exists: str, + store_fill: bool, + separator: str | None, + axis_types: Sequence[str | None] | None, + context: ts.Context, +) -> tuple[str, str, bool]: + """Write the group document, then the array. Returns ``(path, format, adopted)``: + ``adopted`` is True when another process created the array between the probe and the + attempt and ``exists == "adopt"`` (the caller then verifies it). Raises + ``FileExistsError`` when something is already at ``path`` and ``exists == "error"``.""" + path = normalize_path(path) + name, group_path = path.rsplit("/", 1)[-1], path.rsplit("/", 1)[0] + if exists not in ("error", "adopt"): + raise ValueError("exists is 'error' or 'adopt'") + shape = tuple(int(n) for n in shape) + chunks = tuple(int(c) for c in chunks) + if len(chunks) != len(shape): + raise ValueError(f"chunks {chunks} must have one entry per axis of {shape}") + if shards is not None: + shards = tuple(int(s) for s in shards) + if len(shards) != len(shape) or any(s % c for s, c in zip(shards, chunks)): + raise ValueError( + f"shards {shards} must be a whole number of chunks {chunks} on every axis" + ) + axes = tuple(str(a) for a in axes) + if len(axes) != len(shape): + raise ValueError(f"axes {axes} must name every axis of {shape}") + kinds = axis_kinds(axes, axis_types) + vs = per_axis(voxel_size, kinds, 1.0, f"{path}: voxel_size") + tr = per_axis(translation, kinds, 0.0, f"{path}: translation") + + fmt, _missing = group_format(group_path, format) + if fmt == "zarr2" and shards is not None: + raise ValueError("shards are zarr v3 only") + if detect(path, missing_ok=True) is not None: + if exists == "error": + raise FileExistsError( + f"{path} exists; the package never overwrites (exists='adopt' opens an identical array)" + ) + return path, fmt, True + + # the group documents first (Behaviour 11), then the array + ensure_group_entry(group_path, name, vs, tr, axes, units, axis_types, group_attrs, format=fmt) + + def attempt(kvstore: dict) -> ts.TensorStore: + spec = array_spec( + kvstore, + fmt, + shape, + dtype, + chunks, + shards, + fill_value, + codecs, + axes, + attrs, + separator, + store_fill, + ) + return ts.open(spec, create=True, open=False, context=context).result() + + try: + with_access(kvstore_specs(path), attempt) + except Exception as e: # noqa: BLE001 + if "ALREADY_EXISTS" not in str(e) or is_refused(e): + raise + if exists == "error": + raise FileExistsError(f"{path} was created by another process meanwhile") from e + _wait_for(path) + return path, fmt, True + if fmt == "zarr2" and attrs: + Location(path).write_json(".zattrs", dict(attrs)) + return path, fmt, False + + +def _wait_for(path: str) -> None: + """Re-read a path another process is creating: its metadata appears late on NRS.""" + for k in range(1, RACE_ATTEMPTS + 1): + if detect(path, missing_ok=True) is not None: + return + time.sleep(RACE_DELAY * k) + raise FileExistsError( + f"{path}: another process started creating it and its metadata has not appeared" + ) + + +def differences( + meta, + *, + shape: Sequence[int], + dtype, + chunks: Sequence[int], + shards: Sequence[int] | None, + voxel_size: Sequence[float], + translation: Sequence[float], +) -> list[str]: + """How an existing array's ``ArrayMeta`` differs from what a caller wanted to create: + one line per field, empty when identical (placement within ``RTOL``).""" + out = [] + if tuple(meta.shape) != tuple(int(n) for n in shape): + out.append(f"shape: found {tuple(meta.shape)}, wanted {tuple(shape)}") + if np.dtype(meta.dtype) != np.dtype(dtype): + out.append(f"dtype: found {np.dtype(meta.dtype)}, wanted {np.dtype(dtype)}") + if tuple(meta.chunk_shape) != tuple(int(c) for c in chunks): + out.append(f"chunks: found {tuple(meta.chunk_shape)}, wanted {tuple(chunks)}") + found_shards = tuple(meta.shard_shape) if meta.shard_shape is not None else None + wanted_shards = tuple(int(s) for s in shards) if shards is not None else None + if found_shards != wanted_shards: + out.append(f"shards: found {found_shards}, wanted {wanted_shards}") + for what, found, wanted in ( + ("voxel_size", meta.voxel_size, voxel_size), + ("translation", meta.translation, translation), + ): + f, w = np.asarray(found, dtype=float), np.asarray(wanted, dtype=float) + if f.shape != w.shape or not np.allclose(f, w, rtol=RTOL, atol=RTOL): + out.append( + f"{what}: found {tuple(float(x) for x in f)}, wanted {tuple(float(x) for x in w)}" + ) + return out diff --git a/tests/fixtures/write.json b/tests/fixtures/write.json new file mode 100644 index 0000000..5b2f683 --- /dev/null +++ b/tests/fixtures/write.json @@ -0,0 +1,1690 @@ +{ + "_comment": "ome_multiscales(axes, units, levels, version, axis_types=, name=multiscale_name, metadata=, type=multiscale_type) -> expected group attributes (Behaviour 10). Levels are (path, voxel_size, translation) over the space axes or every axis; translation is voxel 0's centre, written always; units are written in OME's names (\"nm\" -> \"nanometer\"), an unknown unit as given. `resolved` is the per-axis form of the inputs (units broadcast to every axis, types from the names), and the same call with resolved.axis_units and resolved.axis_types must give the same document. Compare after JSON normalisation. Shared with chunkmirage's TypeScript builder.", + "cases": [ + { + "name": "zyx_ome_04", + "axes": [ + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 8, + 8, + 8 + ], + [ + 4, + 4, + 4 + ] + ], + [ + "s1", + [ + 16, + 16, + 16 + ], + [ + 8, + 8, + 8 + ] + ] + ], + "version": "0.4", + "resolved": { + "axis_units": [ + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "space", + "space", + "space" + ] + }, + "expect": { + "multiscales": [ + { + "version": "0.4", + "name": "", + "axes": [ + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 8.0, + 8.0, + 8.0 + ] + }, + { + "type": "translation", + "translation": [ + 4.0, + 4.0, + 4.0 + ] + } + ] + }, + { + "path": "s1", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 16.0, + 16.0, + 16.0 + ] + }, + { + "type": "translation", + "translation": [ + 8.0, + 8.0, + 8.0 + ] + } + ] + } + ] + } + ] + } + }, + { + "name": "zyx_ome_05", + "axes": [ + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 8, + 8, + 8 + ], + [ + 4, + 4, + 4 + ] + ], + [ + "s1", + [ + 16, + 16, + 16 + ], + [ + 8, + 8, + 8 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 8.0, + 8.0, + 8.0 + ] + }, + { + "type": "translation", + "translation": [ + 4.0, + 4.0, + 4.0 + ] + } + ] + }, + { + "path": "s1", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 16.0, + 16.0, + 16.0 + ] + }, + { + "type": "translation", + "translation": [ + 8.0, + 8.0, + 8.0 + ] + } + ] + } + ] + } + ] + } + } + }, + { + "name": "czyx_ome_04", + "axes": [ + "c", + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 4, + 4, + 4 + ], + [ + 2, + 2, + 2 + ] + ] + ], + "version": "0.4", + "resolved": { + "axis_units": [ + "", + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "channel", + "space", + "space", + "space" + ] + }, + "expect": { + "multiscales": [ + { + "version": "0.4", + "name": "", + "axes": [ + { + "name": "c", + "type": "channel" + }, + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 4.0, + 4.0, + 4.0 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 2.0, + 2.0, + 2.0 + ] + } + ] + } + ] + } + ] + } + }, + { + "name": "czyx_ome_05", + "axes": [ + "c", + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 4, + 4, + 4 + ], + [ + 2, + 2, + 2 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "", + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "channel", + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "c", + "type": "channel" + }, + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 4.0, + 4.0, + 4.0 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 2.0, + 2.0, + 2.0 + ] + } + ] + } + ] + } + ] + } + } + }, + { + "name": "tczyx_ome_04", + "axes": [ + "t", + "c", + "z", + "y", + "x" + ], + "units": { + "space": "micrometer", + "time": "second" + }, + "levels": [ + [ + "s0", + [ + 1, + 1, + 0.5, + 0.1, + 0.1 + ], + [ + 0, + 0, + 0.25, + 0.05, + 0.05 + ] + ], + [ + "s1", + [ + 1, + 1, + 0.5, + 0.2, + 0.2 + ], + [ + 0, + 0, + 0.25, + 0.1, + 0.1 + ] + ] + ], + "version": "0.4", + "resolved": { + "axis_units": [ + "second", + "", + "micrometer", + "micrometer", + "micrometer" + ], + "axis_types": [ + "time", + "channel", + "space", + "space", + "space" + ] + }, + "expect": { + "multiscales": [ + { + "version": "0.4", + "name": "", + "axes": [ + { + "name": "t", + "type": "time", + "unit": "second" + }, + { + "name": "c", + "type": "channel" + }, + { + "name": "z", + "type": "space", + "unit": "micrometer" + }, + { + "name": "y", + "type": "space", + "unit": "micrometer" + }, + { + "name": "x", + "type": "space", + "unit": "micrometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 1.0, + 0.5, + 0.1, + 0.1 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 0.0, + 0.25, + 0.05, + 0.05 + ] + } + ] + }, + { + "path": "s1", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 1.0, + 0.5, + 0.2, + 0.2 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 0.0, + 0.25, + 0.1, + 0.1 + ] + } + ] + } + ] + } + ] + } + }, + { + "name": "tczyx_ome_05", + "axes": [ + "t", + "c", + "z", + "y", + "x" + ], + "units": { + "space": "micrometer", + "time": "second" + }, + "levels": [ + [ + "s0", + [ + 1, + 1, + 0.5, + 0.1, + 0.1 + ], + [ + 0, + 0, + 0.25, + 0.05, + 0.05 + ] + ], + [ + "s1", + [ + 1, + 1, + 0.5, + 0.2, + 0.2 + ], + [ + 0, + 0, + 0.25, + 0.1, + 0.1 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "second", + "", + "micrometer", + "micrometer", + "micrometer" + ], + "axis_types": [ + "time", + "channel", + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "t", + "type": "time", + "unit": "second" + }, + { + "name": "c", + "type": "channel" + }, + { + "name": "z", + "type": "space", + "unit": "micrometer" + }, + { + "name": "y", + "type": "space", + "unit": "micrometer" + }, + { + "name": "x", + "type": "space", + "unit": "micrometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 1.0, + 0.5, + 0.1, + 0.1 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 0.0, + 0.25, + 0.05, + 0.05 + ] + } + ] + }, + { + "path": "s1", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 1.0, + 0.5, + 0.2, + 0.2 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 0.0, + 0.25, + 0.1, + 0.1 + ] + } + ] + } + ] + } + ] + } + } + }, + { + "name": "non_whole_factor_ome_04", + "axes": [ + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 4, + 4, + 4 + ], + [ + 2, + 2, + 2 + ] + ], + [ + "s1", + [ + 10.48, + 10.48, + 10.48 + ], + [ + 3.7, + 3.7, + 3.7 + ] + ] + ], + "version": "0.4", + "resolved": { + "axis_units": [ + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "space", + "space", + "space" + ] + }, + "expect": { + "multiscales": [ + { + "version": "0.4", + "name": "", + "axes": [ + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 4.0, + 4.0, + 4.0 + ] + }, + { + "type": "translation", + "translation": [ + 2.0, + 2.0, + 2.0 + ] + } + ] + }, + { + "path": "s1", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 10.48, + 10.48, + 10.48 + ] + }, + { + "type": "translation", + "translation": [ + 3.7, + 3.7, + 3.7 + ] + } + ] + } + ] + } + ] + } + }, + { + "name": "non_whole_factor_ome_05", + "axes": [ + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 4, + 4, + 4 + ], + [ + 2, + 2, + 2 + ] + ], + [ + "s1", + [ + 10.48, + 10.48, + 10.48 + ], + [ + 3.7, + 3.7, + 3.7 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 4.0, + 4.0, + 4.0 + ] + }, + { + "type": "translation", + "translation": [ + 2.0, + 2.0, + 2.0 + ] + } + ] + }, + { + "path": "s1", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 10.48, + 10.48, + 10.48 + ] + }, + { + "type": "translation", + "translation": [ + 3.7, + 3.7, + 3.7 + ] + } + ] + } + ] + } + ] + } + } + }, + { + "name": "unknown_axis_untyped_ome_05", + "axes": [ + "q", + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 8, + 8, + 8 + ], + [ + 4, + 4, + 4 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "", + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + null, + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "q" + }, + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 8.0, + 8.0, + 8.0 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 4.0, + 4.0, + 4.0 + ] + } + ] + } + ] + } + ] + } + } + }, + { + "name": "custom_axis_types_ome_04", + "axes": [ + "b", + "z", + "y", + "x" + ], + "axis_types": [ + "time", + "space", + "space", + "space" + ], + "units": { + "space": "nanometer", + "time": "second" + }, + "levels": [ + [ + "s0", + [ + 0.5, + 8, + 8, + 8 + ], + [ + 0, + 4, + 4, + 4 + ] + ] + ], + "version": "0.4", + "resolved": { + "axis_units": [ + "second", + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "time", + "space", + "space", + "space" + ] + }, + "expect": { + "multiscales": [ + { + "version": "0.4", + "name": "", + "axes": [ + { + "name": "b", + "type": "time", + "unit": "second" + }, + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 0.5, + 8.0, + 8.0, + 8.0 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 4.0, + 4.0, + 4.0 + ] + } + ] + } + ] + } + ] + } + }, + { + "name": "time_axis_with_unit_ome_05", + "axes": [ + "t", + "z", + "y", + "x" + ], + "units": { + "space": "nanometer", + "time": "second" + }, + "levels": [ + [ + "s0", + [ + 2, + 8, + 8, + 8 + ], + [ + 0, + 4, + 4, + 4 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "second", + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "time", + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "t", + "type": "time", + "unit": "second" + }, + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 2.0, + 8.0, + 8.0, + 8.0 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 4.0, + 4.0, + 4.0 + ] + } + ] + } + ] + } + ] + } + } + }, + { + "name": "name_type_metadata_ome_04", + "axes": [ + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 8, + 8, + 8 + ], + [ + 4, + 4, + 4 + ] + ], + [ + "s1", + [ + 16, + 16, + 16 + ], + [ + 8, + 8, + 8 + ] + ] + ], + "version": "0.4", + "multiscale_name": "raw", + "multiscale_type": "mean", + "metadata": { + "method": "skimage.transform.downscale_local_mean", + "version": "0.21" + }, + "resolved": { + "axis_units": [ + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "space", + "space", + "space" + ] + }, + "expect": { + "multiscales": [ + { + "version": "0.4", + "name": "raw", + "axes": [ + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 8.0, + 8.0, + 8.0 + ] + }, + { + "type": "translation", + "translation": [ + 4.0, + 4.0, + 4.0 + ] + } + ] + }, + { + "path": "s1", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 16.0, + 16.0, + 16.0 + ] + }, + { + "type": "translation", + "translation": [ + 8.0, + 8.0, + 8.0 + ] + } + ] + } + ], + "type": "mean", + "metadata": { + "method": "skimage.transform.downscale_local_mean", + "version": "0.21" + } + } + ] + } + }, + { + "name": "zero_translation_ome_05", + "axes": [ + "z", + "y", + "x" + ], + "units": "nanometer", + "levels": [ + [ + "s0", + [ + 4, + 4, + 4 + ], + [ + 0, + 0, + 0 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 4.0, + 4.0, + 4.0 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 0.0, + 0.0 + ] + } + ] + } + ] + } + ] + } + } + }, + { + "name": "short_units_ome_04", + "axes": [ + "z", + "y", + "x" + ], + "units": "nm", + "levels": [ + [ + "s0", + [ + 8, + 8, + 8 + ], + [ + 4, + 4, + 4 + ] + ] + ], + "version": "0.4", + "resolved": { + "axis_units": [ + "nanometer", + "nanometer", + "nanometer" + ], + "axis_types": [ + "space", + "space", + "space" + ] + }, + "expect": { + "multiscales": [ + { + "version": "0.4", + "name": "", + "axes": [ + { + "name": "z", + "type": "space", + "unit": "nanometer" + }, + { + "name": "y", + "type": "space", + "unit": "nanometer" + }, + { + "name": "x", + "type": "space", + "unit": "nanometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 8.0, + 8.0, + 8.0 + ] + }, + { + "type": "translation", + "translation": [ + 4.0, + 4.0, + 4.0 + ] + } + ] + } + ] + } + ] + } + }, + { + "name": "short_units_by_kind_ome_05", + "axes": [ + "t", + "z", + "y", + "x" + ], + "units": { + "space": "um", + "time": "s" + }, + "levels": [ + [ + "s0", + [ + 1, + 0.1, + 0.1, + 0.1 + ], + [ + 0, + 0.05, + 0.05, + 0.05 + ] + ] + ], + "version": "0.5", + "resolved": { + "axis_units": [ + "second", + "micrometer", + "micrometer", + "micrometer" + ], + "axis_types": [ + "time", + "space", + "space", + "space" + ] + }, + "expect": { + "ome": { + "version": "0.5", + "multiscales": [ + { + "name": "", + "axes": [ + { + "name": "t", + "type": "time", + "unit": "second" + }, + { + "name": "z", + "type": "space", + "unit": "micrometer" + }, + { + "name": "y", + "type": "space", + "unit": "micrometer" + }, + { + "name": "x", + "type": "space", + "unit": "micrometer" + } + ], + "datasets": [ + { + "path": "s0", + "coordinateTransformations": [ + { + "type": "scale", + "scale": [ + 1.0, + 0.1, + 0.1, + 0.1 + ] + }, + { + "type": "translation", + "translation": [ + 0.0, + 0.05, + 0.05, + 0.05 + ] + } + ] + } + ] + } + ] + } + } + } + ] +} diff --git a/tests/test_write.py b/tests/test_write.py new file mode 100644 index 0000000..d19f383 --- /dev/null +++ b/tests/test_write.py @@ -0,0 +1,754 @@ +"""Writing: Convention 6 and Behaviours 10 to 16 (spec, section 5 "Writing (0.2)").""" + +from __future__ import annotations + +import json +import multiprocessing as mp +import pickle + +import numpy as np +import pytest +import tensorstore as ts + +import worldarray as wa +from conftest import ( + FIXTURES, + ome_multiscales, + pattern, + write_n5, + write_precomputed, + write_zarr2, + write_zarr2_group, + write_zarr3, + write_zarr3_group, +) + +WRITE_CASES = json.loads((FIXTURES / "write.json").read_text())["cases"] + + +# --- Behaviour 10: documents round-trip ------------------------------------------------- + + +def _fixture_call(case, units, axis_types): + return wa.ome_multiscales( + case["axes"], + units, + [tuple(lvl) for lvl in case["levels"]], + version=case["version"], + axis_types=axis_types, + name=case.get("multiscale_name", ""), + metadata=case.get("metadata"), + type=case.get("multiscale_type"), + ) + + +@pytest.mark.parametrize("case", WRITE_CASES, ids=[c["name"] for c in WRITE_CASES]) +def test_ome_multiscales_matches_the_fixture(case): + """The broadcast inputs and their resolved per-axis form give the pinned document.""" + doc = _fixture_call(case, case["units"], case.get("axis_types")) + assert json.loads(json.dumps(doc)) == case["expect"] + resolved = case["resolved"] + per_axis = _fixture_call(case, resolved["axis_units"], resolved["axis_types"]) + assert json.loads(json.dumps(per_axis)) == case["expect"] + + +def test_units_are_written_in_ome_names_and_unknown_ones_as_given(): + for spelling in ("nm", "nanometre", "Nanometers", "nanometer"): + assert wa.units.ome_name(spelling) == "nanometer" + assert wa.units.ome_name("um") == wa.units.ome_name("\u00b5m") == "micrometer" + assert wa.units.ome_name("s") == "second" + with pytest.raises(ValueError, match="unknown unit"): + wa.units.ome_name("pixel") + doc = wa.ome_multiscales(("z", "y", "x"), "pixel", [("s0", (1, 1, 1), (0.5,) * 3)]) + assert [a["unit"] for a in doc["ome"]["multiscales"][0]["axes"]] == ["pixel"] * 3 + doc = wa.ome_multiscales(("z", "y", "x"), "nm", [("s0", (1, 1, 1), (0.5,) * 3)], type="mean") + ms = doc["ome"]["multiscales"][0] + assert [a["unit"] for a in ms["axes"]] == ["nanometer"] * 3 and ms["type"] == "mean" + plain = wa.ome_multiscales(("z", "y", "x"), "nm", [("s0", (1, 1, 1), (0.5,) * 3)]) + assert "type" not in plain["ome"]["multiscales"][0] + + +def test_level_shape_absorbs_float_factor_noise(): + f = 0.3 / 0.1 # 2.9999999999999996, chunkmirage's ratio of float voxel sizes + assert wa.level_shape((9, 9, 9), f) == (3, 3, 3) + assert wa.level_shape((10, 9, 8), 3) == (4, 3, 3) + assert wa.level_shape((9, 9, 9), f, origin=3) == (2, 2, 2) + assert wa.level_shape((9,), (9 + 1e-9,)) == (1,) + + +@pytest.mark.parametrize("fmt", ["zarr2", "zarr3"]) +@pytest.mark.parametrize( + "axes,vs,tr,shape,units", + [ + (("z", "y", "x"), (8, 8, 8), (4, 4, 4), (8, 8, 8), "nanometer"), + (("c", "z", "y", "x"), (8, 8, 8), (4, 4, 4), (2, 8, 8, 8), "nanometer"), # spatial-length + ( + ("t", "c", "z", "y", "x"), + (1, 1, 4, 4, 4), + (0, 0, 2, 2, 2), + (3, 2, 8, 8, 8), + {"space": "nanometer", "time": "second"}, + ), + ], +) +def test_create_round_trips_placement(tmp_path, fmt, axes, vs, tr, shape, units): + path = tmp_path / f"vol.{fmt}" / "s0" + vol = wa.create( + str(path), shape, "uint8", vs, tr, axes=axes, chunks=shape, format=fmt, units=units + ) + assert vol.writable and vol.meta.format == fmt + meta = wa.metadata(str(path)) + kinds = wa.write.axis_kinds(axes) + assert meta.voxel_size == wa.write.per_axis(vs, kinds, 1.0, "vs") + assert meta.translation == wa.write.per_axis(tr, kinds, 0.0, "tr") + assert meta.axes == axes and meta.placed_by == "ome" + attrs = wa.detect(str(path.parent)).attrs + if fmt == "zarr2": + assert attrs["multiscales"][0]["version"] == "0.4" and "ome" not in attrs + axis_docs = attrs["multiscales"][0]["axes"] + else: + assert attrs["ome"]["version"] == "0.5" and "multiscales" not in attrs + axis_docs = attrs["ome"]["multiscales"][0]["axes"] + by_name = {a["name"]: a for a in axis_docs} + assert "unit" not in by_name.get("c", {}) # units on space and time axes only + if "t" in by_name: + assert by_name["t"] == {"name": "t", "type": "time", "unit": "second"} + + +def test_level_placement_and_shape_invert_the_reader_rule(): + # the reader's N5 case: base 4 nm with the corner at the origin (translation 2), factor 2 + assert wa.level_placement((4, 4, 4), (2, 2, 2), 2) == ((8.0, 8.0, 8.0), (4.0, 4.0, 4.0)) + assert wa.level_shape((16, 16, 15), 2) == (8, 8, 8) + # an origin moves the level by that many source voxels (mesh-n-bone's bug cannot be written) + vs, tr = wa.level_placement((4, 4, 4), (2, 2, 2), 2, origin=16) + assert vs == (8.0, 8.0, 8.0) and tr == (68.0, 68.0, 68.0) + assert wa.level_shape((48, 48, 48), 2, origin=16) == (16, 16, 16) + # per-axis factors, non-spatial axes untouched + assert wa.level_placement((1, 4, 4), (0, 2, 2), (1, 2, 2)) == ((1.0, 8.0, 8.0), (0.0, 4.0, 4.0)) + + +def test_level_placement_agrees_with_the_block_mean(tmp_path): + """A level added by factors sits exactly where Behaviour 4's block mean of the source + lands: reading the source over the level's box at the level's voxel size gives the mean + of each 2x2x2 block, with no half-voxel shift.""" + data = pattern((16, 16, 16), "float32") + group = tmp_path / "pyr.zarr" + s0 = wa.create( + str(group / "s0"), data.shape, data.dtype, (4, 4, 4), (2, 2, 2), chunks=(8, 8, 8) + ) + s0.write_index((0, 0, 0), data.shape, data) + s1 = wa.open_pyramid(str(group), mode="r+").add_level("s1", factors=2) + assert (s1.shape, s1.voxel_size, s1.translation) == ( + (8, 8, 8), + (8.0, 8.0, 8.0), + (4.0, 4.0, 4.0), + ) + mean = wa.open(str(group / "s0")).read(s1.box, voxel_size=s1.voxel_size, method="linear") + expected = data.reshape(8, 2, 8, 2, 8, 2).mean(axis=(1, 3, 5)) + np.testing.assert_allclose(mean, expected) + s1.write_index((0, 0, 0), s1.shape, mean.astype(np.float32)) + np.testing.assert_allclose(wa.open(str(group / "s1")).read_index((0, 0, 0), s1.shape), expected) + + +# --- Behaviour 11: create --------------------------------------------------------------- + + +def test_create_follows_the_group_format_and_refuses_a_contradiction(tmp_path): + group = tmp_path / "v2.zarr" + write_zarr2_group(group, ome_multiscales("zyx", [("s0", (4, 4, 4), (2, 2, 2))])) + write_zarr2(group / "s0", pattern((8, 8, 8))) + s1 = wa.create(str(group / "s1"), (4, 4, 4), "uint16", (8, 8, 8), (4, 4, 4), chunks=(4, 4, 4)) + assert s1.meta.format == "zarr2" + attrs = wa.detect(str(group)).attrs + assert [d["path"] for d in attrs["multiscales"][0]["datasets"]] == ["s0", "s1"] + with pytest.raises(ValueError, match="mix zarr versions"): + wa.create( + str(group / "s2"), + (2, 2, 2), + "uint16", + (16, 16, 16), + (8, 8, 8), + chunks=(2, 2, 2), + format="zarr3", + ) + # no group document: the format argument decides, zarr v3 by default + fresh = wa.create( + str(tmp_path / "new.zarr" / "s0"), + (4, 4, 4), + "uint8", + (1, 1, 1), + (0.5, 0.5, 0.5), + chunks=(4, 4, 4), + ) + assert fresh.meta.format == "zarr3" and wa.detect(str(tmp_path / "new.zarr")).format == "zarr3" + + +def test_create_requires_chunks_and_whole_shards(tmp_path): + with pytest.raises(ValueError, match="chunks is required"): + wa.create(str(tmp_path / "a.zarr/s0"), (8, 8, 8), "uint8", (1, 1, 1), (0.5, 0.5, 0.5)) + with pytest.raises(ValueError, match="whole number of chunks"): + wa.create( + str(tmp_path / "b.zarr/s0"), + (8, 8, 8), + "uint8", + (1, 1, 1), + (0.5,) * 3, + chunks=(4, 4, 4), + shards=(6, 6, 6), + ) + with pytest.raises(ValueError, match="zarr v3 only"): + wa.create( + str(tmp_path / "c.zarr/s0"), + (8, 8, 8), + "uint8", + (1, 1, 1), + (0.5,) * 3, + chunks=(4, 4, 4), + shards=(8, 8, 8), + format="zarr2", + ) + + +def test_codecs_default_to_zstd_and_follow_the_dialect(tmp_path): + v3 = wa.create( + str(tmp_path / "v3.zarr/s0"), (8, 8, 8), "uint16", (1, 1, 1), (0.5,) * 3, chunks=(8, 8, 8) + ) + codecs = wa.detect(v3.level_path).docs["zarr.json"]["codecs"] + assert [c["name"] for c in codecs] == ["bytes", "zstd"] and codecs[1]["configuration"][ + "level" + ] == 3 + v2 = wa.create( + str(tmp_path / "v2.zarr/s0"), + (8, 8, 8), + "uint16", + (1, 1, 1), + (0.5,) * 3, + chunks=(8, 8, 8), + format="zarr2", + ) + assert wa.detect(v2.level_path).docs[".zarray"]["compressor"] == {"id": "zstd", "level": 3} + raw = wa.create( + str(tmp_path / "raw.zarr/s0"), + (8, 8, 8), + "uint16", + (1, 1, 1), + (0.5,) * 3, + chunks=(8, 8, 8), + codecs="none", + ) + assert [c["name"] for c in wa.detect(raw.level_path).docs["zarr.json"]["codecs"]] == ["bytes"] + with pytest.raises(ValueError, match="dialect"): + wa.create( + str(tmp_path / "x.zarr/s0"), + (8, 8, 8), + "uint16", + (1, 1, 1), + (0.5,) * 3, + chunks=(8, 8, 8), + codecs={"id": "zstd", "level": 1}, + ) + + +def test_create_never_overwrites_and_adopts_only_an_identical_array(tmp_path): + path = str(tmp_path / "vol.zarr/s0") + args = ((8, 8, 8), "uint16", (4, 4, 4), (2, 2, 2)) + first = wa.create(path, *args, chunks=(4, 4, 4)) + first.write_index((0, 0, 0), (4, 4, 4), np.full((4, 4, 4), 9, np.uint16)) + with pytest.raises(FileExistsError): + wa.create(path, *args, chunks=(4, 4, 4)) + again = wa.create(path, *args, chunks=(4, 4, 4), exists="adopt") + assert ( + again.writable and int(again.read_index((0, 0, 0), (4, 4, 4)).sum()) == 9 * 64 + ) # the data survived + with pytest.raises(ValueError) as e: + wa.create(path, (8, 8, 16), "uint8", (4, 4, 4), (2, 2, 2), chunks=(4, 4, 4), exists="adopt") + message = str(e.value) + assert ( + "shape: found (8, 8, 8), wanted (8, 8, 16)" in message + and "dtype: found uint16, wanted uint8" in message + ) + with pytest.raises(ValueError, match="translation: found"): + wa.create(path, *args[:3], (3, 2, 2), chunks=(4, 4, 4), exists="adopt") + + +def test_create_keeps_everyone_elses_attributes(tmp_path): + group = tmp_path / "shared.zarr" + write_zarr3_group( + group, + { + "fs": [4, 4, 4], + "registry": {"owner": "someone"}, + **ome_multiscales("zyx", [("s0", (4, 4, 4), (2, 2, 2))], version="0.5"), + }, + ) + before = wa.detect(str(group)).attrs + vol = wa.create( + str(group / "s1"), + (4, 4, 4), + "uint16", + (8, 8, 8), + (4, 4, 4), + chunks=(4, 4, 4), + attrs={"voxel_size": [8, 8, 8], "note": "legacy keys pass through"}, + group_attrs={"built_by": "test"}, + ) + after = wa.detect(str(group)).attrs + assert after["fs"] == before["fs"] and after["registry"] == before["registry"] + assert after["built_by"] == "test" + assert [d["path"] for d in after["ome"]["multiscales"][0]["datasets"]] == ["s0", "s1"] + assert ( + after["ome"]["multiscales"][0]["datasets"][0] + == before["ome"]["multiscales"][0]["datasets"][0] + ) + assert vol.meta.attrs["note"] == "legacy keys pass through" and vol.meta.attrs[ + "voxel_size" + ] == [8, 8, 8] + # the array's own attributes agree with the OME placement, so nothing warns or raises + assert wa.metadata(vol.level_path).voxel_size == (8.0, 8.0, 8.0) + + +def test_fill_only_units_are_not_stored_unless_asked(tmp_path): + vol = wa.create( + str(tmp_path / "a.zarr/s0"), (8, 8, 8), "uint8", (1, 1, 1), (0.5,) * 3, chunks=(4, 4, 4) + ) + vol.write_index((0, 0, 0), (4, 4, 4), np.zeros((4, 4, 4), np.uint8)) + vol.write_index((4, 4, 4), (8, 8, 8), np.ones((4, 4, 4), np.uint8)) + assert vol.list_chunks() == [(1, 1, 1)] + kept = wa.create( + str(tmp_path / "b.zarr/s0"), + (8, 8, 8), + "uint8", + (1, 1, 1), + (0.5,) * 3, + chunks=(4, 4, 4), + store_fill=True, + ) + kept.write_index((0, 0, 0), (4, 4, 4), np.zeros((4, 4, 4), np.uint8)) + assert kept.list_chunks() == [(0, 0, 0)] + + +def test_a_lost_creation_race_is_adopted(tmp_path, monkeypatch): + """The probe says absent, the store says present: another process got there first.""" + path = str(tmp_path / "race.zarr/s0") + wa.create(path, (8, 8, 8), "uint16", (4, 4, 4), (2, 2, 2), chunks=(4, 4, 4)) + real_detect, calls = wa.write.detect, {"n": 0} + + def forgetful_detect(p, *, missing_ok=False): + if p.endswith("/s0") and calls["n"] == 0: + calls["n"] += 1 + return None # the probe before the attempt misses it + return real_detect(p, missing_ok=missing_ok) + + monkeypatch.setattr(wa.write, "detect", forgetful_detect) + monkeypatch.setattr(wa.write, "RACE_DELAY", 0.0) + vol = wa.create( + path, (8, 8, 8), "uint16", (4, 4, 4), (2, 2, 2), chunks=(4, 4, 4), exists="adopt" + ) + assert vol.writable and vol.shape == (8, 8, 8) + calls["n"] = 0 + with pytest.raises(FileExistsError, match="another process"): + wa.create(path, (8, 8, 8), "uint16", (4, 4, 4), (2, 2, 2), chunks=(4, 4, 4)) + + +def test_open_for_write_and_pickling(tmp_path): + path = str(tmp_path / "vol.zarr/s0") + wa.create(path, (8, 8, 8), "uint16", (4, 4, 4), (2, 2, 2), chunks=(4, 4, 4)) + ro = wa.open(path) + assert not ro.writable + with pytest.raises(ValueError, match="read-only"): + ro.write_index((0, 0, 0), (4, 4, 4), np.zeros((4, 4, 4), np.uint16)) + rw = wa.open(path, mode="r+") + assert rw.writable and rw.voxel_size == (4.0, 4.0, 4.0) + again = pickle.loads(pickle.dumps(rw)) + assert again.writable and again.translation == rw.translation and again.path == rw.path + again.write_index((4, 4, 4), (8, 8, 8), np.full((4, 4, 4), 3, np.uint16)) + assert int(ro.read_index((4, 4, 4), (8, 8, 8)).sum()) == 3 * 64 + + +# --- Behaviour 13: aligned writes ------------------------------------------------------- + + +def test_write_index_alignment_rules(tmp_path): + vol = wa.create( + str(tmp_path / "vol.zarr/s0"), (10, 8, 8), "uint16", (1, 1, 1), (0.5,) * 3, chunks=(4, 4, 4) + ) + vol.write_index((0, 0, 0), (4, 8, 8), np.ones((4, 8, 8), np.uint16)) # whole units + vol.write_index((8, 0, 0), (10, 4, 4), np.ones((2, 4, 4), np.uint16)) # the clipped last unit + with pytest.raises(ValueError, match=r"axis 1 \(y\) is not whole write units of 4"): + vol.write_index((0, 2, 0), (4, 6, 4), np.ones((4, 4, 4), np.uint16)) + with pytest.raises(ValueError, match="nothing is cast"): + vol.write_index((0, 0, 0), (4, 4, 4), np.ones((4, 4, 4), np.uint8)) + with pytest.raises(ValueError, match="outside the array"): + vol.write_index((8, 0, 0), (12, 4, 4), np.ones((4, 4, 4), np.uint16)) + with pytest.raises(ValueError, match="data has shape"): + vol.write_index((0, 0, 0), (4, 4, 4), np.ones((4, 4, 2), np.uint16)) + vol.write_index( + (1, 1, 1), (3, 3, 3), np.full((2, 2, 2), 5, np.uint16), partial=True + ) # read-modify-write + got = vol.read_index((0, 0, 0), (4, 4, 4)) + assert got[1, 1, 1] == 5 and got[0, 0, 0] == 1 and got[3, 3, 3] == 1 + with pytest.raises(ValueError, match="no resampling on write"): + wa.open( + vol.level_path, mode="r+", voxel_size=(2, 2, 2), on_mismatch="resample" + ).write_index((0, 0, 0), (2, 2, 2), np.zeros((2, 2, 2), np.uint16)) + + +def _write_unit(args): + path, index, value = args + vol = wa.open(path, mode="r+") + begin = tuple(i * c for i, c in zip(index, vol.write_shape)) + end = tuple(b + c for b, c in zip(begin, vol.write_shape)) + vol.write_index(begin, end, np.full(vol.write_shape, value, vol.dtype)) + return index + + +def test_two_processes_write_disjoint_units(tmp_path): + path = str(tmp_path / "vol.zarr/s0") + wa.create(path, (8, 8, 8), "uint16", (1, 1, 1), (0.5,) * 3, chunks=(4, 4, 4)) + jobs = [(path, idx, 1 + i) for i, idx in enumerate(np.ndindex(2, 2, 2))] + with mp.get_context("spawn").Pool(2) as pool: + done = pool.map(_write_unit, jobs) + assert sorted(done) == sorted(j[1] for j in jobs) + got = wa.open(path).read_index((0, 0, 0), (8, 8, 8)) + for i, idx in enumerate(np.ndindex(2, 2, 2)): + sl = tuple(slice(4 * k, 4 * k + 4) for k in idx) + assert np.all(got[sl] == 1 + i), idx + + +def test_write_by_world_box_uses_the_centre_rule(tmp_path): + vol = wa.create( + str(tmp_path / "vol.zarr/s0"), (8, 8, 8), "uint16", (4, 4, 4), (2, 2, 2), chunks=(4, 4, 4) + ) + vol.write( + ((0.0, 0.0, 0.0), (16.0, 16.0, 16.0)), np.full((4, 4, 4), 2, np.uint16) + ) # voxels 0..3 + assert int(vol.read_index((0, 0, 0), (4, 4, 4)).sum()) == 2 * 64 + with pytest.raises(ValueError, match="outside the array"): + vol.write(((16.0, 16.0, 16.0), (40.0, 40.0, 40.0)), np.zeros((6, 6, 6), np.uint16)) + + +# --- Behaviour 14: add_level ------------------------------------------------------------ + + +def test_add_level_places_sizes_orders_and_copies(tmp_path): + group = tmp_path / "pyr.zarr" + s0 = wa.create( + str(group / "s0"), + (16, 16, 16), + "uint32", + (4, 4, 4), + (2, 2, 2), + chunks=(4, 4, 4), + shards=(8, 8, 8), + fill_value=7, + ) + pyr = wa.open_pyramid(str(group), mode="r+") + s2 = pyr.add_level("s2", factors=4) + s1 = pyr.add_level("s1", factors=2) + assert pyr.names == ["s0", "s1", "s2"] # inserted in order of voxel size, not of creation + assert (s1.shape, s1.voxel_size, s1.translation) == ( + (8, 8, 8), + (8.0, 8.0, 8.0), + (4.0, 4.0, 4.0), + ) + assert (s2.shape, s2.voxel_size, s2.translation) == ( + (4, 4, 4), + (16.0, 16.0, 16.0), + (8.0, 8.0, 8.0), + ) + assert ( + s1.dtype == s0.dtype + and s1.meta.chunk_shape == (4, 4, 4) + and s1.meta.shard_shape == (8, 8, 8) + ) + assert s1.meta.fill_value == 7 and s1.list_chunks() == [] # metadata plus an empty array + codecs = wa.detect(s1.level_path).docs["zarr.json"]["codecs"] + assert codecs[0]["name"] == "sharding_indexed" + # an explicit placement, a level from another level, and replacing a same-named entry + s3 = pyr.add_level( + "s3", voxel_size=(32, 32, 32), translation=(16, 16, 16), shape=(2, 2, 2), from_level="s2" + ) + assert s3.translation == (16.0, 16.0, 16.0) + again = pyr.add_level("s3", factors=2, from_level="s2", exists="adopt") + assert again.translation == (16.0, 16.0, 16.0) and pyr.names.count("s3") == 1 + with pytest.raises(ValueError, match="mode='r\\+'"): + wa.open_pyramid(str(group)).add_level("s4", factors=2) + + +# --- Behaviours 15 and 16: chunk keys and write units ----------------------------------- + + +@pytest.mark.parametrize( + "fmt,kwargs,index,key", + [ + ("zarr2", {}, (1, 2, 3), "1.2.3"), + ("zarr2", {"separator": "/"}, (1, 2, 3), "1/2/3"), + ("zarr3", {}, (1, 2, 3), "c/1/2/3"), + ("zarr3", {"separator": "."}, (1, 2, 3), "c.1.2.3"), + ("zarr3", {"encoding": "v2"}, (1, 2, 3), "1.2.3"), + ("n5", {}, (1, 2, 3), "3/2/1"), + ( + "precomputed", + {"chunk_shape": (4, 8, 16), "shape": (10, 16, 32), "origin": (5, 0, 7)}, + (2, 1, 1), + "23-39_8-16_13-15", + ), + ], +) +def test_chunk_key_round_trips(fmt, kwargs, index, key): + assert wa.chunk_key(index, fmt, **kwargs) == key + assert wa.parse_chunk_key(key, fmt, **kwargs) == index + assert wa.parse_chunk_key("zarr.json", fmt, **kwargs) is None + assert wa.parse_chunk_key(".zarray", fmt, **kwargs) is None + + +def test_array_chunk_keys_and_list_chunks_agree(tmp_path): + v2 = wa.create( + str(tmp_path / "v2.zarr/s0"), + (8, 8, 8), + "uint8", + (1, 1, 1), + (0.5,) * 3, + chunks=(4, 4, 4), + format="zarr2", + separator="/", + ) + v2.write_index((4, 0, 4), (8, 4, 8), np.ones((4, 4, 4), np.uint8)) + assert v2.chunk_key((1, 0, 1)) == "1/0/1" and v2.list_chunks() == [(1, 0, 1)] + n5 = tmp_path / "vol.n5" + write_n5(n5 / "s0", pattern((8, 8, 8)), attrs={"resolution": [1, 1, 1]}, chunks=(4, 4, 4)) + arr = wa.open(str(n5 / "s0")) + assert arr.chunk_key((1, 0, 1)) == "1/0/1" and arr.parse_chunk_key("attributes.json") is None + assert len(arr.list_chunks()) == 8 + pre = tmp_path / "pc" + write_precomputed( + pre, [{"data": pattern((8, 8, 8)), "resolution": [4, 4, 4], "voxel_offset": [5, 5, 5]}] + ) + vol = wa.open(str(pre / "s0")) + assert vol.chunk_key((0, 0, 0)) == "5-13_5-13_5-13" and vol.list_chunks() == [(0, 0, 0)] + + +def test_write_shape_is_the_shard_when_sharded(tmp_path): + vol = wa.create( + str(tmp_path / "s.zarr/s0"), + (16, 16, 16), + "uint8", + (1, 1, 1), + (0.5,) * 3, + chunks=(4, 4, 4), + shards=(8, 8, 8), + ) + assert ( + vol.meta.chunk_shape == (4, 4, 4) + and vol.meta.shard_shape == (8, 8, 8) + and vol.write_shape == (8, 8, 8) + ) + assert wa.metadata(vol.level_path, open=False).shard_shape == (8, 8, 8) + assert list(vol.chunks_covering((0, 0, 0), (8, 8, 8), write=True)) == [(0, 0, 0)] + assert len(list(vol.chunks_covering((0, 0, 0), (8, 8, 8)))) == 8 + with pytest.raises(ValueError, match="not whole write units of 8"): + vol.write_index((0, 0, 0), (4, 4, 4), np.ones((4, 4, 4), np.uint8)) + vol.write_index((8, 8, 8), (16, 16, 16), np.ones((8, 8, 8), np.uint8)) + assert vol.chunk_key((1, 1, 1)) == "c/1/1/1" and vol.list_chunks() == [(1, 1, 1)] + plain = wa.create( + str(tmp_path / "p.zarr/s0"), (8, 8, 8), "uint8", (1, 1, 1), (0.5,) * 3, chunks=(4, 4, 4) + ) + assert plain.meta.shard_shape is None and plain.write_shape == (4, 4, 4) + assert ts.open( + {"driver": "zarr3", "kvstore": {"driver": "file", "path": vol.level_path + "/"}} + ).result().shape == (16, 16, 16) + + +# --- Behaviours 11 and 12 as the consumer sessions found them (2026-10-22 review) -------- + + +def _group_doc(path, fmt): + return json.loads((path / ("zarr.json" if fmt == "zarr3" else ".zgroup")).read_text()) + + +def test_create_writes_missing_ancestor_groups_up_to_the_zarr_root(tmp_path): + """cellmap-analyze's hierarchy: every group from the container down has a document.""" + root = tmp_path / "a.zarr" + wa.create(str(root / "sub/mito/s0"), (8, 8, 8), "uint8", (8, 8, 8), (4, 4, 4), chunks=(8, 8, 8)) + for group in (root, root / "sub", root / "sub/mito"): + doc = _group_doc(group, "zarr3") + assert doc["zarr_format"] == 3 and doc["node_type"] == "group" + assert wa.detect(str(group)).kind == "group" + assert "attributes" not in _group_doc(root, "zarr3") # bare: nothing invented above the level + assert "ome" in _group_doc(root / "sub/mito", "zarr3")["attributes"] + assert not (tmp_path / "zarr.json").exists() and not (tmp_path / ".zgroup").exists() + assert wa.open_pyramid(str(root / "sub/mito")).names == ["s0"] + # zarr v2 asked for: .zgroup at every missing ancestor + v2 = tmp_path / "b.zarr" + wa.create( + str(v2 / "mito/s0"), + (8, 8, 8), + "uint8", + (8, 8, 8), + (4, 4, 4), + chunks=(8, 8, 8), + format="zarr2", + ) + assert _group_doc(v2, "zarr2") == {"zarr_format": 2} and _group_doc(v2 / "mito", "zarr2") == { + "zarr_format": 2 + } + assert "multiscales" in json.loads((v2 / "mito/.zattrs").read_text()) + assert not (v2 / ".zattrs").exists() + + +def test_create_follows_the_nearest_existing_ancestors_format(tmp_path): + """A v2 container gets v2 levels and OME 0.4 even when the level's own parent is absent.""" + root = tmp_path / "c.zarr" + write_zarr2_group(root, {"owner": "cellmap-analyze"}) + vol = wa.create( + str(root / "mito/s0"), (8, 8, 8), "uint8", (8, 8, 8), (4, 4, 4), chunks=(8, 8, 8) + ) + assert vol.meta.format == "zarr2" + assert (root / "mito/.zgroup").exists() and (root / "mito/s0/.zarray").exists() + assert "multiscales" in json.loads((root / "mito/.zattrs").read_text()) + assert json.loads((root / ".zattrs").read_text()) == {"owner": "cellmap-analyze"} # untouched + with pytest.raises(ValueError, match="mix zarr versions"): + wa.create( + str(root / "er/s0"), + (8, 8, 8), + "uint8", + (8, 8, 8), + (4, 4, 4), + chunks=(8, 8, 8), + format="zarr3", + ) + assert not (root / "er").exists() + # an array in the way raises, and no document is written above a container with no .zarr name + with pytest.raises(ValueError, match="is an array"): + wa.create( + str(root / "mito/s0/deeper/s0"), + (8, 8, 8), + "uint8", + (8, 8, 8), + (4, 4, 4), + chunks=(8, 8, 8), + ) + wa.create( + str(tmp_path / "out/sub/s0"), (8, 8, 8), "uint8", (8, 8, 8), (4, 4, 4), chunks=(8, 8, 8) + ) + assert (tmp_path / "out/sub/zarr.json").exists() and not (tmp_path / "out/zarr.json").exists() + + +def test_adopting_an_identical_array_writes_no_group_document(tmp_path, monkeypatch): + """registration-station's 20 block jobs each call create on one output: only the first + creation writes; the adopts read, compare and leave the document alone.""" + path = str(tmp_path / "shared.zarr/0") + wa.create( + path, + (8, 8, 8), + "float32", + (0.5, 0.5, 0.5), + (0.25,) * 3, + chunks=(4, 4, 4), + group_attrs={"fs": 2}, + ) + real, writes = wa.write.write_group_attrs, [] + + def counting(loc, fmt, attrs, raw): + writes.append(attrs) + return real(loc, fmt, attrs, raw) + + monkeypatch.setattr(wa.write, "write_group_attrs", counting) + for _ in range(3): + vol = wa.create( + path, + (8, 8, 8), + "float32", + (0.5, 0.5, 0.5), + (0.25 + 1e-9,) * 3, + chunks=(4, 4, 4), + exists="adopt", + group_attrs={"fs": 2}, + ) + assert vol.writable + assert writes == [] + # a different translation for the same level is a replacement, written once + wa.ensure_group_entry( + str(tmp_path / "shared.zarr"), "0", (0.5, 0.5, 0.5), (1.25,) * 3, ("z", "y", "x") + ) + assert len(writes) == 1 and wa.open(path).translation == (1.25, 1.25, 1.25) + + +def test_a_group_entry_lost_to_a_concurrent_writer_is_written_again(tmp_path, monkeypatch): + group = tmp_path / "parts.zarr" + wa.create(str(group / "c000"), (4, 4, 4), "float32", (1, 1, 1), (0.5,) * 3, chunks=(4, 4, 4)) + before = json.loads((group / "zarr.json").read_text()) + real, calls = wa.write.write_group_attrs, {"n": 0} + + def racing(loc, fmt, attrs, raw): + calls["n"] += 1 + real(loc, fmt, attrs, raw) + if calls["n"] == 1: # another process's older document lands after ours + (group / "zarr.json").write_text(json.dumps(before)) + + monkeypatch.setattr(wa.write, "write_group_attrs", racing) + wa.ensure_group_entry(str(group), "c001", (1, 1, 1), (4.5, 0.5, 0.5), ("z", "y", "x")) + assert calls["n"] == 2 and wa.open_pyramid(str(group)).names == ["c000", "c001"] + + def always_lost(loc, fmt, attrs, raw): + real(loc, fmt, attrs, raw) + (group / "zarr.json").write_text(json.dumps(before)) + + monkeypatch.setattr(wa.write, "write_group_attrs", always_lost) + monkeypatch.setattr(wa.write, "RACE_ATTEMPTS", 2) + with pytest.raises(RuntimeError, match="lost 2 times"): + wa.ensure_group_entry(str(group), "c002", (1, 1, 1), (8.5, 0.5, 0.5), ("z", "y", "x")) + + +def test_ensure_group_entry_lists_an_array_the_package_did_not_create(tmp_path, monkeypatch): + """mesh-n-bone's symlinked s0, cellmap-analyze's level listed after it is filled.""" + group = tmp_path / "pyr.zarr" + write_zarr3(group / "s0", pattern((8, 8, 8), np.uint8), chunks=(4, 4, 4)) + assert wa.detect(str(group), missing_ok=True) is not None # a group by its s0, with no document + wa.ensure_group_entry(str(group), "s0", (8, 8, 8), (4, 4, 4), ("z", "y", "x")) + assert wa.open_pyramid(str(group)).names == ["s0"] + assert wa.open(str(group / "s0")).translation == (4.0, 4.0, 4.0) + assert wa.open(str(group / "s0")).meta.units == ("nanometer",) * 3 + writes = [] + monkeypatch.setattr(wa.write, "write_group_attrs", lambda *a: writes.append(a)) + wa.ensure_group_entry(str(group), "s0", (8, 8, 8), (4, 4, 4), ("z", "y", "x")) + assert writes == [] + with pytest.raises(ValueError, match="mix zarr versions"): + wa.ensure_group_entry( + str(group), "s1", (16, 16, 16), (8, 8, 8), ("z", "y", "x"), format="zarr2" + ) + + +def test_a_pickled_handle_keeps_its_context(tmp_path): + path = str(tmp_path / "vol.zarr/s0") + wa.create(path, (8, 8, 8), "uint16", (4, 4, 4), (2, 2, 2), chunks=(4, 4, 4)) + capped = wa.open(path, mode="r+", context=wa.capped_context(concurrency=2)) + again = pickle.loads(pickle.dumps(capped)) + assert again.context.spec.to_json() == capped.context.spec.to_json() + assert again.context is not wa.shared_context(0) + assert ( + pickle.loads(pickle.dumps(capped)).context is again.context + ) # one rebuilt context per spec + shared = pickle.loads(pickle.dumps(wa.open(path))) + assert shared.context is wa.shared_context(0) + + +def test_codecs_of_a_path_and_the_fill_write_that_removes_a_key(tmp_path): + path = str(tmp_path / "vol.zarr/s0") + vol = wa.create( + path, + (8, 8, 8), + "uint8", + (4, 4, 4), + (2, 2, 2), + chunks=(4, 4, 4), + format="zarr2", + codecs={"id": "blosc", "cname": "lz4", "clevel": 5, "shuffle": 1}, + ) + codecs, shards = wa.write.codecs_of(path) + assert codecs["id"] == "blosc" and codecs["cname"] == "lz4" and shards is None + vol.write_index((0, 0, 0), (4, 4, 4), np.ones((4, 4, 4), np.uint8)) + assert vol.list_chunks() == [(0, 0, 0)] + vol.write_index((0, 0, 0), (4, 4, 4), np.zeros((4, 4, 4), np.uint8)) + assert vol.list_chunks() == [] # the one key a write may remove: its own, now all fill + assert int(vol.read_index((0, 0, 0), (4, 4, 4)).sum()) == 0 diff --git a/uv.lock b/uv.lock index 7ef0f56..c73739b 100644 --- a/uv.lock +++ b/uv.lock @@ -1247,7 +1247,7 @@ wheels = [ [[package]] name = "worldarray" -version = "0.1.1" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },