From a37157896f19a6ed88457b5bae62f184c58b8b2e Mon Sep 17 00:00:00 2001 From: Valentin Boussot Date: Sun, 30 Aug 2026 20:44:13 +0200 Subject: [PATCH] feat(transform): stream stored displacement fields and price what they read One squashed commit from feat/stream-stored-fields (12 commits): - feat(geometry): bound a displacement by where it reaches, not how far - feat(transform): read a stored dense field on the region's window - feat(transform): price a stored field at the identity, and stop reading it to plan - fix(transform): fold a stored map's interval through coverage with its sign - fix(reduce): probe a fold with a short region, so the probe cannot be the kill - perf(transform): hold a field in the walk's dtype, not float64 whatever the walk - perf(reduce): select the median by a window of min/max, never by sorting a stack - fix(reduce): a host meter does not charge a scope for the chunk cache it filled - fix(reduce): the probe's allowance leaves the chunk cache its share - fix(resample): the plan's coverage note takes the refusal's exemption - perf(resample): apply a float32 field with sitk.Warp, not a float64 transform - fix(transform): hold a field at the width its store holds it, never wider --- konfai/data/case_reduction.py | 50 +++- konfai/data/patching.py | 20 +- konfai/data/reduction.py | 95 ++++--- konfai/data/transform.py | 300 ++++++++++++++++++----- konfai/utils/ITK.py | 87 ++++++- konfai/utils/ome_zarr.py | 17 ++ tests/unit/test_case_reduction.py | 112 ++++++++- tests/unit/test_itk_transforms.py | 15 +- tests/unit/test_resample_to_reference.py | 26 ++ tests/unit/test_resample_transform.py | 185 +++++++++++++- 10 files changed, 785 insertions(+), 122 deletions(-) diff --git a/konfai/data/case_reduction.py b/konfai/data/case_reduction.py index 4d5bb22c..ae165ab1 100644 --- a/konfai/data/case_reduction.py +++ b/konfai/data/case_reduction.py @@ -604,6 +604,17 @@ def _member_region(self, manager: DatasetManager, region: tuple[slice, ...]) -> # keeps a reserve for the same reason (Predictor._ACCUMULATE_MARGIN): the measurement is of the # region that just ran, and the next one meets an allocator in a different state. _MEASURED_MARGIN = 0.9 + #: The probe's share of the planned height. The probe is the one region that runs BEFORE any + #: measurement can bound it, so it is the one region that must not be able to kill the run on + #: its own. At the planned height it could: a fold over registration fields held 1.42x, 1.47x + #: and 1.50x what its first region was allowed at three budgets, and at an `auto` budget of 77 + #: GiB that first region reached 90 GiB resident on a 122 GiB host, and the host went down + #: before the probe could read anything. The host gives no OutOfMemoryError to catch: the + #: kernel kills. A quarter-height probe overshooting by the same 1.5x holds 0.4 of the budget, + #: which is survivable, and the ratio it measures is the same one -- the halo does not shrink + #: with the region, so a short region over-holds by MORE than a tall one, and a refit from it + #: is conservative. Its price is one extra region: seconds, on a fold of minutes. + _PROBE_SHARE = 0.25 def _folds(self, spatial: list[int], measure: bool = False): """Every region's fold, in order: the loop both passes share. @@ -619,7 +630,11 @@ def _folds(self, spatial: list[int], measure: bool = False): """ start, refitted = 0, not measure while start < int(spatial[0]): - stop = min(start + self.slab_rows, int(spatial[0])) + # The probe is SHORT. Every later region is cut against what it measured; the probe + # itself is cut against nothing, so it is sized so that its own overshoot cannot + # reach the host's limit (_PROBE_SHARE). + rows = self.slab_rows if refitted else max(1, int(self.slab_rows * self._PROBE_SHARE)) + stop = min(start + rows, int(spatial[0])) region = (slice(start, stop), *(slice(0, extent) for extent in spatial[1:])) # Only around the region that is actually the probe. The host meter RESETS the # process's resident high-water mark to take its reading, and that mark is what the @@ -643,26 +658,37 @@ def _refit_to_measurement(self, meter: HeldMeter | None, rows: int, spatial: lis bounds the next one from above, exactly as the predictor's gate reads a forward's transient from the batch that just ran (:meth:`Predictor._accumulate_device`). - Against the BUDGET, not the share the sizing aims at: a share is how a height is chosen, - and what must not be exceeded is the whole declaration. This exists to prevent a kill, not - to shave bytes -- and it cannot shave many, since the region that set the peak has already - run and cutting the rest never undoes it. What it catches is a LATER region holding more - than the first: a case with a wider halo, a region touching more chunks. + Against the whole declaration LESS the chunk cache's share, because that is what the + reading covers. A share is how a height is chosen and what must not be exceeded is the + declaration -- but the meter no longer counts the decoded-chunk cache (it outlives the + region, and charging the region for it cut every region after the probe), so the cache's + bytes have to come off the other side of the comparison too. Judged against the whole + budget, a reading that excludes the cache lets the cache be spent twice: once inside the + allowance, and again by the cache itself. This exists to prevent a kill, not to shave + bytes. The probe is a short region (_PROBE_SHARE), so what it held is scaled to + the planned height before it is judged: a probe that held its share of the budget says the + full region would hold the budget, and a probe that held more says the full region would + be the kill this exists to prevent. Only ever shorter: a probe that came in under its + share does not talk the fold into a taller region than the plan allowed. """ del spatial held = meter.held() if meter is not None else None if held is None or not self._budget_bytes or self._budget_bytes <= 0 or rows <= 0 or held <= 0: return - allowed = float(self._budget_bytes) * self._MEASURED_MARGIN - if held <= allowed: + cache = budget_share("cache", self._budget_bytes) or 0.0 + allowed = (float(self._budget_bytes) - cache) * self._MEASURED_MARGIN + # What the FULL region would hold, from what the probe held: the halo is a fixed cost the + # probe paid in full, so scaling by height over-estimates, which is the safe direction. + projected = held * (self.slab_rows / float(rows)) + if projected <= allowed: return - fitted = max(1, int(rows * allowed / held)) + fitted = max(1, int(self.slab_rows * allowed / projected)) if fitted >= self.slab_rows: return print( - f"[Reduce] '{self.reduce.output}': first region held {format_bytes(held)} of the" - f" {format_bytes(allowed)} its {rows} row(s) may hold --" - f" the rest are cut to {fitted} row(s).", + f"[Reduce] '{self.reduce.output}': a {rows}-row probe held {format_bytes(held)}, so the planned" + f" {self.slab_rows} row(s) would hold {format_bytes(projected)} of the {format_bytes(allowed)}" + f" allowed -- the rest are cut to {fitted} row(s).", flush=True, ) self.slab_rows = fitted diff --git a/konfai/data/patching.py b/konfai/data/patching.py index fdc95b61..4e6f3f74 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -406,7 +406,25 @@ def device_peak() -> int | None: if not reset_resident_peak(): return None resident = resident_bytes() - return None if resident is None else HeldMeter(peak_resident_bytes, int(resident)) + if resident is None: + return None + # THE CACHE IS NOT THE SCOPE'S. A host peak is the whole process's high-water mark, and the + # decoded-chunk cache sits inside it: a scope that reads from a store fills the cache on its + # way, and the cache keeps what it decoded past the scope, for the next one. Charging the scope + # for that is charging it for a budget line that has its own share (BUDGET_SHARES['cache']). + # Measured on a fold's probe over ten native members: 24.4 GiB read, 13.2 of it the cache + # filling from empty, and the fold cut to 78 % of the height its regions actually needed. + from konfai.utils.ome_zarr import chunk_cache_held_bytes + + cache_at_start = chunk_cache_held_bytes() + + def resident_peak_less_cache() -> int | None: + peak = peak_resident_bytes() + if peak is None: + return None + return peak - max(0, chunk_cache_held_bytes() - cache_at_start) + + return HeldMeter(resident_peak_less_cache, int(resident)) def save_destination(save: Save, default_dataset: Dataset, default_group: str) -> tuple[Dataset, str]: diff --git a/konfai/data/reduction.py b/konfai/data/reduction.py index e194e4c9..e2b8590f 100644 --- a/konfai/data/reduction.py +++ b/konfai/data/reduction.py @@ -191,17 +191,29 @@ class Median(Reduction): """ voxel_local = True - # ``torch.stack`` copies the buffer and the sort along the case axis returns values and - # int64 indices over that: measured at 4x the stack it is handed (6 x 16 MiB float32 cases). - # A fold of three, four or five members takes the network instead and costs far less; the - # attribute is the worst case, :meth:`working_multiple_for` is what the plan asks. - working_multiple = 4.0 - - #: What the selection networks below hold beside the members they are handed, measured on a - #: 24 MiB member (float32): the sort's own 4.0 is what anything wider still costs. + # THE MIDDLE IS SELECTED, NEVER SORTED. A sort along the case axis copies the stack and returns + # int64 indices over it -- eight bytes an element whatever the members weigh -- so ten uint16 + # regions of 33 x 1331 x 1775 (1.45 GiB) sorted at 6.0x their own size, and ten float32 ones + # at 4.0x (peak resident above the members, measured). A selection network of element-wise + # min/max holds a WINDOW of the k+1 smallest members seen so far, in the averaging dtype, and + # inserts each member into it: no stack, no indices, and the members stay in the dtype they + # arrived in. Ten uint16 members: 1.8x, against 6.0x. Twice the arithmetic of the sort (3.4 s + # against 1.6 on that region) on a fold whose clock is the disk by 40 to 1, and whose regions + # the planner may now cut two to three times taller. + # + # The attribute is the worst case the plan may see; :meth:`working_multiple_for` prices the + # network for the count it is handed. + working_multiple = 2.5 + #: What the hand-written networks (three to five) hold beside the members they are handed, + #: measured on a 24 MiB float32 member. _NETWORK_MULTIPLE: ClassVar[dict[int, float]] = {1: 1.0, 2: 1.5, 3: 1.0, 4: 2.5, 5: 1.5} + #: Past five, the window: k+1 float32 buffers for k = count // 2, and the two it blends, + #: measured on a 293 MiB uint16 member at ten. + _WINDOW_MULTIPLE = 1.8 def working_multiple_for(self, cases: int) -> float: + if cases > 5: + return self._WINDOW_MULTIPLE return self._NETWORK_MULTIPLE.get(cases, float(self.working_multiple)) @staticmethod @@ -210,23 +222,23 @@ def _median_of_three(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch def __call__(self, tensors: list[torch.Tensor]) -> torch.Tensor: dtype = _averaged_dtype(tensors[0].dtype) - members = [tensor.to(dtype) for tensor in tensors] - if len(members) == 1: - return members[0] - # Three to five members is what a fold has, and there the middle is SELECTED by a network of - # element-wise min/max rather than found by sorting the whole stack: same values to the bit, - # a fraction of the time (CUDA 7.20 -> 0.45 ms at three, 8.42 -> 1.10 at five; CPU 55 -> 26 - # at three), and no stack to hold, which is what lets the planner cut taller slabs. Beyond - # five the sort is simpler and no slower: what torch.quantile computes without its - # interpolation machinery (1.5-2x on CPU, 3.5x on CUDA, measured). - low, high = self._middle_pair(members) + if len(tensors) == 1: + return tensors[0].to(dtype) + # The members are handed over in the dtype they arrived in and widened one at a time as the + # network takes them: torch has no integer min/max kernel on the CPU, and widening ten + # members up front is what put ten float32 copies beside ten uint16 regions. + low, high = self._middle_pair(tensors, dtype) return low if low is high else torch.lerp(low, high, 0.5) - def _middle_pair(self, members: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + def _middle_pair(self, members: list[torch.Tensor], dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: """The one middle member of an odd fold (the same tensor twice), or the two an even fold - averages: by network up to five members, off a sorted stack past it.""" + averages: by a hand-written network up to five members, by the insertion window past it. + ``dtype`` is what the network computes in; the members are widened to it as they enter.""" minimum, maximum = torch.minimum, torch.maximum count = len(members) + if count > 5: + return self._middle_pair_by_window(members, dtype) + members = [member.to(dtype) for member in members] if count == 2: first, second = members return minimum(first, second), maximum(first, second) @@ -239,19 +251,40 @@ def _middle_pair(self, members: list[torch.Tensor]) -> tuple[torch.Tensor, torch c, d = minimum(c, d), maximum(c, d) second, third = maximum(a, c), minimum(b, d) return minimum(second, third), maximum(second, third) - if count == 5: - a, b, c, d, e = members - a, b = minimum(a, b), maximum(a, b) - c, d = minimum(c, d), maximum(c, d) - a, c = minimum(a, c), maximum(a, c) # a is the fold's smallest: out of the running - b, d = minimum(b, d), maximum(b, d) # d is its largest: out too - middle = self._median_of_three(b, c, e) - return middle, middle - ranked = torch.stack(members, dim=0).sort(dim=0).values + a, b, c, d, e = members + a, b = minimum(a, b), maximum(a, b) + c, d = minimum(c, d), maximum(c, d) + a, c = minimum(a, c), maximum(a, c) # a is the fold's smallest: out of the running + b, d = minimum(b, d), maximum(b, d) # d is its largest: out too + middle = self._median_of_three(b, c, e) + return middle, middle + + @staticmethod + def _middle_pair_by_window(members: list[torch.Tensor], dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + """The middle pair of any count, by an insertion window of the ``k + 1`` smallest seen. + + The middle of ``count`` members is rank ``count // 2`` (0-based; the pair ``count // 2 - 1`` + and ``count // 2`` on an even count). A window that keeps the ``k + 1`` smallest members + seen so far, ``k = count // 2``, holds those ranks exactly once every member has passed + through it: a member larger than the whole window can be no smaller than rank ``k + 1`` of + the members seen, so dropping it off the end loses nothing the answer needs. Each insertion + is a chain of element-wise min/max, which is what makes the selection exact -- the same + values a full sort returns, to the bit (pinned against ``torch.sort`` in the tests). + """ + count = len(members) + keep = count // 2 + 1 + window: list[torch.Tensor] = [] + for member in members: + window.append(member.to(dtype)) + for index in range(len(window) - 1, 0, -1): + lower, upper = window[index - 1], window[index] + window[index - 1], window[index] = torch.minimum(lower, upper), torch.maximum(lower, upper) + if len(window) > keep: + window.pop() + middle = window[count // 2] if count % 2: - middle = ranked[count // 2] return middle, middle - return ranked[count // 2 - 1], ranked[count // 2] + return window[count // 2 - 1], middle class Vote(Reduction): diff --git a/konfai/data/transform.py b/konfai/data/transform.py index 40b4cde9..e9eddf0b 100755 --- a/konfai/data/transform.py +++ b/konfai/data/transform.py @@ -47,6 +47,7 @@ Grid, SpatialStages, TransformBound, + WorldBox, bound_of, ) from konfai.data.reduction import Reduction @@ -1405,6 +1406,54 @@ def drop(self) -> None: self._image = self._key = None +def _warp_field_float32(stages: SpatialStages, region: Grid) -> "Any | None": + """The one displacement this whole map is, as a float32 vector image on ``region`` -- or None. + + ``sitk.Warp`` is templated on the field's own type, where ``DisplacementFieldTransform`` is + not: SimpleITK's transform hierarchy is ``TransformBaseTemplate``, so a field handed to + a transform is cast to float64 whatever it was read as, and interleaved into a vector image at + that width. On a 122-row native region that is 6.9 GiB built in 5.0 s, 39% of everything the + member-region costs -- for a field that came off the store as float32 and whose values fit it + exactly. + + Taken only where it changes no value. The field must ALREADY be float32 -- which is what + ``precision: fast`` reads, and what a store written in float32 holds either way -- because + narrowing a genuine float64 field is a different map, not a cheaper one (pinned: on a field + carrying float64 values the two routes disagree). Warp also evaluates the displacement on the + OUTPUT grid rather than interpolating it at the target point, so it stands in only where the + field IS the output grid, and only where the field is the whole map: one order-1 stage, no + affine beside it, nothing to compose. Anything else keeps the composite path. + + Verified bit-identical on two ExaSPIM members at 89% and 98% real coverage (max |diff| 0 over + the region), 1.3x faster end to end, and the field's buffer halves. + """ + from konfai.data.geometry import DisplacementStage + + if len(stages) != 1: + return None + stage = stages[0] + if not isinstance(stage, DisplacementStage) or stage.order != 1: + return None + if stage.values.dtype != np.float32: + return None + grid = stage.grid + if ( + tuple(grid.size_zyx) != tuple(region.size_zyx) + or not np.allclose(grid.origin_xyz, region.origin_xyz, rtol=0.0, atol=1e-9) + or not np.allclose(grid.spacing_xyz, region.spacing_xyz, rtol=0.0, atol=1e-9) + or not np.allclose(grid.direction_xyz, region.direction_xyz, rtol=0.0, atol=1e-9) + ): + return None + components = [ + sitk.GetImageFromArray(np.ascontiguousarray(stage.values[component])) for component in range(grid.rank) + ] + field = sitk.Compose(components) + field.SetOrigin(np.asarray(grid.origin_xyz, dtype=np.float64).tolist()) + field.SetSpacing(np.asarray(grid.spacing_xyz, dtype=np.float64).tolist()) + field.SetDirection(np.asarray(grid.direction_xyz, dtype=np.float64).ravel().tolist()) + return field + + def _resample_with_sitk( payload: torch.Tensor, region: Grid, @@ -1433,7 +1482,14 @@ def _resample_with_sitk( if mode == "cubic": return None # ITK's BSpline is not Keys' Catmull-Rom: the walk keeps its own cubic rank = source.rank - transform = encode_transform_stages(stages) if stages else sitk.Transform(rank, sitk.sitkIdentity) + # A map that is one field ON the output grid is applied by the filter templated on the field's + # own type, never by one that casts it to float64 to hold it (:func:`_warp_field_float32`). + warp_field = _warp_field_float32(stages, region) if stages else None + transform = ( + None + if warp_field is not None + else (encode_transform_stages(stages) if stages else sitk.Transform(rank, sitk.sitkIdentity)) + ) # The window's own origin: the source origin moved by the window's start along each axis. start_index = np.asarray(list(reversed(region_starts)), dtype=np.float64) # (x, y, z) window_origin = source.index_to_world.apply(start_index) @@ -1445,7 +1501,8 @@ def _resample_with_sitk( resampler.SetOutputOrigin(np.asarray(region.origin_xyz, dtype=np.float64).tolist()) resampler.SetOutputSpacing(np.asarray(region.spacing_xyz, dtype=np.float64).tolist()) resampler.SetOutputDirection(np.asarray(region.direction_xyz, dtype=np.float64).ravel().tolist()) - resampler.SetTransform(transform) + if transform is not None: + resampler.SetTransform(transform) resampler.SetInterpolator(interpolator[mode]) resampler.SetDefaultPixelValue(float(fill)) # A blend interpolates in the dtype the walk accumulates in, and torch makes the final cast: @@ -1471,9 +1528,26 @@ def _resample_with_sitk( image.SetOrigin(np.asarray(window_origin, dtype=np.float64).tolist()) image.SetSpacing(np.asarray(source.spacing_xyz, dtype=np.float64).tolist()) image.SetDirection(np.asarray(source.direction_xyz, dtype=np.float64).ravel().tolist()) - resampler.SetOutputPixelType(image.GetPixelID() if mode == "nearest" else blend_pixel_id) + pixel_id = image.GetPixelID() if mode == "nearest" else blend_pixel_id # Held: a view borrows the image's buffer, and a temporary's is freed under it. - resampled = resampler.Execute(image) + if warp_field is not None: + # Warp answers in its INPUT's type where the resampler is asked for an output type, so + # the blend's dtype is carried in rather than requested. ITK interpolates in double + # from either, and float32 carries a uint16 payload exactly, so the values are the + # ones the resampler would have written (pinned in test_resample_transform.py). + resampled = sitk.Warp( + image if image.GetPixelID() == pixel_id else sitk.Cast(image, pixel_id), + warp_field, + interpolator[mode], + [int(e) for e in reversed(region.size_zyx)], + np.asarray(region.origin_xyz, dtype=np.float64).tolist(), + np.asarray(region.spacing_xyz, dtype=np.float64).tolist(), + np.asarray(region.direction_xyz, dtype=np.float64).ravel().tolist(), + float(fill), + ) + else: + resampler.SetOutputPixelType(pixel_id) + resampled = resampler.Execute(image) np.copyto(landing[channel], sitk.GetArrayViewFromImage(resampled), casting="unsafe") return result @@ -1492,6 +1566,9 @@ class _StoredMap: bound: TransformBound affine: bool + #: Whether a dense field member was priced as the identity rather than bounded. The plan cannot + #: bound one from headers, so the run measures its window from the values it samples anyway. + field: bool = False def _stages_bytes(stages: SpatialStages) -> int: @@ -1512,10 +1589,11 @@ def _stages_bytes(stages: SpatialStages) -> int: #: the run. _FIELD_WINDOW_COPIES = 3.0 -#: What one element of a decoded field weighs. It is read as float64 whatever the store holds -#: (:meth:`_DisplacementSource.read`), while the plan counts its volumes at -#: :data:`~konfai.data.patching._SWEEP_ELEMENT_BYTES`: the ratio is what a field window costs in the -#: currency the plan is written in. +#: What one element of a decoded field weighs under the bit-exact walk: read as float64 whatever +#: the store holds (:meth:`_DisplacementSource.read`), while the plan counts its volumes at +#: :data:`~konfai.data.patching._SWEEP_ELEMENT_BYTES`: the ratio is what a field window costs in +#: the currency the plan is written in. Under ``precision: fast`` the field is held in float32 and +#: weighs half (:meth:`Resample._field_element_bytes`). _FIELD_ELEMENT_BYTES = 8 @@ -1627,7 +1705,9 @@ def __init__( "'exact' (the default) walks coordinates in float64, bit-identical to" " sitk.Resample. 'fast' lets the device walk in float32: half the bytes and about" " twice the rows per slab, at ~|world|/2^24 of coordinate error -- for INTENSITY" - " resamples only (on the host ITK's own resampler is used either way)." + " resamples only. On the host it holds a stored field at float32 and applies it" + " with sitk.Warp instead of a float64 transform, which for a field STORED in" + " float32 is the same answer to the bit and half the field's memory." " A nearest pick that lands within that band of a voxel boundary picks the other" " voxel, so a label map must stay 'exact'.", ) @@ -1656,7 +1736,8 @@ def __init__( self._maps: dict[str, _StoredMap] = {} #: The decoded stages of the last cases sampled, most recent last, within #: ``stored_stage_bytes``. Not pickled: a rank decodes what it samples, not the cohort. - self._stored: OrderedDict[str, SpatialStages] = OrderedDict() + # Keyed on (case, box): a field read for one region answers for that region alone. + self._stored: OrderedDict[tuple, SpatialStages] = OrderedDict() #: The last field window read, kept for the sampler: sizing a region's source window reads #: the very field slab the sampler needs next, so one slot makes the two one read. self._field_window: tuple[str, object, DisplacementStage] | None = None @@ -1676,6 +1757,8 @@ def __getstate__(self) -> dict: #: bytes each) and the dense fields of the last two or three, so a fold reading N members per #: region does not decode a member per region. stored_stage_bytes = 512 << 20 + #: Region-keyed entries kept at once: an all-affine map prices zero bytes and would never evict. + stored_stage_slots = 64 @staticmethod def _target_from( @@ -1813,19 +1896,25 @@ def _grids_of(self, name: str) -> tuple[Grid, Grid]: # ------------------------------------------------------------------ the map - def _stored_stages(self, name: str) -> SpatialStages: + def _stored_stages(self, name: str, box: WorldBox | None = None) -> SpatialStages: """This case's stored transforms, decoded and composed, in application order. The last cases' stages are held, most recent last, within ``stored_stage_bytes``: a run samples its cases one after the other, a fold reads its members region by region, and a loader interleaving cases through a dense field decodes at each switch past the bound. + + KEYED ON THE BOX, because a region's stages are not the case's. A field read for one region + answers for that region and no other, so a second region asking with the same case name + would otherwise be handed the first one's window and sample outside it -- the border value, + silently, where the values it needed were on disk all along. """ - stages = self._stored.pop(name, None) + key = (name, None if box is None else (tuple(box.low_xyz), tuple(box.high_xyz))) + stages = self._stored.pop(key, None) if stages is None: - stages = self._decode_stored(name) - self._stored[name] = stages + stages = self._decode_stored(name, box) + self._stored[key] = stages held = sum(_stages_bytes(kept) for kept in self._stored.values()) - while len(self._stored) > 1 and held > self.stored_stage_bytes: + while len(self._stored) > 1 and (held > self.stored_stage_bytes or len(self._stored) > self.stored_stage_slots): held -= _stages_bytes(self._stored.popitem(last=False)[1]) return stages @@ -1833,17 +1922,53 @@ def _stored_map(self, name: str) -> _StoredMap: """The plan's record of this case's stored map, decoded once and kept without its stages.""" stored = self._maps.get(name) if stored is None: - stages = self._stored.get(name) - if stages is None: - stages = self._decode_stored(name) rank = self._source_grid(name).rank + # HEADERS ONLY, always: a cached region's stages are that region's, not the case's, and + # the plan's own read must not pull a field's values for a bound it cannot form from + # them anyway. What comes back is the affine part, exact, with any dense field standing + # as the identity -- the same price the declared route puts on its field. + priced = self._decode_stored(name, headers_only=True) + has_field = self._stored_has_field(name) stored = self._maps[name] = _StoredMap( - bound_of(stages, rank), all(isinstance(stage, AffineStage) for stage in stages) + bound_of(priced, rank), + not has_field and all(isinstance(stage, AffineStage) for stage in priced), + field=has_field, ) return stored - def _decode_stored(self, name: str) -> SpatialStages: - """This case's stored transforms read and decoded, application order, nothing kept.""" + def _stored_has_field(self, name: str) -> bool: + """Whether any member of this case's stored map is a dense field, from headers alone. + + Asked rather than counted: one entry can decode to several stages (a composite of two + affines is two), so a stage count says nothing about how many members there were, let + alone which of them the plan priced away. + """ + from konfai.utils.dataset import DISPLACEMENT_FIELD_ATTRIBUTE + + for group in cast("dict[str, bool]", self.transforms or {}): + for dataset in self.datasets: + if not dataset.is_dataset_exist(group, name): + continue + if getattr(dataset, "read_data", None) is None: + break # a transform-only store serves no field + _shape, header = dataset.get_infos(group, name) + if DISPLACEMENT_FIELD_ATTRIBUTE in header: + return True + break + return False + + def _decode_stored(self, name: str, box: WorldBox | None = None, headers_only: bool = False) -> SpatialStages: + """This case's stored transforms read and decoded, application order, nothing kept. + + ``box`` is the world box the map will be evaluated over, folded through the members already + decoded so each one is read on the box IT sees rather than the one the region started as: a + field applied second is evaluated where the first sent the points, and reading it on the + region's own box would be short by everything the first one moves. Without a box the whole + entry is read, which is the whole-volume route's answer. + + ``headers_only`` is the plan's read: a dense field member decodes to no stage, which is the + identity, and its values are never touched. + """ from konfai.utils.ITK import invert_stages, read_transform_stages _require_simpleitk() @@ -1857,7 +1982,7 @@ def _decode_stored(self, name: str) -> SpatialStages: decoded = None for dataset in self.datasets: if dataset.is_dataset_exist(group, name): - decoded = read_transform_stages(dataset, group, name) + decoded = read_transform_stages(dataset, group, name, box, headers_only, self._field_dtype) break if decoded is None: raise TransformError( @@ -1876,6 +2001,13 @@ def _decode_stored(self, name: str) -> SpatialStages: " invert it where it is written.", ) decoded = inverted + if box is not None: + # The box the NEXT member is read on, which is where this one sends the points it + # was read for: the EFFECTIVE stages, after any inversion. An affine moves it + # exactly; a field grows it by the range of the values just read, and a member that + # only pushes one way moves the box instead of opening it (TransformBound is an + # interval, not a radius). + box = bound_of(decoded, rank).map_box(box) stages.extend(decoded) return tuple(stages) @@ -1896,31 +2028,36 @@ def _field_stage(self, name: str, region: Grid) -> DisplacementStage: spatial = [int(extent) for extent in shape[1:]] grid = Grid.of(spatial, attribute, f"the field for case '{name}'") window = grid.index_window(region.world_box(), margin=1) - values = source.read(name, window, len(spatial)) + values = source.read(name, window, len(spatial), self._field_dtype) stage = DisplacementStage(grid.sub_grid(window), values.numpy(), order=1) self._field_window = (name, key, stage) return stage def stream_abort(self, name: str) -> None: - self._stored.pop(name, None) + for key in [key for key in self._stored if key[0] == name]: + self._stored.pop(key, None) if self._field_window is not None and self._field_window[0] == name: self._field_window = None self._sitk_input.drop() def _stages(self, name: str, region: Grid) -> SpatialStages: - """The whole map over one target region, in application order.""" + """The whole map over one target region, in application order, each stage read on the box + the stages before it send that region to.""" stages: list[AffineStage | DisplacementStage] = [] + box = region.world_box() if self.displacement is not None: - stages.append(self._field_stage(name, region)) + field = self._field_stage(name, region) + stages.append(field) + box = bound_of((field,), self._source_grid(name).rank).map_box(box) if self.transforms is not None: - stages.extend(self._stored_stages(name)) + stages.extend(self._stored_stages(name, box)) return tuple(stages) def _bound(self, name: str) -> TransformBound: """What the map is guaranteed to do, from stored coefficients alone, no voxel read.""" rank = self._source_grid(name).rank folded = TransformBound.exact(AffineMap.identity(rank)) - if self.displacement is not None: + if self.displacement is not None or (self.transforms is not None and self._stored_map(name).field): raise TransformError( "a field's reach is unknown before its values are read; nothing bounds it from headers." ) @@ -1931,12 +2068,12 @@ def _bound(self, name: str) -> TransformBound: def _pricing_bound(self, name: str) -> TransformBound: """The map's bound as the PLAN prices it: headers and declarations, never a voxel. - A field prices as zero displacement. The run never trusts this window: a field's - regions are sized from the values it reads for sampling anyway - (:meth:`measured_region_source`), so the optimism here costs estimate accuracy, not bytes. + A field prices as zero displacement, declared or stored: nothing bounds one from headers, + and reading its values to find out costs a case's worth of memory for a number the run + replaces anyway. The run never trusts this window -- a field's regions are sized from the + values it reads for sampling (:meth:`measured_region_source`) -- so the optimism here costs + estimate accuracy, not bytes. What is left is the affine part, which is exact. """ - if self.displacement is None: - return self._bound(name) rank = self._source_grid(name).rank folded = TransformBound.exact(AffineMap.identity(rank)) if self.transforms is not None: @@ -2052,8 +2189,17 @@ def stream_region_source( @property def measures_at_run(self) -> bool: - """Whether the run sizes this stage's windows from the data it reads: any declared field.""" - return self.displacement is not None + """Whether the run sizes this stage's windows from the data it reads: any field at all. + + Declared or stored: both are priced as the identity by the plan, so both need the run to + say where they actually reach. An affine-only ``transforms`` is not one of them -- its + bound is exact from the coefficients, the plan's window is already the right one, and + measuring it per region would buy nothing and cost a decode. + + Read off the plan's own records, which ``transform_shape`` fills for every case before + anything asks this. + """ + return self.displacement is not None or any(stored.field for stored in self._maps.values()) def case_working_multiple(self, name: str) -> float: """The sampling grid, plus the field window this case's region holds beside it. @@ -2096,10 +2242,28 @@ def case_working_multiple(self, name: str) -> float: # Charging it at one was counting eight bytes as four. from konfai.data.patching import _SWEEP_ELEMENT_BYTES - widening = _FIELD_ELEMENT_BYTES / _SWEEP_ELEMENT_BYTES + widening = self._field_element_bytes / _SWEEP_ELEMENT_BYTES window = max(1, int(shape[0])) * (target_voxel / field_voxel) * widening return base + window * _FIELD_WINDOW_COPIES + @property + def _field_dtype(self) -> type: + """The CEILING a field's values are held at, never the width they are widened to. + + float64 is the bit-exact contract with SimpleITK and narrows nothing; a field stored in + float32 stays float32 under it, losslessly (see + :func:`~konfai.utils.ITK._displacement_stage`). ``precision: fast`` lowers the ceiling to + the float32 its coordinate walk runs in, which is where a genuinely float64 field narrows. + """ + return np.float32 if self.precision == "fast" else np.float64 + + @property + def _field_element_bytes(self) -> int: + """What the PLAN charges a field value, at the ceiling rather than at the width the store + turns out to hold: the plan reads no field header, and over-charging reserves memory a run + then does not need, which is the safe direction to be wrong in.""" + return int(np.dtype(self._field_dtype).itemsize) + def measured_region_source( self, name: str, target_slices: tuple[slice, ...], source_spatial_shape: list[int], cache_attribute: Attribute ) -> list[slice]: @@ -2296,8 +2460,8 @@ def _coverage(cls, source: Grid, target: Grid, bound: TransformBound | None = No Judged THROUGH the declared map's affine part: a stored transform is what makes a cross-frame pair meet (an MR and a CT in different scanner frames with a rigid bridging them), and a coverage judged before applying it would call every such registration - disjoint. The residual (a spline's or a field's sup-norm) only ever moves a sample by a - bounded amount, so it widens the inside band rather than moving the lattice. Counted on a + disjoint. The interval MOVES the lattice too, by the offset a stored map carries, and only + what varies widens the inside band. Counted on a capped lattice rather than solved, because the sampled set is a box only while the grids are axis-aligned and a rotation makes it a polytope. """ @@ -2308,16 +2472,21 @@ def _coverage(cls, source: Grid, target: Grid, bound: TransformBound | None = No lattice = np.stack([axis.ravel() for axis in np.meshgrid(*axes, indexing="ij")], axis=-1) to_world = target.index_to_world if bound is None else target.index_to_world.then(bound.affine) index = to_world.then(source.world_to_index).apply(lattice) - margin_xyz = ( - np.zeros(source.rank) - if bound is None - # A world-space residual box reaches |W2I| @ r in index space, component-wise. - else np.abs(source.world_to_index.matrix) @ np.asarray(bound.residual_xyz, dtype=np.float64) - ) + low = high = np.zeros((1, source.rank)) + if bound is not None: + # The interval, folded into index space the way a world box is: its two ends land + # where the matrix sends them, and a negative entry swaps which end is which. As a + # radius, a map that sent every sample 500 m PAST the case reached back over it just + # as far, and covered it. + matrix = source.world_to_index.matrix + rise, fall = np.maximum(matrix, 0.0), np.minimum(matrix, 0.0) + low = (rise @ bound.low_xyz + fall @ bound.high_xyz)[None, :] + high = (rise @ bound.high_xyz + fall @ bound.low_xyz)[None, :] inside = np.ones(index.shape[0], dtype=bool) for axis in range(source.rank): extent = float(source.size_zyx[source.rank - 1 - axis]) - inside &= (index[:, axis] >= -0.5 - margin_xyz[axis]) & (index[:, axis] < extent - 0.5 + margin_xyz[axis]) + # A probe reaches [index + low, index + high]: inside when that span meets the grid. + inside &= (index[:, axis] + high[:, axis] >= -0.5) & (index[:, axis] + low[:, axis] < extent - 0.5) return float(np.count_nonzero(inside)) / float(inside.size) def _refuse_if_disjoint(self, name: str) -> None: @@ -2328,10 +2497,13 @@ def _refuse_if_disjoint(self, name: str) -> None: 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. - Never with a field configured: its reach is unknown before its values are read, and - bridging two frames is precisely what a field may be for. + Never with a field configured, DECLARED OR STORED: its reach is unknown before its values + are read -- the plan prices both at the identity -- and bridging two frames is precisely + what a field may be for. A cohort registered onto a template it sits 25 mm from covers + exactly nothing until its own field is applied, and refusing it here would drop every + member of the build the stage exists to serve. """ - if self._target_is_own or self.displacement is not None or self.coverage(name) > 0.0: + if self._target_is_own or self._prices_a_field(name) or self.coverage(name) > 0.0: return where = f"case '{name}'" if name else "the case" raise TransformError( @@ -2342,6 +2514,21 @@ def _refuse_if_disjoint(self, name: str) -> None: " actually surrounds, or drop this case with 'subset'.", ) + def _prices_a_field(self, name: str) -> bool: + """Whether this case's map carries a field the plan prices at the identity. + + A field's reach is known only once its values are read, and the plan reads none: it prices + a declared field and a stored one alike as zero displacement. Every geometric judgement + built on that price is then a judgement about two grids sitting bare in world space, not + about where the samples land -- so neither the refusal nor the plan's coverage note may + speak. Measured on a ten-member ExaSPIM build whose fields bridge a 20 mm gap: judged bare, + one member covered 0.0% of the target and the note called everything it wrote fill, while + the run it was describing read that member in full and the template carried its anatomy. + """ + if self.displacement is not None: + return True + return self.transforms is not None and bool(self._stored_map(name).field) + def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribute: Attribute) -> str | None: """What this case covers of the target grid: measured on the header HANDED OVER. @@ -2361,7 +2548,7 @@ def plan_note(self, group_dest: str, name: str, shape: list[int], cache_attribut ) try: source, missing = Grid.from_header([int(extent) for extent in shape], cache_attribute, f"case '{name}'") - if not missing & self._target.needs: + if not missing & self._target.needs and not self._prices_a_field(name): covered = self._coverage(source, self._target.of(source, name), self._map_bound(name)) if covered < self._WORTH_SAYING: notes.append( @@ -3085,18 +3272,21 @@ def probe(self, name: str) -> None: ) from error self._probed.add(name) - def read(self, name: str, region: tuple[slice, ...] | None, channels: int) -> torch.Tensor: + def read( + self, name: str, region: tuple[slice, ...] | None, channels: int, dtype: type = np.float64 + ) -> torch.Tensor: group = self.group_for(name) root = self._root_for(name) if region is None: data, _attributes = root.read_data(group, name) else: data, _attributes = root.read_data_slice(group, name, (slice(None), *region)) - # float64, not .float(): the walk evaluates the field in float64 (DisplacementStage's own - # contract), and a .float() here quantized a float64-stored field before the exact - # arithmetic ever saw it -- a silent sitk divergence for any field an external tool wrote - # in double. A float32 store widens losslessly. - field = torch.from_numpy(np.ascontiguousarray(data)).to(torch.float64) + # The walk's dtype, handed down by the owner. float64 is the bit-exact contract: a .float() + # here would quantise a float64-stored field before the exact arithmetic ever saw it -- a + # silent sitk divergence for any field an external tool wrote in double. float32 is what a + # `precision: fast` walk takes the values in regardless, so widening them first is a copy + # that costs twice the window and buys nothing. + field = torch.from_numpy(np.ascontiguousarray(data)).to(torch.float32 if dtype is np.float32 else torch.float64) if field.shape[0] != channels: raise TransformError( f"The field for case '{name}' has {field.shape[0]} component(s) where the case has" diff --git a/konfai/utils/ITK.py b/konfai/utils/ITK.py index 0ab6e11f..d4998e3f 100644 --- a/konfai/utils/ITK.py +++ b/konfai/utils/ITK.py @@ -31,7 +31,7 @@ from konfai.utils.errors import TransformError if TYPE_CHECKING: - from konfai.data.geometry import AffineMap, AffineStage, DisplacementStage, Grid, SpatialStages + from konfai.data.geometry import AffineMap, AffineStage, DisplacementStage, Grid, SpatialStages, WorldBox def _require_simpleitk() -> None: @@ -205,12 +205,34 @@ def _grid_of_image(image: sitk.Image) -> Grid: ) -def _displacement_stage(grid: Grid, values: np.ndarray, order: int, what: str) -> DisplacementStage: - """A stage over ``values``, component-first ``(rank, *grid)``, contiguous float64: one copy where - they are not that already, none where they are.""" +def _displacement_stage( + grid: Grid, values: np.ndarray, order: int, what: str, dtype: np.dtype | type = np.float64 +) -> DisplacementStage: + """A stage over ``values``, component-first ``(rank, *grid)``, contiguous in ``dtype``: one copy + where they are not that already, none where they are. + + ``dtype`` is a CEILING, not a target: the values are held at the width the STORE holds them at, + never widened past it. A field written in float32 -- which is how ITK writes one, and how a DVF + normally sits on disk -- carries no more information as float64, so widening it buys a copy of + twice the bytes to say exactly the same thing, three times over once the sampler holds its own + (``_FIELD_WINDOW_COPIES``), and quantised straight back by any walk that runs in float32. + + So float64 (the default ceiling, the bit-exact contract with SimpleITK) narrows nothing that + was stored wide, and lets a float32 store stay float32 -- losslessly, with no flag to set and + no precision traded, which is also what lets :func:`~konfai.data.transform._warp_field_float32` + apply such a field without a float64 transform to hold it. A caller whose coordinate walk runs + in float32 lowers the ceiling to float32, and there a genuinely float64 field does narrow: that + is the trade ``precision: fast`` names. + + Anything not already float32 or float64 -- an integer field, a float16 one -- is converted to + the ceiling: those are widths SimpleITK has no pixel type for, and the walk no kernel for. + """ from konfai.data.geometry import DisplacementStage - values = np.ascontiguousarray(values, dtype=np.float64) + ceiling = np.dtype(dtype) + stored = np.asarray(values).dtype + held = stored if stored.kind == "f" and np.float32().itemsize <= stored.itemsize <= ceiling.itemsize else ceiling + values = np.ascontiguousarray(values, dtype=held) if not np.isfinite(values).all(): raise TransformError( f"{what} carries a non-finite displacement value, so no bound on its reach exists.", @@ -264,7 +286,14 @@ def decode_transform_stages(transform: sitk.Transform) -> SpatialStages: ) -def read_transform_stages(dataset: Any, group: str, name: str) -> SpatialStages: +def read_transform_stages( + dataset: Any, + group: str, + name: str, + box: WorldBox | None = None, + headers_only: bool = False, + field_dtype: np.dtype | type = np.float64, +) -> SpatialStages: """The stored transform ``(group, name)`` of ``dataset`` as geometry stages in APPLICATION order. A displacement entry becomes its stage straight from the array the store hands over, on the @@ -273,6 +302,25 @@ def read_transform_stages(dataset: Any, group: str, name: str) -> SpatialStages: once more on the way out). Every other entry decodes through :func:`decode_transform_stages`, as does everything a store that serves transforms alone (``read_transform`` and nothing else) hands over. + + ``box`` is the world box the caller will evaluate the map over, and it is read from the headers + before a voxel is fetched: a field entry then comes back as its own sub-grid over the window + that box falls in, plus the lattice point linear interpolation reaches for. Without it the whole + entry is read, which is what a whole-volume call and the plan's own decode still want. A field + solved at full resolution is gigabytes -- 14.5 GiB per ExaSPIM case, and float64 on the way in + doubles it -- so a region that read the whole one would hold, per case, more than the budget + sizing it was ever told about. + + Only a displacement entry is windowed. An affine is a matrix and a BSpline a coarse control + grid: both are small, and a BSpline's coefficients are not indexed by the box anyway. + + ``headers_only`` is the plan's read: a dense field comes back as NO stage at all, which is the + identity, and its values are never touched. Nothing bounds a field from headers -- so a plan + that read them would pay a case's worth of memory for a number it cannot use, which is exactly + what the declared route already declines to do (:meth:`Resample._pricing_bound`). + + ``field_dtype`` is what a dense field's values are held in: float64 for the bit-exact walk, + float32 for a caller whose walk runs in float32 and would quantise them back anyway. """ from konfai.data.geometry import Grid from konfai.utils.dataset import DISPLACEMENT_FIELD_ATTRIBUTE, data_to_transform @@ -280,12 +328,35 @@ def read_transform_stages(dataset: Any, group: str, name: str) -> SpatialStages: read_data = getattr(dataset, "read_data", None) if read_data is None: return decode_transform_stages(dataset.read_transform(group, name)) + what = f"the displacement field '{group}' of case '{name}'" + if headers_only: + # The PLAN's read: a dense field answers as the identity and its values are never touched. + # Nothing bounds a field from headers, so there is nothing to read them for -- and reading + # them anyway is what put 29 GiB of one native case beside a budget that never saw it. + # Everything else is coefficients, small, and bounded exactly, so it decodes as usual. + shape, header = dataset.get_infos(group, name) + if DISPLACEMENT_FIELD_ATTRIBUTE in header: + return () + data, attribute = read_data(group, name) + return decode_transform_stages(data_to_transform(data, attribute, name)) + window = None + # A backend that cannot serve a slice reads whole, which is what it would do for any window it + # was given: the transform-only stores (read_transform and nothing else) are already out above. + if box is not None and getattr(dataset, "read_data_slice", None) is not None: + shape, header = dataset.get_infos(group, name) + if DISPLACEMENT_FIELD_ATTRIBUTE in header: + grid = Grid.of([int(extent) for extent in shape[1:]], header, what) + window = grid.index_window(box, 1) + if window is not None: + data, attribute = dataset.read_data_slice(group, name, (slice(None), *window)) + return (_displacement_stage(grid.sub_grid(window), data, 1, what, field_dtype),) data, attribute = read_data(group, name) if DISPLACEMENT_FIELD_ATTRIBUTE not in attribute: return decode_transform_stages(data_to_transform(data, attribute, name)) - what = f"the displacement field '{group}' of case '{name}'" return ( - _displacement_stage(Grid.of([int(extent) for extent in np.shape(data)[1:]], attribute, what), data, 1, what), + _displacement_stage( + Grid.of([int(extent) for extent in np.shape(data)[1:]], attribute, what), data, 1, what, field_dtype + ), ) diff --git a/konfai/utils/ome_zarr.py b/konfai/utils/ome_zarr.py index 94363cf4..17759b5b 100644 --- a/konfai/utils/ome_zarr.py +++ b/konfai/utils/ome_zarr.py @@ -390,6 +390,11 @@ def put(self, key: tuple, chunk: np.ndarray) -> None: self._bytes += chunk.nbytes self._trim() + @property + def held_bytes(self) -> int: + """What the cache holds right now: decoded chunks, in the bytes they take resident.""" + return self._bytes + def set_capacity(self, capacity_bytes: int) -> None: """Re-cap the cache, evicting down to the new ceiling.""" with self._lock: @@ -466,6 +471,18 @@ def _steps_to_next_use(self, key: tuple, undeclared: int) -> int: CHUNK_CACHE_FLOOR = 256 << 20 +def chunk_cache_held_bytes() -> int: + """What the decoded-chunk cache holds resident right now, or 0 with no cache. + + For an instrument reading the process's resident memory over one scope of work: the cache + outlives that scope by design (it is what a later region asks for again), so what it gained + during the scope is not the scope's own cost. A fold's probe region read VmHWM over its ten + members and charged the region 24.4 GiB, 13.2 of which was this cache filling from empty -- + and cut every region after it to 78 % of the height that would have fit. + """ + return _CHUNK_CACHE.held_bytes if _CHUNK_CACHE is not None else 0 + + def bound_chunk_cache() -> int: """Resize the decoded-chunk cache to the budget this rank published (:func:`~konfai.utils.budget.set_per_rank_budget`) and answer its capacity from now on.""" diff --git a/tests/unit/test_case_reduction.py b/tests/unit/test_case_reduction.py index c9822a57..53c467b7 100644 --- a/tests/unit/test_case_reduction.py +++ b/tests/unit/test_case_reduction.py @@ -56,7 +56,7 @@ Transform, resolve_operator, ) -from konfai.utils.budget import BUDGET_SHARES +from konfai.utils.budget import BUDGET_SHARES, budget_share from konfai.utils.dataset import Attribute, Dataset from konfai.utils.errors import ReductionError, TransformError @@ -620,7 +620,7 @@ def test_the_peak_is_charged_at_each_side_s_own_width( assert plan.peak_bytes == member_regions * member, why -@pytest.mark.parametrize("cases", [1, 2, 3, 4, 5, 6, 7]) +@pytest.mark.parametrize("cases", [1, 2, 3, 4, 5, 6, 7, 10, 16]) def test_median_selects_the_middle_instead_of_sorting_the_stack(cases: int) -> None: """Up to five members the middle is SELECTED by a network of element-wise min/max; past that a sort finds it. The values are the same to the bit either way -- ``torch.quantile`` is the @@ -635,7 +635,7 @@ def test_median_selects_the_middle_instead_of_sorting_the_stack(cases: int) -> N folded = Median()(members) assert torch.equal(folded, torch.quantile(torch.stack(members, dim=0), 0.5, dim=0)) - assert Median().working_multiple_for(cases) == {1: 1.0, 2: 1.5, 3: 1.0, 4: 2.5, 5: 1.5}.get(cases, 4.0) + assert Median().working_multiple_for(cases) == {1: 1.0, 2: 1.5, 3: 1.0, 4: 2.5, 5: 1.5}.get(cases, 1.8) @pytest.mark.parametrize( @@ -697,6 +697,23 @@ def _vote_by_sorting(tensors: list[torch.Tensor]) -> torch.Tensor: return best +def test_median_keeps_integer_members_narrow_and_holds_a_window_not_a_stack() -> None: + """Ten uint16 regions are folded without ten float32 copies of them and without a sorted stack. + + A sort along the case axis returns int64 indices, eight bytes an element whatever the members + weigh: ten uint16 regions sorted at 6.0x their own size. The window holds k + 1 float32 buffers + and nothing else -- 1.8x, measured on a 293 MiB member -- and the members stay uint16. + """ + torch.manual_seed(5) + members = [torch.randint(0, 60000, (1, 1, 8, 96, 96), dtype=torch.int32).to(torch.uint16) for _ in range(10)] + folded = Median()(members) + ranked = torch.stack([member.float() for member in members], dim=0).sort(dim=0).values + assert torch.equal(folded, torch.lerp(ranked[4], ranked[5], 0.5)), "the window selects what the sort ranks" + assert folded.dtype is torch.float32 + assert all(member.dtype is torch.uint16 for member in members), "the members were not widened in place" + assert Median().working_multiple_for(10) == Median._WINDOW_MULTIPLE < Median.working_multiple + 1 + + @pytest.mark.parametrize("cases", [2, 3, 4, 6, 7]) @pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.int32, torch.float32]) def test_vote_counts_every_candidate_instead_of_sorting_the_stack(cases: int, dtype: torch.dtype) -> None: @@ -894,6 +911,27 @@ def _engine_for_refit(tmp_path: Path, budget: float, rows: int) -> CaseReduction return engine +def test_the_probe_is_short_so_its_own_overshoot_cannot_be_the_kill(tmp_path: Path) -> None: + """The one region no measurement can bound is the probe, because it runs before any. + + At the planned height a probe over registration fields held 1.5x its allowance three times + out of three, and at an `auto` budget of 77 GiB on a 122 GiB host the probe alone reached 90 + GiB and the host went down with nothing measured. The host gives nothing to catch: the kernel + kills. So the probe walks a quarter of the planned height, and every region after it walks + what the projection allows. + """ + engine = _run(tmp_path / "probe", [], Reduce(operator="Mean", output="t"), [])[0] + engine._budget_bytes = 1 << 30 + engine.slab_rows = 100 + walked = [] + engine._fold = lambda region: walked.append(region[0].stop - region[0].start) or torch.zeros(1) # type: ignore[method-assign] + engine._open_meter = lambda: HeldMeter(lambda: 1, 0) # type: ignore[method-assign] + list(engine._folds([1000, 10, 6], measure=True)) + assert walked[0] == int(100 * CaseReduction._PROBE_SHARE), "the probe is a quarter of the planned height" + assert all(rows == 100 for rows in walked[1:-1]), "and the rest walk the planned height once it fits" + assert sum(walked) == 1000, "every row is folded exactly once" + + def test_the_first_region_is_the_probe_and_only_ever_shortens_the_rest(tmp_path: Path) -> None: """What the plan priced is a model; what the first region held is a fact. @@ -906,7 +944,9 @@ def test_the_first_region_is_the_probe_and_only_ever_shortens_the_rest(tmp_path: routes at once: which instrument answered is the meter's business and not the fold's. """ budget = 1 << 30 - allowed = budget * CaseReduction._MEASURED_MARGIN + # Less the cache's share: the meter does not count the decoded-chunk cache, so the comparison + # does not offer it either -- see test_the_allowance_leaves_the_chunk_cache_its_share. + allowed = (budget - (budget_share("cache", budget) or 0.0)) * CaseReduction._MEASURED_MARGIN def engine_holding(name: str, held: int) -> CaseReduction: engine = _run(tmp_path / name, [], Reduce(operator="Mean", output="t"), [])[0] @@ -914,16 +954,27 @@ def engine_holding(name: str, held: int) -> CaseReduction: engine.slab_rows = 100 return engine - # Held twice what the declaration allows: the rest are cut to half the height. + # A full-height probe that held twice what the declaration allows: the rest are cut to half. engine = engine_holding("over", 0) engine._refit_to_measurement(HeldMeter(lambda: int(allowed * 2), 0), 100, [1000, 10, 6]) assert engine.slab_rows == 50 + # A QUARTER-height probe that held half the allowance: scaled to the planned height that is + # twice the allowance, and the rest are cut to half. The probe is short so that its own + # overshoot cannot reach the host; what it measures is projected before it is judged. + engine = engine_holding("short-probe", 0) + engine._refit_to_measurement(HeldMeter(lambda: int(allowed * 0.5), 0), 25, [1000, 10, 6]) + assert engine.slab_rows == 50 + # Held less than the declaration: nothing moves. A measurement is never a licence to spend a # budget the sizing declined to spend, and the region that set the peak has already run. engine = engine_holding("under", 0) engine._refit_to_measurement(HeldMeter(lambda: int(allowed * 0.25), 0), 100, [1000, 10, 6]) assert engine.slab_rows == 100, "a region that fits must not make the next one taller" + # ... and a short probe under ITS share does not either, once projected. + engine = engine_holding("under-short", 0) + engine._refit_to_measurement(HeldMeter(lambda: int(allowed * 0.2), 0), 25, [1000, 10, 6]) + assert engine.slab_rows == 100 # The baseline is subtracted: the same peak over a higher starting point held less. engine = engine_holding("baseline", 0) @@ -964,6 +1015,57 @@ def test_a_host_chain_gets_a_meter_and_a_kernel_without_one_gets_none(monkeypatc assert open_held_meter(None) is None +def test_the_allowance_leaves_the_chunk_cache_its_share(tmp_path: Path) -> None: + """The meter's reading and the figure it is judged against must cover the same bytes. + + The meter stopped counting the decoded-chunk cache (it outlives the region, and charging the + region for it cut every region after the probe). Judged against the whole budget, a reading + that excludes the cache would let the cache be spent twice: once inside the allowance the + regions may fill, and again by the cache itself -- which is how a run held 1.09x what it + declared. The allowance comes down by exactly the cache's share. + """ + budget = 1 << 30 + engine = _run(tmp_path, [], Reduce(operator="Mean", output="t"), [])[0] + engine._budget_bytes = budget + engine.slab_rows = 100 + + cache = budget_share("cache", budget) or 0.0 + assert cache > 0, "the fixture only says anything where the cache has a share" + # A probe holding just under the OLD allowance (the whole budget) and above the new one. + held = int((budget - cache / 2) * CaseReduction._MEASURED_MARGIN) + engine._refit_to_measurement(HeldMeter(lambda: held, 0), 100, [1000, 10, 6]) + assert engine.slab_rows < 100, "a region filling the cache's share as well is cut" + + engine.slab_rows = 100 + fits = int((budget - cache) * CaseReduction._MEASURED_MARGIN) + engine._refit_to_measurement(HeldMeter(lambda: fits, 0), 100, [1000, 10, 6]) + assert engine.slab_rows == 100, "a region inside the allowance keeps the height the plan chose" + + +def test_a_host_meter_does_not_charge_the_scope_for_the_chunk_cache_it_filled(monkeypatch: pytest.MonkeyPatch) -> None: + """The decoded-chunk cache sits inside the process's resident peak and outlives any one scope. + + A fold's probe read VmHWM over its ten members and charged the region 24.4 GiB, 13.2 of which + was the cache filling from empty -- a budget line with a share of its own -- and cut every + region after it to 78 % of the height that would have fit. What the cache gained during the + scope is subtracted; what the scope held on its own is what it is charged. + """ + from konfai.utils import ome_zarr + + cache = ome_zarr._DecodedChunkCache(1 << 30) + monkeypatch.setattr(ome_zarr, "_CHUNK_CACHE", cache) + monkeypatch.setattr(patching_module, "reset_resident_peak", lambda: True) + monkeypatch.setattr(patching_module, "resident_bytes", lambda: 1_000_000) + peak = {"value": 1_000_000} + monkeypatch.setattr(patching_module, "peak_resident_bytes", lambda: peak["value"]) + + meter = open_held_meter(None) + chunk = np.ones((256, 256), np.float32) # 256 KiB decoded, kept by the cache past the scope + cache.put(("store", 0), chunk) + peak["value"] = 1_000_000 + chunk.nbytes + 100_000 # the process grew by the chunk and by the scope's own 100 KB + assert meter is not None and meter.held() == 100_000, "the cache's growth is not the scope's cost" + + def test_the_folds_a_stat_pass_keeps_come_out_of_the_regions_share(tmp_path: Path) -> None: """A kept fold is memory the regions are not holding, so it comes out of the same share. diff --git a/tests/unit/test_itk_transforms.py b/tests/unit/test_itk_transforms.py index c11e45ce..4a0a2067 100644 --- a/tests/unit/test_itk_transforms.py +++ b/tests/unit/test_itk_transforms.py @@ -119,8 +119,14 @@ def test_decoding_a_displacement_field_copies_it_once() -> None: def test_a_stored_displacement_entry_decodes_to_the_stage_the_image_route_gives(tmp_path) -> None: - """Straight from the store's array, on the grid its attributes describe: the same stage, bit - for bit, as reading the entry as a transform and decoding that.""" + """Straight from the store's array, on the grid its attributes describe: the same stage, value + for value, as reading the entry as a transform and decoding that. + + Not byte for byte, because the two routes do not hold it at the same WIDTH and should not. The + direct read holds a field at the width the store holds it at; the transform route has no store + to ask -- SimpleITK's field is float64 whatever was written -- so it answers the same numbers + widened. Equal as maps, and the direct one is the one that fits in half the memory. + """ from konfai.utils.dataset import Attribute, Dataset from konfai.utils.ITK import decode_transform_stages, read_transform_stages @@ -139,7 +145,10 @@ def test_a_stored_displacement_entry_decodes_to_the_stage_the_image_route_gives( assert direct.grid.size_zyx == through_itk.grid.size_zyx for axis in ("origin_xyz", "spacing_xyz", "direction_xyz"): np.testing.assert_array_equal(getattr(direct.grid, axis), getattr(through_itk.grid, axis)) - assert direct.values.tobytes() == through_itk.values.tobytes() + # Exactly equal, not close: widening a float32 is lossless, so no tolerance is called for. + np.testing.assert_array_equal(direct.values.astype(np.float64), through_itk.values) + assert direct.values.dtype == np.float32, "a float32 store is not widened to be held" + assert through_itk.values.dtype == np.float64, "and a SimpleITK field has no other width" def test_a_store_serving_transforms_alone_still_decodes() -> None: diff --git a/tests/unit/test_resample_to_reference.py b/tests/unit/test_resample_to_reference.py index 619134da..f952a841 100644 --- a/tests/unit/test_resample_to_reference.py +++ b/tests/unit/test_resample_to_reference.py @@ -438,6 +438,32 @@ def test_a_case_that_never_meets_the_reference_is_refused(tmp_path: Path) -> Non _stage(dataset).transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), source) +def test_a_case_apart_says_nothing_about_coverage_when_a_field_bridges_it(tmp_path: Path) -> None: + """The plan's note and the refusal read the same coverage, so they take the same exemption. + + A field's reach is known only when its values are read, and the plan reads none -- it prices + the field at the identity. The refusal has always stood aside for that (bridging two frames is + what a field is for). The note did not, and said the opposite of the truth: on a ten-member + ExaSPIM build whose fields bridge a 20 mm gap, it called one member 0.0% covered and everything + it wrote fill, while the run read that member in full and the template carried its anatomy. + """ + 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) + ) + fields = Dataset(tmp_path / "dvf", "mha") + fields.write("DVF", _CASE, np.zeros((3, *_SOURCE_SPATIAL), np.float32), source) + stage = Resample(reference=_CASE, reference_group="Reference", field=f"{tmp_path / 'dvf'}:mha", field_group="DVF") + stage.set_datasets([dataset]) + + # The same geometry the refusal above raises on, now with a field configured. + stage.transform_shape("Case", _CASE, list(_SOURCE_SPATIAL), source) + note = stage.plan_note("Case_out", _CASE, list(_SOURCE_SPATIAL), source) + assert note is None or "covers" not in note, f"a field-bridged case must not be given a coverage: {note}" + + def test_an_unknown_entry_is_refused(dataset: Dataset) -> None: stage = Resample(reference="NOT_THERE", reference_group="Reference") stage.set_datasets([dataset]) diff --git a/tests/unit/test_resample_transform.py b/tests/unit/test_resample_transform.py index b0091d2c..89f34c86 100644 --- a/tests/unit/test_resample_transform.py +++ b/tests/unit/test_resample_transform.py @@ -21,12 +21,14 @@ """ import sys +from dataclasses import replace from pathlib import Path import konfai.data.transform as transform_module import numpy as np import pytest import torch +from konfai.data.geometry import DisplacementStage, Grid from konfai.data.transform import ( LocalityKind, RegionContext, @@ -34,7 +36,7 @@ _optional_image_filler, _SitkInput, ) -from konfai.utils.dataset import Attribute +from konfai.utils.dataset import DISPLACEMENT_FIELD_ATTRIBUTE, Attribute from konfai.utils.errors import TransformError sitk = pytest.importorskip("SimpleITK") @@ -156,6 +158,84 @@ def _golden_resamples() -> dict[str, np.ndarray]: return resampled +@pytest.mark.parametrize("interpolation", ["linear", "nearest"]) +def test_a_float32_field_is_warped_and_says_what_the_float64_transform_says( + monkeypatch: pytest.MonkeyPatch, interpolation: str +) -> None: + """A map that is one float32 field on the output grid does not go through a float64 transform. + + ``sitk.Warp`` is templated on the field's own type; ``DisplacementFieldTransform`` is not -- + SimpleITK's hierarchy is ``TransformBaseTemplate``, so a field read as float32 is cast + to float64 to be held and interleaved into a vector image at that width. On a 122-row native + ExaSPIM region that is 6.9 GiB built in 5.0 s, 39% of what the whole member-region costs. + + The claim is that Warp over the float32 values and the transform over those same values widened + (a lossless widening: every float32 is a float64) write the very same bytes. Both are run over + one captured region, so nothing but the route differs. + """ + image = _image(oblique=False) + volume = _phantom() if interpolation == "linear" else _label_map() + stage = _stage(image, _field(image), interpolation=interpolation) + + captured: list[tuple] = [] + real_resample = transform_module._resample_with_sitk + + def capture(payload, region, source, stages, starts, mode, fill, sitk_input=None): + captured.append((payload, region, source, stages, starts, mode, fill)) + return real_resample(payload, region, source, stages, starts, mode, fill, sitk_input) + + monkeypatch.setattr(transform_module, "_resample_with_sitk", capture) + stage(CASE, torch.from_numpy(volume).unsqueeze(0), Attribute(_attribute(image))) + assert captured, "the host route is what this pins" + payload, region, source, stages, starts, mode, fill = captured[0] + assert len(stages) == 1 and isinstance(stages[0], DisplacementStage) + + narrowed = replace(stages[0], values=stages[0].values.astype(np.float32)) + widened = replace(narrowed, values=narrowed.values.astype(np.float64)) + assert transform_module._warp_field_float32((narrowed,), region) is not None, "Warp is for exactly this" + assert transform_module._warp_field_float32((widened,), region) is None, "and never for float64" + + fast = real_resample(payload, region, source, (narrowed,), starts, mode, fill) + exact = real_resample(payload, region, source, (widened,), starts, mode, fill) + assert torch.equal(fast, exact), "float32 Warp and the float64 transform are the same map" + + +def test_a_float32_store_is_held_at_float32_with_no_flag_to_set() -> None: + """The width a field is held at comes from the STORE, not from a precision the user declares. + + A DVF normally sits on disk in float32 -- it is how ITK writes one -- and widening it to + float64 buys a copy of twice the bytes to say exactly the same thing. Held at its own width it + is lossless, costs no flag, and is what lets ``sitk.Warp`` apply it without a float64 transform + to hold it. The default ceiling narrows nothing: a genuinely float64 field stays float64. + """ + from konfai.utils.ITK import _displacement_stage + + grid = Grid(size_zyx=(4, 5, 6), origin_xyz=np.zeros(3), spacing_xyz=np.ones(3), direction_xyz=np.eye(3)) + stored32 = np.zeros((3, 4, 5, 6), np.float32) + stored64 = np.zeros((3, 4, 5, 6), np.float64) + + assert _displacement_stage(grid, stored32, 1, "f", np.float64).values.dtype == np.float32, "not widened" + assert _displacement_stage(grid, stored64, 1, "f", np.float64).values.dtype == np.float64, "not narrowed" + # 'fast' lowers the ceiling, and there a genuinely float64 field does narrow: the declared trade. + assert _displacement_stage(grid, stored64, 1, "f", np.float32).values.dtype == np.float32, "the trade" + # A width SimpleITK has no pixel type for goes to the ceiling rather than staying as it is. + stored16 = np.zeros((3, 4, 5, 6), np.float16) + assert _displacement_stage(grid, stored16, 1, "f", np.float64).values.dtype == np.float64, "no f16 pixel type" + + +def test_a_float64_field_is_not_narrowed_into_the_warp() -> None: + """The fast route is a cheaper way to the same values, never a cheaper set of values. + + A field carrying float64 values has bits float32 does not, so narrowing it to reach + ``sitk.Warp`` would answer a different map. 'exact' promises bit-identity with + ``sitk.Resample``, and the guard keeps that promise by declining. + """ + + grid = Grid(size_zyx=(4, 5, 6), origin_xyz=np.zeros(3), spacing_xyz=np.ones(3), direction_xyz=np.eye(3)) + values = np.full((3, 4, 5, 6), np.nextafter(1.0, 2.0), dtype=np.float64) # not a float32 + assert transform_module._warp_field_float32((DisplacementStage(grid, values, 1),), grid) is None + + def test_a_stored_resample_reproduces_its_golden_output() -> None: """The values this stage produced when the fixture was stored, on the CPU's exact route. @@ -525,15 +605,16 @@ def planned(cases: int) -> Resample: volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0) first = ten("CASE_003", volume, Attribute(attribute)) - assert set(ten._stored) == {"CASE_003"} + # The cache is keyed on (case, box): a field read for one region answers for that one. + assert {key[0] for key in ten._stored} == {"CASE_003"} ten("CASE_007", volume, Attribute(attribute)) - assert set(ten._stored) == {"CASE_003", "CASE_007"}, "the last cases' stages are held, within the bound" + assert {key[0] for key in ten._stored} == {"CASE_003", "CASE_007"}, "the last cases are held, in bound" ten.stored_stage_bytes = 1 ten("CASE_005", volume, Attribute(attribute)) - assert set(ten._stored) == {"CASE_005"}, "past the bound, the slot follows the case being sampled" + assert {key[0] for key in ten._stored} == {"CASE_005"}, "past the bound, the slot follows the case" # The same values as a stage that decoded the map once and kept it. kept = planned(10) - kept._stored["CASE_003"] = kept._decode_stored("CASE_003") + kept._stored[("CASE_003", None)] = kept._decode_stored("CASE_003") torch.testing.assert_close(first, kept("CASE_003", volume, Attribute(attribute)), rtol=0.0, atol=0.0) # What a rank receives holds no stages either; it decodes what it samples. rank = pickle.loads(pickle.dumps(ten)) @@ -541,11 +622,101 @@ def planned(cases: int) -> Resample: torch.testing.assert_close(rank("CASE_003", volume, Attribute(attribute)), first, rtol=0.0, atol=0.0) +class _SlicedFieldStore: + """A store that serves a dense field group and can serve a slice of it, counting what is asked.""" + + def __init__(self, attribute: Attribute, reference: "sitk.Image | None" = None) -> None: + self.values = np.random.RandomState(4).normal(0.0, 2.0, (3, *SIZE)) + self.header = Attribute(attribute) + self.header[DISPLACEMENT_FIELD_ATTRIBUTE] = "true" + self.reference = reference + self.asked: list[tuple[int, ...]] = [] + self.whole = 0 + + def is_dataset_exist(self, group: str, name: str) -> bool: + del name + return group == "reg" or (group == "Reference" and self.reference is not None) + + def get_infos(self, group: str, name: str) -> tuple[list[int], Attribute]: + del name + if group == "Reference" and self.reference is not None: + return [1, *list(self.reference.GetSize())[::-1]], _attribute(self.reference) + return [int(extent) for extent in self.values.shape], Attribute(self.header) + + def read_data(self, group: str, name: str) -> tuple[np.ndarray, Attribute]: + del group, name + self.whole += 1 + return self.values, Attribute(self.header) + + def read_data_slice(self, group: str, name: str, slices) -> tuple[np.ndarray, Attribute]: + del group, name + block = self.values[tuple(slices)] + self.asked.append(tuple(block.shape[1:])) + return block, Attribute(self.header) + + +def test_a_stored_field_bridging_disjoint_frames_is_not_refused_as_disjoint(): + """The rigid case above, with the bridge stored as a dense field instead of a matrix. + + The plan prices a field at the identity, declared or stored, because nothing bounds one from + headers. So coverage judged through the priced map is zero for exactly the cohort the stage + exists to serve -- ten brains registered onto a template they sit 25 mm from -- and the all-fill + gate has to let a field through the way it already lets a declared one through. It did not, and + the whole build would have been refused before a byte was read. + """ + case = _image(oblique=False) + case.SetOrigin((1000.0, 0.0, 0.0)) + reference = _image(oblique=False) # origin (10, -5, 2): ~1000 mm from the case in x + source = _SlicedFieldStore(_attribute(reference), reference=reference) + source.values[:] = 0.0 + source.values[0] = 990.0 # the offset between the frames, carried in every voxel + + stage = Resample(reference="ref", reference_group="Reference", transforms={"reg": False}) + stage.set_datasets([source]) + # Does not raise, which is the whole test: the gate cannot know where a field reaches, so it + # does not pretend to. transform_shape is where it fires, before a byte is read. + assert stage.transform_shape("", "CASE_000", list(SIZE), Attribute(_attribute(case))) == list(SIZE) + # And the target is NOT the case's own grid, so the gate was really reached: the affine part of + # the priced map is the identity, and coverage judged through it is nothing at all. + assert stage.coverage("CASE_000") == 0.0 + + +def test_a_region_reads_the_window_of_a_stored_field_and_not_the_field(): + """A stored dense field is read on the box its region reaches, not whole. + + The declared route has always done this (``_field_stage``); the stored one read every voxel of + every member, so a chain the plan sized in gigabytes held tens of them beside the budget -- + a cost the sizing never saw, because a decoded transform lives outside the regions it prices. + At full resolution that is 14.5 GiB per case on disk and twice that in memory, float64 on the + way in. + + Counted at the store, not inferred: what a region asks for is the slice it asks for. + """ + attribute = _attribute(_image(oblique=False)) + source = _SlicedFieldStore(attribute) + stage = Resample(transforms={"reg": False}) + stage.set_datasets([source]) + stage.transform_shape("", "CASE_000", list(SIZE), Attribute(attribute)) + + grid = Grid.of(list(SIZE), Attribute(attribute), "the case") + source.asked.clear() + stage._stages("CASE_000", grid.sub_grid((slice(0, 2), slice(0, 2), slice(0, 2)))) + + assert source.asked, "the region read a window rather than the whole entry" + read = source.asked[-1] + assert all(got < full for got, full in zip(read, SIZE, strict=True)), ( + f"a 2x2x2 region read {list(read)} of a {list(SIZE)} field" + ) + # A second region is its own read: the cache is keyed on the box, so it cannot be handed the + # first one's window and sample outside it. + source.asked.clear() + stage._stages("CASE_000", grid.sub_grid((slice(18, 20), slice(24, 26), slice(30, 32)))) + assert (source.asked and source.asked[-1] != read) or len(stage._stored) == 2 + + def test_the_slabbed_walk_lands_each_slab_in_the_one_output(monkeypatch: pytest.MonkeyPatch) -> None: """Above the walk budget the general path gathers slab by slab, and each slab is written into the output as it lands: bit for bit the single pass, with no parts held for a cat.""" - from konfai.data import transform as transform_module - image = _image(oblique=True) attribute = _attribute(image) volume = torch.from_numpy(sitk.GetArrayFromImage(image)).unsqueeze(0)