feat(geometry): bound a displacement by where it reaches, not how far

konfai/data/geometry.py

@@ -120,6 +120,18 @@ def grown(self, radius_xyz: np.ndarray | float) -> WorldBox:
         radius = np.broadcast_to(np.asarray(radius_xyz, dtype=np.float64), self.low_xyz.shape)
         return WorldBox(self.low_xyz - radius, self.high_xyz + radius)
 
+    def extended(self, low_xyz: np.ndarray, high_xyz: np.ndarray) -> WorldBox:
+        """This box plus a per-component interval: each end moved by its own end of it.
+
+        The asymmetric form of :meth:`grown`, and the one a signed displacement bound needs. An
+        interval that does not straddle zero MOVES the box instead of widening it, which is the
+        difference between a field's reach and twice its largest value.
+        """
+        return WorldBox(
+            self.low_xyz + np.asarray(low_xyz, dtype=np.float64),
+            self.high_xyz + np.asarray(high_xyz, dtype=np.float64),
+        )
+
     def image_under(self, affine: AffineMap) -> WorldBox:
         """The axis-aligned hull of this box's image under ``affine``.
 
@@ -361,37 +373,74 @@ def sub_grid(self, region_zyx: tuple[slice, ...]) -> Grid:
 
 @dataclass(frozen=True)
 class TransformBound:
-    """What a stored transform is guaranteed to do: an exact affine part and a bounded residual.
-
-    ``T(p)`` lies in ``affine(p) ± residual_xyz`` for every ``p``, per world component. For a
-    linear transform the residual is zero and the statement is exact; for a BSpline it is the
-    sup-norm of the coefficients (non-negative basis functions summing to one make every
-    displacement a convex combination of them); for a dense field it is the recorded or declared
-    per-component bound. The affine part is read structurally off the transform, never probed:
-    a probe measures a local gradient and extrapolates it, which under-bounds (measured).
+    """What a stored transform is guaranteed to do: an exact affine part and a bounded interval.
+
+    ``T(p)`` lies in ``affine(p) + [low_xyz, high_xyz]`` for every ``p``, per world component. For a
+    linear transform the interval is empty and the statement is exact; for a BSpline it is the range
+    of the coefficients (non-negative basis functions summing to one make every displacement a
+    convex combination of them, so it lies between their smallest and largest); for a dense field it
+    is the range of its values. The affine part is read structurally off the transform, never
+    probed: a probe measures a local gradient and extrapolates it, which under-bounds (measured).
+
+    SIGNED, NOT A RADIUS. A displacement field solved between two frames carries the offset between
+    them in its values, and an interval that does not straddle zero MOVES a region's window rather
+    than widening it. Measured on an ExaSPIM field whose z component runs [-28.1, -22.2] mm on a
+    volume 20.6 mm thick: as a radius it reaches 28.1 mm either way, so every region pulls the whole
+    volume and the fold refuses (23.57 GiB held against a 19.01 GiB budget); as an interval it
+    reaches 5.9 mm, and a 24-row region pulls 175 source rows of 514. The same two reductions
+    produce either (:attr:`DisplacementStage.range_xyz`), so the tighter one is free.
     """
 
     affine: AffineMap
-    residual_xyz: np.ndarray
+    low_xyz: np.ndarray
+    high_xyz: np.ndarray
 
     @staticmethod
     def exact(affine: AffineMap) -> TransformBound:
-        return TransformBound(affine, np.zeros(affine.rank))
+        return TransformBound(affine, np.zeros(affine.rank), np.zeros(affine.rank))
+
+    @staticmethod
+    def interval(low_xyz: np.ndarray, high_xyz: np.ndarray) -> TransformBound:
+        """A pure displacement whose value lies in ``[low_xyz, high_xyz]`` per component."""
+        low = np.asarray(low_xyz, dtype=np.float64)
+        return TransformBound(AffineMap.identity(int(low.size)), low, np.asarray(high_xyz, dtype=np.float64))
 
     @staticmethod
     def shift(residual_xyz: np.ndarray) -> TransformBound:
-        return TransformBound(AffineMap.identity(int(residual_xyz.size)), np.asarray(residual_xyz, dtype=np.float64))
+        """A pure displacement bounded in magnitude only, ``± residual_xyz``.
+
+        For a caller that knows a radius and not a range. Anything that can state both ends should
+        say so with :meth:`interval`: this one is twice as wide wherever the range is one-sided.
+        """
+        radius = np.asarray(residual_xyz, dtype=np.float64)
+        return TransformBound.interval(-radius, radius)
+
+    @property
+    def residual_xyz(self) -> np.ndarray:
+        """The symmetric envelope of the interval, for a caller that wants one number per axis."""
+        return np.maximum(np.abs(self.low_xyz), np.abs(self.high_xyz))
 
     def after(self, inner: TransformBound) -> TransformBound:
-        """The bound of ``self(inner(p))``: interval arithmetic through the outer affine."""
+        """The bound of ``self(inner(p))``: interval arithmetic through the outer affine.
+
+        NOT ``|A| @ residual``. That is right for an interval centred on zero and wrong for one that
+        is not: a negative entry of ``A`` sends the inner interval's low end to the outer's high,
+        and taking absolute values first loses which end went where -- so a rotation folded onto a
+        one-sided field would be bounded by a box that does not contain it. Splitting the matrix
+        into its non-negative and non-positive parts is the same arithmetic written to hold either
+        way, and it reduces to ``|A| @ r`` when ``low = -high``.
+        """
+        matrix = self.affine.matrix
+        rise, fall = np.maximum(matrix, 0.0), np.minimum(matrix, 0.0)
         return TransformBound(
             inner.affine.then(self.affine),
-            np.abs(self.affine.matrix) @ inner.residual_xyz + self.residual_xyz,
+            rise @ inner.low_xyz + fall @ inner.high_xyz + self.low_xyz,
+            rise @ inner.high_xyz + fall @ inner.low_xyz + self.high_xyz,
         )
 
     def map_box(self, box: WorldBox) -> WorldBox:
         """Where the image of ``box`` is guaranteed to lie."""
-        return box.image_under(self.affine).grown(self.residual_xyz)
+        return box.image_under(self.affine).extended(self.low_xyz, self.high_xyz)
 
 
 @dataclass(frozen=True)
@@ -467,28 +516,40 @@ def tensor_by_voxel(self, device: torch.device, dtype: torch.dtype) -> torch.Ten
     def __getstate__(self) -> dict:
         state = dict(self.__dict__)
         state["_tensors"] = {}
+        state.pop("range_xyz", None)
         state.pop("bound_xyz", None)
         return state
 
     @cached_property
-    def bound_xyz(self) -> np.ndarray:
-        """``sup |values|`` per component: one pass over the field, kept for the stage's life.
+    def range_xyz(self) -> tuple[np.ndarray, np.ndarray]:
+        """``(min, max)`` per component: one pass over the field, kept for the stage's life.
 
         Every pull map asks for it (per patch, per plan block, per pushed slab), so it is kept:
         a thousand patches recomputing one constant 3-vector cost 73 s on a 3x160x256x256 field.
         ``cached_property`` writes through ``__dict__``, which a frozen dataclass allows; the entry
         is dropped from the pickle beside ``_tensors``.
 
-        Two reductions and no temporary: ``sup |v| = max(max v, -min v)`` is exact for real values,
-        where ``np.abs(...).max()`` first writes a values-sized copy and then walks it again. On the
-        31 M-voxel field the difference is 72 ms and 250 MiB; on a native ExaSPIM field window
-        (3 x 141 x 1331 x 1775) the copy alone is 4 GiB, and the fold spent 26.4 s of its 179 s here.
+        Two reductions and no temporary, which is what the pair costs and what its magnitude cost
+        before it: ``np.abs(...).max()`` would first write a values-sized copy and then walk it
+        again. On the 31 M-voxel field that copy is 72 ms and 250 MiB; on a native ExaSPIM field
+        window (3 x 141 x 1331 x 1775) it is 4 GiB, and the fold spent 26.4 s of its 179 s here.
+        Keeping both ends instead of the larger magnitude measured 79 ms against 80 on a 184
+        M-value window: the range is free, and it is the one a region can be sized from.
         """
         flat = self.values.reshape(self.values.shape[0], -1)
-        return np.maximum(flat.max(axis=1), -flat.min(axis=1))
+        return flat.min(axis=1), flat.max(axis=1)
+
+    @cached_property
+    def bound_xyz(self) -> np.ndarray:
+        """``sup |values|`` per component, the range's symmetric envelope."""
+        low, high = self.range_xyz
+        return np.maximum(np.abs(low), np.abs(high))
 
     def bound(self) -> TransformBound:
-        return TransformBound.shift(self.bound_xyz)
+        # The interval is clamped to include zero: the stage applies NO displacement outside its
+        # grid, so a target region past the field's edge still needs its identity-mapped samples.
+        low, high = self.range_xyz
+        return TransformBound.interval(np.minimum(low, 0.0), np.maximum(high, 0.0))
 
 
 #: A decoded stored transform: stages in APPLICATION order (first applied first). SimpleITK's

konfai/utils/ome_zarr.py

@@ -39,6 +39,7 @@
 import itertools
 import operator
 import shutil
+import tempfile
 import threading
 from collections import OrderedDict
 from collections.abc import Sequence
@@ -48,6 +49,7 @@
 
 import numpy as np
 
+from konfai.utils import uri
 from konfai.utils.errors import DatasetManagerError
 from konfai.utils.runtime import map_over_rank_pool
 
@@ -177,6 +179,22 @@ def _konfai_attributes(store_path: str) -> dict[str, Any]:
         return {}
 
 
+def _from_ngff_zarr(store_path: str | Path) -> Any:
+    """ngff-zarr's multiscales for ``store_path``. A remote root goes in as a key-to-bytes mapping
+    over its own filesystem (ngff-zarr >= 0.44 reads a remote string as a local path) and as its URL
+    when ngff-zarr refuses the mapping, which older releases resolve themselves."""
+    if not uri.is_uri(store_path):
+        return ngff_zarr.from_ngff_zarr(str(store_path))
+    # Through uri.filesystem, so a missing fsspec backend or configuration is the structured
+    # DatasetManagerError, never a raw dependency error.
+    filesystem = uri.filesystem(store_path)
+    _, target = uri.split_scheme(str(store_path))
+    try:
+        return ngff_zarr.from_ngff_zarr(filesystem.get_mapper(target))
+    except ValueError:
+        return ngff_zarr.from_ngff_zarr(str(store_path))
+
+
 @lru_cache(maxsize=8)
 def _load_image(store_path: str, level: int) -> Any:
     """Return the ``NgffImage`` for ``level`` of an OME-Zarr store, memoised per (store, level).
@@ -194,7 +212,7 @@ def _load_image(store_path: str, level: int) -> Any:
     """
     _require_ngff_zarr()
     try:
-        multiscales = ngff_zarr.from_ngff_zarr(str(store_path))
+        multiscales = _from_ngff_zarr(store_path)
     except (KeyError, IndexError, OSError, TypeError, ValueError) as exc:
         raise DatasetManagerError(
             f"Cannot open OME-Zarr store '{store_path}' (level {level}).",
@@ -258,7 +276,7 @@ def is_displacement_field(store_path: str | Path) -> bool:
     if not _NGFF_ZARR_AVAILABLE:
         return False
     try:
-        metadata = ngff_zarr.from_ngff_zarr(str(store_path)).metadata
+        metadata = _from_ngff_zarr(store_path).metadata
     except Exception:
         # "Not a displacement field" is the only answer this owes: it is asked purely to decide HOW to
         # read an entry, and an absent or unreadable store is not one either.
@@ -608,7 +626,7 @@ def place(coords: tuple) -> None:
 def _level_path(store_path: str, level: int) -> str | None:
     """The zarr path of one level, from the store's multiscales metadata, memoised beside the image."""
     try:
-        datasets = ngff_zarr.from_ngff_zarr(store_path).metadata.datasets
+        datasets = _from_ngff_zarr(store_path).metadata.datasets
         return str(datasets[level if len(datasets) > 1 else 0].path)
     except Exception:
         return None
@@ -624,7 +642,32 @@ def _read_level_window(store_path: str, level: int, image: Any, index: tuple) ->
             array = _level_array(store_path, level_path)
             if tuple(array.shape) == tuple(image.data.shape):
                 return _read_chunked(store_path, level_path, array, index)
-    return np.asarray(image.data[index])
+    return _lazy_window(image.data, index)
+
+
+def _lazy_window(data: Any, index: tuple) -> np.ndarray:
+    """The plain lazy read. A lazy array that refuses a stepped slice (ngff-zarr >= 0.44 wraps the
+    level in an adapter that takes unit steps only) is read over the unit-step span and stepped here."""
+    try:
+        return np.asarray(data[index])
+    except NotImplementedError:
+        # Each slice normalized against the shape, then its ascending unit-step span; a negative
+        # step reads that span backwards from its own end, which lands on the indices the original
+        # slice named.
+        bounds = tuple(
+            slice(*item.indices(size)) if isinstance(item, slice) else item
+            for item, size in zip(index, data.shape, strict=True)
+        )
+        span = tuple(
+            slice(item.stop + 1, item.start + 1)
+            if isinstance(item, slice) and item.step < 0
+            else slice(item.start, item.stop)
+            if isinstance(item, slice)
+            else item
+            for item in bounds
+        )
+        steps = tuple(slice(None, None, item.step) for item in bounds if isinstance(item, slice))
+        return np.asarray(data[span])[steps]
 
 
 def _store_index(
@@ -813,6 +856,11 @@ def write_ome_zarr(
     exist only from 0.6, so a caller passing both could only ever pass them consistently: an
     invariant worth removing rather than documenting.
     """
+    if scale_factors and uri.is_uri(store_path):
+        raise DatasetManagerError(
+            f"Cannot append pyramid levels to the remote store '{store_path}'.",
+            "Levels are derived in place through local paths; write the store locally and upload it.",
+        )
     array_data = np.asarray(data)
     # The one write path: the store described and created empty (ngff-zarr's metadata, the
     # caller's chunking), filled by zarr itself, its levels grafted beside level 0. Handing
@@ -876,6 +924,30 @@ def _type_component_axis(multiscales: Any, axis_type: str) -> None:
                 axis.type = axis_type
 
 
+def _write_skeleton(store_path: str | Path, multiscales: Any, version: str) -> None:
+    """ngff-zarr's metadata for the store, written in place. ngff-zarr (>= 0.44) writes local
+    directories only, so a remote root gets the skeleton written locally and uploaded through the
+    root's own filesystem: a few bytes of metadata, before the array is created underneath it."""
+    if not uri.is_uri(store_path):
+        ngff_zarr.to_ngff_zarr(str(store_path), multiscales, overwrite=True, version=version)
+        return
+    filesystem = uri.filesystem(store_path)
+    _, target = uri.split_scheme(str(store_path))
+    if "/" not in target.strip("/"):
+        raise DatasetManagerError(
+            f"Refusing to create the store at the filesystem root '{store_path}'.",
+            "Name a key under the root (e.g. '.../dataset.ome.zarr'): creating a store replaces what its path holds.",
+        )
+    if filesystem.exists(target):
+        filesystem.rm(target, recursive=True)
+    filesystem.makedirs(target, exist_ok=True)
+    with tempfile.TemporaryDirectory() as scratch:
+        local = Path(scratch) / "skeleton"
+        ngff_zarr.to_ngff_zarr(str(local), multiscales, overwrite=True, version=version)
+        for file in sorted(path for path in local.rglob("*") if path.is_file()):
+            filesystem.put_file(str(file), uri.join(target, file.relative_to(local).as_posix()))
+
+
 def create_ome_zarr_store(
     store_path: str | Path,
     shape: Sequence[int],
@@ -934,7 +1006,7 @@ def create_ome_zarr_store(
         _type_component_axis(multiscales, _DISPLACEMENT_AXIS_TYPE)
         version = _RFC5_VERSION
     # version is explicit because to_ngff_zarr defaults to 0.5, which zarr-python 2 cannot write.
-    ngff_zarr.to_ngff_zarr(str(store_path), multiscales, overwrite=True, version=version)
+    _write_skeleton(store_path, multiscales, version)
 
     # The level-0 key comes from the metadata rather than a literal: ngff-zarr builds it from the
     # image name, so "scale0/image" is its convention to change, not ours to hardcode.
@@ -1077,13 +1149,18 @@ def append_ome_zarr_levels(
     KonfAI attribute sidecar is untouched, being a key beside theirs; a displacement field keeps its
     typed component axis through the same call that types it at creation.
     """
+    if uri.is_uri(store_path):
+        raise DatasetManagerError(
+            f"Cannot append pyramid levels to the remote store '{store_path}'.",
+            "Levels are derived in place through local paths; write the store locally and upload it.",
+        )
     _require_ngff_zarr()
     if not scale_factors:
         return
     store = Path(store_path)
     clear_ome_zarr_cache(store)
     field = is_displacement_field(store)
-    base = ngff_zarr.from_ngff_zarr(str(store)).images[0]
+    base = _from_ngff_zarr(store).images[0]
     stored_chunks = tuple(int(size) for size in base.data.chunksize)
     if downsample_method in (None, "ITKWASM_BIN_SHRINK"):
         multiscales = _bin_shrink_multiscales(base, _level_zero_scale_factors(scale_factors), stored_chunks)
@@ -1152,7 +1229,7 @@ def get_ome_zarr_info(store_path: str | Path, level: int = 0) -> dict[str, Any]:
     image = _load_image(str(store_path), level)
     dims = [str(axis).lower() for axis in image.dims]
     try:
-        n_levels = len(ngff_zarr.from_ngff_zarr(str(store_path)).images)
+        n_levels = len(_from_ngff_zarr(store_path).images)
     except (OSError, TypeError, ValueError):
         n_levels = 1
     return {