worldarray 0.1.0__py3-none-any.whl

worldarray/context.py

@@ -0,0 +1,53 @@
+"""Tensorstore contexts: one shared per cache budget, or a capped one per call.
+
+``shared_context(cache_bytes)`` is one ``tensorstore.Context`` per budget for the whole
+process, so the levels of a pyramid and repeat opens of an array draw on one cache pool.
+The default budget is 0: tensorstore's decoded-chunk cache is off, because every consumer
+either reads each region once or caches decoded chunks itself.
+
+``capped_context`` is mesh-n-bone's worker-process context: every tensorstore concurrency
+resource limited, so dozens of worker processes on one host do not multiply into thousands
+of threads. It returns a **new** context each call, as those workers inline the spec into
+every open today; pass the same object twice to share pools (spec, Behaviour 7).
+"""
+
+from __future__ import annotations
+
+import threading
+
+import tensorstore as ts
+
+_shared: dict[int, ts.Context] = {}
+_lock = threading.Lock()
+
+
+def shared_context(cache_bytes: int = 0) -> ts.Context:
+    with _lock:
+        ctx = _shared.get(int(cache_bytes))
+        if ctx is None:
+            ctx = ts.Context({"cache_pool": {"total_bytes_limit": int(cache_bytes)}})
+            _shared[int(cache_bytes)] = ctx
+        return ctx
+
+
+def capped_context_spec(cache_bytes: int = 0, concurrency: int = 1, network: bool = True) -> dict:
+    """The spec dict behind ``capped_context``, for callers building their own opens.
+    ``network=False`` caps only decoding and local file I/O, leaving tensorstore's request
+    pools alone (cellmap-analyze's choice for remote reads)."""
+    limit = {"limit": int(concurrency)}
+    spec = {
+        "data_copy_concurrency": limit,
+        "file_io_concurrency": limit,
+        "cache_pool": {"total_bytes_limit": int(cache_bytes)},
+    }
+    if network:
+        spec.update(
+            gcs_request_concurrency=limit,
+            s3_request_concurrency=limit,
+            http_request_concurrency=limit,
+        )
+    return spec
+
+
+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))

worldarray/geometry.py

@@ -0,0 +1,173 @@
+"""Boxes, and the one rule that maps a world box to voxel indices.
+
+A world box ``[begin, end)`` is half-open per axis. A voxel is read iff its **centre**
+lies in the box: with voxel size ``s`` and translation ``t`` (the centre of voxel 0), the
+voxels are ``ceil((begin - t) / s) <= i < ceil((end - t) / s)``, after float noise is
+snapped away. A box aligned to voxel corners (every blockwise block) and a box aligned to
+voxel centres (mesh-n-bone's) both map to whole voxels with nothing rounded
+(spec, Behaviour 2).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+SNAP_TOL = 1e-6  # of a voxel: a quotient this close to a whole number is that whole number
+
+
+def snap_integral(values, tol: float = SNAP_TOL) -> np.ndarray:
+    arr = np.asarray(values, dtype=float)
+    rounded = np.round(arr)
+    return np.where(np.abs(arr - rounded) <= tol, rounded, arr)
+
+
+@dataclass(frozen=True)
+class Box:
+    """``[begin, end)`` per axis, in world units (floats) or voxels (whole floats)."""
+
+    begin: tuple[float, ...]
+    end: tuple[float, ...]
+
+    def __post_init__(self) -> None:
+        if len(self.begin) != len(self.end):
+            raise ValueError(f"begin and end differ in length: {self.begin} vs {self.end}")
+        object.__setattr__(self, "begin", tuple(float(b) for b in self.begin))
+        object.__setattr__(self, "end", tuple(float(e) for e in self.end))
+
+    @property
+    def shape(self) -> tuple[float, ...]:
+        return tuple(e - b for b, e in zip(self.begin, self.end))
+
+    @property
+    def ndim(self) -> int:
+        return len(self.begin)
+
+    @property
+    def centre(self) -> tuple[float, ...]:
+        return tuple((b + e) / 2 for b, e in zip(self.begin, self.end))
+
+    @property
+    def empty(self) -> bool:
+        return any(e <= b for b, e in zip(self.begin, self.end))
+
+    def intersect(self, other) -> Box:
+        """The overlap of two boxes; a box of zero extent where there is none."""
+        other = Box.coerce(other)
+        begin = tuple(max(a, b) for a, b in zip(self.begin, other.begin))
+        end = tuple(max(bg, min(a, b)) for bg, a, b in zip(begin, self.end, other.end))
+        return Box(begin, end)
+
+    def union(self, other) -> Box:
+        """The smallest box holding both."""
+        other = Box.coerce(other)
+        return Box(
+            tuple(min(a, b) for a, b in zip(self.begin, other.begin)),
+            tuple(max(a, b) for a, b in zip(self.end, other.end)),
+        )
+
+    def contains(self, other) -> bool:
+        """Whether a point (one coordinate per axis) lies in the box, or another box lies
+        within it. Half-open, so a point on the far edge is outside."""
+        if isinstance(other, Box) or hasattr(other, "begin") or hasattr(other, "offset"):
+            o = Box.coerce(other)
+            return all(
+                b <= ob and oe <= e for b, e, ob, oe in zip(self.begin, self.end, o.begin, o.end)
+            )
+        return all(b <= p < e for b, e, p in zip(self.begin, self.end, other))
+
+    def grow(self, before, after=None) -> Box:
+        """Extended by ``before`` on the low side and ``after`` (default ``before``) on the
+        high side, a scalar or one value per axis; negative values shrink."""
+        before = self._per_axis(before)
+        after = self._per_axis(before if after is None else after)
+        return Box(
+            tuple(b - x for b, x in zip(self.begin, before)),
+            tuple(e + x for e, x in zip(self.end, after)),
+        )
+
+    def shift(self, by) -> Box:
+        by = self._per_axis(by)
+        return Box(
+            tuple(b + x for b, x in zip(self.begin, by)), tuple(e + x for e, x in zip(self.end, by))
+        )
+
+    def snap(self, voxel_size, translation, mode: str = "grow") -> Box:
+        """Aligned to the voxel corners of a grid whose voxel 0 is centred at
+        ``translation``: ``grow`` rounds outward, ``shrink`` inward, ``closest`` to the
+        nearest corner. After ``grow`` the box reads exactly the voxels it touches."""
+        vs = np.asarray(voxel_size, dtype=float)
+        corner0 = np.asarray(translation, dtype=float) - vs / 2
+        qb = snap_integral((np.asarray(self.begin) - corner0) / vs)
+        qe = snap_integral((np.asarray(self.end) - corner0) / vs)
+        if mode == "grow":
+            qb, qe = np.floor(qb), np.ceil(qe)
+        elif mode == "shrink":
+            qb, qe = np.ceil(qb), np.floor(qe)
+        elif mode == "closest":
+            qb, qe = np.round(qb), np.round(qe)
+        else:
+            raise ValueError("mode is 'grow', 'shrink' or 'closest'")
+        return Box(tuple(corner0 + qb * vs), tuple(corner0 + np.maximum(qe, qb) * vs))
+
+    def to_index(self, voxel_size, translation) -> tuple[tuple[int, ...], tuple[int, ...]]:
+        """The voxels whose centres lie in the box (the centre rule)."""
+        return index_range(self.begin, self.end, voxel_size, translation)
+
+    @classmethod
+    def from_index(cls, begin, end, voxel_size, translation) -> Box:
+        """The world box exactly covering voxels ``[begin, end)`` of a grid."""
+        vs = np.asarray(voxel_size, dtype=float)
+        corner0 = np.asarray(translation, dtype=float) - vs / 2
+        return cls(tuple(corner0 + np.asarray(begin) * vs), tuple(corner0 + np.asarray(end) * vs))
+
+    def _per_axis(self, value) -> tuple[float, ...]:
+        if np.ndim(value) == 0:
+            return (float(value),) * self.ndim
+        value = tuple(float(v) for v in value)
+        if len(value) != self.ndim:
+            raise ValueError(f"expected {self.ndim} values, got {value}")
+        return value
+
+    @classmethod
+    def coerce(cls, box: Any) -> Box:
+        """A ``Box`` from a ``(begin, end)`` pair, anything with ``.begin``/``.end`` or
+        ``.offset``/``.shape`` (a funlib ``Roi``, duck-typed), or a ``Box``."""
+        if isinstance(box, Box):
+            return box
+        if hasattr(box, "begin") and hasattr(box, "end"):
+            return cls(tuple(box.begin), tuple(box.end))
+        if hasattr(box, "offset") and hasattr(box, "shape"):
+            offset, shape = tuple(box.offset), tuple(box.shape)
+            return cls(offset, tuple(o + s for o, s in zip(offset, shape)))
+        try:
+            begin, end = box
+            return cls(tuple(begin), tuple(end))
+        except (TypeError, ValueError):
+            raise TypeError(
+                "a box is a (begin, end) pair, an object with .begin/.end or .offset/.shape, or a Box"
+            ) from None
+
+
+def index_range(begin, end, voxel_size, translation) -> tuple[tuple[int, ...], tuple[int, ...]]:
+    """The half-open voxel index range whose centres lie in ``[begin, end)``."""
+    s = np.asarray(voxel_size, dtype=float)
+    t = np.asarray(translation, dtype=float)
+    lo = np.ceil(snap_integral((np.asarray(begin, dtype=float) - t) / s))
+    hi = np.ceil(snap_integral((np.asarray(end, dtype=float) - t) / s))
+    hi = np.maximum(hi, lo)
+    return tuple(int(v) for v in lo), tuple(int(v) for v in hi)
+
+
+def is_aligned(begin, end, voxel_size, translation) -> bool:
+    """Whether every edge sits on a voxel corner or a voxel centre."""
+    s = np.asarray(voxel_size, dtype=float)
+    t = np.asarray(translation, dtype=float)
+    for edges in (begin, end):
+        q = (np.asarray(edges, dtype=float) - t) / s
+        twice = snap_integral(2 * q)
+        if not np.all(twice == np.round(twice)):
+            return False
+    return True

worldarray/levels.py

@@ -0,0 +1,147 @@
+"""The levels of a pyramid, and choosing one for a voxel size.
+
+A group's levels are the paths in its ``multiscales[0].datasets`` (OME-NGFF, COSEM N5), in
+that order; without that, ``s0, s1, ...`` are probed. A precomputed volume's levels are its
+``scales``, in order. Levels are matched by name, never by position (spec, Conventions 4;
+Behaviour 3).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+
+from .placement import ArrayMeta
+from .store import Location, Node, basename, detect, is_root, join, parent
+
+Mode = Literal["floor", "finer", "exact", "nearest", "resample"]
+
+
+@dataclass(frozen=True)
+class LevelRef:
+    name: str
+    path: str  # the array path (for precomputed, the volume root)
+    scale_index: int | None = None  # precomputed only
+
+
+def _group_multiscale_paths(node: Node) -> list[str]:
+    multiscales = node.multiscales
+    if not multiscales:
+        return []
+    return [str(ds["path"]).strip("/") for ds in multiscales[0].get("datasets", []) if "path" in ds]
+
+
+def _probe_levels(path: str) -> list[str]:
+    names: list[str] = []
+    for i in range(64):
+        name = f"s{i}"
+        docs = Location(join(path, name)).read_many(("zarr.json", ".zarray", "attributes.json"))
+        if all(v is None for v in docs.values()):
+            break
+        names.append(name)
+    return names
+
+
+def discover(node: Node) -> list[LevelRef]:
+    """The levels of a group node, or of the group an array node belongs to (just the array
+    itself when its parent is not a multiscale group)."""
+    if node.format == "precomputed":
+        return [
+            LevelRef(str(s.get("key", i)), node.path, i)
+            for i, s in enumerate(node.docs["info"]["scales"])
+        ]
+    if node.kind == "array":
+        ancestor, rel = parent(node.path), basename(node.path)
+        for _ in range(4):  # a dataset path may have more than one component ("seg/s0")
+            if is_root(ancestor):
+                break
+            try:
+                group = detect(ancestor, missing_ok=True)
+            except Exception:  # noqa: BLE001  (a server root answering 301, say)
+                break
+            if (
+                group is not None
+                and group.kind == "group"
+                and rel in _group_multiscale_paths(group)
+            ):
+                return discover(group)
+            if parent(ancestor) is None:
+                break
+            rel = f"{basename(ancestor)}/{rel}"
+            ancestor = parent(ancestor)
+        return [LevelRef(basename(node.path), node.path)]
+    names = _group_multiscale_paths(node) or _probe_levels(node.path)
+    return [LevelRef(name, join(node.path, name)) for name in names]
+
+
+def _spatial_voxel_size(level) -> np.ndarray:
+    """The spatial voxel size of an ``ArrayMeta``, or of a plain per-spatial-axis tuple."""
+    if hasattr(level, "spatial"):
+        return np.asarray([level.voxel_size[i] for i in level.spatial], dtype=float)
+    return np.asarray(level, dtype=float)
+
+
+def same_voxel_size(a, b) -> bool:
+    a, b = np.asarray(a, dtype=float), np.asarray(b, dtype=float)
+    return a.shape == b.shape and bool(np.allclose(a, b, rtol=1e-6, atol=0))
+
+
+def coarser_anywhere(resolution, target) -> bool:
+    r, t = np.asarray(resolution, dtype=float), np.asarray(target, dtype=float)
+    return bool(np.any((r > t) & ~np.isclose(r, t, rtol=1e-6, atol=0)))
+
+
+def select_level(
+    levels: Sequence[tuple[str, ArrayMeta | tuple[float, ...]]], voxel_size, mode: Mode = "floor"
+) -> tuple[str, ArrayMeta | tuple[float, ...]]:
+    """The ``(name, meta)`` to read at ``voxel_size`` (spatial axes, the pyramid's units).
+
+    - ``floor``: the level at ``voxel_size``, else the coarsest level that is no coarser than
+      it on any axis, else the first.
+    - ``finer``: as ``floor``, but ``ValueError`` instead of falling back to the first level
+      when every level is coarser than the target on some axis (cellmap-analyze's rule).
+    - ``exact``: the level at ``voxel_size``, else ``ValueError`` listing the levels.
+    - ``nearest``: the level closest on a log scale summed over the axes; the finer on a tie.
+    - ``resample``: the level to resample from: the level at ``voxel_size`` if there is one;
+      else the coarsest level that is no coarser on any axis; else the finest.
+    """
+    if not levels:
+        raise ValueError("no levels to choose from")
+    target = np.asarray(voxel_size, dtype=float)
+    sizes = [(name, meta, _spatial_voxel_size(meta)) for name, meta in levels]
+    if target.shape != sizes[0][2].shape:
+        raise ValueError(
+            f"voxel_size has {target.size} entries, the levels have {sizes[0][2].size} spatial axes"
+        )
+    for name, meta, vs in sizes:
+        if same_voxel_size(vs, target):
+            return name, meta
+    if mode == "exact":
+        raise ValueError(f"no level at voxel size {target.tolist()}; levels: " + _describe(sizes))
+    if mode in ("floor", "finer", "resample"):
+        candidates = [(name, meta) for name, meta, vs in sizes if not coarser_anywhere(vs, target)]
+        if not candidates:
+            if mode == "finer":
+                raise ValueError(
+                    f"every level is coarser than {target.tolist()} on some axis; levels: "
+                    + _describe(sizes)
+                )
+            return levels[0]
+        # the coarsest of them: largest voxel volume
+        return max(candidates, key=lambda nm: float(np.prod(_spatial_voxel_size(nm[1]))))
+    if mode == "nearest":
+
+        def distance(entry):
+            vs = entry[2]
+            return float(np.sum(np.abs(np.log(vs) - np.log(target)))), float(np.prod(vs))
+
+        name, meta, _ = min(sizes, key=distance)
+        return name, meta
+    raise ValueError(f"unknown mode {mode!r}")
+
+
+def _describe(sizes) -> str:
+    return ", ".join(f"{name} {vs.tolist()}" for name, _, vs in sizes)