From 4c260dedf28f0a1bce0f3e719c190c96f6919da5 Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Tue, 4 Aug 2026 10:52:35 +0200 Subject: [PATCH 1/3] feat(transform): resample a case onto the grid of a declared reference A cohort has to be on one grid before it can be folded, and until now that meant a separate pass before the transform. ResampleToReference puts a case on the grid of a declared reference case -- extents, spacing, origin and direction -- so the members of a Reduce meet on a grid that is one of their own rather than an invented one. --- docs/source/config_guide/transform.md | 71 ++- konfai/data/data_manager.py | 5 + konfai/data/patching.py | 8 +- konfai/data/transform.py | 510 ++++++++++++++++++ konfai/transformer.py | 53 +- tests/unit/test_resample_to_reference.py | 477 ++++++++++++++++ .../unit/test_transform_locality_contract.py | 27 +- 7 files changed, 1133 insertions(+), 18 deletions(-) create mode 100644 tests/unit/test_resample_to_reference.py diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index e6147abd..c7c0fb49 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -192,6 +192,74 @@ One chain changes its cardinality at most once. Composing the two — augment a cohort, then fold it — is two invocations, the second reading the first one's output back. +### `Reduce`: one grid for the cohort, first + +Folding cases together only means something if they are the *same* voxels, so +`Reduce` compares the grid each case's chain **lands on**, and `grid:` says how +strictly: + +| `grid:` | Compares | Use it when | +| --- | --- | --- | +| `strict` | extent **and** `Spacing` / `Origin` / `Direction` | the cases really do share a space | +| `shape_only` | extent alone | you know they share one and the headers disagree | +| `reference:` | extent; that case's header is what the output carries | one member is the one to believe | + +A cohort as acquired rarely passes `strict`: extents differ, and origins can +differ by more than the volumes are wide, because an acquisition's stage +coordinates are not an anatomical frame. `ResampleToReference` is what makes +`strict` true rather than waived — it resamples each case onto the grid of a +**declared reference**, adopting its extent, spacing, origin and direction: + +```yaml +transforms: + ResampleToReference: {entry: case_0, group: CT, fill: 0.0} + Reduce: {operator: Median, output: template, grid: strict} + Write: {dataset: ./Template:mha} +``` + +The reference is a stored image, named by `entry` — and by `group` when the +store holds more than one. It is looked up by entry, not by the case being +processed, because one grid serves the whole cohort: in the run's own +`dataset_filenames`, or in a store of its own. + +```yaml +transforms: + ResampleToReference: {entry: case_1, group: CT, dataset: ./Raw:mha} + Write: {dataset: ./OnTemplate:mha} +``` + +`dataset:` takes the same `path[:format]` spec as everywhere else, so the +reference can live anywhere — which is the atlas loop: point round N+1 at the +store round N wrote its template into. + +Naming an image rather than fifteen numbers is deliberate. A grid is an extent +in array order `(Z, Y, X)` plus an origin, a spacing and a direction in physical +`(x, y, z)` — transcribing those by hand is the mistake that actually gets made, +and a transposed grid resamples perfectly well onto the wrong place. A header +cannot make that mistake. + +```{note} +It streams: a slab of the output reads only the part of the input under it, so a +case never has to fit in memory. The sampler is `sitk.Resample`'s — linear with +taps clamped to the buffer, nearest by round-half-up, and `fill` wherever the +reference grid reaches past the case. +``` + +**What it refuses**, rather than write something plausible and wrong: + +- a case or a reference carrying no `Origin` / `Spacing` / `Direction` — without + geometry there is no physical space to resample in, and a size ratio must not + quietly stand in for one; +- a reference whose `Direction` differs from the case's — the map between them + is then a rotation, not a scale and a shift per axis. Reorient first + (`Canonical`); +- a case that does not meet the reference grid **anywhere** — its output would + be `fill` from edge to edge, and a median would take that as anatomy. + +Partial overlap is legal and ordinary: the rest of the output is `fill`, and the +plan prints the fraction of the grid each case covers. "Most of this template is +background" is then something read before the run rather than after it. + ### `Expand`: one case, N copies `Expand` multiplies, and nothing else. The draws are **ordinary stages of the @@ -434,7 +502,8 @@ refused (cheaper to load the volume), and the plan says so. | a bounded neighbourhood | `HALO`, with `halo=(r,)` | nothing | | the volume flipped/permuted | `ORIENTATION` | `stream_region_source()` | | a translated sub-box | `CROP` | `stream_region_source()` | -| a resampled grid | `RESCALE` | inherit from `Resample` | +| the same box, resampled | `RESCALE` | inherit from `Resample` | +| another grid entirely | `REGRID` | `stream_region_source()` and `stream_region()` | | whole-volume Min/Max/Mean/Std | `GLOBAL_STAT`, with `stat_keys` | nothing | | the same, per component | `GLOBAL_STAT`, with `MinPerChannel`/`MaxPerChannel`/`MeanPerChannel`/`StdPerChannel` | nothing | | genuinely the whole volume | nothing (the default) | nothing | diff --git a/konfai/data/data_manager.py b/konfai/data/data_manager.py index 978d3038..53a08ce3 100755 --- a/konfai/data/data_manager.py +++ b/konfai/data/data_manager.py @@ -193,6 +193,11 @@ def _check_patch_transform_locality(transform: Transform, group_src: str, group_ f"'{name}' resamples its input: applied to one patch it rescales that patch about its own" " extent and changes the patch grid predictions are reassembled onto." ), + LocalityKind.REGRID: ( + f"'{name}' resamples its input onto another grid: applied to one patch it would hand back" + " the whole reference extent, which is neither the patch nor the patch grid predictions" + " are reassembled onto." + ), LocalityKind.WHOLE_VOLUME: f"'{name}' needs the whole volume.", } raise ConfigError( diff --git a/konfai/data/patching.py b/konfai/data/patching.py index e6148d7d..63360ba6 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -208,7 +208,13 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) # The region kinds a composed streamed read (or write) carries between its pointwise stages. -_REGION_KINDS = (LocalityKind.HALO, LocalityKind.ORIENTATION, LocalityKind.CROP, LocalityKind.RESCALE) +_REGION_KINDS = ( + LocalityKind.HALO, + LocalityKind.ORIENTATION, + LocalityKind.CROP, + LocalityKind.RESCALE, + LocalityKind.REGRID, +) # The pull maps are callable dataclasses, not closures, because a plan crosses a process boundary: diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 7d6b4475..c79f23b6 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -61,6 +61,12 @@ class LocalityKind(Enum): - ``GLOBAL_STAT`` -- needs whole-volume stats (``stat_keys`` subset of Min/Max/Mean/Std), obtained once from disk and cached: read the exact patch + the cached stat. - ``RESCALE`` -- resample: source region via the scale mapping + interpolation halo. + - ``REGRID`` -- resample onto ANOTHER grid. ``RESCALE``'s source and target cover the same + physical box and differ only in sampling density, so a size ratio is the whole of its map; + this one's target is a grid in its own right, placed by its own origin, so the map carries an + offset as well as a scale and part of the target may read from outside the source altogether. + The stage owns both halves: it declares the source region a target region pulls + (:meth:`Transform.stream_region_source`) and interpolates it (:meth:`Transform.stream_region`). - ``SLAB`` -- per-voxel value map, plus a side effect that needs the slab's place in the volume (a per-region side write): the streamed-WRITE dispatcher runs it through :meth:`Transform.stream_slab` with region context; the read dispatcher has no such context and @@ -74,6 +80,7 @@ class LocalityKind(Enum): CROP = "crop" GLOBAL_STAT = "global_stat" RESCALE = "rescale" + REGRID = "regrid" SLAB = "slab" WHOLE_VOLUME = "whole_volume" @@ -231,6 +238,23 @@ def prepare(self, konfai_args: str) -> None: nothing: only a stage with a sub-object of its own overrides it. """ + def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribute: Attribute) -> str | None: + """Something about this case the plan should say, beyond its regime and its cost. + + A stage can be correct, stream, fit the budget, and still surprise the reader — a cost the + plan has no column for. The plan is where a run is read before it is trusted, so that is + where the sentence belongs, rather than in a viewer afterwards. + + Answered from headers on the launcher, per (chain, case), under :meth:`patch_locality`'s + rules: read-only, no volume read, and an answer for any case. Identical notes are printed + once, so a note about the STAGE may repeat per case without repeating on the page, while a + note about the CASE stays one line each. + + The base holds nothing: most stages have nothing to add to their regime and their bytes. + """ + del group_dest, name, shape, cache_attribute + return None + def stream_abort(self, name: str) -> None: """Drop whatever ``stream_slab`` holds open for ``name`` after a mid-case failure. @@ -946,6 +970,7 @@ def source_window( scales: list[float], n_in: list[int], halo: int = 1, + offsets: list[float] | None = None, ) -> list[slice]: """The clamped source region a target region reads from, per axis, given the scales. @@ -955,7 +980,15 @@ def source_window( index). Under strong downsampling the nearest voxel of the first output column falls BELOW the linear window's start: the window must include it, or the gather wraps a negative local index onto the far edge. + + ``offsets`` generalises the map to ``source = scale * target + offset``, for a resample whose + target grid is placed by its own origin rather than sharing the source's box (``REGRID``). + Left ``None``, every coordinate below is the one this always computed -- the half-pixel map + is not re-derived through a more general formula that would round differently, because the + paths that must stay bit-identical to ``F.interpolate`` run exactly this code. """ + if offsets is not None: + return Resample._offset_window(target_slices, scales, offsets, n_in, halo) source_slices: list[slice] = [] for k, sl in enumerate(target_slices): smin = int(np.floor(scales[k] * (sl.start + 0.5) - 0.5)) @@ -967,6 +1000,34 @@ def source_window( source_slices.append(slice(max(0, start), min(n_in[k], stop))) return source_slices + @staticmethod + def _offset_window( + target_slices: tuple[slice, ...] | list[slice], + scales: list[float], + offsets: list[float], + n_in: list[int], + halo: int, + ) -> list[slice]: + """``source_window`` for an offset map, clamped to a non-empty region of the source. + + Non-empty even when the target region lies entirely off the source, which is a real place + for a ``REGRID`` and not an error: every sample there is out of bounds and takes the fill, + so what the read returns is never looked at -- but a zero-width read is not something every + backend serves, and one voxel costs nothing. + """ + source_slices: list[slice] = [] + for k, sl in enumerate(target_slices): + first = scales[k] * sl.start + offsets[k] + last = scales[k] * (sl.stop - 1) + offsets[k] + low, high = (first, last) if first <= last else (last, first) + # floor(low) is the linear map's i0 and bounds its nearest pick; +2 reaches i1 past + # floor(high), matching the `smax + 2` the half-pixel window uses for the same reason. + start = int(np.floor(low)) - halo + stop = int(np.floor(high)) + 2 + halo + start = min(max(start, 0), n_in[k] - 1) + source_slices.append(slice(start, min(max(stop, start + 1), n_in[k]))) + return source_slices + def resample_region( self, sub_tensor: torch.Tensor, @@ -974,6 +1035,7 @@ def resample_region( region_starts: list[int], scales: list[float], n_in: list[int], + offsets: list[float] | None = None, ) -> torch.Tensor: """Interpolate a source sub-region to the target patch extent. @@ -981,7 +1043,13 @@ def resample_region( ``region_starts`` are the global source indices of its first voxel per axis. Uses the same global coordinate formula as the whole-volume path, indexing the sub-region as ``sub[i - region_start]``. + + ``offsets`` generalises the map to ``source = scale * target + offset``, as in + :meth:`source_window`, and a sample landing outside the source then takes + :attr:`fill_value`. Left ``None``, this runs the half-pixel code it always did. """ + if offsets is not None: + return self._resample_offset_region(sub_tensor, target_slices, region_starts, scales, n_in, offsets) mode = self._stream_mode(sub_tensor) dev = sub_tensor.device ndim = len(target_slices) @@ -1025,6 +1093,91 @@ def resample_region( out += gathered * weight return out.type(sub_tensor.dtype) + #: What a sample landing outside the source is worth. Only an offset map can land outside at + #: all, so only a ``REGRID`` stage ever reads this, and it sets it from its own configuration. + fill_value: float = 0.0 + + def _resample_offset_region( + self, + sub_tensor: torch.Tensor, + target_slices: tuple[slice, ...], + region_starts: list[int], + scales: list[float], + n_in: list[int], + offsets: list[float], + ) -> torch.Tensor: + """``resample_region`` for an offset map: ITK's sampler, and a fill where the source stops. + + THE SAMPLING RULE IS ``sitk.Resample``'S, deliberately: a sample is inside while its + continuous source index lies in ``[-0.5, n - 0.5)``; inside, the interpolation taps are + clamped to the buffer, so the half-voxel rim beyond the outermost voxel CENTRES reproduces + the border value rather than falling off; outside, the sample is :attr:`fill_value`. Written + against SimpleITK because that is what an independent check of this arithmetic will be, and + a sampler that is only nearly the same as the reference makes every such check a negotiation. + + Coordinates are global (the target index, not its offset within the region), so a region and + the whole volume put the same sample in the same place -- which is what makes the streamed + and whole-volume paths equal by construction here rather than by agreement. + """ + device = sub_tensor.device + ndim = len(target_slices) + window = [int(extent) for extent in sub_tensor.shape[1:]] + # Coordinates in float32 whatever the payload: a float16 volume would otherwise index itself + # with float16 coordinates, which cannot even count the voxels of a large axis. + coordinates = [ + scales[k] * torch.arange(sl.start, sl.stop, device=device, dtype=torch.float32) + offsets[k] + for k, sl in enumerate(target_slices) + ] + inside = [(axis >= -0.5) & (axis < n_in[k] - 0.5) for k, axis in enumerate(coordinates)] + out_shape = [int(sub_tensor.shape[0])] + [sl.stop - sl.start for sl in target_slices] + if not all(bool(axis.any()) for axis in inside): + # Nothing of this region is on the source. Real for a REGRID -- a target grid may reach + # past its case -- and worth its own exit: the gather below would read a window that was + # only ever clamped to something legal, and then be overwritten by the fill anyway. + return torch.full(out_shape, self.fill_value, device=device, dtype=torch.float32).type(sub_tensor.dtype) + + def local(index: torch.Tensor, k: int) -> torch.Tensor: + """A global source index as an offset into the window that was read, kept in range.""" + return torch.clamp(torch.clamp(index, 0, n_in[k] - 1) - region_starts[k], 0, window[k] - 1) + + if self._stream_mode(sub_tensor) == "nearest": + # floor(c + 0.5) is ITK's nearest -- round half up -- not F.interpolate's floor(o * scale), + # which is a statement about a size ratio and says nothing once a grid has its own origin. + picks = [local(torch.floor(axis + 0.5).long(), k) for k, axis in enumerate(coordinates)] + gathered = sub_tensor[(slice(None), *torch.meshgrid(*picks, indexing="ij"))] + out = gathered if gathered.is_floating_point() else gathered.type(torch.float32) + else: + if not sub_tensor.is_floating_point() or ( + sub_tensor.device.type == "cpu" and sub_tensor.dtype in (torch.float16, torch.bfloat16) + ): + work = sub_tensor.type(torch.float32) + else: + work = sub_tensor + taps = [] + for k, axis in enumerate(coordinates): + base = torch.floor(axis) + weight = (axis - base).to(work.dtype) + index = base.long() + taps.append(((local(index, k), 1 - weight), (local(index + 1, k), weight))) + out = torch.zeros([work.shape[0], *out_shape[1:]], device=device, dtype=work.dtype) + for combo in itertools.product(*taps): + gathered = work + weight = torch.ones([1] * (ndim + 1), device=device, dtype=work.dtype) + for k, (index, lam) in enumerate(combo): + gathered = gathered.index_select(k + 1, index) + shape = [1] * (ndim + 1) + shape[k + 1] = -1 + weight = weight * lam.reshape(shape) + out += gathered * weight + # The axes' masks compose by outer product: a sample is inside where every axis of it is. + mask = inside[0] + for axis in inside[1:]: + mask = mask.unsqueeze(-1) & axis + # Filled while still floating, then cast ONCE: torch implements masked_fill for float dtypes + # and not for every integer one -- uint16 is a microscope's native dtype and has no fill at + # all -- so filling after the cast fails on exactly the volumes this stage is built for. + return out.masked_fill(~mask.unsqueeze(0), self.fill_value).type(sub_tensor.dtype) + @abstractmethod def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: """Record the same 'Spacing'/'Size' stack a whole-volume ``__call__`` would. @@ -1134,6 +1287,356 @@ def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatia cache_attribute["Size"] = shape +class ResampleToReference(Resample): + """Resample a case onto the grid of a declared reference — extent, spacing, origin, direction. + + The stage that makes a cohort foldable. ``ResampleToResolution`` lines up SPACINGS and + ``ResampleToShape`` lines up EXTENTS, but both leave each case where it was: the cases still + sit at different origins, so folding them (``Reduce``) can only proceed by looking away + (``grid: shape_only``). This adopts the reference's grid whole, which is what gives + ``grid: strict`` — it compares ``Spacing``, ``Origin`` and ``Direction`` — something true to + check. That is the atlas-template build: bring every case onto one grid, then take the median. + + THE REFERENCE IS A STORED IMAGE, NOT A LIST OF NUMBERS. A grid is fifteen numbers in two axis + orders at once (``shape`` counts (Z, Y, X); ``Origin``, ``Spacing`` and ``Direction`` are + physical (x, y, z)), and transcribing them by hand is the mistake this file's history says is + always made — silently, because a transposed grid resamples perfectly well onto the wrong + place. Naming an image cannot make that mistake: the header IS the declaration. It is also what + an atlas loop needs, where round N+1's reference is round N's own output. + + WHAT IT REFUSES, rather than resample onto a grid it cannot honestly reach: + + - a case, or a reference, whose header carries no ``Origin``/``Spacing``/``Direction`` — with no + geometry there is no physical space to resample IN, and a size ratio would silently stand in + for one; + - a reference whose ``Direction`` differs from the case's — the two grids' axes then do not line + up, and the map is a rotation, not a scale and a shift per axis. ``Canonical`` first; + - a case that does not meet the reference grid at all — the output would be pure ``fill``, and + an all-background member is a plausible, wrong contribution to a median. + + Everything else it declares. A case that reaches only part of the reference is legal and + common — the rest takes ``fill`` — and the plan prints how much of the grid each case covers, + because "most of this template is fill" is not something to discover in a viewer. + """ + + def __init__( + self, + entry: str, + group: str | None = None, + dataset: str | None = None, + fill: float = 0.0, + inverse: bool = True, + ) -> None: + super().__init__(inverse) + if not entry or not str(entry).strip(): + raise TransformError( + "'ResampleToReference' needs an 'entry': the stored image whose grid to adopt.", + "Name it, e.g. ResampleToReference: {entry: 822174, group: Volume}.", + ) + self.entry = str(entry).strip() + self.group = group + self.fill_value = float(fill) + # A root of its own, exactly as Warp takes one for its field. Left out, the reference is + # looked up in the run's own dataset_filenames -- the common case, where the grid to adopt + # is one member of the very cohort being brought together. + self.reference_dataset: Dataset | None = None + if dataset is not None and str(dataset).strip(): + filename, _flag, file_format = split_path_spec(str(dataset), default_format="mha") + self.reference_dataset = Dataset(Path(filename), file_format) + self._grid: tuple[list[int], np.ndarray, np.ndarray, np.ndarray] | None = None + # Each case's map, kept from where its own header was in hand. See _recorded(). + self._maps: dict[str, tuple[list[int], list[float], list[float]]] = {} + + # ------------------------------------------------------------------ the reference + + def _roots(self) -> list[Dataset]: + return [self.reference_dataset] if self.reference_dataset is not None else list(self.datasets) + + def _group_in(self, dataset: Dataset) -> str: + """Which group of ``dataset`` holds the reference — the declared one, or its only one.""" + if self.group is not None: + return self.group + groups = [str(group) for group in dataset.get_group()] + if len(groups) == 1: + return groups[0] + raise TransformError( + f"'ResampleToReference' cannot tell which group of '{dataset.filename}' holds entry" + f" '{self.entry}': it has {len(groups)} ({', '.join(sorted(groups)) or 'none'}).", + "Name it: ResampleToReference: {entry: " + self.entry + ", group: }.", + ) + + def reference_grid(self) -> tuple[list[int], np.ndarray, np.ndarray, np.ndarray]: + """The reference's grid, read from its header once: extent (Z, Y, X) and Origin/Spacing/Direction. + + Headers only, and memoized: a grid is declared once for the stage while a case is one of + many, so re-reading it per case would be the same answer bought again. + """ + if self._grid is not None: + return self._grid + roots = self._roots() + if not roots: + raise TransformError( + f"'ResampleToReference' has no dataset to look entry '{self.entry}' up in.", + "Give the stage a root of its own -- ResampleToReference: {entry: " + + self.entry + + ", dataset: ./Reference:omezarr} -- or run it in a workflow, which hands its" + " dataset_filenames to every stage.", + ) + for dataset in roots: + group = self._group_in(dataset) + if dataset.is_dataset_exist(group, self.entry): + shape, attribute = dataset.get_infos(group, self.entry) + spatial = [int(extent) for extent in shape[1:]] + origin, spacing, direction = self._geometry(attribute, len(spatial), f"reference '{self.entry}'") + self._grid = (spatial, origin, spacing, direction) + return self._grid + raise TransformError( + f"'ResampleToReference' cannot find entry '{self.entry}'" + + (f" in group '{self.group}'" if self.group is not None else "") + + f" in {', '.join(str(dataset.filename) for dataset in roots)}.", + "Check the entry name and its group; the reference is looked up by entry, not by the" + " case being processed, because one grid serves the whole cohort.", + ) + + @staticmethod + def _geometry(attribute: Attribute, rank: int, what: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """``(Origin, Spacing, Direction)`` in physical (x, y, z), or a refusal naming what is missing.""" + missing = [key for key in ("Origin", "Spacing", "Direction") if key not in attribute] + if missing: + raise TransformError( + f"'ResampleToReference' needs the geometry of {what} and its header carries no {', '.join(missing)}.", + "Resampling onto another grid happens in physical space: without an origin, a" + " spacing and a direction there is no space to do it in. Use a source whose" + " geometry is readable (mha, nii, h5, or an OME-Zarr written by KonfAI).", + ) + origin = np.asarray(attribute.get_np_array("Origin"), dtype=np.float64).ravel() + spacing = np.asarray(attribute.get_np_array("Spacing"), dtype=np.float64).ravel() + direction = np.asarray(attribute.get_np_array("Direction"), dtype=np.float64).ravel() + if origin.size != rank or spacing.size != rank or direction.size != rank * rank: + raise TransformError( + f"'ResampleToReference' read a geometry of {what} that does not describe a" + f" {rank}-dimensional grid: Origin {origin.size}, Spacing {spacing.size}," + f" Direction {direction.size} (expected {rank}, {rank}, {rank * rank}).", + ) + if not np.all(spacing > 0.0): + raise TransformError( + f"'ResampleToReference' read a Spacing of {spacing.tolist()} on {what}.", + "A spacing is a physical extent per voxel and must be positive on every axis.", + ) + return origin, spacing, direction.reshape(rank, rank) + + # ------------------------------------------------------------------ the map + + def grid_map( + self, name: str, shape: list[int], cache_attribute: Attribute + ) -> tuple[list[int], list[float], list[float]]: + """``(target extent, scales, offsets)`` in array order — where each target voxel reads from. + + A target voxel ``o`` is the physical point ``O_ref + D (S_ref * o)``; the source index of + that point is ``(D^-1 (p - O_src)) / S_src``. With one shared ``D`` the two compose to + ``scale * o + offset`` per axis, which is the whole map -- and the reason a differing + ``D`` is refused rather than approximated: it would make the map a rotation, whose source + region for a target box is a rotated box no per-axis window can describe. + """ + target, ref_origin, ref_spacing, ref_direction = self.reference_grid() + where = f"case '{name}'" if name else "the case" + if len(target) != len(shape): + raise TransformError( + f"'ResampleToReference' cannot resample {where}, which has {len(shape)} spatial" + f" axis/axes, onto reference '{self.entry}', which has {len(target)}.", + ) + origin, spacing, direction = self._geometry(cache_attribute, len(shape), where) + if not np.allclose(direction, ref_direction, rtol=0.0, atol=1e-6): + raise TransformError( + f"'ResampleToReference' will not resample {where} onto reference '{self.entry}':" + f" their Direction cosines differ ({direction.ravel().tolist()} against" + f" {ref_direction.ravel().tolist()}).", + "The two grids' axes do not line up, so the map between them is a rotation rather" + " than a scale and a shift per axis. Bring them to a common orientation first" + " (Canonical), or make the reference one of the cases as they are stored.", + ) + # (x, y, z) throughout, then reversed once at the end: Origin/Spacing are physical order and + # the scales/offsets a region window is cut with are array order. + scale_xyz = ref_spacing / spacing + offset_xyz = (direction.T @ (ref_origin - origin)) / spacing + scales = [float(value) for value in scale_xyz[::-1]] + offsets = [float(value) for value in offset_xyz[::-1]] + self._refuse_if_disjoint(name, shape, target, scales, offsets) + if name: + self._maps[name] = (target, scales, offsets) + return target, scales, offsets + + def _recorded(self, name: str) -> tuple[list[int], list[float], list[float]]: + """The map computed for this case back when its own header was in hand. + + A REGION read hands back the REGION's ``Origin`` — honest about what it read, and not the + case's. A map recomputed from the attribute a streamed region arrives with would therefore + place the case by the corner of whichever slab is being written, and slide it further with + every slab; every voxel would still be an interpolation of real data, and nothing about the + result would look wrong. + + Where a case's placement is known is where its own header is: :meth:`transform_shape`, which + the manager calls for every case as it is built, before any part of it is read. + """ + recorded = self._maps.get(name) + if recorded is None: + raise TransformError( + f"'ResampleToReference' was asked for a region of case '{name}' before its grid was established.", + "This is a bug if it was reached: transform_shape records the map of every case as" + " its manager is built, and a region is only ever streamed afterwards.", + ) + return recorded + + def _refuse_if_disjoint( + self, name: str, shape: list[int], target: list[int], scales: list[float], offsets: list[float] + ) -> None: + """Refuse a case that does not meet the reference grid anywhere. + + Its output would be ``fill`` from edge to edge. That is not an error the arithmetic can + find -- every voxel of it is exactly what was asked for -- so it is one nothing downstream + would report: a median over the cohort would simply be pulled toward the background by a + member that contributed no anatomy. Counted from the headers, before a byte is read. + """ + covered = self.coverage(shape, target, scales, offsets) + if covered > 0.0: + return + where = f"case '{name}'" if name else "the case" + raise TransformError( + f"'ResampleToReference' would write {where} as nothing but 'fill': it does not overlap" + f" reference '{self.entry}' anywhere, so no voxel of the reference grid reads from it.", + "The two are in different places in physical space. Check that they share a frame" + " (an acquisition's stage coordinates are not an anatomical one), pick a reference the" + " cohort actually surrounds, or drop this case with 'subset'.", + ) + + @staticmethod + def coverage(shape: list[int], target: list[int], scales: list[float], offsets: list[float]) -> float: + """The fraction of the reference grid that reads from inside the case, from headers alone. + + The product of the per-axis fractions, which is exact: the sampled set is a box, so a target + voxel is inside exactly where every one of its axes is. + """ + fraction = 1.0 + for axis, extent in enumerate(target): + index = np.arange(extent) * scales[axis] + offsets[axis] + inside = int(np.count_nonzero((index >= -0.5) & (index < shape[axis] - 0.5))) + fraction *= inside / extent if extent else 0.0 + return float(fraction) + + #: Below this, a case is worth a line in the plan: it reaches only part of the reference grid + #: and the rest of what it writes is fill. Above it, the note would round to "100.0%" and say + #: nothing, and a plan that says nothing on every line is one nobody reads. + _WORTH_SAYING = 0.999 + + def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribute: Attribute) -> str | None: + del group_dest + target, scales, offsets = self.grid_map(name, shape, cache_attribute) + covered = self.coverage(shape, target, scales, offsets) + if covered >= self._WORTH_SAYING: + return None + return ( + f"case '{name}' covers {covered * 100:.1f}% of reference '{self.entry}';" + f" the rest of what it writes is fill ({self.fill_value:g})" + ) + + # ------------------------------------------------------------------ the contract + + def transform_shape(self, group_src: str, name: str, shape: list[int], cache_attribute: Attribute) -> list[int]: + return self.grid_map(name, shape, cache_attribute)[0] + + def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: + # Not RESCALE: that kind's map is a size ratio, computed by the dispatcher from the two + # extents, and this stage's target grid has an origin of its own. REGRID hands the map back + # to the stage -- stream_region_source below -- which is the only place it is known. + return PatchLocality(LocalityKind.REGRID) + + def stream_region_source( + self, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute + ) -> list[slice]: + shape = [int(extent) for extent in source_spatial_shape] + _target, scales, offsets = self.grid_map("", shape, cache_attribute) + return Resample.source_window(target_slices, scales, shape, offsets=offsets) + + def stream_region( + self, name: str, tensor: torch.Tensor, context: RegionContext, cache_attribute: Attribute + ) -> torch.Tensor: + # The recorded map, not one read off `cache_attribute`: what arrives here describes the + # REGION, down to an Origin of its own. See _recorded(). + del cache_attribute + _target, scales, offsets = self._recorded(name) + return self.resample_region( + tensor, + tuple(context.target), + [sl.start for sl in context.source], + scales, + [int(extent) for extent in context.source_shape], + offsets, + ) + + def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: + shape = [int(extent) for extent in tensor.shape[1:]] + target, scales, offsets = self.grid_map(name, shape, cache_attribute) + # The same sampler the streamed path runs, over one region that happens to be the whole + # grid: equality between the two paths is then a property of the code, not a claim about it. + result = self.resample_region( + tensor, tuple(slice(0, extent) for extent in target), [0] * len(shape), scales, shape, offsets + ) + self.write_stream_cache_attribute(cache_attribute, shape) + return result + + def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: + target, origin, spacing, direction = self.reference_grid() + # The case now IS the reference grid, header and all -- which is the point, and what + # Reduce's `grid: strict` reads back. Pushed, not replaced: the source geometry stays + # underneath for inverse() to pop back to. + cache_attribute["Spacing"] = spacing + cache_attribute["Origin"] = origin + cache_attribute["Direction"] = direction.ravel() + cache_attribute["Size"] = np.asarray([int(extent) for extent in source_spatial_shape]) + cache_attribute["Size"] = np.asarray(target) + + # ------------------------------------------------------------------ the inverse + + def _inverse_geometry(self, cache_attribute: Attribute) -> list[int]: + size = super()._inverse_geometry(cache_attribute) + # The forward pushed Origin and Direction as well as Spacing and Size; the parent pops the + # two it knows about, and these are the two it does not. + for key in ("Origin", "Direction"): + cache_attribute.pop_np_array(key) + return size + + def inverse_patch_locality(self, cache_attribute: Attribute) -> PatchLocality: + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=( + "a resample onto a reference grid inverts to a resample back off it, and the region" + " remap for that inverse is not declared -- so a prediction finalize through this" + " stage assembles the volume. The forward direction streams" + ), + ) + + def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: + shape = [int(extent) for extent in tensor.shape[1:]] + target = self._inverse_geometry(cache_attribute) + # _inverse_geometry restored the case's own geometry, so this is the forward map again -- + # read off the same headers -- and the inverse is that map solved for the other index. + _spatial, scales, offsets = self.grid_map(name, target, cache_attribute) + back_scales = [1.0 / scale for scale in scales] + back_offsets = [-offset / scale for offset, scale in zip(offsets, scales, strict=True)] + return self.resample_region( + tensor, tuple(slice(0, extent) for extent in target), [0] * len(shape), back_scales, shape, back_offsets + ) + + def stream_region_target( + self, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute + ) -> list[slice]: + raise TransformError( + "'ResampleToReference' declares a whole-volume inverse and has no target region remap.", + "This is a bug if it was reached: the streamed-write dispatcher should not ask a" + " WHOLE_VOLUME inverse for its regions.", + ) + + class ResampleTransform(TransformInverse): """Resample a volume through stored transforms (a displacement field, an affine). @@ -2430,6 +2933,13 @@ def __init__( # segmentation's patch_size: it degrades the result; the allocator hint below keeps memory in check). self.config_overrides = config_overrides + def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribute: Attribute) -> str | None: + del name, shape, cache_attribute + return ( + f"chain '{group_dest}' runs a NESTED KonfAI inference: its GPU and RAM usage live" + " outside the declared memory_budget, and the plan cannot bound them" + ) + def infer_entry(self, dataset_path: Path, output_path: Path, gpu: list[int]): # Defragment the nested run's CUDA allocator: a heavy model (e.g. a 3D segmentation a metric relies # on) can OOM on a large volume purely from reserved-but-unallocated blocks even though the live diff --git a/konfai/transformer.py b/konfai/transformer.py index 1cea425e..6aca6c49 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -40,7 +40,7 @@ from konfai.data.case_reduction import CaseReduction, split_chain from konfai.data.data_manager import DataTransform, _format_gib, _node_local_ranks from konfai.data.patching import DatasetManager, _save_destination -from konfai.data.transform import Save, split_expand +from konfai.data.transform import Save, Transform, split_expand from konfai.utils.config import apply_config, config from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import ConfigError, TransformerError @@ -98,6 +98,10 @@ class TransformPlan: world_size: int dropped_cases: dict[str, int] dtype_hypothesis: str + #: What the stages themselves asked the plan to say (``Transform.plan_note``) — a cost the + #: columns above have no room for. Part of the plan, not of the run, so ``--plan`` carries it: + #: a note only worth reading after the bytes are written is not worth printing. + notes: tuple[str, ...] = () @property def fallback_entries(self) -> list[TransformPlanEntry]: @@ -135,6 +139,7 @@ def report(self) -> str: f"[Transformer] {dropped} case(s) present in '{group_src}' only are DROPPED:" " several groups_src keep the intersection of their case names." ) + lines.extend(f"[Transformer] NOTE: {note}" for note in self.notes) by_chain: dict[tuple[str, str], list[TransformPlanEntry]] = {} for entry in self.entries: by_chain.setdefault((entry.group_src, entry.group_dest), []).append(entry) @@ -515,7 +520,39 @@ def compute_plan(self, world_size: int = 1, overwrite: bool = False) -> Transfor available.update(dataset.get_names(group_src)) dropped[group_src] = len(available - kept) dtype_hypothesis = f"{'/'.join(sorted(planned_dtypes)) or 'float32'} / source channels" - return TransformPlan(entries, per_rank_budget, budget.description, world_size, dropped, dtype_hypothesis) + return TransformPlan( + entries, + per_rank_budget, + budget.description, + world_size, + dropped, + dtype_hypothesis, + tuple(self._plan_notes()), + ) + + def _plan_notes(self) -> list[str]: + """What the stages themselves ask the plan to say, in chain order, each said once. + + A verdict and a byte count are what the plan can compute ABOUT a chain; this is what the + chain knows about itself — a nested inference whose memory nothing here can bound, a case + that meets only part of the grid it is being resampled onto. Deduplicated because a note + about the stage repeats identically for every case of its chain, while a note about the + case does not: what is printed is the set of distinct things there are to say. + """ + notes: list[str] = [] + for group_dest, managers in self._managers().items(): + for manager in managers: + for stage in manager.transforms: + # A chain's stages are transforms AND draws; only a transform declares a note. + # A draw has nothing to add anyway: what its copies cost is the `regime` column. + if not isinstance(stage, Transform): + continue + note = stage.plan_note( + group_dest, manager.name, list(manager.base_shape[1:]), manager.stored_attributes + ) + if note is not None and note not in notes: + notes.append(note) + return notes def setup(self, world_size: int): """Plan, print, enforce, shard — before any spawn, before any byte.""" @@ -557,18 +594,6 @@ def setup(self, world_size: int): "Use one process, or a directory destination (omezarr, mha, nii.gz).", ) - inference_chains = sorted( - group_dest - for group_dest, managers in self._managers().items() - if managers and any(type(t).__name__ == "KonfAIInference" for t in managers[0].transforms) - ) - if inference_chains: - print( - f"[Transformer] NOTE: chain(s) {', '.join(inference_chains)} run a NESTED KonfAI" - " inference: its GPU and RAM usage live outside the declared memory_budget, and the" - " plan cannot bound them." - ) - violations = plan.budget_violations() if violations: worst = max(violations, key=lambda entry: entry.working_set_bytes) diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py new file mode 100644 index 00000000..953ec27e --- /dev/null +++ b/tests/unit/test_resample_to_reference.py @@ -0,0 +1,477 @@ +# Copyright (c) 2025 Valentin Boussot +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Resampling a case onto a declared reference grid, and what that refuses. + +The streamed-equals-whole-volume property is proven for this stage where it is proven for every +other, in ``test_transform_locality_contract.py``; what is proven here is the half that contract +cannot see. A stage can be perfectly self-consistent -- both its paths agreeing on the same wrong +place -- so the arithmetic is checked against SimpleITK, which resamples in physical space and knows +nothing about this file. And a stage that lands on another grid can be wrong in ways no equality +test reaches: by writing the right voxels under the wrong header, or by writing a case that never +met the reference at all. +""" + +from pathlib import Path + +import numpy as np +import pytest +import torch +from konfai.data.case_reduction import CaseReduction +from konfai.data.data_manager import _check_patch_transform_locality +from konfai.data.patching import DatasetManager, DatasetPatch +from konfai.data.transform import LocalityKind, Reduce, Resample, ResampleToReference, Write +from konfai.utils.dataset import Attribute, Dataset +from konfai.utils.errors import ConfigError, TransformError + +pytest.importorskip("SimpleITK") +import SimpleITK as sitk + +_CASE = "CASE_000" +_SOURCE_SPATIAL = (9, 13, 11) +_REFERENCE_SPATIAL = (7, 10, 15) +# Physical (x, y, z), as a header stores them. Nothing lines up: neither extent, nor spacing, nor +# origin -- and the reference reaches past the case on x, so part of it has no data to read. +_SOURCE_ORIGIN, _SOURCE_SPACING = [-3.0, 5.0, 11.0], [1.5, 1.1, 2.0] +_REFERENCE_ORIGIN, _REFERENCE_SPACING = [-1.25, 4.2, 12.7], [1.9, 1.7, 1.3] +_FILL = -777.0 + + +def _attributes(origin: list[float], spacing: list[float], direction: np.ndarray | None = None) -> Attribute: + attributes = Attribute() + attributes["Origin"] = np.asarray(origin) + attributes["Spacing"] = np.asarray(spacing) + attributes["Direction"] = (np.eye(3) if direction is None else direction).reshape(-1) + return attributes + + +def _volume(shape: tuple[int, ...], seed: int = 0) -> np.ndarray: + # A step, not a smooth field: interpolating a smooth volume onto a shifted grid gives nearly the + # right answer even when the shift is wrong, so a smooth fixture would pass a broken map. + rng = np.random.default_rng(seed) + return (rng.normal(size=shape) * 100).astype(np.float32)[None] + + +@pytest.fixture +def dataset(tmp_path: Path) -> Dataset: + """A case and a reference, on grids that agree about nothing but their direction.""" + dataset = Dataset(tmp_path / "Dataset", "mha") + dataset.write("Case", _CASE, _volume(_SOURCE_SPATIAL), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + dataset.write( + "Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING) + ) + return dataset + + +def _stage(dataset: Dataset, **kwargs: object) -> ResampleToReference: + arguments: dict[str, object] = {"entry": _CASE, "group": "Reference", "fill": _FILL, **kwargs} + stage = ResampleToReference(**arguments) # type: ignore[arg-type] + stage.set_datasets([dataset]) + return stage + + +def _manager(dataset: Dataset, stage: ResampleToReference, group: str = "Case") -> DatasetManager: + return DatasetManager( + index=0, + group_src=group, + group_dest=group, + name=_CASE, + dataset=dataset, + patch=DatasetPatch([4, 4, 4]), + transforms=[stage], + data_augmentations_list=[], + ) + + +def _simpleitk(volume: np.ndarray, source: Attribute, reference: Attribute, nearest: bool = False) -> np.ndarray: + """The same resample, done by SimpleITK — the oracle this stage's arithmetic is checked against.""" + image = sitk.GetImageFromArray(volume[0]) + image.SetOrigin(source.get_np_array("Origin").tolist()) + image.SetSpacing(source.get_np_array("Spacing").tolist()) + image.SetDirection(source.get_np_array("Direction").tolist()) + grid = sitk.Image(*reversed(_REFERENCE_SPATIAL), sitk.sitkFloat32) + grid.SetOrigin(reference.get_np_array("Origin").tolist()) + grid.SetSpacing(reference.get_np_array("Spacing").tolist()) + grid.SetDirection(reference.get_np_array("Direction").tolist()) + interpolator = sitk.sitkNearestNeighbor if nearest else sitk.sitkLinear + return sitk.GetArrayFromImage(sitk.Resample(image, grid, sitk.Transform(), interpolator, _FILL)) + + +# --------------------------------------------------------------------- the arithmetic + + +def test_it_resamples_where_simpleitk_does(dataset: Dataset) -> None: + """The check the streamed-equals-whole-volume contract cannot make: is the place right at all. + + Both of this stage's paths run one sampler, so they agree with each other by construction -- + including on a grid placed in the wrong spot. SimpleITK resamples in physical space through an + implementation that shares no line with this one, so agreeing with it is evidence about the + geometry rather than about the code's self-consistency. + """ + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + volume = dataset.read_data("Case", _CASE)[0] + got = _stage(dataset)(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy()[0] + want = _simpleitk(volume, source, _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING)) + + assert got.shape == want.shape == _REFERENCE_SPATIAL + # float32 weights summed in a different order than ITK's nested lerps: a few ulps of the data's + # own range, not a difference of placement (which would be a whole voxel of gradient). + np.testing.assert_allclose(got, want, rtol=0, atol=64 * float(np.spacing(np.float32(np.abs(volume).max())))) + + +def test_the_edge_of_the_data_is_where_simpleitk_puts_it(dataset: Dataset) -> None: + """The fill boundary, voxel for voxel — the half of a regrid that a tolerance cannot check. + + A map off by one voxel still interpolates real data almost everywhere; where it shows is at the + rim, in which voxels stop having a source at all. Counting them against ITK's own + ``[-0.5, n - 0.5)`` is what pins the convention rather than assuming it. + """ + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + volume = dataset.read_data("Case", _CASE)[0] + got = _stage(dataset)(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy()[0] + want = _simpleitk(volume, source, _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING)) + + assert 0 < int((want == _FILL).sum()) < want.size, "the fixture must have a rim, and not be all rim" + np.testing.assert_array_equal(got == _FILL, want == _FILL) + + +def test_a_label_map_takes_the_nearest_voxel_simpleitk_takes(dataset: Dataset, tmp_path: Path) -> None: + """uint8 resamples by nearest, and nearest here is ITK's round-half-up on the physical index. + + ``F.interpolate``'s nearest is ``floor(o * scale)`` -- a statement about a size ratio, which says + nothing once the target grid has an origin of its own. A label map interpolated by the wrong rule + is still a label map, so nothing downstream would report it. + """ + labels = (np.arange(int(np.prod(_SOURCE_SPATIAL))) % 4).reshape(_SOURCE_SPATIAL).astype(np.uint8)[None] + dataset.write("Labels", _CASE, labels, _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + + got = _stage(dataset, fill=0.0)(_CASE, torch.from_numpy(labels.copy()), Attribute(source)).numpy()[0] + image = sitk.GetImageFromArray(labels[0]) + image.SetOrigin(_SOURCE_ORIGIN) + image.SetSpacing(_SOURCE_SPACING) + grid = sitk.Image(*reversed(_REFERENCE_SPATIAL), sitk.sitkUInt8) + grid.SetOrigin(_REFERENCE_ORIGIN) + grid.SetSpacing(_REFERENCE_SPACING) + want = sitk.GetArrayFromImage(sitk.Resample(image, grid, sitk.Transform(), sitk.sitkNearestNeighbor, 0)) + + assert got.dtype == np.uint8 + np.testing.assert_array_equal(got, want) + + +@pytest.mark.parametrize("dtype", [np.uint16, np.int16, np.float32]) +def test_it_resamples_the_dtypes_a_microscope_and_a_scanner_store(tmp_path: Path, dtype: type) -> None: + """uint16 is what a light-sheet volume IS, and torch fills only some integer dtypes. + + ``masked_fill`` is unimplemented for uint16, so filling after the cast back to the source dtype + raises on precisely the volumes this stage was built for -- and only where the reference grid + reaches past the case, which is to say on the interesting ones. + """ + dataset = Dataset(tmp_path / f"Dtype{np.dtype(dtype).name}", "mha") + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + volume = (np.arange(int(np.prod(_SOURCE_SPATIAL))) % 900).reshape(_SOURCE_SPATIAL).astype(dtype)[None] + dataset.write("Case", _CASE, volume, source) + dataset.write( + "Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING) + ) + + got = _stage(dataset, fill=0.0)(_CASE, torch.from_numpy(volume.copy()), Attribute(source)) + assert got.numpy().dtype == dtype + assert list(got.shape[1:]) == list(_REFERENCE_SPATIAL) + # The rim the reference reaches past the case takes the fill, in the source's own dtype. + assert int((got.numpy() == 0).sum()) > 0 + + +def test_an_oblique_pair_resamples_where_simpleitk_does(tmp_path: Path) -> None: + """A shared non-axis-aligned direction is legal, and the origin shift travels through it. + + The offset is ``D^-1 (O_ref - O_src) / S_src``: drop the ``D^-1`` and an axis-aligned pair still + lands perfectly, because there ``D`` is the identity. Only an oblique pair can tell. + """ + direction = np.linalg.qr(np.asarray([[0.936, -0.352, 0.0], [0.352, 0.936, 0.0], [0.0, 0.0, 1.0]]))[0] + dataset = Dataset(tmp_path / "Oblique", "mha") + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING, direction) + reference = _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING, direction) + volume = _volume(_SOURCE_SPATIAL) + dataset.write("Case", _CASE, volume, source) + dataset.write("Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), reference) + + got = _stage(dataset)(_CASE, torch.from_numpy(volume.copy()), Attribute(source)).numpy()[0] + want = _simpleitk(volume, source, reference) + np.testing.assert_allclose(got, want, rtol=0, atol=64 * float(np.spacing(np.float32(np.abs(volume).max())))) + + +# --------------------------------------------------------------------- the header + + +def test_the_case_lands_on_the_reference_grid(dataset: Dataset) -> None: + """Extent, spacing, origin and direction all adopted — which is the whole point of the stage. + + Right voxels under the wrong header is the failure this stage exists to prevent, and the one a + value comparison cannot see: ``Reduce(grid: strict)`` reads exactly these four. + """ + manager = _manager(dataset, _stage(dataset)) + landed = manager.landed_attributes() + + assert list(manager.spatial_shape) == list(_REFERENCE_SPATIAL) + np.testing.assert_allclose(landed.get_np_array("Origin"), _REFERENCE_ORIGIN) + np.testing.assert_allclose(landed.get_np_array("Spacing"), _REFERENCE_SPACING) + np.testing.assert_allclose(landed.get_np_array("Direction"), np.eye(3).reshape(-1)) + + +def test_a_cohort_on_one_reference_passes_grid_strict(tmp_path: Path) -> None: + """The end this stage is for: heterogeneous cases fold under ``grid: strict``, which is a real check. + + Without the stage the same cohort disagrees on extent AND on geometry, so the reduction has to be + told to look away (``shape_only``). Both halves are asserted, because "strict passes" only means + something if strict would otherwise have refused. + """ + dataset = Dataset(tmp_path / "Cohort", "mha") + origins = [[-3.0, 5.0, 11.0], [-2.0, 5.6, 11.4], [-3.4, 4.7, 10.6]] + shapes = [(9, 13, 11), (8, 12, 12), (10, 13, 10)] + for index, (origin, shape) in enumerate(zip(origins, shapes, strict=True)): + dataset.write("Case", f"C{index}", _volume(shape, index), _attributes(origin, _SOURCE_SPACING)) + dataset.write( + "Reference", "GRID", _volume(_REFERENCE_SPATIAL, 9), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING) + ) + + def managers(with_stage: bool) -> list[DatasetManager]: + built = [] + for index in range(len(shapes)): + stages: list[object] = [] + if with_stage: + stage = ResampleToReference(entry="GRID", group="Reference", fill=_FILL) + stage.set_datasets([dataset]) + stages.append(stage) + built.append( + DatasetManager( + index=index, + group_src="Case", + group_dest="Case", + name=f"C{index}", + dataset=dataset, + patch=None, + transforms=[*stages], + data_augmentations_list=[], + ) + ) + return built + + bare = CaseReduction( + managers(False), Reduce(operator="Median", output="template", grid="strict"), [], dataset, "Case" + ) + assert bare.check_grid() is not None, "the cohort must disagree, or this proves nothing" + + folded = CaseReduction( + managers(True), Reduce(operator="Median", output="template", grid="strict"), [], dataset, "Case" + ) + assert folded.check_grid() is None + assert [list(manager.spatial_shape) for manager in folded.managers] == [list(_REFERENCE_SPATIAL)] * 3 + + +# --------------------------------------------------------------------- the memory bound + + +def test_it_never_assembles_the_volume(dataset: Dataset, tmp_path: Path) -> None: + """The memory bound, asserted as a bound: a regridded case is written without ever being loaded. + + Values alone would pass even if the chain had read everything into RAM first; only forbidding the + whole-volume read proves that a case larger than memory can go through this stage at all. + """ + stage = _stage(dataset) + manager = DatasetManager( + index=0, + group_src="Case", + group_dest="Case", + name=_CASE, + dataset=dataset, + patch=None, + transforms=[stage, Write(str(tmp_path / "Out"))], + data_augmentations_list=[], + ) + assert manager.stream_refusal(0) is None + + def refuse(*args: object, **kwargs: object) -> None: + raise AssertionError("the chain read the whole volume") + + monkeypatched = pytest.MonkeyPatch() + monkeypatched.setattr(Dataset, "read_data", refuse) + try: + assert manager.materialize() is True + finally: + monkeypatched.undo() + + written, attributes = Dataset(tmp_path / "Out", "mha").read_data("Case", _CASE) + assert list(written.shape[1:]) == list(_REFERENCE_SPATIAL) + np.testing.assert_allclose(attributes.get_np_array("Origin"), _REFERENCE_ORIGIN) + np.testing.assert_allclose(attributes.get_np_array("Spacing"), _REFERENCE_SPACING) + + +@pytest.mark.parametrize("offset", [-500.0, 500.0]) +def test_a_region_off_the_source_reads_one_voxel(offset: float) -> None: + """A target region with no source under it pulls one voxel, not the extent it was clamped from. + + A case whose grid meets the reference somewhere still has slabs that do not -- that is the + ordinary shape of a cohort resampled onto one grid -- and every one of them reads a window it + will overwrite with fill. Clamping to a legal-but-empty region is what keeps that read at one + voxel instead of the case's whole cross-section, on every such slab. + """ + window = Resample.source_window( + tuple(slice(0, 4) for _ in _SOURCE_SPATIAL), + [1.0, 1.0, 1.0], + list(_SOURCE_SPATIAL), + offsets=[offset] * len(_SOURCE_SPATIAL), + ) + assert [sl.stop - sl.start for sl in window] == [1, 1, 1] + assert all(0 <= sl.start < extent for sl, extent in zip(window, _SOURCE_SPATIAL, strict=True)) + + +# --------------------------------------------------------------------- what it refuses + + +def test_it_declares_regrid_not_rescale(dataset: Dataset) -> None: + # RESCALE would hand the dispatcher a size ratio and lose the origin entirely, silently. + assert _stage(dataset).patch_locality(Attribute()).kind is LocalityKind.REGRID + + +def test_it_is_refused_as_a_patch_transform(dataset: Dataset, monkeypatch: pytest.MonkeyPatch) -> None: + """Under `patch_transforms:` it would hand back the whole reference extent for every patch. + + Every other region kind is refused there by name, with its own sentence; a kind missing from + that table would raise a KeyError instead of saying what to do about it. + """ + monkeypatch.setenv("KONFAI_ROOT", "Trainer") + with pytest.raises(ConfigError, match="onto another grid"): + _check_patch_transform_locality(_stage(dataset), "CT", "CT") + + +def test_a_differing_direction_is_refused(tmp_path: Path) -> None: + """Axes that do not line up make the map a rotation, which no per-axis window describes.""" + dataset = Dataset(tmp_path / "Rotated", "mha") + turned = np.asarray([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + dataset.write("Case", _CASE, _volume(_SOURCE_SPATIAL), source) + dataset.write( + "Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING, turned) + ) + + with pytest.raises(TransformError, match="Direction cosines differ"): + _stage(dataset).transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), source) + + +def test_a_case_with_no_geometry_is_refused(dataset: Dataset) -> None: + """No origin, no physical space to resample in — and a size ratio must not stand in for one.""" + bare = Attribute() + bare["Spacing"] = np.asarray(_SOURCE_SPACING) + with pytest.raises(TransformError, match="carries no Origin, Direction"): + _stage(dataset).transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), bare) + + +def test_a_case_that_never_meets_the_reference_is_refused(tmp_path: Path) -> None: + """Its output would be fill from edge to edge, and a median would quietly take it as anatomy. + + This is the refusal the real cohort needed: acquisition stage coordinates are not an anatomical + frame, so two brains can be metres apart in physical space and look perfectly normal apart. + """ + dataset = Dataset(tmp_path / "Apart", "mha") + source = _attributes([1000.0, 1000.0, 1000.0], _SOURCE_SPACING) + dataset.write("Case", _CASE, _volume(_SOURCE_SPATIAL), source) + dataset.write( + "Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING) + ) + + with pytest.raises(TransformError, match="nothing but 'fill'"): + _stage(dataset).transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), source) + + +def test_an_unknown_entry_is_refused(dataset: Dataset) -> None: + stage = ResampleToReference(entry="NOT_THERE", group="Reference") + stage.set_datasets([dataset]) + with pytest.raises(TransformError, match="cannot find entry 'NOT_THERE'"): + stage.reference_grid() + + +def test_an_unnamed_group_is_refused_when_the_store_has_several(dataset: Dataset) -> None: + """Warp's rule: guessing which group holds the reference is not a guess worth making.""" + stage = ResampleToReference(entry=_CASE) + stage.set_datasets([dataset]) + with pytest.raises(TransformError, match="cannot tell which group"): + stage.reference_grid() + + +def test_an_empty_entry_is_refused_at_construction() -> None: + with pytest.raises(TransformError, match="needs an 'entry'"): + ResampleToReference(entry=" ") + + +# --------------------------------------------------------------------- what it announces + + +def test_the_plan_is_told_how_much_of_the_grid_the_case_covers(dataset: Dataset) -> None: + """Partial coverage is legal, common, and worth a line: the rest of the output is fill. + + Nothing else in the plan can say it -- the verdict is STREAM and the byte count is the same + either way -- so a template that is mostly background would otherwise be a discovery made in a + viewer. + """ + note = _stage(dataset).plan_note( + "Case_out", _CASE, list(_SOURCE_SPATIAL), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + ) + assert note is not None + assert "covers" in note and "fill" in note + + +def test_a_case_that_fills_the_grid_says_nothing(tmp_path: Path) -> None: + """A note on every line is a note nobody reads: full coverage is the unremarkable case.""" + dataset = Dataset(tmp_path / "Nested", "mha") + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + dataset.write("Case", _CASE, _volume((20, 20, 20)), source) + # A reference well inside the case: every voxel of it has data under it. + dataset.write("Reference", _CASE, _volume((4, 4, 4), 1), _attributes([2.0, 9.0, 15.0], [1.0, 1.0, 1.0])) + inside = ResampleToReference(entry=_CASE, group="Reference") + inside.set_datasets([dataset]) + + assert inside.plan_note("Case_out", _CASE, [20, 20, 20], source) is None + + +# --------------------------------------------------------------------- the way back + + +def test_the_inverse_returns_the_case_to_its_own_grid(dataset: Dataset) -> None: + """A prediction made on the reference grid comes back to the grid the case was stored on. + + The inverse is the same map solved for the other index, so what it must restore is the extent + and the header -- not the values, which an interpolation onto a coarser grid has already lost. + """ + stage = _stage(dataset) + source = _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + attribute = Attribute(source) + volume = dataset.read_data("Case", _CASE)[0] + + forward = stage(_CASE, torch.from_numpy(volume.copy()), attribute) + assert list(forward.shape[1:]) == list(_REFERENCE_SPATIAL) + + back = stage.inverse(_CASE, forward, attribute) + assert list(back.shape[1:]) == list(_SOURCE_SPATIAL) + np.testing.assert_allclose(attribute.get_np_array("Origin"), _SOURCE_ORIGIN) + np.testing.assert_allclose(attribute.get_np_array("Spacing"), _SOURCE_SPACING) + + +def test_the_inverse_declares_the_whole_volume_and_says_why(dataset: Dataset) -> None: + """Declared, not discovered: the write-side region remap is not implemented, so it says so.""" + locality = _stage(dataset).inverse_patch_locality(Attribute()) + assert locality.kind is LocalityKind.WHOLE_VOLUME + assert locality.reason is not None and "reference grid" in locality.reason diff --git a/tests/unit/test_transform_locality_contract.py b/tests/unit/test_transform_locality_contract.py index 9198f7c2..007f3a67 100644 --- a/tests/unit/test_transform_locality_contract.py +++ b/tests/unit/test_transform_locality_contract.py @@ -62,6 +62,7 @@ PatchLocality, Percentage, Reduce, + ResampleToReference, ResampleToResolution, ResampleToShape, ResampleTransform, @@ -87,6 +88,14 @@ _PATCH_SIZE = [4, 4, 4] _SPACING = [1.5, 1.5, 2.0] +# The grid the "Reference" group is stored on, for a stage that resamples onto a declared reference. +# Chosen so the two boxes OVERLAP WITHOUT NESTING: in physical x the reference runs to 13.4 where the +# case stops at 12.0, so its last columns read from outside the case and take the fill. A reference +# contained in its case would prove the sampler and never the boundary, which is the half that +# differs between the streamed and whole-volume paths. +_REFERENCE_SPATIAL = (7, 8, 9) +_REFERENCE_SPACING = [1.8, 1.2, 2.5] + # No extent is a multiple of the patch size, so the last patch of every axis is a border patch the read # plan has to pad: the grid is 3x3x3 and 19 of its 27 patches touch a border. _PEAK = 450.0 @@ -158,6 +167,14 @@ class _Case: _Case(ResampleToResolution([2.0, 1.0, 3.0]), group="Labels"), ], "ResampleToShape": [_Case(ResampleToShape([12, 8, 14]), atol=_RESCALE_ATOL)], + # Onto a grid of its own, so part of the target reads from outside the case and takes the fill. + # atol is 0: unlike the scale-only resamples -- whose whole-volume path is F.interpolate and whose + # streamed path is resample_region, two implementations that agree to a rounding -- both paths of + # this one run the SAME sampler over global coordinates, so they agree bit for bit or not at all. + "ResampleToReference": [ + _Case(ResampleToReference(entry=_CASE_NAME, group="Reference")), + _Case(ResampleToReference(entry=_CASE_NAME, group="Reference"), group="Labels"), + ], "ResampleTransform": [_Case(ResampleTransform({"transform": True}))], "Save": [_Case(Save("Dataset"))], # Warp needs a field on disk to run, which this registry cannot build: with no declared @@ -226,14 +243,17 @@ def _volumes() -> dict[str, np.ndarray]: # Stored with the foreground box already on it, which is how a crop is a translation rather # than a question about the voxels. "Boxed": intensity.astype(np.float32)[None], + # A grid of its OWN -- other extent, other spacing, other origin -- for a stage that resamples + # onto a reference rather than about the case's own extent. Only its header is ever read. + "Reference": rng.standard_normal(_REFERENCE_SPATIAL).astype(np.float32)[None], } def _attributes(group: str) -> Attribute: """The metadata a group is stored with -- and so what a declaration about it is handed.""" attributes = Attribute() - attributes["Origin"] = np.asarray([-3.0, 5.0, 11.0]) - attributes["Spacing"] = np.asarray(_SPACING) + attributes["Origin"] = np.asarray([-1.0, 6.0, 12.0] if group == "Reference" else [-3.0, 5.0, 11.0]) + attributes["Spacing"] = np.asarray(_REFERENCE_SPACING if group == "Reference" else _SPACING) attributes["Direction"] = {"Oblique": _OBLIQUE, "Permuting": _PERMUTING}.get(group, _AXIS_ALIGNED).reshape(-1) if group == "Ensemble": # What a `combine: Concat` reduction writes: the per-model channel counts MergeLabels and @@ -299,6 +319,9 @@ def _streamable_cases() -> list[_Case]: def _manager(dataset: Dataset, case: _Case) -> DatasetManager: + # What a run does before it builds a manager (Data.prepare): a stage that reads a SECOND entry -- + # a mask, a field, a reference grid -- is handed the roots to find it in. + case.transform.set_datasets([dataset]) return DatasetManager( index=0, group_src=case.group, From 16ee5b32f6c643ca4881b562644a168730dbdebe Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Tue, 4 Aug 2026 10:52:36 +0200 Subject: [PATCH 2/3] feat(transform): compose the grid change and the warp into one pass The stage takes an optional field, so two passes over a volume become one and the intermediate -- a case resampled onto the reference but not yet warped -- never has to exist on disk. ShapeUpdate goes: it was never generic, and the shared stage covers what it did. --- docs/source/config_guide/transform.md | 84 +- konfai/data/transform.py | 760 +++++++++++++----- .../test_transform_doc_examples.py | 21 + ...ate.py => test_per_component_statistic.py} | 50 +- tests/unit/test_resample_to_reference.py | 312 +++++++ .../unit/test_transform_locality_contract.py | 29 +- 6 files changed, 1000 insertions(+), 256 deletions(-) rename tests/unit/{test_shape_update.py => test_per_component_statistic.py} (55%) diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index c7c0fb49..c1fe2681 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -232,6 +232,54 @@ transforms: reference can live anywhere — which is the atlas loop: point round N+1 at the store round N wrote its template into. +#### Through a displacement field, in one interpolation + +Add `field:` and the stage becomes the whole of a registration's apply step — +`sitk.Resample(image, reference_grid, DisplacementFieldTransform(field))`. For +each voxel of the **target** grid the field is read at that voxel's world +position, added, and the source sampled **once** at the displaced point: + +```yaml +transforms: + ResampleToReference: + entry: case_0 + group: CT + field: ./Fields:mha + field_group: DVF + max_displacement: 4.0 + Write: {dataset: ./Registered:mha} +``` + +Fields stored *beside* the cases — one entry per case, in the same roots — need +no path at all: `field_group: DVF` on its own finds them. + +Doing this as two stages instead (resample onto the grid, then `Warp`) costs +**two** interpolations, and the second cannot restore the detail the first +smoothed away. That is not a small effect: on a high-frequency volume the +second pass moves voxels by a large fraction of the range, which is exactly why +an atlas rebuilds its appearance from native volumes rather than from resampled +ones. + +The field lives on **its own grid**, usually coarser than either the source or +the target, and is read where it is asked — it is defined in world units, so a +field solved at 120 µm moves a volume stored at 30 µm without being upsampled +first. Outside its own extent the displacement is zero: the transform is the +identity where the field says nothing, as SimpleITK has it. + +`max_displacement` sizes the source region each target region must read, and is +**checked against every field region actually read** — a field that exceeds +what it declared raises rather than sampling zeros, which would show up as a +dark rim around the moved anatomy and nothing else. It takes `auto`, reading +the bound KonfAI records on a field it writes. With no bound at all the stage +declares `WHOLE_VOLUME` and says so in the plan, exactly as `Warp` does. + +```{note} +`Warp` still exists and is not this: it adds a displacement on the case's *own* +grid, which is the shape update of an atlas build, and it neither changes the +grid nor needs a reference. Use it when the field was solved on the very grid +it is applied to. +``` + Naming an image rather than fifteen numbers is deliberate. A grid is an extent in array order `(Z, Y, X)` plus an origin, a spacing and a direction in physical `(x, y, z)` — transcribing those by hand is the mistake that actually gets made, @@ -409,32 +457,22 @@ not. `Clip` then `Normalize` therefore takes the whole-volume path, and the plan says so. Reorder the chain, or cut it with a `Save`. ``` -### `ShapeUpdate`: the shape residual of a displacement field +### Statistics a stage may ask for -The shape update of an atlas build: `output = -step * (field - t)`, where `t` is -the field's per-component spatial mean in world units. Resampling a template -through the result moves it along the cohort's shape residual, at ANTs' -gradient step. +A stage that needs a figure over the WHOLE volume does not have to assemble it. +Declaring the statistic in `PatchLocality(LocalityKind.GLOBAL_STAT, stat_keys=…)` +tells the planner to read it once from the stored volume; the stage is then a +value map, and a volume of any size runs region by region. -```yaml -transforms: - ShapeUpdate: - step: 0.25 - Write: - dataset: ./Update:omezarr -``` +Alongside the pooled `Mean`, `Min`, `Max` and `Std`, a **per-component** mean is +available as `MeanPerChannel`. It exists because a per-channel quantity has as +many parts as the volume has components, and the pooled mean of all of them +describes none of them — a three-component displacement field centred by one +number is centred on no axis. -`t` is **stripped, not applied**, and that is the whole point of the stage. A -total field maps template coordinates into each specimen's OWN world frame, so -its spatial mean is dominated by the frame-to-frame offset, not by any pose error -of the template. Applying it in full translates the template out of its own grid -and clips the anatomy; stripping it is what keeps the template anchored — the -same thing ANTs' `AverageAffineTransformNoRigid` is for. - -Because the statistic is declared rather than recomputed per region, a field of -any size runs region by region: the volume is never assembled. Handed the whole -volume anyway — a chain that fell back for another reason — the stage takes the -statistic from the tensor in hand, so both paths leave the same state behind. +Handed the whole volume anyway — a chain that fell back for another reason — a +stage should take the statistic from the tensor in hand and record it, so both +paths leave the same state behind. ## Writing your own transform diff --git a/konfai/data/transform.py b/konfai/data/transform.py index c79f23b6..ba9125a1 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -24,7 +24,7 @@ from enum import Enum from multiprocessing import current_process, get_context from pathlib import Path -from typing import Any +from typing import Any, cast import numpy as np import torch @@ -1105,6 +1105,7 @@ def _resample_offset_region( scales: list[float], n_in: list[int], offsets: list[float], + fill: float | None = None, ) -> torch.Tensor: """``resample_region`` for an offset map: ITK's sampler, and a fill where the source stops. @@ -1118,7 +1119,12 @@ def _resample_offset_region( Coordinates are global (the target index, not its offset within the region), so a region and the whole volume put the same sample in the same place -- which is what makes the streamed and whole-volume paths equal by construction here rather than by agreement. + + ``fill`` overrides :attr:`fill_value` for a caller sampling something that is not the case's + own voxels: a displacement field says nothing outside its extent, and what it means there is + zero -- the identity -- not the image's background. """ + outside = self.fill_value if fill is None else fill device = sub_tensor.device ndim = len(target_slices) window = [int(extent) for extent in sub_tensor.shape[1:]] @@ -1134,7 +1140,7 @@ def _resample_offset_region( # Nothing of this region is on the source. Real for a REGRID -- a target grid may reach # past its case -- and worth its own exit: the gather below would read a window that was # only ever clamped to something legal, and then be overwritten by the fill anyway. - return torch.full(out_shape, self.fill_value, device=device, dtype=torch.float32).type(sub_tensor.dtype) + return torch.full(out_shape, outside, device=device, dtype=torch.float32).type(sub_tensor.dtype) def local(index: torch.Tensor, k: int) -> torch.Tensor: """A global source index as an offset into the window that was read, kept in range.""" @@ -1176,7 +1182,7 @@ def local(index: torch.Tensor, k: int) -> torch.Tensor: # Filled while still floating, then cast ONCE: torch implements masked_fill for float dtypes # and not for every integer one -- uint16 is a microscope's native dtype and has no fill at # all -- so filling after the cast fails on exactly the volumes this stage is built for. - return out.masked_fill(~mask.unsqueeze(0), self.fill_value).type(sub_tensor.dtype) + return out.masked_fill(~mask.unsqueeze(0), outside).type(sub_tensor.dtype) @abstractmethod def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: @@ -1287,6 +1293,29 @@ def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatia cache_attribute["Size"] = shape +@dataclass(frozen=True) +class _ReferenceMap: + """Everything a case needs to be read onto a reference grid, computed once from the headers. + + All in array order (Z, Y, X) except ``source_spacing`` and ``direction``, which stay in the + physical (x, y, z) their headers are written in — a displacement's components are physical too, + and converting them is the one place the two orders meet. + """ + + target: list[int] + #: Where a target voxel reads from: ``source_index = scale * target_index + offset``. + scales: list[float] + offsets: list[float] + source_spacing: list[float] + direction: np.ndarray + #: The same affine map from the target grid to the FIELD's grid, when one is declared. The field + #: has a grid of its own -- typically coarser than both -- and is interpolated where it is asked, + #: which is what lets a field solved at 120 um move a volume stored at 30 um. + field_scales: list[float] | None = None + field_offsets: list[float] | None = None + field_shape: list[int] | None = None + + class ResampleToReference(Resample): """Resample a case onto the grid of a declared reference — extent, spacing, origin, direction. @@ -1304,19 +1333,38 @@ class ResampleToReference(Resample): place. Naming an image cannot make that mistake: the header IS the declaration. It is also what an atlas loop needs, where round N+1's reference is round N's own output. + ``field`` MAKES IT A RECALAGE, IN ONE INTERPOLATION. With a displacement field declared, this is + ``sitk.Resample(image, reference_grid, DisplacementFieldTransform(field))`` — for each voxel of + the TARGET grid, the field is read at that voxel's world position, added, and the source sampled + once at the displaced point. Doing it as two stages instead (this one, then ``Warp``) resamples + twice, and a volume that has been interpolated twice has lost detail the second pass invented no + more of — which is the whole reason the appearance of an atlas is rebuilt from native volumes. + + The field lives on ITS OWN grid, and is read where it is asked: it is defined in world units, so + a field solved at 120 um moves a volume stored at 30 um without being upsampled first. Outside + its own extent the displacement is zero, as SimpleITK has it — the transform is the identity + where the field says nothing. + + ``max_displacement`` sizes the source region a target region needs, exactly as on ``Warp``, and + is CHECKED against every field region actually read. + WHAT IT REFUSES, rather than resample onto a grid it cannot honestly reach: - - a case, or a reference, whose header carries no ``Origin``/``Spacing``/``Direction`` — with no - geometry there is no physical space to resample IN, and a size ratio would silently stand in - for one; - - a reference whose ``Direction`` differs from the case's — the two grids' axes then do not line + - a case, a reference or a field whose header carries no ``Origin``/``Spacing``/``Direction`` — + with no geometry there is no physical space to resample IN, and a size ratio would silently + stand in for one; + - a reference or a field whose ``Direction`` differs from the case's — the axes then do not line up, and the map is a rotation, not a scale and a shift per axis. ``Canonical`` first; - a case that does not meet the reference grid at all — the output would be pure ``fill``, and - an all-background member is a plausible, wrong contribution to a median. + an all-background member is a plausible, wrong contribution to a median; + - a field that displaces further than ``max_displacement`` declared, because the region read was + sized from that number and sampling past it returns zeros — a dark rim around the anatomy and + nothing else to see. Everything else it declares. A case that reaches only part of the reference is legal and common — the rest takes ``fill`` — and the plan prints how much of the grid each case covers, - because "most of this template is fill" is not something to discover in a viewer. + because "most of this template is fill" is not something to discover in a viewer. A field with + no bound at all declares ``WHOLE_VOLUME`` and says so, as ``Warp`` does. """ def __init__( @@ -1324,6 +1372,9 @@ def __init__( entry: str, group: str | None = None, dataset: str | None = None, + field: str | None = None, + field_group: str | None = None, + max_displacement: float | str = 0.0, fill: float = 0.0, inverse: bool = True, ) -> None: @@ -1343,9 +1394,30 @@ def __init__( if dataset is not None and str(dataset).strip(): filename, _flag, file_format = split_path_spec(str(dataset), default_format="mha") self.reference_dataset = Dataset(Path(filename), file_format) + # The transform argument, in SimpleITK's sense: absent, this resamples onto the grid and + # nothing more. Its field lookup, its bound and its refusals are Warp's, shared. Either + # spelling declares one -- a store of its own, or a group beside the cases -- and a bound + # declared without a field is refused rather than quietly doing nothing. + declared = (field is not None and str(field).strip()) or field_group is not None + if not declared and _is_declared_displacement(max_displacement): + raise TransformError( + f"'ResampleToReference' was given a max_displacement of {max_displacement!r} and no field to apply.", + "Name the field the displacement belongs to — field: ./DVF:omezarr, or field_group:" + " DVF for fields stored beside the cases — or drop max_displacement: without a field" + " this stage resamples onto the grid and nothing more.", + ) + self.displacement: _DisplacementSource | None = ( + _DisplacementSource("ResampleToReference", field, field_group, max_displacement) if declared else None + ) self._grid: tuple[list[int], np.ndarray, np.ndarray, np.ndarray] | None = None # Each case's map, kept from where its own header was in hand. See _recorded(). - self._maps: dict[str, tuple[list[int], list[float], list[float]]] = {} + self._maps: dict[str, _ReferenceMap] = {} + + def set_datasets(self, datasets: list[Dataset]) -> None: + super().set_datasets(datasets) + # A field declared by group alone lives beside the cases, so it looks in the same roots. + if self.displacement is not None: + self.displacement.roots = list(datasets) # ------------------------------------------------------------------ the reference @@ -1427,16 +1499,32 @@ def _geometry(attribute: Attribute, rank: int, what: str) -> tuple[np.ndarray, n # ------------------------------------------------------------------ the map - def grid_map( - self, name: str, shape: list[int], cache_attribute: Attribute - ) -> tuple[list[int], list[float], list[float]]: - """``(target extent, scales, offsets)`` in array order — where each target voxel reads from. + @staticmethod + def _affine_between( + from_origin: np.ndarray, + from_spacing: np.ndarray, + onto_origin: np.ndarray, + onto_spacing: np.ndarray, + direction: np.ndarray, + ) -> tuple[list[float], list[float]]: + """``(scales, offsets)`` in array order taking an index of the FROM grid to one of the ONTO grid. + + An index ``o`` of the from-grid is the physical point ``O_from + D (S_from * o)``; the index + of that point on the onto-grid is ``(D^-1 (p - O_onto)) / S_onto``. With one shared ``D`` + those compose to ``scale * o + offset`` per axis, which is the whole map. Used twice per + case -- target to source, and target to the field's own grid -- because it is the same + question asked of two grids. + """ + scale_xyz = from_spacing / onto_spacing + offset_xyz = (direction.T @ (from_origin - onto_origin)) / onto_spacing + return [float(value) for value in scale_xyz[::-1]], [float(value) for value in offset_xyz[::-1]] + + def grid_map(self, name: str, shape: list[int], cache_attribute: Attribute) -> _ReferenceMap: + """Where each target voxel reads from — and, with a field, where it reads the displacement. - A target voxel ``o`` is the physical point ``O_ref + D (S_ref * o)``; the source index of - that point is ``(D^-1 (p - O_src)) / S_src``. With one shared ``D`` the two compose to - ``scale * o + offset`` per axis, which is the whole map -- and the reason a differing - ``D`` is refused rather than approximated: it would make the map a rotation, whose source - region for a target box is a rotated box no per-axis window can describe. + The reason a differing ``Direction`` is refused rather than approximated: it would make the + map a rotation, whose source region for a target box is a rotated box no per-axis window can + describe. """ target, ref_origin, ref_spacing, ref_direction = self.reference_grid() where = f"case '{name}'" if name else "the case" @@ -1446,27 +1534,57 @@ def grid_map( f" axis/axes, onto reference '{self.entry}', which has {len(target)}.", ) origin, spacing, direction = self._geometry(cache_attribute, len(shape), where) - if not np.allclose(direction, ref_direction, rtol=0.0, atol=1e-6): - raise TransformError( - f"'ResampleToReference' will not resample {where} onto reference '{self.entry}':" - f" their Direction cosines differ ({direction.ravel().tolist()} against" - f" {ref_direction.ravel().tolist()}).", - "The two grids' axes do not line up, so the map between them is a rotation rather" - " than a scale and a shift per axis. Bring them to a common orientation first" - " (Canonical), or make the reference one of the cases as they are stored.", - ) - # (x, y, z) throughout, then reversed once at the end: Origin/Spacing are physical order and - # the scales/offsets a region window is cut with are array order. - scale_xyz = ref_spacing / spacing - offset_xyz = (direction.T @ (ref_origin - origin)) / spacing - scales = [float(value) for value in scale_xyz[::-1]] - offsets = [float(value) for value in offset_xyz[::-1]] + self._refuse_differing_direction(direction, ref_direction, where, f"reference '{self.entry}'") + scales, offsets = self._affine_between(ref_origin, ref_spacing, origin, spacing, direction) self._refuse_if_disjoint(name, shape, target, scales, offsets) + field_scales, field_offsets, field_shape = self._field_map(name, direction, ref_origin, ref_spacing, where) + recorded = _ReferenceMap( + target=target, + scales=scales, + offsets=offsets, + source_spacing=[float(value) for value in spacing], + direction=direction, + field_scales=field_scales, + field_offsets=field_offsets, + field_shape=field_shape, + ) if name: - self._maps[name] = (target, scales, offsets) - return target, scales, offsets + self._maps[name] = recorded + return recorded + + @staticmethod + def _refuse_differing_direction(direction: np.ndarray, other: np.ndarray, where: str, what: str) -> None: + if np.allclose(direction, other, rtol=0.0, atol=1e-6): + return + raise TransformError( + f"'ResampleToReference' will not resample {where} through {what}: their Direction" + f" cosines differ ({direction.ravel().tolist()} against {other.ravel().tolist()}).", + "The grids' axes do not line up, so the map between them is a rotation rather than a" + " scale and a shift per axis. Bring them to a common orientation first (Canonical), or" + " use grids as they were stored.", + ) + + def _field_map( + self, name: str, direction: np.ndarray, ref_origin: np.ndarray, ref_spacing: np.ndarray, where: str + ) -> tuple[list[float] | None, list[float] | None, list[int] | None]: + """The affine map from the TARGET grid onto the field's own grid, from headers alone. - def _recorded(self, name: str) -> tuple[list[int], list[float], list[float]]: + The field is a second image with a geometry of its own, so asking "where on the field does + this target voxel land" is the same question as "where on the source", asked of another + grid -- and separable for the same reason. That is what lets the field be interpolated onto + the target region with the ordinary sampler, filled with zero where it does not reach, which + is precisely what ``DisplacementFieldTransform`` does outside its own extent. + """ + if self.displacement is None or not name: + return None, None, None + shape, attribute = self.displacement.infos(name) + spatial = [int(extent) for extent in shape[1:]] + field_origin, field_spacing, field_direction = self._geometry(attribute, len(spatial), f"the field for {where}") + self._refuse_differing_direction(direction, field_direction, where, f"the field for {where}") + scales, offsets = self._affine_between(ref_origin, ref_spacing, field_origin, field_spacing, direction) + return scales, offsets, spatial + + def _recorded(self, name: str) -> _ReferenceMap: """The map computed for this case back when its own header was in hand. A REGION read hands back the REGION's ``Origin`` — honest about what it read, and not the @@ -1530,8 +1648,8 @@ def coverage(shape: list[int], target: list[int], scales: list[float], offsets: def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribute: Attribute) -> str | None: del group_dest - target, scales, offsets = self.grid_map(name, shape, cache_attribute) - covered = self.coverage(shape, target, scales, offsets) + recorded = self.grid_map(name, shape, cache_attribute) + covered = self.coverage(shape, recorded.target, recorded.scales, recorded.offsets) if covered >= self._WORTH_SAYING: return None return ( @@ -1542,20 +1660,54 @@ def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribut # ------------------------------------------------------------------ the contract def transform_shape(self, group_src: str, name: str, shape: list[int], cache_attribute: Attribute) -> list[int]: - return self.grid_map(name, shape, cache_attribute)[0] + return self.grid_map(name, shape, cache_attribute).target def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: # Not RESCALE: that kind's map is a size ratio, computed by the dispatcher from the two # extents, and this stage's target grid has an origin of its own. REGRID hands the map back # to the stage -- stream_region_source below -- which is the only place it is known. + if self.displacement is None: + return PatchLocality(LocalityKind.REGRID) + # Through a field, the source region is the affine box GROWN by the displacement, so the + # same missing bound that makes Warp whole-volume makes this one whole-volume too -- and for + # the same reason: what it reaches is unknown, not unbounded-in-principle. + if self.displacement.component_bound() is None: + return PatchLocality(LocalityKind.WHOLE_VOLUME, reason=self.displacement.undeclared_reason()) + if _array_order_spacing(cache_attribute) is None: + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=( + "the case carries no 'Spacing', so a displacement in world units cannot be" + " turned into a halo in voxels" + ), + ) return PatchLocality(LocalityKind.REGRID) + def _displacement_halo(self, source_spacing: list[float]) -> int: + """How far past the affine box a target region reaches, in SOURCE voxels. + + One number rather than one per axis, because ``source_window`` grows every axis by the same + margin; the largest per-axis halo is the only safe collapse of a per-component bound. + """ + if self.displacement is None: + return 1 + bound = self.displacement.component_bound() + if bound is None: + return 1 + return 1 + max(_halo_from_bound(bound, list(reversed(source_spacing)))) + def stream_region_source( self, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute ) -> list[slice]: shape = [int(extent) for extent in source_spatial_shape] - _target, scales, offsets = self.grid_map("", shape, cache_attribute) - return Resample.source_window(target_slices, scales, shape, offsets=offsets) + recorded = self.grid_map("", shape, cache_attribute) + return Resample.source_window( + target_slices, + recorded.scales, + shape, + halo=self._displacement_halo(recorded.source_spacing), + offsets=recorded.offsets, + ) def stream_region( self, name: str, tensor: torch.Tensor, context: RegionContext, cache_attribute: Attribute @@ -1563,27 +1715,156 @@ def stream_region( # The recorded map, not one read off `cache_attribute`: what arrives here describes the # REGION, down to an Origin of its own. See _recorded(). del cache_attribute - _target, scales, offsets = self._recorded(name) - return self.resample_region( + return self._sample_target_region( + name, tensor, + self._recorded(name), tuple(context.target), [sl.start for sl in context.source], - scales, [int(extent) for extent in context.source_shape], - offsets, ) def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: shape = [int(extent) for extent in tensor.shape[1:]] - target, scales, offsets = self.grid_map(name, shape, cache_attribute) - # The same sampler the streamed path runs, over one region that happens to be the whole - # grid: equality between the two paths is then a property of the code, not a claim about it. - result = self.resample_region( - tensor, tuple(slice(0, extent) for extent in target), [0] * len(shape), scales, shape, offsets + recorded = self.grid_map(name, shape, cache_attribute) + # The same call the streamed path makes, over one region that happens to be the whole grid: + # equality between the two paths is then a property of the code, not a claim about it. + result = self._sample_target_region( + name, tensor, recorded, tuple(slice(0, extent) for extent in recorded.target), [0] * len(shape), shape ) self.write_stream_cache_attribute(cache_attribute, shape) return result + # ------------------------------------------------------------------ the sampling + + def _sample_target_region( + self, + name: str, + sub_tensor: torch.Tensor, + recorded: _ReferenceMap, + target_slices: tuple[slice, ...], + region_starts: list[int], + n_in: list[int], + ) -> torch.Tensor: + """One region of the target grid, read from the source in a SINGLE interpolation. + + Without a field this is the separable map and the ordinary sampler. With one, the + displacement is added to the target voxel's source coordinate BEFORE anything is sampled -- + so the source is read once, at the displaced point, and never resampled onto an intermediate + grid first. Resampling then warping is two interpolations of the same voxels, and the second + cannot restore what the first smoothed away. + """ + if self.displacement is None: + return self.resample_region( + sub_tensor, target_slices, region_starts, recorded.scales, n_in, recorded.offsets + ) + displacement = self._displacement_in_source_voxels(name, recorded, target_slices, sub_tensor.device) + coordinates = [ + recorded.scales[axis] + * torch.arange(sl.start, sl.stop, device=sub_tensor.device, dtype=torch.float32).reshape( + [-1 if other == axis else 1 for other in range(len(target_slices))] + ) + + recorded.offsets[axis] + + displacement[axis] + for axis, sl in enumerate(target_slices) + ] + return self._sample_at(sub_tensor, coordinates, region_starts, n_in) + + def _displacement_in_source_voxels( + self, name: str, recorded: _ReferenceMap, target_slices: tuple[slice, ...], device: torch.device + ) -> list[torch.Tensor]: + """The field over this target region, per ARRAY axis, in source voxels. + + Three conversions meet here, and getting any of them wrong gives a warp that looks entirely + plausible and moves the anatomy somewhere else: + + - the field is read on ITS grid and interpolated onto the TARGET grid (a separable map, so + the ordinary sampler does it), filled with ZERO outside — a displacement field says + nothing outside its own extent, and SimpleITK reads that as the identity; + - its components are physical (x, y, z) and the axes here are array order (z, y, x); + - and they are world units, which become source voxels through ``D^-1`` and the source + spacing — not through the reference's, which is a different grid. + + One place this and ITK can legitimately differ: a target voxel landing EXACTLY on the + field's last half-voxel (index ``n - 0.5``, the open end of the domain). The map here is + computed exactly, while ITK reaches the same index through world coordinates and loses the + last bits against a large origin, so the two can disagree about which side of the boundary + it falls on. It takes a field grid that is an exact multiple of the target's AND shares its + origin to arrange — which a solved field is not — and it costs one boundary column. + """ + source = self.displacement + if source is None: # unreachable: only a declared field reaches this + raise TransformError(f"'ResampleToReference' has no field to read for case '{name}'.") + window = Resample.source_window( + target_slices, + cast("list[float]", recorded.field_scales), + cast("list[int]", recorded.field_shape), + offsets=cast("list[float]", recorded.field_offsets), + ) + rank = len(target_slices) + field = source.read(name, tuple(window), rank).to(device) + source.check_bound(field, name) + # Zero outside: the transform is the identity where the field does not reach. + on_target = self._resample_offset_region( + field, + target_slices, + [sl.start for sl in window], + cast("list[float]", recorded.field_scales), + cast("list[int]", recorded.field_shape), + cast("list[float]", recorded.field_offsets), + fill=0.0, + ) + # (x, y, z) world -> (x, y, z) source index, then reversed to array order. + rotated = torch.from_numpy(np.ascontiguousarray(recorded.direction.T)).to(device=device, dtype=on_target.dtype) + spacing = torch.tensor(recorded.source_spacing, device=device, dtype=on_target.dtype).reshape(-1, 1) + flat = (rotated @ on_target.reshape(rank, -1)) / spacing + return list(reversed(list(flat.reshape(on_target.shape)))) + + def _sample_at( + self, + sub_tensor: torch.Tensor, + coordinates: list[torch.Tensor], + region_starts: list[int], + n_in: list[int], + ) -> torch.Tensor: + """Trilinear gather at an ARBITRARY coordinate per voxel — ITK's rules, non-separably. + + The sibling of :meth:`Resample._resample_offset_region`, and deliberately not a + generalisation of it: there the coordinate is a product of per-axis maps, so no coordinate + volume is ever built and one ``index_select`` per axis does the work. A displacement is not + separable, so here a coordinate exists per voxel and the eight corners are gathered flat. + Both obey the same convention -- inside is ``[-0.5, n - 0.5)``, taps clamp to the buffer, + outside takes ``fill_value`` -- which is what makes them interchangeable at the seam. + """ + rank = len(coordinates) + window = [int(extent) for extent in sub_tensor.shape[1:]] + extent = list(torch.broadcast_shapes(*[axis.shape for axis in coordinates])) + inside = torch.ones(extent, dtype=torch.bool, device=sub_tensor.device) + bases: list[torch.Tensor] = [] + weights: list[torch.Tensor] = [] + for axis, coordinate in enumerate(coordinates): + inside &= (coordinate >= -0.5) & (coordinate < n_in[axis] - 0.5) + base = torch.floor(coordinate) + bases.append(base.to(torch.int32).expand(extent)) + weights.append((coordinate - base).expand(extent)) + out_shape = [int(sub_tensor.shape[0]), *extent] + if not bool(inside.any()): + return torch.full(out_shape, self.fill_value, device=sub_tensor.device, dtype=torch.float32).type( + sub_tensor.dtype + ) + work = sub_tensor if sub_tensor.is_floating_point() else sub_tensor.type(torch.float32) + flat_source = work.reshape(int(work.shape[0]), -1) + out = torch.zeros(out_shape, device=sub_tensor.device, dtype=work.dtype) + for corner in itertools.product((0, 1), repeat=rank): + flat_index = torch.zeros(extent, dtype=torch.long, device=sub_tensor.device) + weight = torch.ones(extent, device=sub_tensor.device, dtype=work.dtype) + for axis, step in enumerate(corner): + index = torch.clamp(bases[axis].to(torch.long) + step, 0, n_in[axis] - 1) - region_starts[axis] + flat_index = flat_index * window[axis] + torch.clamp(index, 0, window[axis] - 1) + weight = weight * (weights[axis] if step else 1 - weights[axis]).to(work.dtype) + out += flat_source.index_select(1, flat_index.reshape(-1)).reshape(out_shape) * weight + return out.masked_fill(~inside.unsqueeze(0), self.fill_value).type(sub_tensor.dtype) + def write_stream_cache_attribute(self, cache_attribute: Attribute, source_spatial_shape: list[int]) -> None: target, origin, spacing, direction = self.reference_grid() # The case now IS the reference grid, header and all -- which is the point, and what @@ -1619,10 +1900,12 @@ def inverse(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) - shape = [int(extent) for extent in tensor.shape[1:]] target = self._inverse_geometry(cache_attribute) # _inverse_geometry restored the case's own geometry, so this is the forward map again -- - # read off the same headers -- and the inverse is that map solved for the other index. - _spatial, scales, offsets = self.grid_map(name, target, cache_attribute) - back_scales = [1.0 / scale for scale in scales] - back_offsets = [-offset / scale for offset, scale in zip(offsets, scales, strict=True)] + # read off the same headers -- and the inverse is that map solved for the other index. The + # displacement is NOT undone: inverting a field is a solve, not an algebraic step, so a + # warped forward inverts to the grid change alone and says so below. + recorded = self.grid_map(name, target, cache_attribute) + back_scales = [1.0 / scale for scale in recorded.scales] + back_offsets = [-offset / scale for offset, scale in zip(recorded.offsets, recorded.scales, strict=True)] return self.resample_region( tensor, tuple(slice(0, extent) for extent in target), [0] * len(shape), back_scales, shape, back_offsets ) @@ -2030,73 +2313,54 @@ def __init__( super().__init__(dataset, group, scale_factors, downsample_method) -class Warp(Transform): - """Resample a case through a displacement field defined on the same grid. - - ``output(p) = input(p + d(p))``, with ``d`` read from ``field`` in world units. Field and case - must share a grid: this is the shape update of an atlas build, where the field was solved on the - very grid it is applied to. Warping ONTO A DIFFERENT grid is a resample as well as a warp, and - is not this stage. - - THE HALO IS DECLARED, AND VERIFIED. How far a target voxel reaches into the source is the - displacement itself, so the region this needs is the target enlarged by the largest - displacement. That bound is then CHECKED against every region actually read: a field that - exceeds it raises, rather than quietly sampling zeros — which would look like a dark rim around - the moved anatomy and nothing else. +class _DisplacementSource: + """A displacement field on disk: where it is, how far it reaches, and how to read a region of it. - ``max_displacement`` is in the same world units as ``Spacing`` (micrometres for these stores), - and takes ``auto``: a field records its own per-component bound when KonfAI writes it, so - ``auto`` reads it back from the headers instead of asking you for a number you would have to - measure. It is the cohort's bound, not this case's — a locality is declared once for the stage, - while the field is per case — so it over-reads for a gentle case and is never short for a wild - one. Give a number when you know one and want the tightest halo. + Shared by the two stages that resample THROUGH one. What they do with the displacement differs + entirely — :class:`Warp` adds it on the case's own grid, :class:`ResampleToReference` composes it + with a change of grid — but where the field comes from, what bounds it, and the refusal when it + exceeds that bound are the same questions, and were answered twice before this existed. - With no bound at all — ``0.0``, an ``auto`` the headers cannot answer, or a case with no - ``Spacing`` — the stage declares ``WHOLE_VOLUME`` and says which, because a `Warp` that silently - costs the whole volume is this stage's expensive failure: the result is right, so nothing looks - wrong except the peak memory the reader was streaming to control. + ``owner`` is the stage's name, so a refusal reads as coming from the stage the user declared and + not from a helper they have never heard of. """ - def __init__( - self, - field: str, - group: str | None = None, - max_displacement: float | str = 0.0, - interpolation: str = "linear", - ) -> None: - super().__init__() - if not field or not str(field).strip(): - raise TransformError( - "'Warp' needs a 'field': the displacement field to resample through.", - "Declare it, e.g. Warp: {field: ./DVF:omezarr, max_displacement: 250.0}.", - ) - if interpolation not in ("linear", "nearest"): + def __init__(self, owner: str, field: str | None, group: str | None, max_displacement: float | str) -> None: + self.owner = owner + # A root of its own, or none: with no ``field`` path the fields are a GROUP of the run's own + # dataset_filenames, one entry per case — which is how a cohort registered in place stores + # them, beside the volumes they were solved on. + self.dataset: Dataset | None = None + if field is not None and str(field).strip(): + filename, _flag, file_format = split_path_spec(str(field), default_format="mha") + self.dataset = Dataset(Path(filename), file_format) + elif group is None: raise TransformError( - f"'Warp' has an unknown interpolation '{interpolation}'.", - "Use 'linear' for an image or 'nearest' for a label map.", + f"'{owner}' has neither a 'field' path nor a group to find the fields in.", + f"Name the store — {owner}: {{field: ./DVF:omezarr}} — or, for fields stored beside" + f" the cases, the group they are in: {owner}: {{field_group: DVF}}.", ) - filename, _flag, file_format = split_path_spec(str(field), default_format="mha") - self.field_dataset = Dataset(Path(filename), file_format) - self.field_group = group - self.auto_displacement = isinstance(max_displacement, str) and max_displacement.strip().lower() == "auto" - if isinstance(max_displacement, str) and not self.auto_displacement: + self.group = group + #: The run's own roots, handed over by the owner; only consulted when there is no path. + self.roots: list[Dataset] = [] + self.auto = isinstance(max_displacement, str) and max_displacement.strip().lower() == "auto" + if isinstance(max_displacement, str) and not self.auto: try: max_displacement = float(max_displacement) except ValueError: raise TransformError( - f"'Warp' has a max_displacement of '{max_displacement}', which is neither a number nor 'auto'.", + f"'{owner}' has a max_displacement of '{max_displacement}', which is neither a number nor 'auto'.", "Give a distance in the case's world units (max_displacement: 250.0), or 'auto'" " to read the bound the fields recorded when they were written.", ) from None # Per component, in the field's own (x, y, z) order. A scalar bound broadcasts to all three; # `auto` fills this from the headers on first use. Per component and not one number, because # these grids are anisotropic: one collapsed maximum over-reads the fine axes. - self.max_displacement = 0.0 if self.auto_displacement else float(max_displacement) + self.max_displacement = 0.0 if self.auto else float(max_displacement) self._auto_bound: list[float] | None = None self._auto_resolved = False - self.interpolation = interpolation - def _component_bound(self) -> list[float] | None: + def component_bound(self) -> list[float] | None: """The per-component bound this stage warps within, or ``None`` when it has none. For ``auto``, the largest bound any field in the group recorded, read from headers alone and @@ -2104,7 +2368,7 @@ def _component_bound(self) -> list[float] | None: others would be a bound for them and a guess for that one, and this number is what sizes the region every read depends on. """ - if not self.auto_displacement: + if not self.auto: return [self.max_displacement] * 3 if self.max_displacement > 0.0 else None if self._auto_resolved: return self._auto_bound @@ -2113,11 +2377,12 @@ def _component_bound(self) -> list[float] | None: bound: list[float] = [] try: - group = self._group_for(None) + group = self.group_for(None) + root = self._root_for(None) # The header reads belong inside: a directory store lists its entries from the filesystem # alone, so an unreadable field can only surface here, one entry at a time. - for entry in self.field_dataset.get_names(group): - _shape, attribute = self.field_dataset.get_infos(group, entry) + for entry in root.get_names(group): + _shape, attribute = root.get_infos(group, entry) if DISPLACEMENT_BOUND_ATTRIBUTE not in attribute: return self._auto_bound recorded = [float(value) for value in attribute.get_np_array(DISPLACEMENT_BOUND_ATTRIBUTE).ravel()] @@ -2128,63 +2393,63 @@ def _component_bound(self) -> list[float] | None: self._auto_bound = bound return self._auto_bound - def _spacing(self, cache_attribute: Attribute) -> list[float] | None: - """The case's spacing in array order (z, y, x); ``Spacing`` is stored (x, y, z).""" - if "Spacing" not in cache_attribute: - return None - spacing = [float(value) for value in np.asarray(cache_attribute.get_np_array("Spacing")).ravel()] - return list(reversed(spacing)) - - def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: - bound = self._component_bound() - if bound is None: - return PatchLocality( - LocalityKind.WHOLE_VOLUME, - reason=( - "max_displacement is 'auto' and the fields carry no recorded bound to read" - " (KonfAI records one on an OME-Zarr field it writes; other formats and other" - " producers do not)" - if self.auto_displacement - else "no 'max_displacement' is declared" - ) - + " -- how far this warp reaches into its source is unknown and the region it must" - " read is unbounded. Declare it in the case's world units (e.g." - " max_displacement: 250.0) to stream with a halo", - ) - spacing = self._spacing(cache_attribute) - if spacing is None: - return PatchLocality( - LocalityKind.WHOLE_VOLUME, - reason=( - "the case carries no 'Spacing', so a displacement in world units cannot be" - " turned into a halo in voxels" - ), - ) - # The bound is per component in (x, y, z); a halo is per array axis in (z, y, x). - per_axis = list(reversed(bound))[-len(spacing) :] if len(bound) >= len(spacing) else [max(bound)] * len(spacing) - halo = tuple( - int(np.ceil(value / extent)) if extent > 0 else 0 for value, extent in zip(per_axis, spacing, strict=False) + def undeclared_reason(self) -> str: + """Why there is no bound, in the words the plan prints.""" + return ( + "max_displacement is 'auto' and the fields carry no recorded bound to read" + " (KonfAI records one on an OME-Zarr field it writes; other formats and other" + " producers do not)" + if self.auto + else "no 'max_displacement' is declared" + ) + ( + " -- how far this reaches into its source is unknown and the region it must read is" + " unbounded. Declare it in the case's world units (e.g. max_displacement: 250.0) to" + " stream with a halo" ) - return PatchLocality(LocalityKind.HALO, halo=halo) - def _group_for(self, name: str | None) -> str: - if self.field_group is not None: - return self.field_group - groups = [str(group) for group in self.field_dataset.get_group()] + def group_for(self, name: str | None) -> str: + if self.group is not None: + return self.group + if self.dataset is None: # unreachable: a source with no path was given a group to use + raise TransformError( + f"'{self.owner}' has no field store of its own and no group to look for one in.", + f"Name the group the fields are in: {self.owner}: {{field_group: DVF}}.", + ) + groups = [str(group) for group in self.dataset.get_group()] if len(groups) == 1: return groups[0] where = f"the field for case '{name}'" if name is not None else "the fields" raise TransformError( - f"'Warp' cannot tell which group of '{self.field_dataset.filename}' holds {where}: it has {len(groups)}.", - "Name it: Warp: {field: ./DVF:omezarr, group: DVF}.", + f"'{self.owner}' cannot tell which group of '{self.dataset.filename}' holds {where}: it has {len(groups)}.", + f"Name it: {self.owner}: {{field: ./DVF:omezarr, group: DVF}}.", + ) + + def _root_for(self, name: str | None) -> Dataset: + """The store this case's field is in: the declared one, or whichever run root holds it.""" + if self.dataset is not None: + return self.dataset + group = self.group_for(name) + for root in self.roots: + if name is None or root.is_dataset_exist(group, name): + return root + raise TransformError( + f"'{self.owner}' cannot find a field for case '{name}' in group '{group}' of" + f" {', '.join(str(root.filename) for root in self.roots) or 'any dataset'}.", + "A field declared by group alone is looked up beside the cases, one entry per case." + f" Give the store a path of its own instead: {self.owner}: {{field: ./DVF:omezarr}}.", ) - def _read_field(self, name: str, region: tuple[slice, ...] | None, channels: int) -> torch.Tensor: - group = self._group_for(name) + def infos(self, name: str) -> tuple[list[int], Attribute]: + """The field entry's shape and header, without reading a voxel of it.""" + return self._root_for(name).get_infos(self.group_for(name), name) + + def read(self, name: str, region: tuple[slice, ...] | None, channels: int) -> torch.Tensor: + group = self.group_for(name) + root = self._root_for(name) if region is None: - data, _attributes = self.field_dataset.read_data(group, name) + data, _attributes = root.read_data(group, name) else: - data, _attributes = self.field_dataset.read_data_slice(group, name, (slice(None), *region)) + data, _attributes = root.read_data_slice(group, name, (slice(None), *region)) field = torch.from_numpy(np.ascontiguousarray(data)).float() if field.shape[0] != channels: raise TransformError( @@ -2194,6 +2459,119 @@ def _read_field(self, name: str, region: tuple[slice, ...] | None, channels: int ) return field + def check_bound(self, field: torch.Tensor, name: str) -> None: + """The declaration is a promise about the region that was read; check it against the samples. + + Per component, matching how the halo was derived: a field that stays under the collapsed + maximum can still exceed the bound on one axis, which is the axis whose halo was too small. + """ + bound = self.component_bound() + if bound is None or not field.numel(): + return + for component in range(field.shape[0]): + declared = bound[component] if component < len(bound) else max(bound) + largest = float(field[component].abs().max()) + if largest > declared: + raise TransformError( + f"The field for case '{name}' displaces up to {largest:.3f} on component" + f" {component}, beyond the {declared:.3f} '{self.owner}' sized its region from.", + "Raise max_displacement to at least the field's true maximum, or use" + " max_displacement: auto: the region read is sized from that number, so a larger" + " displacement samples outside what was read.", + ) + + +def _halo_from_bound(bound: list[float], spacing: list[float]) -> tuple[int, ...]: + """A world-unit displacement bound as a per-axis halo in voxels. + + ``bound`` is per component in (x, y, z); ``spacing`` is per array axis in (z, y, x). Reversing + one of them is the whole of this function, and the reason it is one. + """ + per_axis = list(reversed(bound))[-len(spacing) :] if len(bound) >= len(spacing) else [max(bound)] * len(spacing) + return tuple( + int(np.ceil(value / extent)) if extent > 0 else 0 for value, extent in zip(per_axis, spacing, strict=False) + ) + + +def _is_declared_displacement(max_displacement: float | str) -> bool: + """Whether a ``max_displacement`` was actually asked for, rather than left at its default.""" + if isinstance(max_displacement, str): + return bool(max_displacement.strip()) + return float(max_displacement) != 0.0 + + +def _array_order_spacing(cache_attribute: Attribute) -> list[float] | None: + """A case's spacing in array order (z, y, x); ``Spacing`` is stored (x, y, z).""" + if "Spacing" not in cache_attribute: + return None + return list(reversed([float(value) for value in np.asarray(cache_attribute.get_np_array("Spacing")).ravel()])) + + +class Warp(Transform): + """Resample a case through a displacement field defined on the same grid. + + ``output(p) = input(p + d(p))``, with ``d`` read from ``field`` in world units. Field and case + must share a grid: this is the shape update of an atlas build, where the field was solved on the + very grid it is applied to. Warping ONTO A DIFFERENT grid is a resample as well as a warp, and + is not this stage. + + THE HALO IS DECLARED, AND VERIFIED. How far a target voxel reaches into the source is the + displacement itself, so the region this needs is the target enlarged by the largest + displacement. That bound is then CHECKED against every region actually read: a field that + exceeds it raises, rather than quietly sampling zeros — which would look like a dark rim around + the moved anatomy and nothing else. + + ``max_displacement`` is in the same world units as ``Spacing`` (micrometres for these stores), + and takes ``auto``: a field records its own per-component bound when KonfAI writes it, so + ``auto`` reads it back from the headers instead of asking you for a number you would have to + measure. It is the cohort's bound, not this case's — a locality is declared once for the stage, + while the field is per case — so it over-reads for a gentle case and is never short for a wild + one. Give a number when you know one and want the tightest halo. + + With no bound at all — ``0.0``, an ``auto`` the headers cannot answer, or a case with no + ``Spacing`` — the stage declares ``WHOLE_VOLUME`` and says which, because a `Warp` that silently + costs the whole volume is this stage's expensive failure: the result is right, so nothing looks + wrong except the peak memory the reader was streaming to control. + """ + + def __init__( + self, + field: str, + group: str | None = None, + max_displacement: float | str = 0.0, + interpolation: str = "linear", + ) -> None: + super().__init__() + if not field or not str(field).strip(): + # Required here and not merely by the shared source: a Warp's field is on the case's own + # grid, so there is no group-beside-the-cases shorthand for it to fall back to. + raise TransformError( + "'Warp' needs a 'field': the displacement field to resample through.", + "Declare it, e.g. Warp: {field: ./DVF:omezarr, max_displacement: 250.0}.", + ) + if interpolation not in ("linear", "nearest"): + raise TransformError( + f"'Warp' has an unknown interpolation '{interpolation}'.", + "Use 'linear' for an image or 'nearest' for a label map.", + ) + self.displacement = _DisplacementSource("Warp", field, group, max_displacement) + self.interpolation = interpolation + + def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: + bound = self.displacement.component_bound() + if bound is None: + return PatchLocality(LocalityKind.WHOLE_VOLUME, reason=self.displacement.undeclared_reason()) + spacing = _array_order_spacing(cache_attribute) + if spacing is None: + return PatchLocality( + LocalityKind.WHOLE_VOLUME, + reason=( + "the case carries no 'Spacing', so a displacement in world units cannot be" + " turned into a halo in voxels" + ), + ) + return PatchLocality(LocalityKind.HALO, halo=_halo_from_bound(bound, spacing)) + def _sample(self, tensor: torch.Tensor, field: torch.Tensor, spacing: list[float]) -> torch.Tensor: """``output(p) = input(p + d(p))`` over the block handed in, sampled with grid_sample. @@ -2221,46 +2599,25 @@ def _sample(self, tensor: torch.Tensor, field: torch.Tensor, spacing: list[float ) return moved.squeeze(0).to(tensor.dtype) - def _check_declared_bound(self, field: torch.Tensor, name: str) -> None: - """The declaration is a promise about the region that was read; check it against the samples. - - Per component, matching how the halo was derived: a field that stays under the collapsed - maximum can still exceed the bound on one axis, which is the axis whose halo was too small. - """ - bound = self._component_bound() - if bound is None or not field.numel(): - return - for component in range(field.shape[0]): - declared = bound[component] if component < len(bound) else max(bound) - largest = float(field[component].abs().max()) - if largest > declared: - raise TransformError( - f"The field for case '{name}' displaces up to {largest:.3f} on component" - f" {component}, beyond the {declared:.3f} this stage sized its region from.", - "Raise max_displacement to at least the field's true maximum, or use" - " max_displacement: auto: the region read is sized from that number, so a larger" - " displacement samples outside what was read.", - ) - def stream_region( self, name: str, tensor: torch.Tensor, context: RegionContext, cache_attribute: Attribute ) -> torch.Tensor: - spacing = self._spacing(cache_attribute) + spacing = _array_order_spacing(cache_attribute) if spacing is None: return self(name, tensor, cache_attribute) - field = self._read_field(name, context.source, len(tensor.shape) - 1) - self._check_declared_bound(field, name) + field = self.displacement.read(name, context.source, len(tensor.shape) - 1) + self.displacement.check_bound(field, name) return self._sample(tensor, field, spacing) def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - spacing = self._spacing(cache_attribute) + spacing = _array_order_spacing(cache_attribute) if spacing is None: raise TransformError( f"'Warp' needs the case's Spacing to turn a world displacement into voxels, and" f" case '{name}' declares none.", "Use a source whose geometry is readable (mha, nii, h5 or omezarr written by KonfAI).", ) - field = self._read_field(name, None, len(tensor.shape) - 1) + field = self.displacement.read(name, None, len(tensor.shape) - 1) return self._sample(tensor, field, spacing) @@ -2336,49 +2693,6 @@ def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) ) -class ShapeUpdate(Transform): - """Pull a template toward the mean SHAPE a displacement field carries, leaving its pose alone. - - ``output = -step * (field - t)``, where ``t`` is the field's per-component spatial mean in world - units. Resampling a template through the result moves it along the cohort's shape residual at - ANTs' gradient step — the shape update of an atlas build. - - WHY ``t`` IS STRIPPED RATHER THAN APPLIED. A total field maps template coordinates into each - specimen's OWN world frame, so its spatial mean is dominated by the frame-to-frame offset, not by - any pose error of the template. Applying it translates the template out of its own grid and clips - the anatomy; stripping it is what keeps the template anchored. - - Per-component on purpose: a translation has as many parts as the field has components, and the - pooled mean of all of them describes nothing. That is why this declares ``MeanPerChannel`` — the - statistic is read once from the stored volume, and the stage is then a value map, so a field of - any size runs region by region. - """ - - def __init__(self, step: float = 0.25) -> None: - super().__init__() - self.step = float(step) - - def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: - return PatchLocality(LocalityKind.GLOBAL_STAT, stat_keys=frozenset({"MeanPerChannel"})) - - def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: - if "MeanPerChannel" not in cache_attribute: - # Whole volume in hand: the statistic IS this tensor's, so take it rather than demand a - # seed. Recorded on the case, as the streamed path records its own, so both leave the - # same state behind and an inverse finds it where it expects. - cache_attribute["MeanPerChannel"] = tensor.reshape(int(tensor.shape[0]), -1).to(torch.float32).mean(dim=1) - mean = torch.as_tensor(cache_attribute.get_np_array("MeanPerChannel"), dtype=torch.float32) - if mean.numel() != tensor.shape[0]: - raise TransformError( - f"'ShapeUpdate' was handed {mean.numel()} component mean(s) for a" - f" {tensor.shape[0]}-component field on '{name}'.", - "The statistic is read per channel from the stored entry: check that the entry is the" - " displacement field itself and not a derived volume.", - ) - centred = tensor.to(torch.float32) - mean.reshape(-1, *([1] * (tensor.dim() - 1))) - return -self.step * centred - - class Expand(Transform): """Turn one case into ``nb`` copies, at a declared point of the chain — ``Reduce``'s mirror. diff --git a/tests/integration/test_transform_doc_examples.py b/tests/integration/test_transform_doc_examples.py index 402577ae..953f9128 100644 --- a/tests/integration/test_transform_doc_examples.py +++ b/tests/integration/test_transform_doc_examples.py @@ -121,6 +121,26 @@ def _write_dataset(root: Path, groups: list[str], cases: int = 2, shape=(8, 12, sitk.WriteImage(image, str(case / f"{group}.mha")) +def _write_fields(root: Path, cases: int = 2, shape=(4, 6, 6)) -> None: + """One displacement field per case, in a store of its own, on a grid COARSER than the images. + + A page that documents resampling THROUGH a field has to be able to show one, and an example is + only runnable if the fixture holds what it names. Coarse on purpose: a field is read where it is + asked rather than resampled to match first, so the example exercises that and not a same-grid + special case. + """ + for index in range(cases): + case = root / f"case_{index}" + case.mkdir(parents=True, exist_ok=True) + # KonfAI stores a field component-first; SimpleITK wants the components last. + field = np.zeros((3, *shape), dtype=np.float32) + field[0], field[1], field[2] = 0.5, -0.25, 0.75 + image = sitk.GetImageFromArray(np.moveaxis(field, 0, -1), isVector=True) + image.SetSpacing((1.0, 1.0, 4.0)) + image.SetOrigin((1.0, 2.0, 3.0)) + sitk.WriteImage(image, str(case / "DVF.mha")) + + def _doc_examples() -> list[tuple[int, str]]: if not DOC.is_file(): return [] @@ -134,6 +154,7 @@ def test_a_documented_example_plans_and_runs(line: int, config: str, tmp_path: P workdir = tmp_path / f"block_{line}" workdir.mkdir() _write_dataset(workdir / "Raw", _source_groups(config)) + _write_fields(workdir / "Fields") (workdir / "Transform.yml").write_text(config, encoding="utf-8") for filename, source in _python_modules(DOC.read_text(encoding="utf-8")).items(): (workdir / filename).write_text(source, encoding="utf-8") diff --git a/tests/unit/test_shape_update.py b/tests/unit/test_per_component_statistic.py similarity index 55% rename from tests/unit/test_shape_update.py rename to tests/unit/test_per_component_statistic.py index 456e8ff9..8feb3bee 100644 --- a/tests/unit/test_shape_update.py +++ b/tests/unit/test_per_component_statistic.py @@ -14,19 +14,53 @@ # # SPDX-License-Identifier: Apache-2.0 -"""ShapeUpdate: the shape residual of a displacement field, streamed.""" +"""The per-component statistic, exercised through a transform defined HERE rather than shipped. + +The stage under test is deliberately a user's: KonfAI ships generic transforms, and centring a +displacement field per component then scaling it is one step of an atlas build, which belongs to the +pipeline that needs it. What the framework owes such a stage is that declaring a locality contract is +enough to stream it -- and that is exactly what these tests pin. + +So this doubles as the worked example: forty lines below, written against the public surface only, +is a user transform that streams a volume it never assembles. +""" from pathlib import Path import numpy as np import pytest +import torch +from konfai.data import Attribute, LocalityKind, PatchLocality, Transform, Write from konfai.data.patching import DatasetManager -from konfai.data.transform import ShapeUpdate, Write -from konfai.utils.dataset import Attribute, Dataset +from konfai.utils.dataset import Dataset pytest.importorskip("SimpleITK") +class _CentreAndScale(Transform): + """``-step * (field - t)``, t being the field's per-component spatial mean. + + Per-component on purpose: a translation has as many parts as the field has components, and the + pooled mean of all of them describes nothing. Declaring ``MeanPerChannel`` is what lets the + statistic be read once from the stored volume and the stage then be a value map, so a field of + any size runs region by region. + """ + + def __init__(self, step: float = 0.25) -> None: + super().__init__() + self.step = float(step) + + def patch_locality(self, cache_attribute: Attribute) -> PatchLocality: + return PatchLocality(LocalityKind.GLOBAL_STAT, stat_keys=frozenset({"MeanPerChannel"})) + + def __call__(self, name: str, tensor: torch.Tensor, cache_attribute: Attribute) -> torch.Tensor: + if "MeanPerChannel" not in cache_attribute: + cache_attribute["MeanPerChannel"] = tensor.reshape(int(tensor.shape[0]), -1).to(torch.float32).mean(dim=1) + mean = torch.as_tensor(cache_attribute.get_np_array("MeanPerChannel"), dtype=torch.float32) + centred = tensor.to(torch.float32) - mean.reshape(-1, *([1] * (tensor.dim() - 1))) + return -self.step * centred + + def _field(tmp_path: Path) -> Dataset: rng = np.random.default_rng(0) attributes = Attribute() @@ -63,9 +97,9 @@ def test_the_statistic_is_per_component(tmp_path: Path) -> None: np.testing.assert_allclose(stats["mean"], volume.mean(), rtol=0, atol=1e-4) -def test_shape_update_streams_and_equals_the_whole_volume(tmp_path: Path) -> None: +def test_a_user_stage_streams_and_equals_the_whole_volume(tmp_path: Path) -> None: source = _field(tmp_path) - manager = _manager(source, [ShapeUpdate(step=0.25), Write(f"{tmp_path / 'out'}:h5")]) + manager = _manager(source, [_CentreAndScale(step=0.25), Write(f"{tmp_path / 'out'}:h5")]) assert manager.stream_refusal(0, apply_augmentations=False) is None assert manager.materialize() is True, "a per-component centring is a value map: it must stream" @@ -77,12 +111,12 @@ def test_shape_update_streams_and_equals_the_whole_volume(tmp_path: Path) -> Non np.testing.assert_allclose(got.reshape(3, -1).mean(axis=1), np.zeros(3), rtol=0, atol=1e-4) -def test_a_streamed_shape_update_never_reads_the_volume(tmp_path: Path, monkeypatch) -> None: +def test_a_streamed_user_stage_never_reads_the_volume(tmp_path: Path, monkeypatch) -> None: source = _field(tmp_path) - manager = _manager(source, [ShapeUpdate(step=0.25), Write(f"{tmp_path / 'out'}:h5")]) + manager = _manager(source, [_CentreAndScale(step=0.25), Write(f"{tmp_path / 'out'}:h5")]) def refuse(*args, **kwargs): - raise AssertionError("the shape update assembled the field") + raise AssertionError("the user stage assembled the volume") monkeypatch.setattr(Dataset, "read_data", refuse) assert manager.materialize() is True diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 953ec27e..070b2bec 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -475,3 +475,315 @@ def test_the_inverse_declares_the_whole_volume_and_says_why(dataset: Dataset) -> locality = _stage(dataset).inverse_patch_locality(Attribute()) assert locality.kind is LocalityKind.WHOLE_VOLUME assert locality.reason is not None and "reference grid" in locality.reason + + +# --------------------------------------------------------------------- through a field + +# The field's own grid: coarser than the case AND than the reference, with an origin of its own -- +# which is the point. A field solved at one resolution moves a volume stored at another, because it +# is defined in world units and read where it is asked. +_FIELD_SPATIAL = (5, 6, 7) +_FIELD_ORIGIN, _FIELD_SPACING = [-4.0, 3.0, 9.0], [4.0, 4.5, 5.0] +_BOUND = 20.0 + + +def _displacement(shape: tuple[int, ...] = _FIELD_SPATIAL) -> np.ndarray: + """A field whose three components are DIFFERENT functions of DIFFERENT axes. + + Component-first, in physical (x, y, z). Every component varies, and none is a multiple of + another, so swapping two of them — or reversing the component axis against the array axes — + moves the anatomy somewhere a comparison against SimpleITK cannot miss. + """ + z, y, x = np.meshgrid(*[np.arange(n) for n in shape], indexing="ij") + return np.stack([2.0 + 0.5 * x, -1.5 + 0.3 * y, 0.8 * np.sin(z * 0.9)]).astype(np.float32) + + +def _high_frequency(shape: tuple[int, ...] = _SOURCE_SPATIAL) -> np.ndarray: + """A source whose detail a second interpolation would visibly smooth away.""" + z, y, x = np.meshgrid(*[np.arange(n) for n in shape], indexing="ij") + return (100 * np.sin(z * 1.7) * np.cos(y * 2.1) + 80 * np.sin(x * 2.9)).astype(np.float32)[None] + + +@pytest.fixture +def warped(tmp_path: Path) -> tuple[Dataset, Dataset, np.ndarray]: + """A case, a reference grid and a field — three grids that agree about nothing.""" + images = Dataset(tmp_path / "Images", "h5") + volume = _high_frequency() + images.write("Case", _CASE, volume, _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + images.write("Reference", _CASE, _volume(_REFERENCE_SPATIAL, 1), _attributes(_REFERENCE_ORIGIN, _REFERENCE_SPACING)) + fields = Dataset(tmp_path / "Fields", "h5") + fields.write("DVF", _CASE, _displacement(), _attributes(_FIELD_ORIGIN, _FIELD_SPACING)) + return images, fields, volume + + +def _warping(images: Dataset, fields: Dataset, **kwargs: object) -> ResampleToReference: + arguments: dict[str, object] = { + "entry": _CASE, + "group": "Reference", + "field": f"{fields.filename}:h5", + "field_group": "DVF", + "max_displacement": _BOUND, + "fill": _FILL, + **kwargs, + } + stage = ResampleToReference(**arguments) # type: ignore[arg-type] + stage.set_datasets([images]) + return stage + + +def _simpleitk_warp(volume: np.ndarray, field: np.ndarray | None = None) -> np.ndarray: + """``sitk.Resample(image, grid, DisplacementFieldTransform(field))`` — the one-pass authority.""" + image = sitk.GetImageFromArray(volume[0]) + image.SetOrigin(_SOURCE_ORIGIN) + image.SetSpacing(_SOURCE_SPACING) + grid = sitk.Image(*reversed(_REFERENCE_SPATIAL), sitk.sitkFloat32) + grid.SetOrigin(_REFERENCE_ORIGIN) + grid.SetSpacing(_REFERENCE_SPACING) + transform: sitk.Transform = sitk.Transform() + if field is not None: + # sitk wants a vector image, (z, y, x, component), where KonfAI stores component-first. + vector = sitk.GetImageFromArray(np.moveaxis(field, 0, -1).astype(np.float64), isVector=True) + vector.SetOrigin(_FIELD_ORIGIN) + vector.SetSpacing(_FIELD_SPACING) + transform = sitk.DisplacementFieldTransform(sitk.Cast(vector, sitk.sitkVectorFloat64)) + return sitk.GetArrayFromImage(sitk.Resample(image, grid, transform, sitk.sitkLinear, _FILL)) + + +def test_it_warps_onto_the_reference_where_simpleitk_does(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """The whole operation, against the one call that defines it. + + Three grids, none agreeing on extent, spacing or origin, and a field with a different function + per component: a mistake in any of the axis-order conversions lands the anatomy elsewhere, and + a comparison against SimpleITK is the only thing that would notice. + """ + images, fields, volume = warped + got = _warping(images, fields)(_CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + want = _simpleitk_warp(volume, _displacement()) + + assert list(got.shape[1:]) == list(_REFERENCE_SPATIAL) + span = float(want.max() - want.min()) + np.testing.assert_allclose(got.numpy()[0], want, rtol=0, atol=64 * float(np.spacing(np.float32(span)))) + + +def test_the_displaced_edge_is_where_simpleitk_puts_it(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """Which voxels have no source is decided by the DISPLACED coordinate, so it tests the whole map. + + A composition that is off by a fraction of a voxel still interpolates real data almost + everywhere; the rim is where it stops having any, and matching it voxel for voxel is a much + sharper statement than any tolerance on the values. + """ + images, fields, volume = warped + got = _warping(images, fields)(_CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + want = _simpleitk_warp(volume, _displacement()) + + assert 0 < int((want == _FILL).sum()) < want.size, "the fixture must have a rim, and not be all rim" + np.testing.assert_array_equal(got.numpy()[0] == _FILL, want == _FILL) + + +def test_the_field_components_are_not_reversed(tmp_path: Path) -> None: + """The (x, y, z) / (z, y, x) test, built so that reversing the two orders cannot pass. + + One component is non-zero and the displacement is a whole number of source voxels on that axis, + so the answer is the source shifted by an exact voxel count along ONE array axis. Put the + displacement on z instead of x — which is what reversing the component axis does — and the + result is shifted along the wrong axis, by a different number of voxels, because the spacings + differ too. + """ + images = Dataset(tmp_path / "Images", "h5") + fields = Dataset(tmp_path / "Fields", "h5") + geometry = _attributes([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]) + volume = _high_frequency((6, 7, 8)) + images.write("Case", _CASE, volume, geometry) + # The reference IS the source grid here: the only thing moving anything is the field. + images.write("Reference", _CASE, volume, geometry) + # +2 world units on x alone, uniform: with unit spacing that is exactly two voxels on array axis 2. + uniform = np.zeros((3, 6, 7, 8), dtype=np.float32) + uniform[0] = 2.0 + fields.write("DVF", _CASE, uniform, geometry) + + stage = ResampleToReference( + entry=_CASE, + group="Reference", + field=f"{fields.filename}:h5", + field_group="DVF", + max_displacement=4.0, + fill=0.0, + ) + stage.set_datasets([images]) + got = stage(_CASE, torch.from_numpy(volume.copy()), Attribute(geometry)).numpy()[0] + + # output(o) = input(o + 2) along the LAST array axis, which is physical x. + np.testing.assert_allclose(got[:, :, :-2], volume[0][:, :, 2:], rtol=0, atol=1e-4) + # And nothing moved along z: a reversed component axis would have shifted this one instead. + assert not np.allclose(got[:-2, :, :], volume[0][2:, :, :], atol=1e-4) + + +def test_it_interpolates_once_not_twice(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """The reason this is one stage: two resamples cost detail that the second cannot put back. + + The two-pass baseline is built entirely in SimpleITK — resample onto the grid, then warp on it — + so what it costs is measured independently of anything here. On a high-frequency source that + cost is a large fraction of the range, while this stage sits at float rounding from the one-pass + result. Asserted as an ORDER OF MAGNITUDE, not a number: the point is the gap, not its digits. + """ + images, fields, volume = warped + got = _warping(images, fields)( + _CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + ).numpy()[0] + one_pass = _simpleitk_warp(volume, _displacement()) + + intermediate = sitk.GetImageFromArray(_simpleitk_warp(volume, None)) + intermediate.SetOrigin(_REFERENCE_ORIGIN) + intermediate.SetSpacing(_REFERENCE_SPACING) + vector = sitk.GetImageFromArray(np.moveaxis(_displacement(), 0, -1).astype(np.float64), isVector=True) + vector.SetOrigin(_FIELD_ORIGIN) + vector.SetSpacing(_FIELD_SPACING) + grid = sitk.Image(*reversed(_REFERENCE_SPATIAL), sitk.sitkFloat32) + grid.SetOrigin(_REFERENCE_ORIGIN) + grid.SetSpacing(_REFERENCE_SPACING) + two_pass = sitk.GetArrayFromImage( + sitk.Resample( + intermediate, + grid, + sitk.DisplacementFieldTransform(sitk.Cast(vector, sitk.sitkVectorFloat64)), + sitk.sitkLinear, + _FILL, + ) + ) + + both = (one_pass != _FILL) & (two_pass != _FILL) + second_pass_costs = float(np.abs(one_pass[both] - two_pass[both]).max()) + this_stage_costs = float(np.abs(got[both] - one_pass[both]).max()) + assert second_pass_costs > 1.0, "the fixture must be one a second interpolation actually damages" + assert this_stage_costs < second_pass_costs / 1000.0, ( + f"this stage is {this_stage_costs:.3g} from the one-pass result where a second interpolation" + f" costs {second_pass_costs:.3g}: it is not resampling once" + ) + + +def test_the_streamed_warp_equals_the_whole_volume(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """Region by region through a field, against the same chain run whole.""" + images, fields, volume = warped + + def manager() -> DatasetManager: + return DatasetManager( + index=0, + group_src="Case", + group_dest="Case", + name=_CASE, + dataset=images, + patch=None, + transforms=[_warping(images, fields)], + data_augmentations_list=[], + ) + + streaming = manager() + assert streaming.stream_refusal(0) is None + whole = _warping(images, fields)( + _CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + ).numpy() + extent = streaming.spatial_shape + got = np.empty((1, *extent), dtype=np.float32) + for start in range(0, extent[0], 3): + stop = min(start + 3, extent[0]) + got[:, start:stop] = streaming.read_region( + (slice(start, stop), slice(0, extent[1]), slice(0, extent[2])) + ).numpy() + np.testing.assert_array_equal(got, whole) + + +def test_the_warped_run_never_assembles_the_volume(warped: tuple[Dataset, Dataset, np.ndarray], tmp_path: Path) -> None: + """The memory bound, through a field: neither the case nor the field is ever read whole.""" + images, fields, _volume = warped + manager = DatasetManager( + index=0, + group_src="Case", + group_dest="Case", + name=_CASE, + dataset=images, + patch=None, + transforms=[_warping(images, fields), Write(str(tmp_path / "Warped"))], + data_augmentations_list=[], + ) + assert manager.stream_refusal(0) is None + + def refuse(*args: object, **kwargs: object) -> None: + raise AssertionError("the chain read a whole volume") + + monkeypatched = pytest.MonkeyPatch() + monkeypatched.setattr(Dataset, "read_data", refuse) + try: + assert manager.materialize() is True + finally: + monkeypatched.undo() + written = Dataset(tmp_path / "Warped", "mha").read_data("Case", _CASE)[0] + assert list(written.shape[1:]) == list(_REFERENCE_SPATIAL) + + +def test_a_field_beyond_its_declared_bound_is_refused(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """The halo is a promise about what was read; a field that breaks it must not sample zeros. + + Sampling past the region that was read gives a dark rim around the moved anatomy and nothing + else to see, which is the shape of a mistake nobody finds. + """ + images, fields, volume = warped + stage = _warping(images, fields, max_displacement=0.5) + with pytest.raises(TransformError, match="displaces up to"): + stage(_CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + + +def test_a_field_with_no_bound_declares_the_whole_volume(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """Warp's rule, and for the same reason: an unbounded reach cannot size a region.""" + images, fields, _volume = warped + locality = _warping(images, fields, max_displacement=0.0).patch_locality( + _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + ) + assert locality.kind is LocalityKind.WHOLE_VOLUME + assert locality.reason is not None and "max_displacement" in locality.reason + # And without a field it is a region stage whatever the bound says. + assert _stage_regrid_kind(images) is LocalityKind.REGRID + + +def _stage_regrid_kind(images: Dataset) -> LocalityKind: + stage = ResampleToReference(entry=_CASE, group="Reference") + stage.set_datasets([images]) + return stage.patch_locality(_attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)).kind + + +def test_a_field_on_another_direction_is_refused(warped: tuple[Dataset, Dataset, np.ndarray], tmp_path: Path) -> None: + """A field whose axes do not line up with the case's would need a rotation, not a per-axis map.""" + images, _fields, _volume = warped + turned = Dataset(tmp_path / "Turned", "h5") + rotated = np.asarray([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + turned.write("DVF", _CASE, _displacement(), _attributes(_FIELD_ORIGIN, _FIELD_SPACING, rotated)) + stage = _warping(images, turned) + with pytest.raises(TransformError, match="Direction"): + stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + + +def test_fields_can_live_beside_the_cases(warped: tuple[Dataset, Dataset, np.ndarray]) -> None: + """No 'field' path: the fields are a group of the run's own roots, one entry per case. + + That is how a cohort registered in place stores them, and it is the same answer as naming the + store explicitly — which is what makes it a shorthand rather than a second code path. + """ + images, fields, volume = warped + images.write("DVF", _CASE, _displacement(), _attributes(_FIELD_ORIGIN, _FIELD_SPACING)) + beside = ResampleToReference(entry=_CASE, group="Reference", field_group="DVF", max_displacement=_BOUND, fill=_FILL) + beside.set_datasets([images]) + + got = beside(_CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING)) + want = _warping(images, fields)( + _CASE, torch.from_numpy(volume.copy()), _attributes(_SOURCE_ORIGIN, _SOURCE_SPACING) + ) + np.testing.assert_array_equal(got.numpy(), want.numpy()) + + +def test_a_bound_with_no_field_is_refused() -> None: + """A max_displacement and nothing to apply it to is a chain that silently does not warp. + + The stage would resample onto the grid perfectly well and the declaration would simply have no + effect — which is the failure that leaves a plausible volume and no error. + """ + with pytest.raises(TransformError, match="no field to apply"): + ResampleToReference(entry=_CASE, group="Reference", max_displacement=1.0) diff --git a/tests/unit/test_transform_locality_contract.py b/tests/unit/test_transform_locality_contract.py index 007f3a67..5a7333bf 100644 --- a/tests/unit/test_transform_locality_contract.py +++ b/tests/unit/test_transform_locality_contract.py @@ -96,6 +96,13 @@ _REFERENCE_SPATIAL = (7, 8, 9) _REFERENCE_SPACING = [1.8, 1.2, 2.5] +# A displacement field on a grid COARSER than either, which is how one is actually solved: the field +# is in world units, so it is read where it is asked rather than resampled to match anything first. +# Its displacements stay inside _FIELD_BOUND, which is what the halo is sized from and checked against. +_FIELD_SPATIAL = (4, 5, 5) +_FIELD_SPACING = [3.0, 2.4, 4.0] +_FIELD_BOUND = 3.0 + # No extent is a multiple of the patch size, so the last patch of every axis is a border patch the read # plan has to pad: the grid is 3x3x3 and 19 of its 27 patches touch a border. _PEAK = 450.0 @@ -174,6 +181,13 @@ class _Case: "ResampleToReference": [ _Case(ResampleToReference(entry=_CASE_NAME, group="Reference")), _Case(ResampleToReference(entry=_CASE_NAME, group="Reference"), group="Labels"), + # Onto the same grid THROUGH a field, which is the whole operation in one stage: the source + # region a target region pulls is the affine box grown by the declared displacement, and the + # sampling is no longer separable. Same contract, and the same atol: one sampler, global + # coordinates. + _Case( + ResampleToReference(entry=_CASE_NAME, group="Reference", field_group="Field", max_displacement=_FIELD_BOUND) + ), ], "ResampleTransform": [_Case(ResampleTransform({"transform": True}))], "Save": [_Case(Save("Dataset"))], @@ -246,14 +260,25 @@ def _volumes() -> dict[str, np.ndarray]: # A grid of its OWN -- other extent, other spacing, other origin -- for a stage that resamples # onto a reference rather than about the case's own extent. Only its header is ever read. "Reference": rng.standard_normal(_REFERENCE_SPATIAL).astype(np.float32)[None], + # A displacement field, component-first in physical (x, y, z), each component a different + # function so a reversed axis order cannot pass unnoticed. Bounded by _FIELD_BOUND. + "Field": np.stack( + [ + 2.0 * np.cos(np.arange(_FIELD_SPATIAL[2]))[None, None, :] * np.ones(_FIELD_SPATIAL), + 1.5 * np.sin(np.arange(_FIELD_SPATIAL[1]))[None, :, None] * np.ones(_FIELD_SPATIAL), + 1.0 * np.cos(np.arange(_FIELD_SPATIAL[0]))[:, None, None] * np.ones(_FIELD_SPATIAL), + ] + ).astype(np.float32), } def _attributes(group: str) -> Attribute: """The metadata a group is stored with -- and so what a declaration about it is handed.""" attributes = Attribute() - attributes["Origin"] = np.asarray([-1.0, 6.0, 12.0] if group == "Reference" else [-3.0, 5.0, 11.0]) - attributes["Spacing"] = np.asarray(_REFERENCE_SPACING if group == "Reference" else _SPACING) + origins = {"Reference": [-1.0, 6.0, 12.0], "Field": [-4.0, 4.0, 10.0]} + spacings = {"Reference": _REFERENCE_SPACING, "Field": _FIELD_SPACING} + attributes["Origin"] = np.asarray(origins.get(group, [-3.0, 5.0, 11.0])) + attributes["Spacing"] = np.asarray(spacings.get(group, _SPACING)) attributes["Direction"] = {"Oblique": _OBLIQUE, "Permuting": _PERMUTING}.get(group, _AXIS_ALIGNED).reshape(-1) if group == "Ensemble": # What a `combine: Concat` reduction writes: the per-model channel counts MergeLabels and From 001ecac1453ad453b4e97607527026f7f5dfeb0c Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Tue, 4 Aug 2026 10:52:36 +0200 Subject: [PATCH 3/3] fix(transform): what reviewing the new stage turned up The auto bound stopped at the first root that answers, but the bound is the cohort's and a field declared by group alone is looked up beside the cases -- which a run may spread over several stores. The plan asked each stage about the case as stored: past a resample a case that covered the whole reference covers 67.5% of it, and the plan said nothing about the third that is fill. And a statistics section contradicted the table above it. --- docs/source/config_guide/transform.md | 13 ---------- konfai/data/transform.py | 17 +++++++------ konfai/transformer.py | 15 +++++++++--- tests/unit/test_resample_to_reference.py | 22 +++++++++++++++++ tests/unit/test_transformer_workflow.py | 31 ++++++++++++++++++++++++ 5 files changed, 75 insertions(+), 23 deletions(-) diff --git a/docs/source/config_guide/transform.md b/docs/source/config_guide/transform.md index c1fe2681..1b615dd8 100644 --- a/docs/source/config_guide/transform.md +++ b/docs/source/config_guide/transform.md @@ -457,19 +457,6 @@ not. `Clip` then `Normalize` therefore takes the whole-volume path, and the plan says so. Reorder the chain, or cut it with a `Save`. ``` -### Statistics a stage may ask for - -A stage that needs a figure over the WHOLE volume does not have to assemble it. -Declaring the statistic in `PatchLocality(LocalityKind.GLOBAL_STAT, stat_keys=…)` -tells the planner to read it once from the stored volume; the stage is then a -value map, and a volume of any size runs region by region. - -Alongside the pooled `Mean`, `Min`, `Max` and `Std`, a **per-component** mean is -available as `MeanPerChannel`. It exists because a per-channel quantity has as -many parts as the volume has components, and the pooled mean of all of them -describes none of them — a three-component displacement field centred by one -number is centred on no axis. - Handed the whole volume anyway — a chain that fell back for another reason — a stage should take the statistic from the tensor in hand and record it, so both paths leave the same state behind. diff --git a/konfai/data/transform.py b/konfai/data/transform.py index ba9125a1..e7a1b335 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -2378,15 +2378,18 @@ def component_bound(self) -> list[float] | None: bound: list[float] = [] try: group = self.group_for(None) - root = self._root_for(None) + # Every root, not the first that answers: this is the COHORT's bound, and a field declared + # by group alone is looked up beside the cases, which a run may spread over several stores. + roots = [self.dataset] if self.dataset is not None else list(self.roots) # The header reads belong inside: a directory store lists its entries from the filesystem # alone, so an unreadable field can only surface here, one entry at a time. - for entry in root.get_names(group): - _shape, attribute = root.get_infos(group, entry) - if DISPLACEMENT_BOUND_ATTRIBUTE not in attribute: - return self._auto_bound - recorded = [float(value) for value in attribute.get_np_array(DISPLACEMENT_BOUND_ATTRIBUTE).ravel()] - bound = recorded if not bound else [max(a, b) for a, b in zip(bound, recorded, strict=False)] + for root in roots: + for entry in root.get_names(group): + _shape, attribute = root.get_infos(group, entry) + if DISPLACEMENT_BOUND_ATTRIBUTE not in attribute: + return self._auto_bound + recorded = [float(value) for value in attribute.get_np_array(DISPLACEMENT_BOUND_ATTRIBUTE).ravel()] + bound = recorded if not bound else [max(a, b) for a, b in zip(bound, recorded, strict=False)] except Exception: # an unreadable field dataset is a whole-volume answer, not a crash return self._auto_bound if bound and max(bound) > 0.0: diff --git a/konfai/transformer.py b/konfai/transformer.py index 6aca6c49..c5bc92e1 100644 --- a/konfai/transformer.py +++ b/konfai/transformer.py @@ -542,16 +542,25 @@ def _plan_notes(self) -> list[str]: notes: list[str] = [] for group_dest, managers in self._managers().items(): for manager in managers: + # Each stage is asked about ITS OWN input, not about the case as stored: past a + # Resample or a Crop the grid a stage meets is no longer the one on disk. Folded the + # way the streamed planner folds it, over a copy of the stored state. + shape = [int(extent) for extent in manager.base_shape[1:]] + attributes = manager.stored_attributes for stage in manager.transforms: # A chain's stages are transforms AND draws; only a transform declares a note. # A draw has nothing to add anyway: what its copies cost is the `regime` column. if not isinstance(stage, Transform): continue - note = stage.plan_note( - group_dest, manager.name, list(manager.base_shape[1:]), manager.stored_attributes - ) + note = stage.plan_note(group_dest, manager.name, list(shape), Attribute(attributes)) if note is not None and note not in notes: notes.append(note) + source = list(shape) + shape = [ + int(extent) + for extent in stage.transform_shape(manager.group_src, manager.name, source, attributes) + ] + stage.write_stream_cache_attribute(attributes, source) return notes def setup(self, world_size: int): diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 070b2bec..47b99f7f 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -36,6 +36,7 @@ from konfai.data.transform import LocalityKind, Reduce, Resample, ResampleToReference, Write from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import ConfigError, TransformError +from konfai.utils.ome_zarr import DISPLACEMENT_BOUND_ATTRIBUTE pytest.importorskip("SimpleITK") import SimpleITK as sitk @@ -779,6 +780,27 @@ def test_fields_can_live_beside_the_cases(warped: tuple[Dataset, Dataset, np.nda np.testing.assert_array_equal(got.numpy(), want.numpy()) +def test_auto_reads_the_bound_from_every_root_not_the_first(tmp_path: Path) -> None: + """The bound is the cohort's, and a cohort declared by group alone can span the run's roots. + + Stopping at the first root that answers gives a halo sized for part of the cohort: the cases + living in the other stores are then read through a region too small for their own displacement. + """ + first, second = Dataset(tmp_path / "first", "mha"), Dataset(tmp_path / "second", "mha") + for root, shift in ((first, 1.0), (second, 9.0)): + field = np.zeros((3, 4, 4, 4), dtype=np.float32) + field[:] = shift + attributes = _attributes(_FIELD_ORIGIN, _FIELD_SPACING) + attributes[DISPLACEMENT_BOUND_ATTRIBUTE] = np.array([shift, shift, shift]) + root.write("DVF", f"CASE_{int(shift)}", field, attributes) + + stage = ResampleToReference(entry="CASE_1", group="Reference", field_group="DVF", max_displacement="auto") + stage.set_datasets([first, second]) + + # 9.0 from the second root, not 1.0 from the first. + assert stage.displacement.component_bound() == [9.0, 9.0, 9.0] + + def test_a_bound_with_no_field_is_refused() -> None: """A max_displacement and nothing to apply it to is a chain that silently does not warp. diff --git a/tests/unit/test_transformer_workflow.py b/tests/unit/test_transformer_workflow.py index 462d8b5e..67535575 100644 --- a/tests/unit/test_transformer_workflow.py +++ b/tests/unit/test_transformer_workflow.py @@ -355,6 +355,37 @@ def test_the_plan_reports_the_dtype_it_probed_the_destinations_with(tmp_path: Pa assert "assumed uint8 / source channels" in plan.report() +_RESAMPLED_THEN_REFERENCED = """\ + ResampleToResolution: + spacing: [2.0, 2.0, 2.0] + ResampleToReference: + entry: CASE_000 + group: CT + Write: + dataset: {out}:h5 +""" + + +def test_a_stage_is_asked_about_its_own_input_not_the_case_as_stored(tmp_path: Path) -> None: + """The note a stage declares describes the grid it MEETS, which the stages before it decide. + + Here the resample takes 12x10x8 down to 12x7x2, so the case covers 67.5% of the reference. Asked + about the case as stored it covers all of it and says nothing -- the plan stays silent about the + remaining third of the output, which will be fill. + """ + _write_source(tmp_path) + _write_config(tmp_path, _RESAMPLED_THEN_REFERENCED.format(out=tmp_path / "out")) + workflow = _build(tmp_path) + manager = workflow._managers()["CT_out"][0] + stored = [int(extent) for extent in manager.base_shape[1:]] + reference = manager.transforms[1] + + notes = workflow._plan_notes() + + assert reference.plan_note("CT_out", "CASE_000", stored, manager.stored_attributes) is None + assert notes and all("covers 67.5%" in note for note in notes) + + def test_unknown_key_is_refused_with_its_path(tmp_path: Path) -> None: """The strict mode: a typo'd key is a parse error, never a silently-used default.""" _write_source(tmp_path)